356 lines
11 KiB
Go
356 lines
11 KiB
Go
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)
|
||
}
|
||
}
|