feat: add a translate stack

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, " "))
+}