Kharec
839e9291d6
Downloading the 300 MB model restarted from zero after any failure and hung forever on a stalled connection. Download into a .partial file that survives interruptions, resume it with a Range request (falling back to a fresh download when the server cannot), skip files already verified, and abort when nothing arrives for 30 s. A model directory only counts as installed when every file has its pinned size, so a truncated model is offered for download again instead of failing to load. The model load held the translation mutex, so Cancel did nothing for the seconds it takes to read 300 MB. Load without holding the lock, make marian.Load context-aware, and let CancelTranslation wait for the task to wind down so a new one can start at once. The idle release timer re-arms when it fires while a download holds the task slot, instead of never freeing the model.
384 lines
11 KiB
Go
384 lines
11 KiB
Go
// 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. The load can take seconds for a large
|
||
// model: it stops early when ctx is cancelled.
|
||
func Load(ctx context.Context, dir string) (*Model, error) {
|
||
if err := ctx.Err(); err != nil {
|
||
return nil, err
|
||
}
|
||
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>"]
|
||
|
||
if err := ctx.Err(); err != nil {
|
||
return nil, err
|
||
}
|
||
tensors, err := readSafetensors(ctx, 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
|
||
}
|