feat: add a translate stack
This commit is contained in:
Kharec committed
1 parent
5cf16588c1
commit
cb356ff8c5
10 files changed
+2996
No files matched your search
@@ -0,0 +1,234 @@
|
||||
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
|
||||
}
|
||||
Reference in new issue
Block a user