feat: add a translate stack

This commit is contained in:
Kharec committed 2026-10-01 10:15:28 +02:00
1 parent 5cf16588c1
commit cb356ff8c5
10 files changed
+2996

No files matched your search

+234
View File
@@ -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
}
+109
View File
@@ -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
}
+355
View File
@@ -0,0 +1,355 @@
package marian
import (
"context"
"crypto/sha256"
"encoding/binary"
"encoding/hex"
"encoding/json"
"fmt"
"math"
"math/rand"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"reflect"
"strings"
"testing"
)
// ---------------------------------------------------------- test helpers
func protoKey(field, wire int) []byte { return binary.AppendUvarint(nil, uint64(field<<3|wire)) }
func protoBytes(field int, b []byte) []byte {
out := append(protoKey(field, 2), binary.AppendUvarint(nil, uint64(len(b)))...)
return append(out, b...)
}
func protoVarint(field int, v uint64) []byte {
return append(protoKey(field, 0), binary.AppendUvarint(nil, v)...)
}
func protoFloat(field int, f float32) []byte {
return binary.LittleEndian.AppendUint32(protoKey(field, 5), math.Float32bits(f))
}
type testPiece struct {
piece string
score float32
typ int
}
// spmModel serialises a SentencePiece ModelProto of the given type.
func spmModel(modelType int, pieces []testPiece) []byte {
var out []byte
for _, p := range pieces {
var msg []byte
msg = append(msg, protoBytes(1, []byte(p.piece))...)
msg = append(msg, protoFloat(2, p.score)...)
msg = append(msg, protoVarint(3, uint64(p.typ))...)
out = append(out, protoBytes(1, msg)...)
}
out = append(out, protoBytes(2, protoVarint(3, uint64(modelType)))...)
return out
}
var pieces = []testPiece{
{"<unk>", 0, pieceUnknown}, {"</s>", 0, pieceControl},
{"▁hello", -1, pieceNormal}, {"▁he", -2, pieceNormal}, {"llo", -2, pieceNormal}, {"▁", -3, pieceNormal},
{"h", -5, pieceNormal}, {"e", -5, pieceNormal}, {"l", -5, pieceNormal}, {"o", -5, pieceNormal},
}
func writeFile(t *testing.T, path string, data []byte) {
t.Helper()
if err := os.WriteFile(path, data, 0o644); err != nil {
t.Fatal(err)
}
}
type testTensor struct {
name string
shape []int
data []float32
}
func writeSafetensors(t *testing.T, path string, tensors []testTensor) {
t.Helper()
header := map[string]any{"__metadata__": map[string]string{"format": "pt"}}
var data []byte
for _, x := range tensors {
start := len(data)
for _, f := range x.data {
data = binary.LittleEndian.AppendUint32(data, math.Float32bits(f))
}
header[x.name] = map[string]any{"dtype": "F32", "shape": x.shape, "data_offsets": []int{start, len(data)}}
}
h, _ := json.Marshal(header)
out := binary.LittleEndian.AppendUint64(nil, uint64(len(h)))
writeFile(t, path, append(append(out, h...), data...))
}
// ----------------------------------------------------------------- tests
func TestSentencePiece(t *testing.T) {
path := filepath.Join(t.TempDir(), "m.spm")
writeFile(t, path, spmModel(1, pieces))
sp, err := loadSentencePiece(path)
if err != nil {
t.Fatal(err)
}
tests := map[string][]string{
"hello": {"▁hello"},
" hello hel ": {"▁hello", "▁he", "l"},
"helloz": {"▁hello", "z"},
"": nil,
}
for in, want := range tests {
if got := sp.encode(in); !reflect.DeepEqual(got, want) {
t.Errorf("encode(%q) = %q, want %q", in, got, want)
}
}
if got := decodePieces([]string{"▁Hello", ",", "▁world"}); got != "Hello, world" {
t.Errorf("decodePieces = %q", got)
}
writeFile(t, path, spmModel(2, pieces))
if _, err := loadSentencePiece(path); err == nil || !strings.Contains(err.Error(), "unigram") {
t.Errorf("BPE model: err = %v", err)
}
writeFile(t, path, []byte{0x0a, 0xff})
if _, err := loadSentencePiece(path); err == nil {
t.Error("malformed model accepted")
}
}
func TestSafetensors(t *testing.T) {
path := filepath.Join(t.TempDir(), "m.safetensors")
writeSafetensors(t, path, []testTensor{{"w", []int{2, 2}, []float32{1, -2, 3.5, 0}}})
got, err := readSafetensors(path)
if err != nil {
t.Fatal(err)
}
if w := got["w"]; !reflect.DeepEqual(w.shape, []int{2, 2}) || !reflect.DeepEqual(w.data, []float32{1, -2, 3.5, 0}) {
t.Errorf("w = %+v", w)
}
// F16 and BF16 tensors are widened.
h := []byte(`{"a":{"dtype":"F16","shape":[2],"data_offsets":[0,4]},"b":{"dtype":"BF16","shape":[1],"data_offsets":[4,6]}}`)
data := binary.LittleEndian.AppendUint64(nil, uint64(len(h)))
data = append(data, h...)
data = append(data, 0x00, 0x3c, 0x00, 0xc0, 0x80, 0x3f) // 1, -2 (F16), 1 (BF16)
writeFile(t, path, data)
got, err = readSafetensors(path)
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(got["a"].data, []float32{1, -2}) || !reflect.DeepEqual(got["b"].data, []float32{1}) {
t.Errorf("a = %v, b = %v", got["a"].data, got["b"].data)
}
h = []byte(`{"a":{"dtype":"F32","shape":[3],"data_offsets":[0,4]}}`)
writeFile(t, path, append(binary.LittleEndian.AppendUint64(nil, uint64(len(h))), append(h, 0, 0, 0, 0)...))
if _, err := readSafetensors(path); err == nil {
t.Error("shape mismatch accepted")
}
}
func TestHalf(t *testing.T) {
tests := map[uint16]float32{0x3c00: 1, 0xc000: -2, 0x7bff: 65504, 0x0001: 1.0 / (1 << 24), 0x8000: float32(math.Copysign(0, -1))}
for h, want := range tests {
if got := half(h); got != want {
t.Errorf("half(%#04x) = %v, want %v", h, got, want)
}
}
if !math.IsInf(float64(half(0x7c00)), 1) {
t.Error("half(0x7c00) is not +Inf")
}
}
func TestTopK(t *testing.T) {
if got := topK([]float32{0.1, 0.9, 0.3, 0.7, 0.9}, 3); !reflect.DeepEqual(got, []int{1, 4, 3}) {
t.Errorf("topK = %v", got)
}
if got := topK([]float32{1}, 3); !reflect.DeepEqual(got, []int{0}) {
t.Errorf("topK of one = %v", got)
}
}
// tinyModel writes a random model of the real architecture, small enough
// to run in a test.
func tinyModel(t *testing.T) string {
t.Helper()
dir := t.TempDir()
const d, ffn, vocab, positions = 8, 16, 12, 32
cfg := fmt.Sprintf(`{"d_model":%d,"encoder_layers":1,"decoder_layers":1,"encoder_attention_heads":2,
"decoder_attention_heads":2,"activation_function":"swish","scale_embedding":true,"normalize_before":false,
"max_position_embeddings":%d,"vocab_size":%d,"pad_token_id":%d,"eos_token_id":0,"decoder_start_token_id":%d,"num_beams":2}`,
d, positions, vocab, vocab-1, vocab-1)
writeFile(t, filepath.Join(dir, "config.json"), []byte(cfg))
writeFile(t, filepath.Join(dir, "source.spm"), spmModel(1, pieces))
writeFile(t, filepath.Join(dir, "target.spm"), spmModel(1, pieces))
ids := map[string]int{"</s>": 0, "<unk>": 1, "<pad>": vocab - 1}
for i, p := range pieces[2:] {
ids[p.piece] = i + 2
}
v, _ := json.Marshal(ids)
writeFile(t, filepath.Join(dir, "vocab.json"), v)
rng := rand.New(rand.NewSource(1))
random := func(n int) []float32 {
out := make([]float32, n)
for i := range out {
out[i] = float32(rng.NormFloat64() * 0.5)
}
return out
}
ones := func(n int) []float32 {
out := make([]float32, n)
for i := range out {
out[i] = 1
}
return out
}
var tensors []testTensor
add := func(name string, data []float32, shape ...int) {
tensors = append(tensors, testTensor{name, shape, data})
}
lin := func(p string, in, out int) {
add(p+".weight", random(in*out), out, in)
add(p+".bias", random(out), out)
}
norm := func(p string) {
add(p+".weight", ones(d), d)
add(p+".bias", make([]float32, d), d)
}
attn := func(p string) {
for _, x := range []string{"q_proj", "k_proj", "v_proj", "out_proj"} {
lin(p+"."+x, d, d)
}
}
add("model.shared.weight", random(vocab*d), vocab, d)
add("model.encoder.embed_positions.weight", random(positions*d), positions, d)
add("model.decoder.embed_positions.weight", random(positions*d), positions, d)
add("final_logits_bias", make([]float32, vocab), 1, vocab)
attn("model.encoder.layers.0.self_attn")
norm("model.encoder.layers.0.self_attn_layer_norm")
lin("model.encoder.layers.0.fc1", d, ffn)
lin("model.encoder.layers.0.fc2", ffn, d)
norm("model.encoder.layers.0.final_layer_norm")
attn("model.decoder.layers.0.self_attn")
norm("model.decoder.layers.0.self_attn_layer_norm")
attn("model.decoder.layers.0.encoder_attn")
norm("model.decoder.layers.0.encoder_attn_layer_norm")
lin("model.decoder.layers.0.fc1", d, ffn)
lin("model.decoder.layers.0.fc2", ffn, d)
norm("model.decoder.layers.0.final_layer_norm")
writeSafetensors(t, filepath.Join(dir, "model.safetensors"), tensors)
return dir
}
func TestTinyModel(t *testing.T) {
dir := tinyModel(t)
if !Installed(dir) {
t.Fatal("tiny model not installed")
}
m, err := Load(dir)
if err != nil {
t.Fatal(err)
}
first, err := m.Translate(context.Background(), "hello hello")
if err != nil {
t.Fatal(err)
}
// Random weights: the output means nothing, but it is deterministic
// and made of target pieces.
again, _ := m.Translate(context.Background(), "hello hello")
if first != again {
t.Errorf("translations differ: %q, %q", first, again)
}
if strings.ContainsAny(first, "▁<>") {
t.Errorf("translation %q contains special pieces", first)
}
if out, _ := m.Translate(context.Background(), " "); out != "" {
t.Errorf("blank text translated to %q", out)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
if _, err := m.Translate(ctx, "hello"); err == nil {
t.Error("cancelled translation succeeded")
}
os.Remove(filepath.Join(dir, "vocab.json"))
if Installed(dir) {
t.Error("incomplete model reported installed")
}
if _, err := Load(dir); err == nil {
t.Error("incomplete model loaded")
}
}
// TestRealModel runs the OPUS-MT model when PLOTZILLA_MT_MODEL names its
// directory (~/.local/share/plotzilla/models/opus-mt-fr-en once
// downloaded by the app).
func TestRealModel(t *testing.T) {
dir := os.Getenv("PLOTZILLA_MT_MODEL")
if dir == "" {
t.Skip("PLOTZILLA_MT_MODEL not set")
}
m, err := Load(dir)
if err != nil {
t.Fatal(err)
}
tests := map[string]string{
"Bonjour, comment allez-vous ?": "Hello, how are you?",
"Je t'avais dit de ne jamais revenir ici.": "I told you never to come back here.",
}
for in, want := range tests {
if got, err := m.Translate(context.Background(), in); err != nil || got != want {
t.Errorf("Translate(%q) = %q, %v, want %q", in, got, err, want)
}
}
}
func TestDownload(t *testing.T) {
content := []byte("model weights")
sum := sha256.Sum256(content)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/missing" {
http.NotFound(w, r)
return
}
w.Write(content)
}))
defer srv.Close()
dir := filepath.Join(t.TempDir(), "model")
file := RemoteFile{Name: "f.bin", URL: srv.URL + "/f", Size: int64(len(content)), SHA256: hex.EncodeToString(sum[:])}
var last int64
if err := Download(context.Background(), srv.Client(), dir, []RemoteFile{file}, func(done, total int64) { last = done }); err != nil {
t.Fatal(err)
}
if got, _ := os.ReadFile(filepath.Join(dir, "f.bin")); string(got) != string(content) || last != file.Size {
t.Errorf("downloaded %q, progress %d", got, last)
}
bad := file
bad.Name, bad.SHA256 = "bad.bin", strings.Repeat("0", 64)
missing := file
missing.Name, missing.URL = "missing.bin", srv.URL+"/missing"
short := file
short.Name, short.Size = "short.bin", file.Size+5
for _, f := range []RemoteFile{bad, missing, short} {
if err := Download(context.Background(), srv.Client(), dir, []RemoteFile{f}, nil); err == nil {
t.Errorf("%s: no error", f.Name)
}
if _, err := os.Stat(filepath.Join(dir, f.Name)); err == nil {
t.Errorf("%s: a failed download left the file", f.Name)
}
}
entries, _ := os.ReadDir(dir)
if len(entries) != 1 {
t.Errorf("temporary files left: %v", entries)
}
}
+376
View File
@@ -0,0 +1,376 @@
// Package marian runs MarianMT translation models (the OPUS-MT models of
// the University of Helsinki, in their Hugging Face form) on the CPU, in
// plain Go: a SentencePiece tokenizer, the transformer encoder-decoder and
// beam search. A model directory holds config.json, model.safetensors,
// source.spm, target.spm and vocab.json.
package marian
import (
"context"
"encoding/json"
"fmt"
"math"
"os"
"path/filepath"
)
// Files are the files a model directory must contain.
var Files = []string{"config.json", "model.safetensors", "source.spm", "target.spm", "vocab.json"}
type config struct {
DModel int `json:"d_model"`
EncoderLayers int `json:"encoder_layers"`
DecoderLayers int `json:"decoder_layers"`
EncoderHeads int `json:"encoder_attention_heads"`
DecoderHeads int `json:"decoder_attention_heads"`
Activation string `json:"activation_function"`
ScaleEmbedding bool `json:"scale_embedding"`
NormalizeBefore bool `json:"normalize_before"`
MaxPositions int `json:"max_position_embeddings"`
VocabSize int `json:"vocab_size"`
PadToken int `json:"pad_token_id"`
EOSToken int `json:"eos_token_id"`
DecoderStart int `json:"decoder_start_token_id"`
NumBeams int `json:"num_beams"`
}
type linear struct {
w, b []float32
in, out int
}
type layerNorm struct{ g, b []float32 }
type attention struct {
q, k, v, o linear
heads int
}
type encoderLayer struct {
self attention
selfNorm layerNorm
fc1, fc2 linear
finalLayerNorm layerNorm
}
type decoderLayer struct {
self, cross attention
selfNorm, crossNorm layerNorm
fc1, fc2 linear
finalLayerNorm layerNorm
}
// Model is a loaded translation model. It is read-only once loaded:
// Translate may be called from several goroutines at once.
type Model struct {
cfg config
embed []float32 // vocab × d, shared by encoder, decoder and output
encPos []float32 // positions × d
decPos []float32
logitsBias []float32
encoder []encoderLayer
decoder []decoderLayer
source *sentencePiece
target *sentencePiece
ids map[string]int
pieces []string
unk int
embedScale float32
activation func(float32) float32
}
// Load reads the model in dir.
func Load(dir string) (*Model, error) {
var cfg config
data, err := os.ReadFile(filepath.Join(dir, "config.json"))
if err != nil {
return nil, err
}
if err := json.Unmarshal(data, &cfg); err != nil {
return nil, fmt.Errorf("reading config.json: %w", err)
}
if cfg.DModel == 0 || cfg.EncoderHeads == 0 || cfg.DecoderHeads == 0 || cfg.DModel%cfg.EncoderHeads != 0 || cfg.DModel%cfg.DecoderHeads != 0 {
return nil, fmt.Errorf("unsupported model dimensions in config.json")
}
if cfg.NormalizeBefore {
return nil, fmt.Errorf("unsupported model: pre-norm layers")
}
m := &Model{cfg: cfg, embedScale: 1}
if cfg.ScaleEmbedding {
m.embedScale = float32(math.Sqrt(float64(cfg.DModel)))
}
switch cfg.Activation {
case "swish", "silu":
m.activation = func(x float32) float32 { return x / (1 + float32(math.Exp(float64(-x)))) }
case "relu":
m.activation = func(x float32) float32 { return max(x, 0) }
case "gelu":
m.activation = func(x float32) float32 { return 0.5 * x * (1 + float32(math.Erf(float64(x)/math.Sqrt2))) }
default:
return nil, fmt.Errorf("unsupported activation %q", cfg.Activation)
}
if m.source, err = loadSentencePiece(filepath.Join(dir, "source.spm")); err != nil {
return nil, err
}
if m.target, err = loadSentencePiece(filepath.Join(dir, "target.spm")); err != nil {
return nil, err
}
if data, err = os.ReadFile(filepath.Join(dir, "vocab.json")); err != nil {
return nil, err
}
if err := json.Unmarshal(data, &m.ids); err != nil {
return nil, fmt.Errorf("reading vocab.json: %w", err)
}
m.pieces = make([]string, cfg.VocabSize)
for p, id := range m.ids {
if id >= 0 && id < len(m.pieces) {
m.pieces[id] = p
}
}
m.unk = m.ids["<unk>"]
tensors, err := readSafetensors(filepath.Join(dir, "model.safetensors"))
if err != nil {
return nil, fmt.Errorf("reading model.safetensors: %w", err)
}
if err := m.bind(tensors); err != nil {
return nil, fmt.Errorf("reading model.safetensors: %w", err)
}
return m, nil
}
// bind takes the weights the model needs from the tensors, checking their
// shapes.
func (m *Model) bind(t map[string]tensor) error {
d := m.cfg.DModel
var missing error
get := func(name string, shape ...int) []float32 {
x, ok := t[name]
if !ok {
if missing == nil {
missing = fmt.Errorf("missing tensor %s", name)
}
return nil
}
n := 1
for _, s := range shape {
n *= s
}
if len(x.data) != n && missing == nil {
missing = fmt.Errorf("tensor %s has shape %v", name, x.shape)
}
return x.data
}
first := func(names ...string) string {
for _, n := range names {
if _, ok := t[n]; ok {
return n
}
}
return names[0]
}
lin := func(prefix string, in, out int) linear {
return linear{w: get(prefix+".weight", out, in), b: get(prefix+".bias", out), in: in, out: out}
}
norm := func(prefix string) layerNorm {
return layerNorm{g: get(prefix+".weight", d), b: get(prefix+".bias", d)}
}
attn := func(prefix string, heads int) attention {
return attention{
q: lin(prefix+".q_proj", d, d), k: lin(prefix+".k_proj", d, d),
v: lin(prefix+".v_proj", d, d), o: lin(prefix+".out_proj", d, d), heads: heads,
}
}
ffn := func(prefix string) int {
if x, ok := t[prefix+".fc1.weight"]; ok && len(x.shape) == 2 {
return x.shape[0]
}
return 4 * d
}
v := m.cfg.VocabSize
m.embed = get(first("model.shared.weight", "model.decoder.embed_tokens.weight", "model.encoder.embed_tokens.weight"), v, d)
m.encPos = get("model.encoder.embed_positions.weight", m.cfg.MaxPositions, d)
m.decPos = get("model.decoder.embed_positions.weight", m.cfg.MaxPositions, d)
m.logitsBias = get("final_logits_bias", 1, v)
for i := range m.cfg.EncoderLayers {
p := fmt.Sprintf("model.encoder.layers.%d", i)
f := ffn(p)
m.encoder = append(m.encoder, encoderLayer{
self: attn(p+".self_attn", m.cfg.EncoderHeads), selfNorm: norm(p + ".self_attn_layer_norm"),
fc1: lin(p+".fc1", d, f), fc2: lin(p+".fc2", f, d), finalLayerNorm: norm(p + ".final_layer_norm"),
})
}
for i := range m.cfg.DecoderLayers {
p := fmt.Sprintf("model.decoder.layers.%d", i)
f := ffn(p)
m.decoder = append(m.decoder, decoderLayer{
self: attn(p+".self_attn", m.cfg.DecoderHeads), selfNorm: norm(p + ".self_attn_layer_norm"),
cross: attn(p+".encoder_attn", m.cfg.DecoderHeads), crossNorm: norm(p + ".encoder_attn_layer_norm"),
fc1: lin(p+".fc1", d, f), fc2: lin(p+".fc2", f, d), finalLayerNorm: norm(p + ".final_layer_norm"),
})
}
return missing
}
// maxSourceTokens leaves room for the end-of-sentence token.
func (m *Model) maxSourceTokens() int { return m.cfg.MaxPositions - 1 }
// Translate translates one text. Long texts are better split into
// sentences first: the models are trained on sentences.
func (m *Model) Translate(ctx context.Context, text string) (string, error) {
pieces := m.source.encode(text)
if len(pieces) == 0 {
return "", nil
}
ids := make([]int, 0, len(pieces)+1)
for _, p := range pieces {
id, ok := m.ids[p]
if !ok {
id = m.unk
}
ids = append(ids, id)
}
if len(ids) > m.maxSourceTokens() {
ids = ids[:m.maxSourceTokens()]
}
ids = append(ids, m.cfg.EOSToken)
out, err := m.beamSearch(ctx, m.encode(ids), len(ids))
if err != nil {
return "", err
}
words := make([]string, 0, len(out))
for _, id := range out {
if id == m.cfg.EOSToken || id == m.cfg.PadToken || id == m.unk || id < 0 || id >= len(m.pieces) {
continue
}
words = append(words, m.pieces[id])
}
return decodePieces(words), nil
}
// encode runs the encoder: one d-vector per source token.
func (m *Model) encode(ids []int) []float32 {
d := m.cfg.DModel
n := len(ids)
x := make([]float32, n*d)
for i, id := range ids {
row := x[i*d : (i+1)*d]
emb := m.embed[id*d : (id+1)*d]
pos := m.encPos[i*d : (i+1)*d]
for j := range row {
row[j] = emb[j]*m.embedScale + pos[j]
}
}
q := make([]float32, n*d)
k := make([]float32, n*d)
v := make([]float32, n*d)
ctxv := make([]float32, n*d)
tmp := make([]float32, n*d)
scores := make([]float32, n)
for _, l := range m.encoder {
l.self.q.rows(q, x, n)
l.self.k.rows(k, x, n)
l.self.v.rows(v, x, n)
for i := range n {
attend(ctxv[i*d:(i+1)*d], q[i*d:(i+1)*d], k, v, n, l.self.heads, scores)
}
l.self.o.rows(tmp, ctxv, n)
addNorm(x, tmp, l.selfNorm, n)
hidden := make([]float32, n*l.fc1.out)
l.fc1.rows(hidden, x, n)
for i := range hidden {
hidden[i] = m.activation(hidden[i])
}
l.fc2.rows(tmp, hidden, n)
addNorm(x, tmp, l.finalLayerNorm, n)
}
return x
}
// rows applies the layer to n row vectors: dst = x·Wᵀ + b.
func (l linear) rows(dst, x []float32, n int) {
for i := range n {
l.apply(dst[i*l.out:(i+1)*l.out], x[i*l.in:(i+1)*l.in])
}
}
func (l linear) apply(dst, x []float32) {
for o := range l.out {
dst[o] = dot(x, l.w[o*l.in:(o+1)*l.in]) + l.b[o]
}
}
// attend computes multi-head scaled dot-product attention of one query
// over n keys and values (n × d each).
func attend(dst, q, k, v []float32, n, heads int, scores []float32) {
d := len(q)
hd := d / heads
scale := float32(1 / math.Sqrt(float64(hd)))
clear(dst)
for h := range heads {
lo, hi := h*hd, (h+1)*hd
maxScore := float32(math.Inf(-1))
for j := range n {
s := dot(q[lo:hi], k[j*d+lo:j*d+hi]) * scale
scores[j] = s
maxScore = max(maxScore, s)
}
var sum float32
for j := range n {
scores[j] = float32(math.Exp(float64(scores[j] - maxScore)))
sum += scores[j]
}
out := dst[lo:hi]
for j := range n {
p := scores[j] / sum
row := v[j*d+lo : j*d+hi]
for c := range out {
out[c] += p * row[c]
}
}
}
}
// addNorm sets x = LayerNorm(x + y) for n rows.
func addNorm(x, y []float32, ln layerNorm, n int) {
d := len(ln.g)
for i := range n {
row := x[i*d : (i+1)*d]
add := y[i*d : (i+1)*d]
var mean float64
for j := range row {
row[j] += add[j]
mean += float64(row[j])
}
mean /= float64(d)
var variance float64
for _, val := range row {
diff := float64(val) - mean
variance += diff * diff
}
inv := 1 / math.Sqrt(variance/float64(d)+1e-5)
for j := range row {
row[j] = float32((float64(row[j])-mean)*inv)*ln.g[j] + ln.b[j]
}
}
}
// dot is the hot loop of the whole model: unrolled, with independent
// accumulators so the CPU can overlap the additions.
func dot(a, b []float32) float32 {
b = b[:len(a)]
var s0, s1, s2, s3 float32
i := 0
for ; i+8 <= len(a); i += 8 {
s0 += a[i]*b[i] + a[i+4]*b[i+4]
s1 += a[i+1]*b[i+1] + a[i+5]*b[i+5]
s2 += a[i+2]*b[i+2] + a[i+6]*b[i+6]
s3 += a[i+3]*b[i+3] + a[i+7]*b[i+7]
}
for ; i < len(a); i++ {
s0 += a[i] * b[i]
}
return s0 + s1 + s2 + s3
}
+137
View File
@@ -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)
}
+225
View File
@@ -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, " "))
}
+136
View File
@@ -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
}
+160
View File
@@ -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)) }
+957
View File
@@ -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])
}
+307
View File
@@ -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)
}
}
}