diff --git a/internal/marian/decode.go b/internal/marian/decode.go new file mode 100644 index 0000000..8f50db9 --- /dev/null +++ b/internal/marian/decode.go @@ -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 +} diff --git a/internal/marian/download.go b/internal/marian/download.go new file mode 100644 index 0000000..85ad758 --- /dev/null +++ b/internal/marian/download.go @@ -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 +} diff --git a/internal/marian/marian_test.go b/internal/marian/marian_test.go new file mode 100644 index 0000000..dd5e057 --- /dev/null +++ b/internal/marian/marian_test.go @@ -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{ + {"", 0, pieceUnknown}, {"", 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{"": 0, "": 1, "": 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) + } +} diff --git a/internal/marian/model.go b/internal/marian/model.go new file mode 100644 index 0000000..ec0c91b --- /dev/null +++ b/internal/marian/model.go @@ -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[""] + + 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 +} diff --git a/internal/marian/safetensors.go b/internal/marian/safetensors.go new file mode 100644 index 0000000..e12d092 --- /dev/null +++ b/internal/marian/safetensors.go @@ -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) +} diff --git a/internal/marian/sentencepiece.go b/internal/marian/sentencepiece.go new file mode 100644 index 0000000..34f7f6a --- /dev/null +++ b/internal/marian/sentencepiece.go @@ -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, " ")) +} diff --git a/internal/translate/detect.go b/internal/translate/detect.go new file mode 100644 index 0000000..9864d5c --- /dev/null +++ b/internal/translate/detect.go @@ -0,0 +1,136 @@ +package translate + +import ( + "strings" + "unicode" +) + +// Lexicon tells which languages know a word, from spelling dictionaries. +// It refines the detection of words that are not function words. +type Lexicon interface { + Knows(word string) (french, english bool) +} + +// Function words of one language that are not words of the other: they +// carry most of the signal in a short line. +var frenchWords = set(`le la les un une des du de au aux et est sont ce cette ces cet il ils elle elles je tu nous vous + ne pas que qui quoi dans sur avec pour par mais où donc ou ni sa ses leur leurs mon ma mes ton ta tes notre votre nos vos lui eux + se ça cela ceci très bien aussi alors encore toujours jamais rien tout tous toute toutes quand puis sans sous chez vers entre + être avoir fait faire va vais suis es était étaient ont avait avaient sera oui non merci voilà voici déjà ici là dont + comme comment pourquoi peut peux veux veut faut quelque quelques chose rien personne moi toi soi y en + après avant pendant depuis jusque contre derrière devant près loin enfin ensuite puis peu beaucoup trop assez`) + +var englishWords = set(`the and is are was were be been being this that these those he she they we you it his her him their + its my your our of to in at with for from by not no yes do does did have has had will would can could should shall + what who whom where when why how there here then than just into out up down over about back all some any very + only also again never always nothing everything something someone anyone everyone because before after while through + i me us them which been get got go goes going went come came look looks looking says said tell told know think`) + +// elisions are French contractions before a vowel: "l'homme", "qu'il". +var elisions = []string{"l'", "d'", "j'", "qu'", "n'", "s'", "c'", "m'", "t'", "jusqu'", "lorsqu'", "puisqu'"} + +// contractions end English words: "don't", "we're", "Mary's". +var contractions = []string{"n't", "'re", "'ll", "'ve", "'m", "'d"} + +func set(words string) map[string]bool { + m := map[string]bool{} + for _, w := range strings.Fields(words) { + m[w] = true + } + return m +} + +const frenchLetters = "éèêëàâçîïôûùüœæ" + +// words splits text into lower-case words, apostrophes included +// ("l'homme", "don't"). +func words(text string) []string { + var out []string + var w []rune + flush := func() { + if len(w) > 0 { + out = append(out, strings.Trim(string(w), "'")) + w = w[:0] + } + } + for _, r := range text { + switch { + case unicode.IsLetter(r) || unicode.IsMark(r): + w = append(w, unicode.ToLower(r)) + case (r == '\'' || r == '’') && len(w) > 0: + w = append(w, '\'') + default: + flush() + } + } + flush() + return out +} + +// scores counts the evidence that text is French and that it is English. +// names are words to ignore (the script's characters). +func scores(text string, lex Lexicon, names map[string]bool) (fr, en int) { + for _, w := range words(text) { + if w == "" || names[w] { + continue + } + switch { + case frenchWords[w]: + fr += 2 + continue + case englishWords[w]: + en += 2 + continue + } + if hasPrefix(w, elisions) { + fr += 2 + continue + } + if hasSuffix(w, contractions) { + en += 2 + continue + } + if strings.HasSuffix(w, "'s") { + en++ + continue + } + if strings.ContainsAny(w, frenchLetters) { + fr++ + continue + } + if lex != nil && len([]rune(w)) > 2 { + switch f, e := lex.Knows(w); { + case f && !e: + fr++ + case e && !f: + en++ + } + } + } + return fr, en +} + +func hasPrefix(w string, prefixes []string) bool { + for _, p := range prefixes { + if strings.HasPrefix(w, p) && len(w) > len(p) { + return true + } + } + return false +} + +func hasSuffix(w string, suffixes []string) bool { + for _, s := range suffixes { + if strings.HasSuffix(w, s) && len(w) > len(s) { + return true + } + } + return false +} + +// notEnglish reports whether text is not clearly English: the parts of a +// line already known to be French (the location of a French slugline). +func (s *session) notEnglish(text string) bool { + fr, en := scores(text, s.lex, s.nameWords) + return en <= fr +} diff --git a/internal/translate/rules.go b/internal/translate/rules.go new file mode 100644 index 0000000..373397d --- /dev/null +++ b/internal/translate/rules.go @@ -0,0 +1,160 @@ +package translate + +import ( + "strings" + "unicode" +) + +// The standard terms of French scripts and their Hollywood equivalents. +// Keys are folded (see fold): upper case, no accents, single spaces. + +var frenchTimes = table(map[string]string{ + "JOUR": "DAY", "NUIT": "NIGHT", "SOIR": "EVENING", "SOIREE": "EVENING", "MATIN": "MORNING", + "MATINEE": "MORNING", "APRES-MIDI": "AFTERNOON", "APRES MIDI": "AFTERNOON", "AUBE": "DAWN", + "AURORE": "DAWN", "CREPUSCULE": "DUSK", "COUCHER DE SOLEIL": "SUNSET", "COUCHER DU SOLEIL": "SUNSET", + "LEVER DE SOLEIL": "SUNRISE", "LEVER DU SOLEIL": "SUNRISE", "MIDI": "NOON", "MINUIT": "MIDNIGHT", + "PLUS TARD": "LATER", "UN PEU PLUS TARD": "A LITTLE LATER", "PEU APRES": "MOMENTS LATER", + "INSTANTS PLUS TARD": "MOMENTS LATER", "QUELQUES INSTANTS PLUS TARD": "MOMENTS LATER", + "QUELQUES MINUTES PLUS TARD": "MINUTES LATER", "CONTINU": "CONTINUOUS", "CONTINUE": "CONTINUOUS", + "CONTINUITE": "CONTINUOUS", "EN CONTINU": "CONTINUOUS", "SUITE": "CONTINUOUS", + "MEME MOMENT": "SAME TIME", "AU MEME MOMENT": "SAME TIME", "MEME TEMPS": "SAME TIME", + "LE LENDEMAIN": "THE NEXT DAY", "LENDEMAIN": "NEXT DAY", "PETIT MATIN": "EARLY MORNING", + "TOT LE MATIN": "EARLY MORNING", "FIN D'APRES-MIDI": "LATE AFTERNOON", "FIN DE JOURNEE": "LATE AFTERNOON", + "DEBUT DE SOIREE": "EARLY EVENING", "TARD LE SOIR": "LATE NIGHT", "TARD DANS LA NUIT": "LATE NIGHT", + "PLUS TOT": "EARLIER", "NUIT AMERICAINE": "DAY FOR NIGHT", +}) + +var englishTimes = table(map[string]string{ + "DAY": "", "NIGHT": "", "MORNING": "", "EVENING": "", "AFTERNOON": "", "DAWN": "", "DUSK": "", + "SUNSET": "", "SUNRISE": "", "NOON": "", "MIDNIGHT": "", "LATER": "", "CONTINUOUS": "", + "MOMENTS LATER": "", "SAME TIME": "", "EARLY MORNING": "", "LATE NIGHT": "", "EARLIER": "", + "MAGIC HOUR": "", "LATE AFTERNOON": "", "EARLY EVENING": "", "SAME": "", "CONT'D": "", +}) + +// sceneNotes are the parenthetical notes that follow a slugline. +var sceneNotes = table(map[string]string{ + "FLASH-BACK": "FLASHBACK", "FLASHBACK": "FLASHBACK", "RETOUR EN ARRIERE": "FLASHBACK", + "FIN DU FLASH-BACK": "END FLASHBACK", "FIN DU FLASHBACK": "END FLASHBACK", "REVE": "DREAM", + "CAUCHEMAR": "NIGHTMARE", "SOUVENIR": "MEMORY", "PRESENT": "PRESENT DAY", "RETOUR AU PRESENT": "PRESENT DAY", + "EN MOUVEMENT": "MOVING", "ROULANT": "MOVING", "A L'ARRET": "STATIONARY", "IMAGINAIRE": "IMAGINED", + "FANTASME": "FANTASY", "STOCK SHOT": "STOCK FOOTAGE", "IMAGES D'ARCHIVE": "STOCK FOOTAGE", +}) + +// extensions are the character extensions, folded without dots, spaces, +// hyphens or apostrophes: "V.O.", "VOIX OFF" and "voix-off" all give "VOIXOFF". +var extensions = map[string]string{ + "VO": "V.O.", "VOIXOFF": "V.O.", "OFF": "V.O.", "ENOFF": "V.O.", "ENPENSEE": "V.O.", "PENSEE": "V.O.", + "PENSEES": "V.O.", "HC": "O.S.", "HORSCHAMP": "O.S.", "HORSCHAMPS": "O.S.", "SUITE": "CONT'D", + "CONT": "CONT'D", "ENSUITE": "CONT'D", "AUTELEPHONE": "ON PHONE", "TELEPHONE": "ON PHONE", + "AUTEL": "ON PHONE", "TEL": "ON PHONE", "ALARADIO": "ON RADIO", "RADIO": "ON RADIO", + "ALATELE": "ON TV", "ALATELEVISION": "ON TV", "TELE": "ON TV", "ALATV": "ON TV", + "SOUSTITRE": "SUBTITLED", "SOUSTITRES": "SUBTITLED", "SOUSTITREE": "SUBTITLED", "PRELAP": "PRE-LAP", + "CHUCHOTE": "WHISPERING", "CHUCHOTANT": "WHISPERING", "ENCHUCHOTANT": "WHISPERING", + "ENREGISTREMENT": "RECORDED", "ENREGISTRE": "RECORDED", "SURECRAN": "ON SCREEN", "ALECRAN": "ON SCREEN", +} + +var transitions = table(map[string]string{ + "COUPE": "CUT TO:", "COUPE SUR": "CUT TO:", "COUPURE": "CUT TO:", "CUT SUR": "CUT TO:", + "FONDU ENCHAINE": "DISSOLVE TO:", "ENCHAINE": "DISSOLVE TO:", "ENCHAINE SUR": "DISSOLVE TO:", + "FONDU ENCHAINE SUR": "DISSOLVE TO:", "FONDU ENCHAINE A": "DISSOLVE TO:", + "FONDU AU NOIR": "FADE TO BLACK.", "FONDU EN NOIR": "FADE TO BLACK.", "FONDU NOIR": "FADE TO BLACK.", + "FONDU AU BLANC": "FADE TO WHITE.", "FONDU EN FERMETURE": "FADE OUT.", "FERMETURE AU NOIR": "FADE OUT.", + "FERMETURE EN FONDU": "FADE OUT.", "FONDU DE FERMETURE": "FADE OUT.", "FONDU DE SORTIE": "FADE OUT.", + "FONDU EN OUVERTURE": "FADE IN:", "OUVERTURE AU NOIR": "FADE IN:", "OUVERTURE EN FONDU": "FADE IN:", + "FONDU D'OUVERTURE": "FADE IN:", "OUVERTURE EN FONDU AU NOIR": "FADE IN:", "OUVERTURE": "FADE IN:", + "COUPE FRANCHE": "SMASH CUT TO:", "CUT FRANC": "SMASH CUT TO:", "COUPE SECHE": "SMASH CUT TO:", + "CUT SEC": "SMASH CUT TO:", "RACCORD": "MATCH CUT TO:", "RACCORD SUR": "MATCH CUT TO:", + "RACCORD MOUVEMENT": "MATCH CUT TO:", "RACCORD REGARD": "MATCH CUT TO:", "RACCORD DANS L'AXE": "MATCH CUT TO:", + "COUPE SAUTEE": "JUMP CUT TO:", "VOLET": "WIPE TO:", "VOLET SUR": "WIPE TO:", "FONDU": "FADE TO:", + "FONDU SUR": "FADE TO:", "FONDU A": "FADE TO:", "FIN": "THE END", +}) + +// shots are matched as the first words of an insert / shot element; the +// English forms start like the shots the Fountain format recognises. +var shots = []struct{ french, english string }{ + {"TRES GROS PLAN SUR", "EXTREME CLOSE ON"}, + {"TRES GROS PLAN", "EXTREME CLOSE-UP"}, + {"GROS PLAN SUR", "CLOSE ON"}, + {"GROS PLAN", "CLOSE UP"}, + {"PLAN RAPPROCHE SUR", "CLOSE ON"}, + {"PLAN SERRE SUR", "CLOSE ON"}, + {"PLAN LARGE SUR", "WIDE ON"}, + {"PLAN LARGE", "WIDE SHOT"}, + {"PLAN D'ENSEMBLE", "WIDE SHOT"}, + {"PLAN GENERAL", "WIDE SHOT"}, + {"PLAN DE SITUATION", "WIDE SHOT"}, + {"PLAN AERIEN", "AERIAL SHOT"}, + {"VUE AERIENNE", "AERIAL SHOT"}, + {"PLAN EN PLONGEE", "OVERHEAD SHOT"}, + {"PLONGEE", "OVERHEAD SHOT"}, + {"VUE DU DESSUS", "OVERHEAD SHOT"}, + {"POINT DE VUE DE", "POV"}, + {"POINT DE VUE", "POV"}, + {"P.D.V.", "POV"}, + {"PDV", "POV"}, + {"PLAN SUBJECTIF", "POV"}, + {"VUE SUBJECTIVE", "POV"}, + {"CONTRECHAMP", "REVERSE ANGLE"}, + {"CONTRE-CHAMP", "REVERSE ANGLE"}, + {"NOUVEL ANGLE", "NEW ANGLE"}, + {"ANGLE SUR", "ANGLE ON"}, + {"RETOUR A LA SCENE", "BACK TO SCENE"}, + {"RETOUR SUR LA SCENE", "BACK TO SCENE"}, + {"INSERT SUR", "INSERT"}, + {"TRAVELLING SUR", "TRACKING ON"}, + {"TRAVELLING", "TRACKING SHOT"}, + {"PLAN MOYEN", "MEDIUM SHOT"}, + {"PLAN AMERICAIN", "MEDIUM SHOT"}, + {"PLAN TAILLE", "MEDIUM SHOT"}, + {"SERIE DE PLANS", "SERIES OF SHOTS"}, + {"FIN DU MONTAGE", "END MONTAGE"}, +} + +// parentheticals are the stock wrylies, folded to lower case. +var parentheticals = map[string]string{ + "un temps": "beat", "temps": "beat", "un temps d'arret": "beat", "un silence": "beat", + "suite": "continuing", "continuant": "continuing", "a voix basse": "quietly", "tout bas": "quietly", + "bas": "quietly", "a part": "aside", "ironique": "wryly", "pour lui-meme": "to himself", + "pour elle-meme": "to herself", "pour soi": "to themself", "off": "O.S.", "en off": "V.O.", + "hors champ": "O.S.", "hors-champ": "O.S.", "voix off": "V.O.", +} + +// credits are the usual title page credits, folded to lower case. +var credits = map[string]string{ + "ecrit par": "Written by", "scenario de": "Screenplay by", "scenario": "Screenplay by", + "scenario et dialogues de": "Screenplay by", "scenario et dialogues": "Screenplay by", + "un film de": "A film by", "realise par": "Directed by", "d'apres une idee de": "Story by", + "histoire originale de": "Story by", +} + +func table(m map[string]string) map[string]string { + out := make(map[string]string, len(m)) + for k, v := range m { + out[fold(k)] = v + } + return out +} + +var accents = strings.NewReplacer( + "À", "A", "Â", "A", "Ä", "A", "Ç", "C", "É", "E", "È", "E", "Ê", "E", "Ë", "E", "Î", "I", "Ï", "I", + "Ô", "O", "Ö", "O", "Ù", "U", "Û", "U", "Ü", "U", "Œ", "OE", "Æ", "AE", "Ÿ", "Y", "’", "'", +) + +// fold upper-cases, removes accents and collapses spaces, to look a term +// up whatever its spelling. +func fold(s string) string { + return strings.Join(strings.Fields(accents.Replace(strings.ToUpper(s))), " ") +} + +// foldTight also drops dots, spaces, hyphens and apostrophes ("V.O." → "VO"). +func foldTight(s string) string { + return strings.Map(func(r rune) rune { + if r == '.' || r == '-' || r == '\'' || unicode.IsSpace(r) { + return -1 + } + return r + }, fold(s)) +} + +// foldLower is fold in lower case, for parentheticals and credits. +func foldLower(s string) string { return strings.ToLower(fold(s)) } diff --git a/internal/translate/translate.go b/internal/translate/translate.go new file mode 100644 index 0000000..e7ce40c --- /dev/null +++ b/internal/translate/translate.go @@ -0,0 +1,957 @@ +// Package translate translates a screenplay from French into English, +// element by element, with the conventions of each element type: a French +// slugline becomes "INT. KITCHEN - NIGHT", "(V.O.)" and "(SUITE)" become +// "(V.O.)" and "(CONT'D)", "FONDU AU NOIR" becomes "FADE TO BLACK.", and so +// on. Free text (action, dialogue, parentheticals, locations) goes through +// a machine translation model, one sentence at a time. Text that does not +// read as French is left untouched, so a script can be translated piece by +// piece, and translating it again changes nothing. +package translate + +import ( + "context" + "regexp" + "slices" + "strings" + "sync" + "unicode" + + "plotzilla/internal/screenplay" +) + +// Translator translates one French sentence into English. It must be +// safe for concurrent use. +type Translator interface { + Translate(ctx context.Context, text string) (string, error) +} + +// Options says what to translate. +type Options struct { + // Elements are the indexes of the elements to translate (flat order, + // as everywhere else); nil translates every element. + Elements []int + // TitlePage also translates the title page credit, source, draft date + // and notes. The title itself is never translated. + TitlePage bool + // Workers is how many sentences are translated at once (default 1). + Workers int + // Progress is called after each sentence, from any goroutine. + Progress func(done, total int) +} + +// Change is a translated element, to put in place of element Index. +type Change struct { + Index int `json:"index"` + Element screenplay.Element `json:"element"` +} + +// Result lists what changed. TitlePage is nil when it did not change. +type Result struct { + Changes []Change `json:"changes"` + TitlePage *screenplay.TitlePage `json:"titlePage,omitempty"` +} + +// Document translates the French text of doc. It does not modify doc. +func Document(ctx context.Context, doc *screenplay.Document, tr Translator, lex Lexicon, opts Options) (Result, error) { + s := newSession(doc, lex) + indexes := opts.Elements + if indexes == nil { + indexes = make([]int, len(doc.Elements)) + for i := range indexes { + indexes[i] = i + } + } + // First pass: collect the sentences to translate. Second pass, once + // they are translated: build the elements from them. + s.collect = true + s.run(doc, indexes, opts.TitlePage) + done, err := s.translateAll(ctx, tr, max(opts.Workers, 1), opts.Progress) + if err != nil { + return Result{}, err + } + s.collect, s.done = false, done + return s.run(doc, indexes, opts.TitlePage), nil +} + +type session struct { + lex Lexicon + // names are the script's characters, lower case, longest first: they + // are kept as written instead of being translated ("Marie" would + // become "Mary"). + names [][]rune + nameWords map[string]bool + // frenchAt says, for each element, whether it or (when it is + // ambiguous) its nearest unambiguous neighbour is French; + // frenchContext is that value for the element being translated. + frenchAt []bool + frenchContext bool + collect bool + wanted []string + seen map[string]bool + done map[string]string +} + +// determiners start a character cue that is a description to translate +// ("LE POLICIER") rather than a name. +var determiners = set("le la les l un une des mon ma mes ton ta tes son sa ses notre votre leur") + +func newSession(doc *screenplay.Document, lex Lexicon) *session { + s := &session{lex: lex, nameWords: map[string]bool{}, seen: map[string]bool{}} + seen := map[string]bool{} + for _, el := range doc.Elements { + if el.Type != screenplay.Character { + continue + } + name := strings.ToLower(screenplay.CharacterName(el.PrintedText())) + ws := words(name) + if name == "" || seen[name] || len(ws) == 0 || determiners[strings.TrimSuffix(firstElision(ws[0]), "'")] { + continue + } + seen[name] = true + s.names = append(s.names, []rune(name)) + for _, w := range ws { + s.nameWords[w] = true + } + } + s.frenchAt = s.languages(doc) + // Longest first: "Marie-Claire" before "Marie". + for i := 1; i < len(s.names); i++ { + for j := i; j > 0 && len(s.names[j]) > len(s.names[j-1]); j-- { + s.names[j], s.names[j-1] = s.names[j-1], s.names[j] + } + } + return s +} + +// languages decides the language of each element: French, English, or, +// for an element without clear evidence ("Attends !", "Oui ?" are short; +// a character cue is a name), the language of its nearest decided +// neighbour. +func (s *session) languages(doc *screenplay.Document) []bool { + n := len(doc.Elements) + lang := make([]int, n) // 1 French, -1 English, 0 undecided + for i, el := range doc.Elements { + fr, en := scores(el.Text(), s.lex, s.nameWords) + switch { + case fr > en: + lang[i] = 1 + case en > fr: + lang[i] = -1 + } + } + out := make([]bool, n) + for i := range n { + if lang[i] != 0 { + out[i] = lang[i] > 0 + continue + } + for d := 1; d < n; d++ { + if i-d >= 0 && lang[i-d] != 0 { + out[i] = lang[i-d] > 0 + break + } + if i+d < n && lang[i+d] != 0 { + out[i] = lang[i+d] > 0 + break + } + } + } + return out +} + +// firstElision returns "l'" for "l'inspecteur", the word itself otherwise. +func firstElision(w string) string { + if i := strings.IndexByte(w, '\''); i > 0 { + return w[:i+1] + } + return w +} + +func (s *session) run(doc *screenplay.Document, indexes []int, titlePage bool) Result { + res := Result{Changes: []Change{}} + for _, i := range indexes { + if i < 0 || i >= len(doc.Elements) { + continue + } + s.frenchContext = s.frenchAt[i] + if el, changed := s.element(doc.Elements[i]); changed { + res.Changes = append(res.Changes, Change{Index: i, Element: el}) + } + } + if titlePage { + s.frenchContext = false + tp := doc.TitlePage + tp.Credit = s.lines(tp.Credit, s.credit) + tp.Source = s.lines(tp.Source, s.prose) + tp.DraftDate = s.lines(tp.DraftDate, s.prose) + tp.Notes = s.lines(tp.Notes, s.prose) + if tp != doc.TitlePage { + res.TitlePage = &tp + } + } + return res +} + +// ------------------------------------------------------------- elements + +// element translates el. Notes are translated on their own; the rest of +// the text as a whole, so sentences keep their meaning. Emphasis that +// covers only part of a translated stretch cannot follow the words: the +// translation keeps the styles common to the whole stretch. +func (s *session) element(el screenplay.Element) (screenplay.Element, bool) { + var runs []screenplay.Run + changed := false + if t := s.retype(el); t != el.Type { + el.Type, changed = t, true + } + rule := s.rule(el.Type) + for _, chunk := range chunks(el.Runs) { + var b strings.Builder + for _, r := range chunk { + b.WriteString(r.Text) + } + text := b.String() + var out string + if chunk[0].Note { + out = s.lines(text, s.prose) + } else { + out = s.lines(text, rule) + } + if out == text { + runs = append(runs, chunk...) + continue + } + changed = true + style := commonStyle(chunk) + style.Text = out + runs = append(runs, style) + } + if !changed { + return el, false + } + oldLen := len([]rune(el.Text())) + el.Runs = screenplay.NormalizeRuns(runs) + el.Locked = remapLocks(el.Locked, oldLen, len([]rune(el.Text()))) + return el, true +} + +// retype recognises a French transition or shot typed as action, as the +// Fountain format reads "FONDU AU NOIR." or "GROS PLAN SUR la lettre" (it +// only knows the English forms): the element becomes a transition or a +// shot, as the English text will be read back. +func (s *session) retype(el screenplay.Element) screenplay.ElementType { + if el.Type != screenplay.Action || len(el.Runs) == 0 || slices.ContainsFunc(el.Runs, func(r screenplay.Run) bool { return r.Note }) { + return el.Type + } + line := strings.TrimSpace(el.Text()) + if strings.Contains(line, "\n") { + return el.Type + } + if key := fold(strings.TrimRight(line, ":. \u202f")); transitions[key] != "" && key != "FIN" && isCaps(line) { + return screenplay.Transition + } + fields := strings.Fields(line) + for _, sh := range shots { + n := len(strings.Fields(sh.french)) + if len(fields) >= n && isCaps(strings.Join(fields[:n], " ")) && fold(strings.Join(fields[:n], " ")) == sh.french { + return screenplay.Shot + } + } + return el.Type +} + +func (s *session) rule(t screenplay.ElementType) func(string) string { + switch t { + case screenplay.Scene: + return s.slugline + case screenplay.Character: + return s.cue + case screenplay.Parenthetical: + return s.parenthetical + case screenplay.Transition: + return s.transition + case screenplay.Shot: + return s.shot + case screenplay.Centered: + return s.centered + } + return s.prose +} + +// chunks groups runs into stretches of text: each note alone, the text +// between notes together. +func chunks(runs []screenplay.Run) [][]screenplay.Run { + var out [][]screenplay.Run + for i, r := range runs { + if i == 0 || r.Note || runs[i-1].Note { + out = append(out, nil) + } + out[len(out)-1] = append(out[len(out)-1], r) + } + return out +} + +func commonStyle(runs []screenplay.Run) screenplay.Run { + c := runs[0] + c.Text = "" + for _, r := range runs[1:] { + c.Bold = c.Bold && r.Bold + c.Italic = c.Italic && r.Italic + c.Underline = c.Underline && r.Underline + if r.Rev != c.Rev { + c.Rev = 0 + } + } + return c +} + +// remapLocks keeps the locked pages of a translated element in proportion +// to the text: the page that started halfway through starts halfway +// through the translation. +func remapLocks(locks []screenplay.PageLock, oldLen, newLen int) []screenplay.PageLock { + if len(locks) == 0 { + return locks + } + out := make([]screenplay.PageLock, len(locks)) + prev := 0 + for i, l := range locks { + off := 0 + if oldLen > 0 { + off = (l.Offset*newLen + oldLen/2) / oldLen + } + off = min(max(off, prev), newLen) + out[i] = screenplay.PageLock{Offset: off, Page: l.Page} + prev = off + } + return out +} + +// lines applies fn to each line of text, keeping the line breaks and the +// spaces around each line. +func (s *session) lines(text string, fn func(string) string) string { + parts := strings.Split(text, "\n") + for i, line := range parts { + core := strings.TrimSpace(line) + if core == "" { + continue + } + lead := line[:strings.Index(line, core)] + trail := line[len(lead)+len(core):] + parts[i] = lead + fn(core) + trail + } + return strings.Join(parts, "\n") +} + +// ---------------------------------------------------------------- rules + +// prose translates free text sentence by sentence, leaving the sentences +// that are not French. +// +// A sentence with as much French as English evidence ("Attends !": to +// attend is English too, "Marie ?") is translated when the text around it +// is French: other sentences of the text, or the neighbouring elements. +func (s *session) prose(text string) string { + list := sentences(text) + type verdict struct{ fr, en int } + verdicts := make([]verdict, len(list)) + frenchAround := s.frenchContext + for i, sentence := range list { + fr, en := scores(sentence, s.lex, s.nameWords) + verdicts[i] = verdict{fr, en} + frenchAround = frenchAround || fr > en + } + var b strings.Builder + for i, sentence := range list { + core := strings.TrimRightFunc(sentence, unicode.IsSpace) + v := verdicts[i] + if v.fr > v.en || (v.fr == v.en && frenchAround && hasLetter(core)) { + b.WriteString(s.mt(core, modeOf(core))) + } else { + b.WriteString(core) + } + b.WriteString(sentence[len(core):]) + } + return b.String() +} + +// french reports whether a line reads as French, or could be French and +// sits among French elements. +func (s *session) french(text string) bool { + fr, en := scores(text, s.lex, s.nameWords) + return fr > en || (fr == en && s.frenchContext && hasLetter(text)) +} + +func hasLetter(text string) bool { return strings.IndexFunc(text, unicode.IsLetter) >= 0 } + +var ( + slugPrefix = regexp.MustCompile(`(?i)^((?:INT\.?\s*/\s*EXT|EXT\.?\s*/\s*INT|I\s*/\s*E|INT|EXT|EST|INT[ÉE]RIEUR|EXT[ÉE]RIEUR)\.?)\s+(.*)$`) + slugDash = regexp.MustCompile(`\s+[-–—]+\s+`) + trailNote = regexp.MustCompile(`^(.*?)\s*\(([^()]*)\)$`) +) + +// slugline translates a French scene heading: "INT. CUISINE DE MARIE - +// NUIT" → "INT. MARIE'S KITCHEN - NIGHT". Times of day and the usual notes +// come from a table, locations from the model. +func (s *session) slugline(line string) string { + prefix, rest := "", line + if m := slugPrefix.FindStringSubmatch(line); m != nil { + prefix, rest = m[1], m[2] + } + note := "" + if m := trailNote.FindStringSubmatch(rest); m != nil && m[1] != "" { + rest, note = m[1], m[2] + } + parts := slugDash.Split(rest, -1) + locations, time := parts, "" + if len(parts) > 1 { + locations, time = parts[:len(parts)-1], parts[len(parts)-1] + } + frTime, isFrTime := frenchTimes[fold(time)] + _, isEnTime := englishTimes[fold(time)] + _, isFrNote := sceneNotes[fold(note)] + fp := fold(prefix) + french := isFrTime || strings.HasPrefix(fp, "INTERIEUR") || strings.HasPrefix(fp, "EXTERIEUR") || + (!isEnTime && (s.french(strings.Join(locations, " ")) || isFrNote && time == "")) + if !french { + return line + } + + switch { + case strings.HasPrefix(fp, "INTERIEUR"): + prefix = "INT." + case strings.HasPrefix(fp, "EXTERIEUR"): + prefix = "EXT." + } + out := make([]string, 0, len(parts)) + for _, loc := range locations { + if s.notEnglish(loc) { + loc = s.mt(loc, location) + } + out = append(out, loc) + } + switch { + case isFrTime: + out = append(out, frTime) + case time != "" && !isEnTime && s.notEnglish(time): + out = append(out, s.mt(time, caps)) + case time != "": + out = append(out, time) + } + heading := strings.Join(out, " - ") + if prefix != "" { + heading = prefix + " " + heading + } + if note != "" { + if en, ok := sceneNotes[fold(note)]; ok { + note = en + } else if s.french(note) { + note = s.mt(note, caps) + } + heading += " (" + note + ")" + } + return heading +} + +var ( + parens = regexp.MustCompile(`\(([^()]*)\)`) + shotHead = regexp.MustCompile(`^([A-Z' .]+?)(\s+[-–—:]\s+)(.+)$`) +) + +// cue translates a character cue's extensions, "MARIE (VOIX OFF)" → +// "MARIE (V.O.)". The name is kept, unless it is a description starting +// with an article ("LE POLICIER" → "THE POLICEMAN"). +func (s *session) cue(line string) string { + name, rest := line, "" + if i := strings.IndexByte(line, '('); i >= 0 { + name, rest = line[:i], line[i:] + } + core := strings.TrimRightFunc(name, unicode.IsSpace) + gap := name[len(core):] + if ws := words(core); len(ws) > 0 && determiners[strings.TrimSuffix(firstElision(ws[0]), "'")] && s.notEnglish(core) { + core = s.mt(core, caps) + } + rest = parens.ReplaceAllStringFunc(rest, func(p string) string { + inner := strings.TrimSpace(p[1 : len(p)-1]) + if en, ok := extensions[foldTight(inner)]; ok { + return "(" + en + ")" + } + if s.french(inner) { + return "(" + s.mt(inner, caps) + ")" + } + return p + }) + return core + gap + rest +} + +// parenthetical translates a wryly, "(un temps)" → "(beat)", "(à Marie)" +// → "(to Marie)", in lower case as is the custom. +func (s *session) parenthetical(line string) string { + open, close := strings.HasPrefix(line, "("), strings.HasSuffix(line, ")") + inner := line + if open { + inner = inner[1:] + } + if close { + inner = inner[:len(inner)-1] + } + inner = strings.TrimSpace(inner) + var out string + if en, ok := parentheticals[foldLower(inner)]; ok { + out = en + } else if s.french(inner) { + out = s.lowerFirst(s.mt(inner, plain)) + } else { + return line + } + if open { + out = "(" + out + } + if close { + out += ")" + } + return out +} + +// lowerFirst lower-cases the first letter unless it starts a name, an +// acronym or "I". +func (s *session) lowerFirst(text string) string { + r := []rune(text) + if len(r) < 2 || !unicode.IsUpper(r[0]) || !unicode.IsLower(r[1]) { + return text + } + if ws := words(text); len(ws) > 0 && s.nameWords[ws[0]] { + return text + } + r[0] = unicode.ToLower(r[0]) + return string(r) +} + +// transition translates "FONDU ENCHAÎNÉ :" → "DISSOLVE TO:". +func (s *session) transition(line string) string { + key := fold(strings.TrimRight(line, ":.  ")) + if en, ok := transitions[key]; ok { + return en + } + if s.french(line) { + return s.mt(line, caps) + } + return line +} + +// shot translates a shot or insert: "GROS PLAN SUR LA LETTRE" → "CLOSE ON +// THE LETTER", "INSERT - UNE PHOTO" → "INSERT - A PHOTO". +func (s *session) shot(line string) string { + fields := strings.Fields(line) + for _, sh := range shots { + n := len(strings.Fields(sh.french)) + if len(fields) < n || fold(strings.Join(fields[:n], " ")) != sh.french { + continue + } + rest := afterWords(line, n) + trimmed := strings.TrimLeft(rest, " \t-–—:") + if strings.TrimSpace(trimmed) == "" { + return sh.english + } + sep := " " + if strings.ContainsAny(rest[:len(rest)-len(trimmed)], "-–—:") { + sep = " - " + } + if s.notEnglish(trimmed) { + trimmed = s.mt(trimmed, modeOf(trimmed)) + } + return sh.english + sep + trimmed + } + // "INSERT - UNE PHOTO": an English shot with French text after it. + if s.french(line) { + if m := shotHead.FindStringSubmatch(line); m != nil && !s.french(m[1]) { + return m[1] + " - " + s.mt(m[3], caps) + } + return s.mt(line, caps) + } + return line +} + +// afterWords returns line after its first n words. +func afterWords(line string, n int) string { + i := 0 + for range n { + for i < len(line) && line[i] == ' ' { + i++ + } + for i < len(line) && line[i] != ' ' { + i++ + } + } + return line[i:] +} + +func (s *session) centered(line string) string { + if fold(strings.Trim(line, ". ")) == "FIN" { + return "THE END" + } + return s.prose(line) +} + +func (s *session) credit(line string) string { + if en, ok := credits[foldLower(strings.TrimRight(line, ": "))]; ok { + return en + } + return s.prose(line) +} + +// isCaps reports whether text is written in capitals. +func isCaps(text string) bool { + letters := 0 + for _, r := range text { + if unicode.IsLower(r) { + return false + } + if unicode.IsLetter(r) { + letters++ + } + } + return letters > 1 +} + +// ------------------------------------------------------------ sentences + +var abbreviations = set("m mm mme mmes mlle mlles dr pr me mgr st ste etc cf p vol ch av bd no mr mrs ms jr sr vs env") + +// sentences splits text into sentences, each with the spaces that follow +// it. A sentence ends with . ! ? or … (and closing quotes or brackets) +// followed by a space and a capital, a digit or an opening quote or dash; +// "M. Dupont" and initials do not end one. +func sentences(text string) []string { + r := []rune(text) + var out []string + start := 0 + for i := 0; i < len(r); i++ { + if !strings.ContainsRune(".!?…", r[i]) { + continue + } + j := closers(r, i+1) + if j >= len(r) || !unicode.IsSpace(r[j]) { + i = j - 1 + continue + } + k := j + for k < len(r) && unicode.IsSpace(r[k]) { + k++ + } + if k == len(r) { + break + } + next := r[k] + if !unicode.IsUpper(next) && !unicode.IsDigit(next) && !strings.ContainsRune("«\"“—–-¿¡([", next) { + i = j - 1 + continue + } + if r[i] == '.' && j == i+1 && abbreviated(r[:i]) { + i = j - 1 + continue + } + out = append(out, string(r[start:k])) + start, i = k, k-1 + } + if start < len(r) { + out = append(out, string(r[start:])) + } + return out +} + +// closers skips what may close a sentence after its final punctuation: +// more punctuation, quotes, brackets, and a French closing guillemet after +// its space ("Non. »"). +func closers(r []rune, j int) int { + for j < len(r) { + switch { + case strings.ContainsRune(".!?…»\"')]’”", r[j]): + j++ + case unicode.IsSpace(r[j]): + k := j + for k < len(r) && unicode.IsSpace(r[k]) && r[k] != '\n' { + k++ + } + if k < len(r) && r[k] == '»' { + j = k + 1 + continue + } + return j + default: + return j + } + } + return j +} + +// abbreviated reports whether text ends with an abbreviation or an +// initial (the word before a full stop). +func abbreviated(text []rune) bool { + i := len(text) + for i > 0 && unicode.IsLetter(text[i-1]) { + i-- + } + w := text[i:] + if len(w) == 0 { + return false + } + return (len(w) == 1 && unicode.IsUpper(w[0])) || abbreviations[strings.ToLower(string(w))] +} + +// ---------------------------------------------------------- translation + +// mode says how a text is put to the model. +type mode byte + +const ( + plain mode = iota + // caps: text in capitals, translated in lower case (the model reads + // capitals badly) and put back in capitals. + caps + // location: a slugline's location, in capitals. It is given an + // article, which the model needs to read it as a noun phrase + // ("CUISINE DE MARIE" alone gives "MARIE CUISINE", "La cuisine de + // Marie" gives "Marie's kitchen"), removed from the translation. + location +) + +func modeOf(text string) mode { + if isCaps(text) { + return caps + } + return plain +} + +// mt returns the translation of a French sentence (during the first pass: +// records it and returns it as is). +func (s *session) mt(text string, m mode) string { + key := string(rune('0'+m)) + "\x00" + text + if s.collect { + if !s.seen[key] { + s.seen[key] = true + s.wanted = append(s.wanted, key) + } + return text + } + if out, ok := s.done[key]; ok && out != "" { + return out + } + return text +} + +func (s *session) translateAll(ctx context.Context, tr Translator, workers int, progress func(done, total int)) (map[string]string, error) { + done := make(map[string]string, len(s.wanted)) + if len(s.wanted) == 0 { + return done, nil + } + ctx, cancel := context.WithCancel(ctx) + defer cancel() + var ( + mu sync.Mutex + firstErr error + wg sync.WaitGroup + ) + jobs := make(chan string) + for range min(workers, len(s.wanted)) { + wg.Add(1) + go func() { + defer wg.Done() + for key := range jobs { + out, err := s.translateOne(ctx, tr, key) + mu.Lock() + if err != nil && firstErr == nil { + firstErr = err + cancel() + } + done[key] = out + n := len(done) + mu.Unlock() + if err == nil && progress != nil { + progress(n, len(s.wanted)) + } + } + }() + } +feed: + for _, key := range s.wanted { + select { + case jobs <- key: + case <-ctx.Done(): + break feed + } + } + close(jobs) + wg.Wait() + if firstErr != nil { + return nil, firstErr + } + if err := ctx.Err(); err != nil { + return nil, err + } + return done, nil +} + +// placeholders stand for the characters' names while a sentence is +// translated: invented names the model copies as they are (checked with +// OPUS-MT fr-en), so it cannot turn "Marie" into "Mary" or "Pierre" into +// "stone". +var placeholders = []string{ + "Zorbin", "Dravok", "Yelmir", "Brontiq", "Quorvel", "Velkor", "Strivan", "Xavek", + "Kwilam", "Tazmik", "Moxtel", "Gorlex", "Pilvam", "Wexol", "Jandrix", +} + +func (s *session) translateOne(ctx context.Context, tr Translator, key string) (string, error) { + m, text := mode(key[0]-'0'), key[2:] + upper := m != plain + src, names := s.protect(text, upper) + article := m == location && !startsWithDeterminer(src) + if article { + src = articleFor(src) + src + } + if upper { + src = capitalize(src) + } + out, err := tr.Translate(ctx, src) + if err != nil { + return "", err + } + if article { + out = leadingArticle.ReplaceAllString(out, "") + } + for i, name := range names { + out = placeholderRes[i].ReplaceAllLiteralString(out, name) + } + return polish(text, out, upper), nil +} + +var ( + leadingArticle = regexp.MustCompile(`(?i)^\s*(the|an?)\s+`) + placeholderRes = func() []*regexp.Regexp { + out := make([]*regexp.Regexp, len(placeholders)) + for i, p := range placeholders { + out[i] = regexp.MustCompile(`(?i)\b` + p + `\b`) + } + return out + }() +) + +func startsWithDeterminer(text string) bool { + ws := words(text) + return len(ws) > 0 && (determiners[strings.TrimSuffix(firstElision(ws[0]), "'")] || ws[0] == "du" || ws[0] == "au" || ws[0] == "aux") +} + +// articleFor picks a French article for a noun phrase: its gender is +// unknown, but any article will do for the model. +func articleFor(text string) string { + first := strings.ToLower(fold(strings.Fields(text)[0])) + switch { + case strings.ContainsAny(first[:1], "aeiouyh"): + return "l'" + case strings.HasSuffix(first, "s") || strings.HasSuffix(first, "x"): + return "les " + } + return "le " +} + +// capitalize upper-cases the first letter: a sentence starts with a +// capital, or the model takes it for a fragment. +func capitalize(text string) string { + r := []rune(text) + for i, c := range r { + if unicode.IsLetter(c) { + r[i] = unicode.ToUpper(c) + break + } + } + return string(r) +} + +// protect replaces the characters' names written with a capital by +// placeholders and, if lowerCase, lower-cases the rest. It returns the +// names in placeholder order. +func (s *session) protect(text string, lowerCase bool) (string, []string) { + r := []rune(text) + lower := make([]rune, len(r)) + for i, c := range r { + lower[i] = unicode.ToLower(c) + } + isWord := func(i int) bool { return i >= 0 && i < len(r) && (unicode.IsLetter(r[i]) || unicode.IsDigit(r[i])) } + var b strings.Builder + var names []string + index := map[string]int{} + for i := 0; i < len(r); { + matched := false + if !isWord(i - 1) { + for _, name := range s.names { + end := i + len(name) + if end > len(r) || isWord(end) || string(lower[i:end]) != string(name) || !hasUpper(r[i:end]) { + continue + } + key := string(name) + n, ok := index[key] + if !ok { + if len(names) == len(placeholders) { + break + } + n = len(names) + index[key] = n + names = append(names, string(r[i:end])) + } + b.WriteString(placeholders[n]) + i, matched = end, true + break + } + } + if matched { + continue + } + if lowerCase { + b.WriteRune(lower[i]) + } else { + b.WriteRune(r[i]) + } + i++ + } + return b.String(), names +} + +func hasUpper(r []rune) bool { + for _, c := range r { + if unicode.IsUpper(c) { + return true + } + } + return false +} + +var ( + openQuote = regexp.MustCompile(`«[\s\x{a0}\x{202f}]*`) + closeQuote = regexp.MustCompile(`[\s\x{a0}\x{202f}]*»`) + spaceMark = regexp.MustCompile(`[\s\x{a0}\x{202f}]+([!?;:])`) +) + +// polish applies English typography to a translation: no space before +// ! ? ; :, straight quotes instead of guillemets, no full stop the source +// did not have (a location, a wryly), capitals if the source had them. +func polish(src, out string, caps bool) string { + out = openQuote.ReplaceAllString(out, `"`) + out = closeQuote.ReplaceAllString(out, `"`) + out = spaceMark.ReplaceAllString(out, "$1") + out = strings.TrimSpace(out) + if !strings.ContainsAny(lastRune(strings.TrimSpace(src)), ".!?…") && + strings.HasSuffix(out, ".") && !strings.HasSuffix(out, "..") { + out = strings.TrimSuffix(out, ".") + } + if caps { + out = screenplay.UpperRunes(out) + } + return out +} + +func lastRune(s string) string { + r := []rune(s) + if len(r) == 0 { + return "" + } + return string(r[len(r)-1]) +} diff --git a/internal/translate/translate_test.go b/internal/translate/translate_test.go new file mode 100644 index 0000000..9c64e86 --- /dev/null +++ b/internal/translate/translate_test.go @@ -0,0 +1,307 @@ +package translate + +import ( + "context" + "errors" + "reflect" + "slices" + "sync" + "testing" + + "plotzilla/internal/screenplay" +) + +// fakeMT translates from a fixed table; a sentence it does not know comes +// back in angle brackets, so the test shows what was asked. +type fakeMT struct { + mu sync.Mutex + dict map[string]string + calls []string + err error +} + +func (f *fakeMT) Translate(ctx context.Context, text string) (string, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.calls = append(f.calls, text) + if f.err != nil { + return "", f.err + } + if out, ok := f.dict[text]; ok { + return out, nil + } + return "<" + text + ">", nil +} + +type fakeLexicon map[string][2]bool + +func (l fakeLexicon) Knows(w string) (bool, bool) { return l[w][0], l[w][1] } + +var lexicon = fakeLexicon{ + "attends": {true, true}, "cuisine": {true, true}, "laughing": {false, true}, "phone": {false, true}, + "rings": {false, true}, "nobody": {false, true}, "moves": {false, true}, "kitchen": {false, true}, + "street": {false, true}, "pleut": {true, false}, "entre": {true, false}, "lettre": {true, false}, "octobre": {true, false}, +} + +var dict = map[string]string{ + "Le cuisine de Zorbin": "Zorbin's kitchen", + "Le rue": "The street", + "Zorbin ?": "Zorbin?", + "Tu es là ?": "Are you there?", + "Zorbin décroche.": "Zorbin picks up.", + "Le policier": "The policeman", + "à Zorbin": "To Zorbin", + "Attends !": "Wait!", + "Elle entre.": "She comes in.", + "Il pleut.": "It's raining.", + "Vérifier la date": "Check the date.", + "la lettre.": "the letter.", + "Je l'ai lue.": "I read it.", + "1er octobre 2026": "1 October 2026", + "Il pleut. Elle entre.": "unused", + "Bonjour, Pierre.": "Hello, Pierre.", + "Bonjour, « Zorbin » !": "Hello, « Zorbin » !", + "Ça marche, Zorbin.": "That works, Zorbin.", + "Zorbin et Dravok sortent.": "Zorbin and Dravok go out.", +} + +func el(t screenplay.ElementType, text string) screenplay.Element { + return screenplay.NewElement(t, text) +} + +func translateDoc(t *testing.T, doc *screenplay.Document, opts Options) (Result, *fakeMT) { + t.Helper() + mt := &fakeMT{dict: dict} + res, err := Document(context.Background(), doc, mt, lexicon, opts) + if err != nil { + t.Fatal(err) + } + return res, mt +} + +// texts applies the changes and returns every element's text. +func texts(doc *screenplay.Document, res Result) []string { + out := make([]string, len(doc.Elements)) + for i, e := range doc.Elements { + out[i] = e.Text() + } + for _, c := range res.Changes { + out[c.Index] = c.Element.Text() + } + return out +} + +func TestScreenplayElements(t *testing.T) { + doc := &screenplay.Document{Elements: []screenplay.Element{ + el(screenplay.Scene, "INT. CUISINE DE MARIE - NUIT"), + el(screenplay.Character, "PIERRE (VOIX OFF)"), + el(screenplay.Dialogue, "Marie ? Tu es là ?"), + el(screenplay.Character, "MARIE"), + el(screenplay.Parenthetical, "(un temps)"), + el(screenplay.Dialogue, "Attends !"), + el(screenplay.Parenthetical, "(à Marie)"), + el(screenplay.Character, "LE POLICIER"), + el(screenplay.Character, "MARIE (SUITE)"), + el(screenplay.Transition, "FONDU ENCHAÎNÉ :"), + el(screenplay.Scene, "EXTÉRIEUR RUE - AUBE (FLASH-BACK)"), + el(screenplay.Action, "The phone rings. Marie décroche."), + el(screenplay.Centered, "FIN"), + }} + res, _ := translateDoc(t, doc, Options{}) + want := []string{ + "INT. MARIE'S KITCHEN - NIGHT", + "PIERRE (V.O.)", + "Marie? Are you there?", + "MARIE", + "(beat)", + "Wait!", + "(to Marie)", + "THE POLICEMAN", + "MARIE (CONT'D)", + "DISSOLVE TO:", + "EXT. STREET - DAWN (FLASHBACK)", + "The phone rings. Marie picks up.", + "THE END", + } + if got := texts(doc, res); !reflect.DeepEqual(got, want) { + for i := range want { + if got[i] != want[i] { + t.Errorf("element %d = %q, want %q", i, got[i], want[i]) + } + } + } + if slices.ContainsFunc(res.Changes, func(c Change) bool { return c.Index == 3 }) { + t.Error("a character name was changed") + } +} + +func TestEnglishLeftAlone(t *testing.T) { + doc := &screenplay.Document{Elements: []screenplay.Element{ + el(screenplay.Scene, "INT. KITCHEN - DAY"), + el(screenplay.Action, "Nobody moves. The phone rings."), + el(screenplay.Character, "PIERRE (O.S.)"), + el(screenplay.Parenthetical, "(laughing)"), + el(screenplay.Dialogue, "I told you, I'm not leaving."), + el(screenplay.Transition, "CUT TO:"), + el(screenplay.Shot, "CLOSE ON THE LETTER"), + el(screenplay.Action, "OK."), + }} + res, mt := translateDoc(t, doc, Options{TitlePage: true}) + if len(res.Changes) != 0 || res.TitlePage != nil || len(mt.calls) != 0 { + t.Errorf("English script changed: %+v, model asked %q", res, mt.calls) + } +} + +func TestActionRetyped(t *testing.T) { + doc := &screenplay.Document{Elements: []screenplay.Element{ + el(screenplay.Action, "FONDU EN OUVERTURE :"), + el(screenplay.Action, "GROS PLAN SUR la lettre."), + el(screenplay.Action, "COUPE SUR :"), + el(screenplay.Action, "Il pleut."), + }} + res, _ := translateDoc(t, doc, Options{}) + want := map[int]screenplay.Element{ + 0: el(screenplay.Transition, "FADE IN:"), + 1: el(screenplay.Shot, "CLOSE ON the letter."), + 2: el(screenplay.Transition, "CUT TO:"), + 3: el(screenplay.Action, "It's raining."), + } + for _, c := range res.Changes { + if w := want[c.Index]; c.Element.Type != w.Type || c.Element.Text() != w.Text() { + t.Errorf("element %d = %s %q, want %s %q", c.Index, c.Element.Type, c.Element.Text(), w.Type, w.Text()) + } + } + if len(res.Changes) != len(want) { + t.Errorf("%d changes, want %d", len(res.Changes), len(want)) + } +} + +func TestNotesStylesAndLocks(t *testing.T) { + doc := &screenplay.Document{Elements: []screenplay.Element{{ + Type: screenplay.Action, + Runs: []screenplay.Run{ + {Text: "Elle entre. ", Bold: true}, + {Text: "Vérifier la date", Note: true}, + {Text: " Il pleut.", Rev: 1}, + }, + Locked: []screenplay.PageLock{{Offset: 0, Page: "2"}, {Offset: 28, Page: "3"}}, + }}} + res, _ := translateDoc(t, doc, Options{}) + if len(res.Changes) != 1 { + t.Fatalf("changes: %+v", res.Changes) + } + got := res.Changes[0].Element + want := []screenplay.Run{ + {Text: "She comes in. ", Bold: true}, + // No full stop: the note had none. + {Text: "Check the date", Note: true}, + {Text: " It's raining.", Rev: 1}, + } + if !reflect.DeepEqual(got.Runs, want) { + t.Errorf("runs = %+v, want %+v", got.Runs, want) + } + // 28 of 38 runes → the same proportion of the 42 translated ones. + if want := []screenplay.PageLock{{Offset: 0, Page: "2"}, {Offset: 31, Page: "3"}}; !reflect.DeepEqual(got.Locked, want) { + t.Errorf("locks = %+v, want %+v", got.Locked, want) + } +} + +func TestSelectionAndTitlePage(t *testing.T) { + doc := &screenplay.Document{ + TitlePage: screenplay.TitlePage{Title: "La Dernière Bobine", Credit: "Écrit par", DraftDate: "1er octobre 2026"}, + Elements: []screenplay.Element{ + el(screenplay.Action, "Il pleut."), + el(screenplay.Action, "Elle entre."), + }, + } + res, _ := translateDoc(t, doc, Options{Elements: []int{1}}) + if len(res.Changes) != 1 || res.Changes[0].Index != 1 || res.TitlePage != nil { + t.Errorf("selection: %+v", res) + } + res, _ = translateDoc(t, doc, Options{TitlePage: true}) + want := screenplay.TitlePage{Title: "La Dernière Bobine", Credit: "Written by", DraftDate: "1 October 2026"} + if res.TitlePage == nil || *res.TitlePage != want { + t.Errorf("title page = %+v, want %+v", res.TitlePage, want) + } +} + +func TestNamesAndTypography(t *testing.T) { + doc := &screenplay.Document{Elements: []screenplay.Element{ + el(screenplay.Character, "MARIE"), + el(screenplay.Character, "PIERRE"), + el(screenplay.Dialogue, "Bonjour, « Marie » !"), + el(screenplay.Action, "Ça marche, PIERRE."), + el(screenplay.Action, "Marie et Pierre sortent."), + }} + res, mt := translateDoc(t, doc, Options{}) + got := texts(doc, res) + want := []string{"MARIE", "PIERRE", `Hello, "Marie"!`, "That works, PIERRE.", "Marie and Pierre go out."} + if !reflect.DeepEqual(got, want) { + t.Errorf("got %q, want %q (model asked %q)", got, want, mt.calls) + } +} + +func TestIdempotent(t *testing.T) { + doc := &screenplay.Document{Elements: []screenplay.Element{ + el(screenplay.Scene, "INT. CUISINE DE MARIE - NUIT"), + el(screenplay.Character, "MARIE (SUITE)"), + el(screenplay.Dialogue, "Je l'ai lue."), + }} + res, _ := translateDoc(t, doc, Options{}) + for _, c := range res.Changes { + doc.Elements[c.Index] = c.Element + } + res, mt := translateDoc(t, doc, Options{}) + if len(res.Changes) != 0 || len(mt.calls) != 0 { + t.Errorf("second translation changed %+v, asked %q", res.Changes, mt.calls) + } +} + +func TestModelError(t *testing.T) { + doc := &screenplay.Document{Elements: []screenplay.Element{el(screenplay.Action, "Il pleut.")}} + boom := errors.New("boom") + if _, err := Document(context.Background(), doc, &fakeMT{err: boom}, nil, Options{Workers: 4}); !errors.Is(err, boom) { + t.Errorf("err = %v, want boom", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := Document(ctx, doc, &fakeMT{dict: dict}, nil, Options{}); !errors.Is(err, context.Canceled) { + t.Errorf("err = %v, want context.Canceled", err) + } +} + +func TestSentences(t *testing.T) { + tests := map[string][]string{ + "Il pleut. Elle entre.": {"Il pleut. ", "Elle entre."}, + "Quoi ? Non !": {"Quoi ? ", "Non !"}, + "M. Dupont entre. J. Doe aussi.": {"M. Dupont entre. ", "J. Doe aussi."}, + "Attends... il revient.": {"Attends... il revient."}, + "« Non. » Il sort.": {"« Non. » ", "Il sort."}, + } + for in, want := range tests { + if got := sentences(in); !reflect.DeepEqual(got, want) { + t.Errorf("sentences(%q) = %q, want %q", in, got, want) + } + } +} + +func TestDetection(t *testing.T) { + tests := []struct { + text string + french bool + }{ + {"Il pleut sur la ville.", true}, + {"It's raining on the city.", false}, + {"Qu'est-ce que tu fais ?", true}, + {"Allô ?", true}, + {"OK.", false}, + {"Don't move.", false}, + {"Marie", false}, + } + s := &session{lex: lexicon, nameWords: map[string]bool{"marie": true}} + for _, tt := range tests { + if got := s.french(tt.text); got != tt.french { + t.Errorf("french(%q) = %v, want %v", tt.text, got, tt.french) + } + } +}