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
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package marian
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
// RemoteFile is a model file to download, checked against its size and
|
||||
// SHA-256 digest.
|
||||
type RemoteFile struct {
|
||||
Name string
|
||||
URL string
|
||||
Size int64
|
||||
SHA256 string
|
||||
}
|
||||
|
||||
// Installed reports whether dir holds every file of a model.
|
||||
func Installed(dir string) bool {
|
||||
for _, f := range Files {
|
||||
if st, err := os.Stat(filepath.Join(dir, f)); err != nil || st.Size() == 0 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// Download fetches files into dir. Each file is written next to its
|
||||
// final name and renamed once its digest is checked, so an interrupted
|
||||
// download never leaves a half-written model behind. progress receives
|
||||
// the bytes downloaded so far and the total.
|
||||
func Download(ctx context.Context, client *http.Client, dir string, files []RemoteFile, progress func(done, total int64)) error {
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
var total, done int64
|
||||
for _, f := range files {
|
||||
total += f.Size
|
||||
}
|
||||
for _, f := range files {
|
||||
n, err := download(ctx, client, dir, f, func(n int64) {
|
||||
if progress != nil {
|
||||
progress(done+n, total)
|
||||
}
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("downloading %s: %w", f.Name, err)
|
||||
}
|
||||
done += n
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func download(ctx context.Context, client *http.Client, dir string, f RemoteFile, progress func(int64)) (int64, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, f.URL, nil)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return 0, fmt.Errorf("server answered %s", resp.Status)
|
||||
}
|
||||
tmp, err := os.CreateTemp(dir, "."+f.Name+".*")
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer os.Remove(tmp.Name()) // no-op once renamed
|
||||
hash := sha256.New()
|
||||
w := &countingWriter{w: io.MultiWriter(tmp, hash), progress: progress}
|
||||
_, err = io.Copy(w, io.LimitReader(resp.Body, f.Size+1))
|
||||
if cerr := tmp.Close(); err == nil {
|
||||
err = cerr
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if w.n != f.Size {
|
||||
return 0, fmt.Errorf("got %d bytes, expected %d", w.n, f.Size)
|
||||
}
|
||||
if sum := hex.EncodeToString(hash.Sum(nil)); sum != f.SHA256 {
|
||||
return 0, fmt.Errorf("checksum mismatch")
|
||||
}
|
||||
if err := os.Chmod(tmp.Name(), 0o644); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return w.n, os.Rename(tmp.Name(), filepath.Join(dir, f.Name))
|
||||
}
|
||||
|
||||
type countingWriter struct {
|
||||
w io.Writer
|
||||
n int64
|
||||
progress func(int64)
|
||||
}
|
||||
|
||||
func (c *countingWriter) Write(p []byte) (int, error) {
|
||||
n, err := c.w.Write(p)
|
||||
c.n += int64(n)
|
||||
c.progress(c.n)
|
||||
return n, err
|
||||
}
|
||||
@@ -0,0 +1,355 @@
|
||||
package marian
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"math/rand"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------- test helpers
|
||||
|
||||
func protoKey(field, wire int) []byte { return binary.AppendUvarint(nil, uint64(field<<3|wire)) }
|
||||
|
||||
func protoBytes(field int, b []byte) []byte {
|
||||
out := append(protoKey(field, 2), binary.AppendUvarint(nil, uint64(len(b)))...)
|
||||
return append(out, b...)
|
||||
}
|
||||
|
||||
func protoVarint(field int, v uint64) []byte {
|
||||
return append(protoKey(field, 0), binary.AppendUvarint(nil, v)...)
|
||||
}
|
||||
|
||||
func protoFloat(field int, f float32) []byte {
|
||||
return binary.LittleEndian.AppendUint32(protoKey(field, 5), math.Float32bits(f))
|
||||
}
|
||||
|
||||
type testPiece struct {
|
||||
piece string
|
||||
score float32
|
||||
typ int
|
||||
}
|
||||
|
||||
// spmModel serialises a SentencePiece ModelProto of the given type.
|
||||
func spmModel(modelType int, pieces []testPiece) []byte {
|
||||
var out []byte
|
||||
for _, p := range pieces {
|
||||
var msg []byte
|
||||
msg = append(msg, protoBytes(1, []byte(p.piece))...)
|
||||
msg = append(msg, protoFloat(2, p.score)...)
|
||||
msg = append(msg, protoVarint(3, uint64(p.typ))...)
|
||||
out = append(out, protoBytes(1, msg)...)
|
||||
}
|
||||
out = append(out, protoBytes(2, protoVarint(3, uint64(modelType)))...)
|
||||
return out
|
||||
}
|
||||
|
||||
var pieces = []testPiece{
|
||||
{"<unk>", 0, pieceUnknown}, {"</s>", 0, pieceControl},
|
||||
{"▁hello", -1, pieceNormal}, {"▁he", -2, pieceNormal}, {"llo", -2, pieceNormal}, {"▁", -3, pieceNormal},
|
||||
{"h", -5, pieceNormal}, {"e", -5, pieceNormal}, {"l", -5, pieceNormal}, {"o", -5, pieceNormal},
|
||||
}
|
||||
|
||||
func writeFile(t *testing.T, path string, data []byte) {
|
||||
t.Helper()
|
||||
if err := os.WriteFile(path, data, 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
type testTensor struct {
|
||||
name string
|
||||
shape []int
|
||||
data []float32
|
||||
}
|
||||
|
||||
func writeSafetensors(t *testing.T, path string, tensors []testTensor) {
|
||||
t.Helper()
|
||||
header := map[string]any{"__metadata__": map[string]string{"format": "pt"}}
|
||||
var data []byte
|
||||
for _, x := range tensors {
|
||||
start := len(data)
|
||||
for _, f := range x.data {
|
||||
data = binary.LittleEndian.AppendUint32(data, math.Float32bits(f))
|
||||
}
|
||||
header[x.name] = map[string]any{"dtype": "F32", "shape": x.shape, "data_offsets": []int{start, len(data)}}
|
||||
}
|
||||
h, _ := json.Marshal(header)
|
||||
out := binary.LittleEndian.AppendUint64(nil, uint64(len(h)))
|
||||
writeFile(t, path, append(append(out, h...), data...))
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------- tests
|
||||
|
||||
func TestSentencePiece(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "m.spm")
|
||||
writeFile(t, path, spmModel(1, pieces))
|
||||
sp, err := loadSentencePiece(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tests := map[string][]string{
|
||||
"hello": {"▁hello"},
|
||||
" hello hel ": {"▁hello", "▁he", "l"},
|
||||
"helloz": {"▁hello", "z"},
|
||||
"": nil,
|
||||
}
|
||||
for in, want := range tests {
|
||||
if got := sp.encode(in); !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("encode(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
if got := decodePieces([]string{"▁Hello", ",", "▁world"}); got != "Hello, world" {
|
||||
t.Errorf("decodePieces = %q", got)
|
||||
}
|
||||
|
||||
writeFile(t, path, spmModel(2, pieces))
|
||||
if _, err := loadSentencePiece(path); err == nil || !strings.Contains(err.Error(), "unigram") {
|
||||
t.Errorf("BPE model: err = %v", err)
|
||||
}
|
||||
writeFile(t, path, []byte{0x0a, 0xff})
|
||||
if _, err := loadSentencePiece(path); err == nil {
|
||||
t.Error("malformed model accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSafetensors(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "m.safetensors")
|
||||
writeSafetensors(t, path, []testTensor{{"w", []int{2, 2}, []float32{1, -2, 3.5, 0}}})
|
||||
got, err := readSafetensors(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w := got["w"]; !reflect.DeepEqual(w.shape, []int{2, 2}) || !reflect.DeepEqual(w.data, []float32{1, -2, 3.5, 0}) {
|
||||
t.Errorf("w = %+v", w)
|
||||
}
|
||||
|
||||
// F16 and BF16 tensors are widened.
|
||||
h := []byte(`{"a":{"dtype":"F16","shape":[2],"data_offsets":[0,4]},"b":{"dtype":"BF16","shape":[1],"data_offsets":[4,6]}}`)
|
||||
data := binary.LittleEndian.AppendUint64(nil, uint64(len(h)))
|
||||
data = append(data, h...)
|
||||
data = append(data, 0x00, 0x3c, 0x00, 0xc0, 0x80, 0x3f) // 1, -2 (F16), 1 (BF16)
|
||||
writeFile(t, path, data)
|
||||
got, err = readSafetensors(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !reflect.DeepEqual(got["a"].data, []float32{1, -2}) || !reflect.DeepEqual(got["b"].data, []float32{1}) {
|
||||
t.Errorf("a = %v, b = %v", got["a"].data, got["b"].data)
|
||||
}
|
||||
|
||||
h = []byte(`{"a":{"dtype":"F32","shape":[3],"data_offsets":[0,4]}}`)
|
||||
writeFile(t, path, append(binary.LittleEndian.AppendUint64(nil, uint64(len(h))), append(h, 0, 0, 0, 0)...))
|
||||
if _, err := readSafetensors(path); err == nil {
|
||||
t.Error("shape mismatch accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHalf(t *testing.T) {
|
||||
tests := map[uint16]float32{0x3c00: 1, 0xc000: -2, 0x7bff: 65504, 0x0001: 1.0 / (1 << 24), 0x8000: float32(math.Copysign(0, -1))}
|
||||
for h, want := range tests {
|
||||
if got := half(h); got != want {
|
||||
t.Errorf("half(%#04x) = %v, want %v", h, got, want)
|
||||
}
|
||||
}
|
||||
if !math.IsInf(float64(half(0x7c00)), 1) {
|
||||
t.Error("half(0x7c00) is not +Inf")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTopK(t *testing.T) {
|
||||
if got := topK([]float32{0.1, 0.9, 0.3, 0.7, 0.9}, 3); !reflect.DeepEqual(got, []int{1, 4, 3}) {
|
||||
t.Errorf("topK = %v", got)
|
||||
}
|
||||
if got := topK([]float32{1}, 3); !reflect.DeepEqual(got, []int{0}) {
|
||||
t.Errorf("topK of one = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// tinyModel writes a random model of the real architecture, small enough
|
||||
// to run in a test.
|
||||
func tinyModel(t *testing.T) string {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
const d, ffn, vocab, positions = 8, 16, 12, 32
|
||||
cfg := fmt.Sprintf(`{"d_model":%d,"encoder_layers":1,"decoder_layers":1,"encoder_attention_heads":2,
|
||||
"decoder_attention_heads":2,"activation_function":"swish","scale_embedding":true,"normalize_before":false,
|
||||
"max_position_embeddings":%d,"vocab_size":%d,"pad_token_id":%d,"eos_token_id":0,"decoder_start_token_id":%d,"num_beams":2}`,
|
||||
d, positions, vocab, vocab-1, vocab-1)
|
||||
writeFile(t, filepath.Join(dir, "config.json"), []byte(cfg))
|
||||
writeFile(t, filepath.Join(dir, "source.spm"), spmModel(1, pieces))
|
||||
writeFile(t, filepath.Join(dir, "target.spm"), spmModel(1, pieces))
|
||||
ids := map[string]int{"</s>": 0, "<unk>": 1, "<pad>": vocab - 1}
|
||||
for i, p := range pieces[2:] {
|
||||
ids[p.piece] = i + 2
|
||||
}
|
||||
v, _ := json.Marshal(ids)
|
||||
writeFile(t, filepath.Join(dir, "vocab.json"), v)
|
||||
|
||||
rng := rand.New(rand.NewSource(1))
|
||||
random := func(n int) []float32 {
|
||||
out := make([]float32, n)
|
||||
for i := range out {
|
||||
out[i] = float32(rng.NormFloat64() * 0.5)
|
||||
}
|
||||
return out
|
||||
}
|
||||
ones := func(n int) []float32 {
|
||||
out := make([]float32, n)
|
||||
for i := range out {
|
||||
out[i] = 1
|
||||
}
|
||||
return out
|
||||
}
|
||||
var tensors []testTensor
|
||||
add := func(name string, data []float32, shape ...int) {
|
||||
tensors = append(tensors, testTensor{name, shape, data})
|
||||
}
|
||||
lin := func(p string, in, out int) {
|
||||
add(p+".weight", random(in*out), out, in)
|
||||
add(p+".bias", random(out), out)
|
||||
}
|
||||
norm := func(p string) {
|
||||
add(p+".weight", ones(d), d)
|
||||
add(p+".bias", make([]float32, d), d)
|
||||
}
|
||||
attn := func(p string) {
|
||||
for _, x := range []string{"q_proj", "k_proj", "v_proj", "out_proj"} {
|
||||
lin(p+"."+x, d, d)
|
||||
}
|
||||
}
|
||||
add("model.shared.weight", random(vocab*d), vocab, d)
|
||||
add("model.encoder.embed_positions.weight", random(positions*d), positions, d)
|
||||
add("model.decoder.embed_positions.weight", random(positions*d), positions, d)
|
||||
add("final_logits_bias", make([]float32, vocab), 1, vocab)
|
||||
attn("model.encoder.layers.0.self_attn")
|
||||
norm("model.encoder.layers.0.self_attn_layer_norm")
|
||||
lin("model.encoder.layers.0.fc1", d, ffn)
|
||||
lin("model.encoder.layers.0.fc2", ffn, d)
|
||||
norm("model.encoder.layers.0.final_layer_norm")
|
||||
attn("model.decoder.layers.0.self_attn")
|
||||
norm("model.decoder.layers.0.self_attn_layer_norm")
|
||||
attn("model.decoder.layers.0.encoder_attn")
|
||||
norm("model.decoder.layers.0.encoder_attn_layer_norm")
|
||||
lin("model.decoder.layers.0.fc1", d, ffn)
|
||||
lin("model.decoder.layers.0.fc2", ffn, d)
|
||||
norm("model.decoder.layers.0.final_layer_norm")
|
||||
writeSafetensors(t, filepath.Join(dir, "model.safetensors"), tensors)
|
||||
return dir
|
||||
}
|
||||
|
||||
func TestTinyModel(t *testing.T) {
|
||||
dir := tinyModel(t)
|
||||
if !Installed(dir) {
|
||||
t.Fatal("tiny model not installed")
|
||||
}
|
||||
m, err := Load(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
first, err := m.Translate(context.Background(), "hello hello")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Random weights: the output means nothing, but it is deterministic
|
||||
// and made of target pieces.
|
||||
again, _ := m.Translate(context.Background(), "hello hello")
|
||||
if first != again {
|
||||
t.Errorf("translations differ: %q, %q", first, again)
|
||||
}
|
||||
if strings.ContainsAny(first, "▁<>") {
|
||||
t.Errorf("translation %q contains special pieces", first)
|
||||
}
|
||||
if out, _ := m.Translate(context.Background(), " "); out != "" {
|
||||
t.Errorf("blank text translated to %q", out)
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
if _, err := m.Translate(ctx, "hello"); err == nil {
|
||||
t.Error("cancelled translation succeeded")
|
||||
}
|
||||
|
||||
os.Remove(filepath.Join(dir, "vocab.json"))
|
||||
if Installed(dir) {
|
||||
t.Error("incomplete model reported installed")
|
||||
}
|
||||
if _, err := Load(dir); err == nil {
|
||||
t.Error("incomplete model loaded")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRealModel runs the OPUS-MT model when PLOTZILLA_MT_MODEL names its
|
||||
// directory (~/.local/share/plotzilla/models/opus-mt-fr-en once
|
||||
// downloaded by the app).
|
||||
func TestRealModel(t *testing.T) {
|
||||
dir := os.Getenv("PLOTZILLA_MT_MODEL")
|
||||
if dir == "" {
|
||||
t.Skip("PLOTZILLA_MT_MODEL not set")
|
||||
}
|
||||
m, err := Load(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tests := map[string]string{
|
||||
"Bonjour, comment allez-vous ?": "Hello, how are you?",
|
||||
"Je t'avais dit de ne jamais revenir ici.": "I told you never to come back here.",
|
||||
}
|
||||
for in, want := range tests {
|
||||
if got, err := m.Translate(context.Background(), in); err != nil || got != want {
|
||||
t.Errorf("Translate(%q) = %q, %v, want %q", in, got, err, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownload(t *testing.T) {
|
||||
content := []byte("model weights")
|
||||
sum := sha256.Sum256(content)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/missing" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Write(content)
|
||||
}))
|
||||
defer srv.Close()
|
||||
dir := filepath.Join(t.TempDir(), "model")
|
||||
file := RemoteFile{Name: "f.bin", URL: srv.URL + "/f", Size: int64(len(content)), SHA256: hex.EncodeToString(sum[:])}
|
||||
|
||||
var last int64
|
||||
if err := Download(context.Background(), srv.Client(), dir, []RemoteFile{file}, func(done, total int64) { last = done }); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, _ := os.ReadFile(filepath.Join(dir, "f.bin")); string(got) != string(content) || last != file.Size {
|
||||
t.Errorf("downloaded %q, progress %d", got, last)
|
||||
}
|
||||
|
||||
bad := file
|
||||
bad.Name, bad.SHA256 = "bad.bin", strings.Repeat("0", 64)
|
||||
missing := file
|
||||
missing.Name, missing.URL = "missing.bin", srv.URL+"/missing"
|
||||
short := file
|
||||
short.Name, short.Size = "short.bin", file.Size+5
|
||||
for _, f := range []RemoteFile{bad, missing, short} {
|
||||
if err := Download(context.Background(), srv.Client(), dir, []RemoteFile{f}, nil); err == nil {
|
||||
t.Errorf("%s: no error", f.Name)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, f.Name)); err == nil {
|
||||
t.Errorf("%s: a failed download left the file", f.Name)
|
||||
}
|
||||
}
|
||||
entries, _ := os.ReadDir(dir)
|
||||
if len(entries) != 1 {
|
||||
t.Errorf("temporary files left: %v", entries)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package marian
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"os"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// tensor is a weight matrix or vector read from a safetensors file, in
|
||||
// row-major order.
|
||||
type tensor struct {
|
||||
shape []int
|
||||
data []float32
|
||||
}
|
||||
|
||||
type tensorInfo struct {
|
||||
DType string `json:"dtype"`
|
||||
Shape []int `json:"shape"`
|
||||
Offsets [2]int64 `json:"data_offsets"`
|
||||
}
|
||||
|
||||
// readSafetensors loads every tensor of a safetensors file (an 8-byte
|
||||
// little-endian header length, a JSON header, then the raw data) as
|
||||
// float32. F32 tensors are read straight into their slice: the file is
|
||||
// large and a second copy would double the memory needed to load it.
|
||||
func readSafetensors(path string) (map[string]tensor, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer f.Close()
|
||||
var size uint64
|
||||
if err := binary.Read(f, binary.LittleEndian, &size); err != nil {
|
||||
return nil, fmt.Errorf("reading header size: %w", err)
|
||||
}
|
||||
if size > 100<<20 {
|
||||
return nil, fmt.Errorf("header too large (%d bytes)", size)
|
||||
}
|
||||
header := make([]byte, size)
|
||||
if _, err := io.ReadFull(f, header); err != nil {
|
||||
return nil, fmt.Errorf("reading header: %w", err)
|
||||
}
|
||||
var infos map[string]json.RawMessage
|
||||
if err := json.Unmarshal(header, &infos); err != nil {
|
||||
return nil, fmt.Errorf("parsing header: %w", err)
|
||||
}
|
||||
base := int64(8 + size)
|
||||
out := make(map[string]tensor, len(infos))
|
||||
for name, raw := range infos {
|
||||
if name == "__metadata__" {
|
||||
continue
|
||||
}
|
||||
var info tensorInfo
|
||||
if err := json.Unmarshal(raw, &info); err != nil {
|
||||
return nil, fmt.Errorf("tensor %s: %w", name, err)
|
||||
}
|
||||
n := 1
|
||||
for _, d := range info.Shape {
|
||||
n *= d
|
||||
}
|
||||
t, err := readTensor(f, base, info, n)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("tensor %s: %w", name, err)
|
||||
}
|
||||
out[name] = tensor{shape: info.Shape, data: t}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func readTensor(f *os.File, base int64, info tensorInfo, n int) ([]float32, error) {
|
||||
start, end := info.Offsets[0], info.Offsets[1]
|
||||
width := map[string]int64{"F32": 4, "F16": 2, "BF16": 2}[info.DType]
|
||||
if width == 0 {
|
||||
return nil, fmt.Errorf("unsupported type %s", info.DType)
|
||||
}
|
||||
if end-start != int64(n)*width {
|
||||
return nil, fmt.Errorf("size %d does not match shape %v", end-start, info.Shape)
|
||||
}
|
||||
data := make([]float32, n)
|
||||
if n == 0 {
|
||||
return data, nil
|
||||
}
|
||||
if info.DType == "F32" {
|
||||
buf := unsafe.Slice((*byte)(unsafe.Pointer(&data[0])), n*4)
|
||||
if _, err := f.ReadAt(buf, base+start); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !littleEndian {
|
||||
for i := range data {
|
||||
data[i] = math.Float32frombits(binary.LittleEndian.Uint32(buf[i*4:]))
|
||||
}
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
buf := make([]byte, n*2)
|
||||
if _, err := f.ReadAt(buf, base+start); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for i := range data {
|
||||
h := binary.LittleEndian.Uint16(buf[i*2:])
|
||||
if info.DType == "BF16" {
|
||||
data[i] = math.Float32frombits(uint32(h) << 16)
|
||||
} else {
|
||||
data[i] = half(h)
|
||||
}
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
var littleEndian = func() bool {
|
||||
x := uint16(1)
|
||||
return *(*byte)(unsafe.Pointer(&x)) == 1
|
||||
}()
|
||||
|
||||
// half converts an IEEE 754 half-precision float.
|
||||
func half(h uint16) float32 {
|
||||
sign := uint32(h>>15) << 31
|
||||
exp := uint32(h>>10) & 0x1f
|
||||
frac := uint32(h) & 0x3ff
|
||||
switch {
|
||||
case exp == 0 && frac == 0:
|
||||
return math.Float32frombits(sign)
|
||||
case exp == 0: // subnormal
|
||||
v := float32(frac) / 1024 / 16384
|
||||
if sign != 0 {
|
||||
return -v
|
||||
}
|
||||
return v
|
||||
case exp == 0x1f:
|
||||
return math.Float32frombits(sign | 0xff<<23 | frac<<13)
|
||||
}
|
||||
return math.Float32frombits(sign | (exp+112)<<23 | frac<<13)
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
package marian
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"strings"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
// space is SentencePiece's visible word boundary ("▁").
|
||||
const space = "▁"
|
||||
|
||||
// sentencePiece is a SentencePiece unigram model: it splits text into the
|
||||
// pieces of the vocabulary that maximise the summed piece scores.
|
||||
type sentencePiece struct {
|
||||
scores map[string]float32
|
||||
maxRunes int
|
||||
unkScore float32
|
||||
// addDummyPrefix starts the text with "▁", as if preceded by a space.
|
||||
addDummyPrefix bool
|
||||
}
|
||||
|
||||
// SentencePiece piece types (sentencepiece_model.proto).
|
||||
const (
|
||||
pieceNormal = 1
|
||||
pieceUnknown = 2
|
||||
pieceControl = 3
|
||||
pieceUserDefined = 4
|
||||
pieceUnused = 5
|
||||
pieceByte = 6
|
||||
)
|
||||
|
||||
// loadSentencePiece reads a serialized ModelProto. Only the fields needed
|
||||
// to encode are decoded: the pieces (field 1), the model type (trainer
|
||||
// spec, field 2) and the dummy-prefix flag (normalizer spec, field 3).
|
||||
func loadSentencePiece(path string) (*sentencePiece, error) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sp := &sentencePiece{scores: map[string]float32{}, addDummyPrefix: true}
|
||||
minScore := float32(0)
|
||||
err = protoFields(data, func(field int, wire int, v uint64, b []byte) error {
|
||||
switch field {
|
||||
case 1:
|
||||
var piece string
|
||||
var score float32
|
||||
typ := pieceNormal
|
||||
err := protoFields(b, func(field, wire int, v uint64, b []byte) error {
|
||||
switch field {
|
||||
case 1:
|
||||
piece = string(b)
|
||||
case 2:
|
||||
score = math.Float32frombits(uint32(v))
|
||||
case 3:
|
||||
typ = int(v)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if typ == pieceNormal || typ == pieceUserDefined {
|
||||
sp.scores[piece] = score
|
||||
sp.maxRunes = max(sp.maxRunes, len([]rune(piece)))
|
||||
minScore = min(minScore, score)
|
||||
}
|
||||
case 2:
|
||||
return protoFields(b, func(field, wire int, v uint64, b []byte) error {
|
||||
if field == 3 && v != 1 {
|
||||
return fmt.Errorf("unsupported SentencePiece model type %d (only unigram)", v)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
case 3:
|
||||
return protoFields(b, func(field, wire int, v uint64, b []byte) error {
|
||||
if field == 3 {
|
||||
sp.addDummyPrefix = v != 0
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("reading %s: %w", path, err)
|
||||
}
|
||||
if len(sp.scores) == 0 {
|
||||
return nil, fmt.Errorf("reading %s: no pieces", path)
|
||||
}
|
||||
// SentencePiece gives unknown characters a score below every piece.
|
||||
sp.unkScore = minScore - 10
|
||||
return sp, nil
|
||||
}
|
||||
|
||||
// protoFields calls fn for each field of a protobuf message: v holds
|
||||
// varints and fixed-width values, b length-delimited ones.
|
||||
func protoFields(data []byte, fn func(field, wire int, v uint64, b []byte) error) error {
|
||||
for len(data) > 0 {
|
||||
key, n := binary.Uvarint(data)
|
||||
if n <= 0 {
|
||||
return errors.New("malformed protobuf key")
|
||||
}
|
||||
data = data[n:]
|
||||
field, wire := int(key>>3), int(key&7)
|
||||
var v uint64
|
||||
var b []byte
|
||||
switch wire {
|
||||
case 0:
|
||||
v, n = binary.Uvarint(data)
|
||||
if n <= 0 {
|
||||
return errors.New("malformed protobuf varint")
|
||||
}
|
||||
data = data[n:]
|
||||
case 1:
|
||||
if len(data) < 8 {
|
||||
return errors.New("truncated protobuf")
|
||||
}
|
||||
v, data = binary.LittleEndian.Uint64(data), data[8:]
|
||||
case 2:
|
||||
l, n := binary.Uvarint(data)
|
||||
if n <= 0 || uint64(len(data)-n) < l {
|
||||
return errors.New("truncated protobuf")
|
||||
}
|
||||
b, data = data[n:n+int(l)], data[n+int(l):]
|
||||
case 5:
|
||||
if len(data) < 4 {
|
||||
return errors.New("truncated protobuf")
|
||||
}
|
||||
v, data = uint64(binary.LittleEndian.Uint32(data)), data[4:]
|
||||
default:
|
||||
return fmt.Errorf("unsupported protobuf wire type %d", wire)
|
||||
}
|
||||
if err := fn(field, wire, v, b); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// normalize approximates SentencePiece's nmt_nfkc normalisation for the
|
||||
// characters scripts actually contain: every kind of space becomes a
|
||||
// plain one, runs of spaces collapse, control characters go away and a
|
||||
// few compatibility characters are decomposed.
|
||||
func normalize(text string) string {
|
||||
var b strings.Builder
|
||||
pendingSpace := false
|
||||
for _, r := range text {
|
||||
switch {
|
||||
case unicode.IsSpace(r):
|
||||
pendingSpace = b.Len() > 0
|
||||
continue
|
||||
case unicode.IsControl(r) || r == '\u200b' || r == '\ufeff':
|
||||
continue
|
||||
}
|
||||
if pendingSpace {
|
||||
b.WriteByte(' ')
|
||||
pendingSpace = false
|
||||
}
|
||||
switch r {
|
||||
case '…':
|
||||
b.WriteString("...")
|
||||
case 'fi':
|
||||
b.WriteString("fi")
|
||||
case 'fl':
|
||||
b.WriteString("fl")
|
||||
case 'ff':
|
||||
b.WriteString("ff")
|
||||
default:
|
||||
b.WriteRune(r)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// encode splits text into pieces (Viterbi search over the unigram scores).
|
||||
func (sp *sentencePiece) encode(text string) []string {
|
||||
text = normalize(text)
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
if sp.addDummyPrefix {
|
||||
text = " " + text
|
||||
}
|
||||
runes := []rune(strings.ReplaceAll(text, " ", space))
|
||||
n := len(runes)
|
||||
best := make([]float64, n+1)
|
||||
from := make([]int, n+1)
|
||||
for i := 1; i <= n; i++ {
|
||||
best[i] = math.Inf(-1)
|
||||
}
|
||||
for end := 1; end <= n; end++ {
|
||||
for start := max(0, end-sp.maxRunes); start < end; start++ {
|
||||
if math.IsInf(best[start], -1) {
|
||||
continue
|
||||
}
|
||||
score, ok := sp.scores[string(runes[start:end])]
|
||||
if !ok {
|
||||
if end-start != 1 {
|
||||
continue
|
||||
}
|
||||
score = sp.unkScore
|
||||
}
|
||||
if s := best[start] + float64(score); s > best[end] {
|
||||
best[end], from[end] = s, start
|
||||
}
|
||||
}
|
||||
}
|
||||
var pieces []string
|
||||
for end := n; end > 0; end = from[end] {
|
||||
pieces = append(pieces, string(runes[from[end]:end]))
|
||||
}
|
||||
for i, j := 0, len(pieces)-1; i < j; i, j = i+1, j-1 {
|
||||
pieces[i], pieces[j] = pieces[j], pieces[i]
|
||||
}
|
||||
return pieces
|
||||
}
|
||||
|
||||
// decodePieces joins target pieces back into text.
|
||||
func decodePieces(pieces []string) string {
|
||||
return strings.TrimSpace(strings.ReplaceAll(strings.Join(pieces, ""), space, " "))
|
||||
}
|
||||
Reference in new issue
Block a user