Skip to content

Commit

Permalink
apply the suggestions of #56
Browse files Browse the repository at this point in the history
  • Loading branch information
ynqa committed Jul 26, 2022
1 parent dc41a6e commit 2dcde9e
Showing 1 changed file with 15 additions and 6 deletions.
21 changes: 15 additions & 6 deletions pkg/model/word2vec/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -77,15 +77,23 @@ func (mod *skipGram) trainOne(
}
}

type cbowToken struct {
agg []float64
tmp []float64
}

type cbow struct {
ch chan []float64
ch chan cbowToken
window int
}

func newCbow(opts Options) mod {
ch := make(chan []float64, opts.Goroutines*2)
ch := make(chan cbowToken, opts.Goroutines)
for i := 0; i < opts.Goroutines; i++ {
ch <- make([]float64, opts.Dim)
ch <- cbowToken{
agg: make([]float64, opts.Dim),
tmp: make([]float64, opts.Dim),
}
}
return &cbow{
ch: ch,
Expand All @@ -100,10 +108,11 @@ func (mod *cbow) trainOne(
param *matrix.Matrix,
optimizer optimizer,
) {
agg, tmp := <-mod.ch, <-mod.ch
token := <-mod.ch
agg, tmp := token.agg, token.tmp
defer func() {
mod.ch <- agg
mod.ch <- tmp
token := cbowToken{agg, tmp}
mod.ch <- token
}()
for i := 0; i < len(agg); i++ {
agg[i], tmp[i] = 0, 0
Expand Down

0 comments on commit 2dcde9e

Please sign in to comment.