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,376 @@
|
||||
// Package marian runs MarianMT translation models (the OPUS-MT models of
|
||||
// the University of Helsinki, in their Hugging Face form) on the CPU, in
|
||||
// plain Go: a SentencePiece tokenizer, the transformer encoder-decoder and
|
||||
// beam search. A model directory holds config.json, model.safetensors,
|
||||
// source.spm, target.spm and vocab.json.
|
||||
package marian
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
// Files are the files a model directory must contain.
|
||||
var Files = []string{"config.json", "model.safetensors", "source.spm", "target.spm", "vocab.json"}
|
||||
|
||||
type config struct {
|
||||
DModel int `json:"d_model"`
|
||||
EncoderLayers int `json:"encoder_layers"`
|
||||
DecoderLayers int `json:"decoder_layers"`
|
||||
EncoderHeads int `json:"encoder_attention_heads"`
|
||||
DecoderHeads int `json:"decoder_attention_heads"`
|
||||
Activation string `json:"activation_function"`
|
||||
ScaleEmbedding bool `json:"scale_embedding"`
|
||||
NormalizeBefore bool `json:"normalize_before"`
|
||||
MaxPositions int `json:"max_position_embeddings"`
|
||||
VocabSize int `json:"vocab_size"`
|
||||
PadToken int `json:"pad_token_id"`
|
||||
EOSToken int `json:"eos_token_id"`
|
||||
DecoderStart int `json:"decoder_start_token_id"`
|
||||
NumBeams int `json:"num_beams"`
|
||||
}
|
||||
|
||||
type linear struct {
|
||||
w, b []float32
|
||||
in, out int
|
||||
}
|
||||
|
||||
type layerNorm struct{ g, b []float32 }
|
||||
|
||||
type attention struct {
|
||||
q, k, v, o linear
|
||||
heads int
|
||||
}
|
||||
|
||||
type encoderLayer struct {
|
||||
self attention
|
||||
selfNorm layerNorm
|
||||
fc1, fc2 linear
|
||||
finalLayerNorm layerNorm
|
||||
}
|
||||
|
||||
type decoderLayer struct {
|
||||
self, cross attention
|
||||
selfNorm, crossNorm layerNorm
|
||||
fc1, fc2 linear
|
||||
finalLayerNorm layerNorm
|
||||
}
|
||||
|
||||
// Model is a loaded translation model. It is read-only once loaded:
|
||||
// Translate may be called from several goroutines at once.
|
||||
type Model struct {
|
||||
cfg config
|
||||
embed []float32 // vocab × d, shared by encoder, decoder and output
|
||||
encPos []float32 // positions × d
|
||||
decPos []float32
|
||||
logitsBias []float32
|
||||
encoder []encoderLayer
|
||||
decoder []decoderLayer
|
||||
source *sentencePiece
|
||||
target *sentencePiece
|
||||
ids map[string]int
|
||||
pieces []string
|
||||
unk int
|
||||
embedScale float32
|
||||
activation func(float32) float32
|
||||
}
|
||||
|
||||
// Load reads the model in dir.
|
||||
func Load(dir string) (*Model, error) {
|
||||
var cfg config
|
||||
data, err := os.ReadFile(filepath.Join(dir, "config.json"))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := json.Unmarshal(data, &cfg); err != nil {
|
||||
return nil, fmt.Errorf("reading config.json: %w", err)
|
||||
}
|
||||
if cfg.DModel == 0 || cfg.EncoderHeads == 0 || cfg.DecoderHeads == 0 || cfg.DModel%cfg.EncoderHeads != 0 || cfg.DModel%cfg.DecoderHeads != 0 {
|
||||
return nil, fmt.Errorf("unsupported model dimensions in config.json")
|
||||
}
|
||||
if cfg.NormalizeBefore {
|
||||
return nil, fmt.Errorf("unsupported model: pre-norm layers")
|
||||
}
|
||||
m := &Model{cfg: cfg, embedScale: 1}
|
||||
if cfg.ScaleEmbedding {
|
||||
m.embedScale = float32(math.Sqrt(float64(cfg.DModel)))
|
||||
}
|
||||
switch cfg.Activation {
|
||||
case "swish", "silu":
|
||||
m.activation = func(x float32) float32 { return x / (1 + float32(math.Exp(float64(-x)))) }
|
||||
case "relu":
|
||||
m.activation = func(x float32) float32 { return max(x, 0) }
|
||||
case "gelu":
|
||||
m.activation = func(x float32) float32 { return 0.5 * x * (1 + float32(math.Erf(float64(x)/math.Sqrt2))) }
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported activation %q", cfg.Activation)
|
||||
}
|
||||
|
||||
if m.source, err = loadSentencePiece(filepath.Join(dir, "source.spm")); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if m.target, err = loadSentencePiece(filepath.Join(dir, "target.spm")); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if data, err = os.ReadFile(filepath.Join(dir, "vocab.json")); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := json.Unmarshal(data, &m.ids); err != nil {
|
||||
return nil, fmt.Errorf("reading vocab.json: %w", err)
|
||||
}
|
||||
m.pieces = make([]string, cfg.VocabSize)
|
||||
for p, id := range m.ids {
|
||||
if id >= 0 && id < len(m.pieces) {
|
||||
m.pieces[id] = p
|
||||
}
|
||||
}
|
||||
m.unk = m.ids["<unk>"]
|
||||
|
||||
tensors, err := readSafetensors(filepath.Join(dir, "model.safetensors"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("reading model.safetensors: %w", err)
|
||||
}
|
||||
if err := m.bind(tensors); err != nil {
|
||||
return nil, fmt.Errorf("reading model.safetensors: %w", err)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// bind takes the weights the model needs from the tensors, checking their
|
||||
// shapes.
|
||||
func (m *Model) bind(t map[string]tensor) error {
|
||||
d := m.cfg.DModel
|
||||
var missing error
|
||||
get := func(name string, shape ...int) []float32 {
|
||||
x, ok := t[name]
|
||||
if !ok {
|
||||
if missing == nil {
|
||||
missing = fmt.Errorf("missing tensor %s", name)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
n := 1
|
||||
for _, s := range shape {
|
||||
n *= s
|
||||
}
|
||||
if len(x.data) != n && missing == nil {
|
||||
missing = fmt.Errorf("tensor %s has shape %v", name, x.shape)
|
||||
}
|
||||
return x.data
|
||||
}
|
||||
first := func(names ...string) string {
|
||||
for _, n := range names {
|
||||
if _, ok := t[n]; ok {
|
||||
return n
|
||||
}
|
||||
}
|
||||
return names[0]
|
||||
}
|
||||
lin := func(prefix string, in, out int) linear {
|
||||
return linear{w: get(prefix+".weight", out, in), b: get(prefix+".bias", out), in: in, out: out}
|
||||
}
|
||||
norm := func(prefix string) layerNorm {
|
||||
return layerNorm{g: get(prefix+".weight", d), b: get(prefix+".bias", d)}
|
||||
}
|
||||
attn := func(prefix string, heads int) attention {
|
||||
return attention{
|
||||
q: lin(prefix+".q_proj", d, d), k: lin(prefix+".k_proj", d, d),
|
||||
v: lin(prefix+".v_proj", d, d), o: lin(prefix+".out_proj", d, d), heads: heads,
|
||||
}
|
||||
}
|
||||
ffn := func(prefix string) int {
|
||||
if x, ok := t[prefix+".fc1.weight"]; ok && len(x.shape) == 2 {
|
||||
return x.shape[0]
|
||||
}
|
||||
return 4 * d
|
||||
}
|
||||
|
||||
v := m.cfg.VocabSize
|
||||
m.embed = get(first("model.shared.weight", "model.decoder.embed_tokens.weight", "model.encoder.embed_tokens.weight"), v, d)
|
||||
m.encPos = get("model.encoder.embed_positions.weight", m.cfg.MaxPositions, d)
|
||||
m.decPos = get("model.decoder.embed_positions.weight", m.cfg.MaxPositions, d)
|
||||
m.logitsBias = get("final_logits_bias", 1, v)
|
||||
for i := range m.cfg.EncoderLayers {
|
||||
p := fmt.Sprintf("model.encoder.layers.%d", i)
|
||||
f := ffn(p)
|
||||
m.encoder = append(m.encoder, encoderLayer{
|
||||
self: attn(p+".self_attn", m.cfg.EncoderHeads), selfNorm: norm(p + ".self_attn_layer_norm"),
|
||||
fc1: lin(p+".fc1", d, f), fc2: lin(p+".fc2", f, d), finalLayerNorm: norm(p + ".final_layer_norm"),
|
||||
})
|
||||
}
|
||||
for i := range m.cfg.DecoderLayers {
|
||||
p := fmt.Sprintf("model.decoder.layers.%d", i)
|
||||
f := ffn(p)
|
||||
m.decoder = append(m.decoder, decoderLayer{
|
||||
self: attn(p+".self_attn", m.cfg.DecoderHeads), selfNorm: norm(p + ".self_attn_layer_norm"),
|
||||
cross: attn(p+".encoder_attn", m.cfg.DecoderHeads), crossNorm: norm(p + ".encoder_attn_layer_norm"),
|
||||
fc1: lin(p+".fc1", d, f), fc2: lin(p+".fc2", f, d), finalLayerNorm: norm(p + ".final_layer_norm"),
|
||||
})
|
||||
}
|
||||
return missing
|
||||
}
|
||||
|
||||
// maxSourceTokens leaves room for the end-of-sentence token.
|
||||
func (m *Model) maxSourceTokens() int { return m.cfg.MaxPositions - 1 }
|
||||
|
||||
// Translate translates one text. Long texts are better split into
|
||||
// sentences first: the models are trained on sentences.
|
||||
func (m *Model) Translate(ctx context.Context, text string) (string, error) {
|
||||
pieces := m.source.encode(text)
|
||||
if len(pieces) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
ids := make([]int, 0, len(pieces)+1)
|
||||
for _, p := range pieces {
|
||||
id, ok := m.ids[p]
|
||||
if !ok {
|
||||
id = m.unk
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if len(ids) > m.maxSourceTokens() {
|
||||
ids = ids[:m.maxSourceTokens()]
|
||||
}
|
||||
ids = append(ids, m.cfg.EOSToken)
|
||||
out, err := m.beamSearch(ctx, m.encode(ids), len(ids))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
words := make([]string, 0, len(out))
|
||||
for _, id := range out {
|
||||
if id == m.cfg.EOSToken || id == m.cfg.PadToken || id == m.unk || id < 0 || id >= len(m.pieces) {
|
||||
continue
|
||||
}
|
||||
words = append(words, m.pieces[id])
|
||||
}
|
||||
return decodePieces(words), nil
|
||||
}
|
||||
|
||||
// encode runs the encoder: one d-vector per source token.
|
||||
func (m *Model) encode(ids []int) []float32 {
|
||||
d := m.cfg.DModel
|
||||
n := len(ids)
|
||||
x := make([]float32, n*d)
|
||||
for i, id := range ids {
|
||||
row := x[i*d : (i+1)*d]
|
||||
emb := m.embed[id*d : (id+1)*d]
|
||||
pos := m.encPos[i*d : (i+1)*d]
|
||||
for j := range row {
|
||||
row[j] = emb[j]*m.embedScale + pos[j]
|
||||
}
|
||||
}
|
||||
q := make([]float32, n*d)
|
||||
k := make([]float32, n*d)
|
||||
v := make([]float32, n*d)
|
||||
ctxv := make([]float32, n*d)
|
||||
tmp := make([]float32, n*d)
|
||||
scores := make([]float32, n)
|
||||
for _, l := range m.encoder {
|
||||
l.self.q.rows(q, x, n)
|
||||
l.self.k.rows(k, x, n)
|
||||
l.self.v.rows(v, x, n)
|
||||
for i := range n {
|
||||
attend(ctxv[i*d:(i+1)*d], q[i*d:(i+1)*d], k, v, n, l.self.heads, scores)
|
||||
}
|
||||
l.self.o.rows(tmp, ctxv, n)
|
||||
addNorm(x, tmp, l.selfNorm, n)
|
||||
hidden := make([]float32, n*l.fc1.out)
|
||||
l.fc1.rows(hidden, x, n)
|
||||
for i := range hidden {
|
||||
hidden[i] = m.activation(hidden[i])
|
||||
}
|
||||
l.fc2.rows(tmp, hidden, n)
|
||||
addNorm(x, tmp, l.finalLayerNorm, n)
|
||||
}
|
||||
return x
|
||||
}
|
||||
|
||||
// rows applies the layer to n row vectors: dst = x·Wᵀ + b.
|
||||
func (l linear) rows(dst, x []float32, n int) {
|
||||
for i := range n {
|
||||
l.apply(dst[i*l.out:(i+1)*l.out], x[i*l.in:(i+1)*l.in])
|
||||
}
|
||||
}
|
||||
|
||||
func (l linear) apply(dst, x []float32) {
|
||||
for o := range l.out {
|
||||
dst[o] = dot(x, l.w[o*l.in:(o+1)*l.in]) + l.b[o]
|
||||
}
|
||||
}
|
||||
|
||||
// attend computes multi-head scaled dot-product attention of one query
|
||||
// over n keys and values (n × d each).
|
||||
func attend(dst, q, k, v []float32, n, heads int, scores []float32) {
|
||||
d := len(q)
|
||||
hd := d / heads
|
||||
scale := float32(1 / math.Sqrt(float64(hd)))
|
||||
clear(dst)
|
||||
for h := range heads {
|
||||
lo, hi := h*hd, (h+1)*hd
|
||||
maxScore := float32(math.Inf(-1))
|
||||
for j := range n {
|
||||
s := dot(q[lo:hi], k[j*d+lo:j*d+hi]) * scale
|
||||
scores[j] = s
|
||||
maxScore = max(maxScore, s)
|
||||
}
|
||||
var sum float32
|
||||
for j := range n {
|
||||
scores[j] = float32(math.Exp(float64(scores[j] - maxScore)))
|
||||
sum += scores[j]
|
||||
}
|
||||
out := dst[lo:hi]
|
||||
for j := range n {
|
||||
p := scores[j] / sum
|
||||
row := v[j*d+lo : j*d+hi]
|
||||
for c := range out {
|
||||
out[c] += p * row[c]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// addNorm sets x = LayerNorm(x + y) for n rows.
|
||||
func addNorm(x, y []float32, ln layerNorm, n int) {
|
||||
d := len(ln.g)
|
||||
for i := range n {
|
||||
row := x[i*d : (i+1)*d]
|
||||
add := y[i*d : (i+1)*d]
|
||||
var mean float64
|
||||
for j := range row {
|
||||
row[j] += add[j]
|
||||
mean += float64(row[j])
|
||||
}
|
||||
mean /= float64(d)
|
||||
var variance float64
|
||||
for _, val := range row {
|
||||
diff := float64(val) - mean
|
||||
variance += diff * diff
|
||||
}
|
||||
inv := 1 / math.Sqrt(variance/float64(d)+1e-5)
|
||||
for j := range row {
|
||||
row[j] = float32((float64(row[j])-mean)*inv)*ln.g[j] + ln.b[j]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// dot is the hot loop of the whole model: unrolled, with independent
|
||||
// accumulators so the CPU can overlap the additions.
|
||||
func dot(a, b []float32) float32 {
|
||||
b = b[:len(a)]
|
||||
var s0, s1, s2, s3 float32
|
||||
i := 0
|
||||
for ; i+8 <= len(a); i += 8 {
|
||||
s0 += a[i]*b[i] + a[i+4]*b[i+4]
|
||||
s1 += a[i+1]*b[i+1] + a[i+5]*b[i+5]
|
||||
s2 += a[i+2]*b[i+2] + a[i+6]*b[i+6]
|
||||
s3 += a[i+3]*b[i+3] + a[i+7]*b[i+7]
|
||||
}
|
||||
for ; i < len(a); i++ {
|
||||
s0 += a[i] * b[i]
|
||||
}
|
||||
return s0 + s1 + s2 + s3
|
||||
}
|
||||
Reference in new issue
Block a user