235 lines
6.1 KiB
Go
235 lines
6.1 KiB
Go
package marian
|
|
|
|
import (
|
|
"context"
|
|
"math"
|
|
"runtime"
|
|
"slices"
|
|
"sync"
|
|
)
|
|
|
|
// beam is a partial translation being extended by beam search.
|
|
type beam struct {
|
|
tokens []int // starts with the decoder start token
|
|
score float64
|
|
// Keys and values of the decoder's self-attention, one slice per layer
|
|
// holding len(tokens) rows: each step only computes its own row.
|
|
keys, values [][]float32
|
|
}
|
|
|
|
type hypothesis struct {
|
|
tokens []int
|
|
score float64 // length-normalised
|
|
}
|
|
|
|
type candidate struct {
|
|
parent, token int
|
|
score float64
|
|
}
|
|
|
|
// beamSearch generates the translation, as Hugging Face's generate() does
|
|
// with the model's defaults: num_beams beams, length penalty 1, the
|
|
// padding token never produced.
|
|
func (m *Model) beamSearch(ctx context.Context, enc []float32, srcLen int) ([]int, error) {
|
|
d := m.cfg.DModel
|
|
numBeams := max(m.cfg.NumBeams, 1)
|
|
crossK := make([][]float32, len(m.decoder))
|
|
crossV := make([][]float32, len(m.decoder))
|
|
for i, l := range m.decoder {
|
|
crossK[i] = make([]float32, srcLen*d)
|
|
crossV[i] = make([]float32, srcLen*d)
|
|
l.cross.k.rows(crossK[i], enc, srcLen)
|
|
l.cross.v.rows(crossV[i], enc, srcLen)
|
|
}
|
|
beams := []beam{{
|
|
tokens: []int{m.cfg.DecoderStart},
|
|
keys: make([][]float32, len(m.decoder)),
|
|
values: make([][]float32, len(m.decoder)),
|
|
}}
|
|
maxLen := min(m.cfg.MaxPositions-1, 2*srcLen+10)
|
|
var done []hypothesis
|
|
for len(beams[0].tokens) <= maxLen {
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
hidden := make([][]float32, len(beams))
|
|
for i := range beams {
|
|
hidden[i] = m.decodeStep(&beams[i], crossK, crossV, srcLen)
|
|
}
|
|
logProbs := m.logProbs(hidden)
|
|
var cands []candidate
|
|
for i, lp := range logProbs {
|
|
for _, t := range topK(lp, 2*numBeams) {
|
|
cands = append(cands, candidate{parent: i, token: t, score: beams[i].score + float64(lp[t])})
|
|
}
|
|
}
|
|
slices.SortStableFunc(cands, func(a, b candidate) int {
|
|
switch {
|
|
case a.score > b.score:
|
|
return -1
|
|
case a.score < b.score:
|
|
return 1
|
|
}
|
|
return 0
|
|
})
|
|
var next []beam
|
|
for rank, c := range cands {
|
|
parent := beams[c.parent]
|
|
if c.token == m.cfg.EOSToken {
|
|
if rank < numBeams {
|
|
done = append(done, hypothesis{tokens: slices.Clone(parent.tokens), score: c.score / float64(len(parent.tokens))})
|
|
}
|
|
continue
|
|
}
|
|
child := beam{
|
|
tokens: append(slices.Clone(parent.tokens), c.token),
|
|
score: c.score,
|
|
keys: make([][]float32, len(parent.keys)),
|
|
values: make([][]float32, len(parent.values)),
|
|
}
|
|
for l := range parent.keys {
|
|
child.keys[l] = slices.Clone(parent.keys[l])
|
|
child.values[l] = slices.Clone(parent.values[l])
|
|
}
|
|
next = append(next, child)
|
|
if len(next) == numBeams {
|
|
break
|
|
}
|
|
}
|
|
slices.SortStableFunc(done, func(a, b hypothesis) int {
|
|
switch {
|
|
case a.score > b.score:
|
|
return -1
|
|
case a.score < b.score:
|
|
return 1
|
|
}
|
|
return 0
|
|
})
|
|
if len(done) > numBeams {
|
|
done = done[:numBeams]
|
|
}
|
|
if len(next) == 0 {
|
|
break
|
|
}
|
|
// Finished: no running beam can still beat the worst kept
|
|
// hypothesis (scores only decrease as tokens are added).
|
|
if len(done) == numBeams && next[0].score/float64(len(next[0].tokens)) <= done[numBeams-1].score {
|
|
break
|
|
}
|
|
beams = next
|
|
}
|
|
if len(done) == 0 {
|
|
return beams[0].tokens[1:], nil
|
|
}
|
|
return done[0].tokens[1:], nil
|
|
}
|
|
|
|
// decodeStep runs the decoder for the beam's last token and returns its
|
|
// output vector.
|
|
func (m *Model) decodeStep(b *beam, crossK, crossV [][]float32, srcLen int) []float32 {
|
|
d := m.cfg.DModel
|
|
pos := len(b.tokens) - 1
|
|
tok := b.tokens[pos]
|
|
x := make([]float32, d)
|
|
emb := m.embed[tok*d : (tok+1)*d]
|
|
p := m.decPos[pos*d : (pos+1)*d]
|
|
for j := range x {
|
|
x[j] = emb[j]*m.embedScale + p[j]
|
|
}
|
|
q := make([]float32, d)
|
|
k := make([]float32, d)
|
|
v := make([]float32, d)
|
|
ctxv := make([]float32, d)
|
|
tmp := make([]float32, d)
|
|
scores := make([]float32, max(pos+1, srcLen))
|
|
for i, l := range m.decoder {
|
|
l.self.q.apply(q, x)
|
|
l.self.k.apply(k, x)
|
|
l.self.v.apply(v, x)
|
|
b.keys[i] = append(b.keys[i], k...)
|
|
b.values[i] = append(b.values[i], v...)
|
|
attend(ctxv, q, b.keys[i], b.values[i], pos+1, l.self.heads, scores)
|
|
l.self.o.apply(tmp, ctxv)
|
|
addNorm(x, tmp, l.selfNorm, 1)
|
|
|
|
l.cross.q.apply(q, x)
|
|
attend(ctxv, q, crossK[i], crossV[i], srcLen, l.cross.heads, scores)
|
|
l.cross.o.apply(tmp, ctxv)
|
|
addNorm(x, tmp, l.crossNorm, 1)
|
|
|
|
hidden := make([]float32, l.fc1.out)
|
|
l.fc1.apply(hidden, x)
|
|
for j := range hidden {
|
|
hidden[j] = m.activation(hidden[j])
|
|
}
|
|
l.fc2.apply(tmp, hidden)
|
|
addNorm(x, tmp, l.finalLayerNorm, 1)
|
|
}
|
|
return x
|
|
}
|
|
|
|
// logProbs projects each beam's vector on the vocabulary (the embedding
|
|
// matrix, shared with the output layer) and returns log-softmax scores.
|
|
// This is the largest product of a step, so it is split across CPUs.
|
|
func (m *Model) logProbs(hidden [][]float32) [][]float32 {
|
|
d, vocab := m.cfg.DModel, m.cfg.VocabSize
|
|
out := make([][]float32, len(hidden))
|
|
for i := range out {
|
|
out[i] = make([]float32, vocab)
|
|
}
|
|
workers := runtime.GOMAXPROCS(0)
|
|
chunk := (vocab + workers - 1) / workers
|
|
var wg sync.WaitGroup
|
|
for lo := 0; lo < vocab; lo += chunk {
|
|
hi := min(lo+chunk, vocab)
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
for t := lo; t < hi; t++ {
|
|
row := m.embed[t*d : (t+1)*d]
|
|
for i, h := range hidden {
|
|
out[i][t] = dot(h, row) + m.logitsBias[t]
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
for _, logits := range out {
|
|
if m.cfg.PadToken >= 0 && m.cfg.PadToken < vocab {
|
|
logits[m.cfg.PadToken] = float32(math.Inf(-1))
|
|
}
|
|
maxLogit := float32(math.Inf(-1))
|
|
for _, l := range logits {
|
|
maxLogit = max(maxLogit, l)
|
|
}
|
|
var sum float64
|
|
for _, l := range logits {
|
|
sum += math.Exp(float64(l - maxLogit))
|
|
}
|
|
norm := maxLogit + float32(math.Log(sum))
|
|
for t := range logits {
|
|
logits[t] -= norm
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// topK returns the indexes of the k largest values, largest first.
|
|
func topK(values []float32, k int) []int {
|
|
idx := make([]int, 0, k+1)
|
|
for i, v := range values {
|
|
if len(idx) == k && v <= values[idx[k-1]] {
|
|
continue
|
|
}
|
|
at := len(idx)
|
|
for at > 0 && values[idx[at-1]] < v {
|
|
at--
|
|
}
|
|
idx = slices.Insert(idx, at, i)
|
|
if len(idx) > k {
|
|
idx = idx[:k]
|
|
}
|
|
}
|
|
return idx
|
|
}
|