Maven/internal/router/onnxembedder.go

package router

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"strings"
	"sync"

	ort "github.com/yalue/onnxruntime_go"
	"golang.org/x/text/unicode/norm"
)

// The deployed model is multilingual-e5-small. e5 was trained with these two
// words glued to the front of every text, and it scores badly without them —
// they are part of the model, not a style choice. Swapping back to a symmetric
// paraphrase model means dropping them again.
const (
	queryPrefix   = "query: "
	passagePrefix = "passage: "
)

const (
	padTokenID = 1
	unkTokenID = 3
	clsTokenID = 0
	sepTokenID = 2
	maxLength  = 128
	embedDim   = 384
)

type onnxEmbedder struct {
	tokenizer *unigramTokenizer
	session   *ort.DynamicSession[int64, float32]
	runtime   *ONNXRuntimeLease
	id        string
	closeOnce sync.Once
	closeErr  error
}

func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, error) {
	runtime, err := AcquireONNXRuntime(libPath)
	if err != nil {
		return nil, err
	}

	tok, err := newUnigramTokenizer(tokenizerPath)
	if err != nil {
		_ = runtime.Close()
		return nil, fmt.Errorf("tokenizer: %w", err)
	}

	session, err := ort.NewDynamicSession[int64, float32](
		modelPath,
		[]string{"input_ids", "attention_mask", "token_type_ids"},
		[]string{"last_hidden_state"},
	)
	if err != nil {
		_ = runtime.Close()
		return nil, fmt.Errorf("onnx: create session: %w", err)
	}

	return &onnxEmbedder{
		tokenizer: tok,
		session:   session,
		runtime:   runtime,
		id:        modelIDFromPath(modelPath),
	}, nil
}

func (e *onnxEmbedder) Dim() int { return embedDim }

// ID names the loaded model for the DB marker (Vikunja #378): the model file's
// own name, the dimension, and the tokenizer revision, so pointing the config
// at another model changes the string on its own.
func (e *onnxEmbedder) ID() string { return e.id }

// tokenizerRev — bumped whenever the tokenizer changes what it emits for the
// same text, because that changes every vector while the model file's name
// stays put. Rev 2 is the fix for the reversed word pieces (V-664): stored
// passages embedded under rev 1 no longer sit in the same space as a query
// embedded now, and ReembedAll rewrites them because this string moved.
const tokenizerRev = 2

// modelIDFromPath turns /opt/.../multilingual-e5-small.onnx into
// "multilingual-e5-small@384/tok2".
func modelIDFromPath(modelPath string) string {
	name := modelPath
	if i := strings.LastIndexAny(name, "/\\"); i >= 0 {
		name = name[i+1:]
	}
	name = strings.TrimSuffix(name, ".onnx")
	if name == "" {
		name = "onnx"
	}
	return fmt.Sprintf("%s@%d/tok%d", name, embedDim, tokenizerRev)
}

// Embed treats the text as a query. The classifier compares one short
// utterance to another short seed phrase, so both sides get the same prefix
// and the comparison stays fair. The recall path must call EmbedQuery and
// EmbedPassage instead.
func (e *onnxEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
	return e.embed(ctx, queryPrefix+text)
}

// EmbedQuery — the question the user just asked.
func (e *onnxEmbedder) EmbedQuery(ctx context.Context, text string) ([]float32, error) {
	return e.embed(ctx, queryPrefix+text)
}

// EmbedPassage — a note or fact being stored, or re-scored at lookup time.
func (e *onnxEmbedder) EmbedPassage(ctx context.Context, text string) ([]float32, error) {
	return e.embed(ctx, passagePrefix+text)
}

func (e *onnxEmbedder) embed(ctx context.Context, text string) ([]float32, error) {
	inputIDs, attentionMask, _ := e.tokenizer.Encode(text)

	inputShape := ort.NewShape(1, int64(maxLength))
	inputT, err := ort.NewTensor(inputShape, inputIDs)
	if err != nil {
		return nil, fmt.Errorf("onnx: input tensor: %w", err)
	}
	defer inputT.Destroy()
	maskT, err := ort.NewTensor(inputShape, attentionMask)
	if err != nil {
		return nil, fmt.Errorf("onnx: mask tensor: %w", err)
	}
	defer maskT.Destroy()
	typeT, err := ort.NewTensor(inputShape, make([]int64, maxLength))
	if err != nil {
		return nil, fmt.Errorf("onnx: type tensor: %w", err)
	}
	defer typeT.Destroy()

	outputShape := ort.NewShape(1, int64(maxLength), embedDim)
	outputT, err := ort.NewTensor(outputShape, make([]float32, maxLength*embedDim))
	if err != nil {
		return nil, fmt.Errorf("onnx: output tensor: %w", err)
	}
	defer outputT.Destroy()

	if err := e.session.Run(
		[]*ort.Tensor[int64]{inputT, maskT, typeT},
		[]*ort.Tensor[float32]{outputT},
	); err != nil {
		return nil, fmt.Errorf("onnx: run: %w", err)
	}

	emb := meanPool(outputT.GetData(), attentionMask, maxLength, embedDim)
	return emb, nil
}

func (e *onnxEmbedder) Close() error {
	if e == nil {
		return nil
	}
	e.closeOnce.Do(func() {
		if e.session != nil {
			e.session.Destroy()
			e.session = nil
		}
		e.closeErr = e.runtime.Close()
	})
	return e.closeErr
}

func meanPool(hidden []float32, mask []int64, seqLen, dim int) []float32 {
	out := make([]float32, dim)
	var maskSum float32
	for i := 0; i < seqLen; i++ {
		if mask[i] == 0 {
			continue
		}
		maskSum++
		for j := 0; j < dim; j++ {
			out[j] += hidden[i*dim+j]
		}
	}
	if maskSum > 0 {
		for j := 0; j < dim; j++ {
			out[j] /= maskSum
		}
	}
	var sumSq float64
	for _, v := range out {
		sumSq += float64(v) * float64(v)
	}
	if sumSq > 0 {
		inv := float32(1.0 / math.Sqrt(sumSq))
		for i := range out {
			out[i] *= inv
		}
	}
	return out
}

type unigramTokenizer struct {
	vocab    map[string]vocabEntry
	unkScore float64
}

type vocabEntry struct {
	id    int64
	score float64
}

func newUnigramTokenizer(path string) (*unigramTokenizer, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read tokenizer.json: %w", err)
	}

	var raw struct {
		Model struct {
			Type  string          `json:"type"`
			Vocab json.RawMessage `json:"vocab"`
		} `json:"model"`
	}
	if err := json.Unmarshal(data, &raw); err != nil {
		return nil, fmt.Errorf("parse tokenizer.json: %w", err)
	}
	if raw.Model.Type != "Unigram" {
		return nil, fmt.Errorf("unsupported tokenizer type: %s", raw.Model.Type)
	}

	var rawVocab [][]json.RawMessage
	if err := json.Unmarshal(raw.Model.Vocab, &rawVocab); err != nil {
		return nil, fmt.Errorf("parse vocab: %w", err)
	}

	vocab := make(map[string]vocabEntry, len(rawVocab))
	var unkScore float64
	for _, pair := range rawVocab {
		if len(pair) < 2 {
			continue
		}
		var token string
		if err := json.Unmarshal(pair[0], &token); err != nil {
			continue
		}
		var score float64
		if err := json.Unmarshal(pair[1], &score); err != nil {
			continue
		}
		vocab[token] = vocabEntry{score: score}
	}
	// Assign IDs based on order
	i := int64(0)
	for _, pair := range rawVocab {
		var token string
		if err := json.Unmarshal(pair[0], &token); err != nil {
			continue
		}
		e := vocab[token]
		e.id = i
		vocab[token] = e
		if i == unkTokenID {
			unkScore = e.score
		}
		i++
	}

	return &unigramTokenizer{vocab: vocab, unkScore: unkScore}, nil
}

func (t *unigramTokenizer) Encode(text string) (inputIDs, attentionMask, tokenTypeIDs []int64) {
	tokens := t.tokenize(text)

	tokens = append([]int64{clsTokenID}, tokens...)
	tokens = append(tokens, sepTokenID)

	if len(tokens) > maxLength {
		tokens = tokens[:maxLength-1]
		tokens = append(tokens, sepTokenID)
	}

	inputIDs = make([]int64, maxLength)
	attentionMask = make([]int64, maxLength)
	tokenTypeIDs = make([]int64, maxLength)
	for i, id := range tokens {
		inputIDs[i] = id
		attentionMask[i] = 1
	}
	return
}

func (t *unigramTokenizer) tokenize(text string) []int64 {
	words := preTokenize(text)
	var ids []int64
	for _, word := range words {
		wordIDs := t.encodeWord(word)
		ids = append(ids, wordIDs...)
	}
	return ids
}

type cand struct {
	start int
	end   int
	id    int64
	score float64
}

func (t *unigramTokenizer) encodeWord(word string) []int64 {
	runes := []rune(word)
	n := len(runes)
	if n == 0 {
		return nil
	}

	var candidates []cand
	for i := 0; i < n; i++ {
		for j := i + 1; j <= n && j-i <= 50; j++ {
			sub := string(runes[i:j])
			if e, ok := t.vocab[sub]; ok {
				candidates = append(candidates, cand{
					start: i, end: j, id: e.id, score: e.score,
				})
			}
		}
	}

	dp := make([]float64, n+1)
	prev := make([]int, n+1)
	bestID := make([]int64, n+1)
	filled := make([]bool, n+1)
	dp[0] = 0
	filled[0] = true

	for i := 1; i <= n; i++ {
		bestScore := math.Inf(-1)
		bestPrev := -1
		bestTokenID := int64(unkTokenID)

		for _, c := range candidates {
			if c.end == i && filled[c.start] {
				candScore := dp[c.start] + c.score
				if candScore > bestScore {
					bestScore = candScore
					bestPrev = c.start
					bestTokenID = c.id
				}
			}
		}

		if bestScore == math.Inf(-1) {
			if filled[i-1] {
				dp[i] = dp[i-1] + t.unkScore
				prev[i] = i - 1
				bestID[i] = unkTokenID
				filled[i] = true
			}
		} else {
			dp[i] = bestScore
			prev[i] = bestPrev
			bestID[i] = bestTokenID
			filled[i] = true
		}
	}

	// Backtracking walks the word from its end, and prepending each piece puts
	// it back in reading order. There used to be a second reverse after this
	// loop, which undid it: every multi-piece word came out backwards, and
	// "query: вода" tokenized to [0 12 1294 41 12489 2] where the reference
	// tokenizer gives [0 41 1294 12 12489 2] (V-664). A transformer reads
	// position, so the pieces of a long Russian word were being read in the
	// wrong order on every turn.
	var result []int64
	for i := n; i > 0; i = prev[i] {
		result = append([]int64{bestID[i]}, result...)
	}
	return result
}

func preTokenize(text string) []string {
	text = norm.NFKC.String(text)
	text = strings.ToLower(text)
	pieces := strings.Fields(text)
	out := make([]string, 0, len(pieces))
	for _, p := range pieces {
		out = append(out, "\u2581"+p)
	}
	return out
}

var _ AsymmetricEmbedder = (*onnxEmbedder)(nil)