diff --git a/internal/router/heads.go b/internal/router/heads.go new file mode 100644 index 0000000..acb2b99 --- /dev/null +++ b/internal/router/heads.go @@ -0,0 +1,243 @@ +package router + +import ( + "context" + "encoding/json" + "fmt" + "math" + "os" + "path/filepath" + + ort "github.com/yalue/onnxruntime_go" +) + +// The routing heads (V-546, V-661, V-664). Four linear heads over one masked +// mean pool of a fine-tuned copy of multilingual-e5-small: intent, +// destination, BIO slot tags and clarify. Trained on workpc, exported to ONNX, +// and read here. +// +// Why this is not the classifier. The classifier compares one utterance to +// frozen seed phrases by cosine. A head is a softmax over the label set, so it +// cannot name a value that does not exist, and its max is a calibratable +// confidence where Confidence: 1.0 was a hardcode. +// +// Why it is not the resident model either. It answers in single-digit +// milliseconds against the model's p50 of 1.19s, and it names a destination +// the classifier arm never names at all. +// +// The body is a COPY of the embedder weights, fine-tuned. It must never +// replace models/embedder/multilingual-e5-small — memory recall depends on +// that file scoring what it scored. +// +// The slot head is exported and deliberately not read. Slots already come from +// the stage-2 extractor, and mapping BIO tags back to text needs character +// offsets the unigram tokenizer does not keep. Reading it is separate work. +const ( + // headsSeq — the sequence length the heads were trained at. Padding is + // masked out of both attention and the pool, so this changes nothing but + // truncation, and truncation is what training did at 64. + headsSeq = 64 + + // headsThreshold — max softmax over the intent head, below which the heads + // decline and the cascade carries on to the resident model. + // + // 0.6 is the knee measured on the 88-case intent fixture + // (docs/evals/2026-08-08-routing-heads-in-go.md). It keeps 81 of 88 cases + // at 97.5% accuracy. Every higher value up to 0.9 drops right answers and + // keeps the same two wrong ones, so it buys nothing. + headsThreshold = 0.6 +) + +// RouterHeads runs the exported graph. Nil is a working value everywhere: a +// deployment with no weights file routes exactly as it did before this +// existed. +type RouterHeads struct { + tokenizer *unigramTokenizer + session *ort.DynamicSession[int64, float32] + intents []Intent + sources []Source + threshold float64 +} + +// headsMeta — router_heads.json, written beside the weights by the exporter. +// The label order is the head's output order and cannot be inferred from Go. +type headsMeta struct { + Intents []string `json:"intents"` + Sources []string `json:"sources"` + Prefix string `json:"prefix"` +} + +// NewRouterHeads loads the graph and its label order. modelPath points at the +// .onnx; the external weights and router_heads.json sit beside it. +// +// It assumes the ONNX environment is already initialised, because the embedder +// does that at startup and the runtime allows it once. +func NewRouterHeads(modelPath, tokenizerPath string) (*RouterHeads, error) { + metaPath := filepath.Join(filepath.Dir(modelPath), "router_heads.json") + raw, err := os.ReadFile(metaPath) + if err != nil { + return nil, fmt.Errorf("heads: read %s: %w", metaPath, err) + } + var meta headsMeta + if err := json.Unmarshal(raw, &meta); err != nil { + return nil, fmt.Errorf("heads: parse %s: %w", metaPath, err) + } + if meta.Prefix != queryPrefix { + return nil, fmt.Errorf("heads: trained with prefix %q, this build uses %q", + meta.Prefix, queryPrefix) + } + + intents := make([]Intent, len(meta.Intents)) + for i, s := range meta.Intents { + intents[i] = Intent(s) + } + sources := make([]Source, len(meta.Sources)) + for i, s := range meta.Sources { + // SourceUnknown is not in Sources, because it is the absence of a + // choice. It is a class the head can emit, and the one it should emit + // often, so it is allowed here and nowhere else. + if s != string(SourceUnknown) && !ValidSource(Source(s)) { + return nil, fmt.Errorf("heads: unknown destination %q in %s", s, metaPath) + } + sources[i] = Source(s) + } + + tok, err := newUnigramTokenizer(tokenizerPath) + if err != nil { + return nil, fmt.Errorf("heads: tokenizer: %w", err) + } + session, err := ort.NewDynamicSession[int64, float32]( + modelPath, + []string{"input_ids", "attention_mask"}, + []string{"intent", "source", "slots", "clarify"}, + ) + if err != nil { + return nil, fmt.Errorf("heads: create session: %w", err) + } + + return &RouterHeads{ + tokenizer: tok, + session: session, + intents: intents, + sources: sources, + threshold: headsThreshold, + }, nil +} + +func (h *RouterHeads) Close() error { + if h == nil { + return nil + } + h.session.Destroy() + return nil +} + +// headsResult — one forward pass, read back. +type headsResult struct { + Intent Intent + Source Source + Confidence float64 + Clarify bool +} + +// Route runs the heads and reports whether they are confident enough to answer. +// A false second return is a decline, not an error: the cascade goes on to the +// resident model, which is what happens today. +func (h *RouterHeads) Route(ctx context.Context, utterance string) (headsResult, bool, error) { + if h == nil { + return headsResult{}, false, nil + } + ids, mask, _ := h.tokenizer.Encode(queryPrefix + utterance) + ids, mask = ids[:headsSeq], mask[:headsSeq] + // The tokenizer pads and truncates to its own length, which is longer than + // this one. Cutting the tail can cut the separator with it, so put it back. + if mask[headsSeq-1] == 1 { + ids[headsSeq-1] = sepTokenID + } + + shape := ort.NewShape(1, headsSeq) + idsT, err := ort.NewTensor(shape, ids) + if err != nil { + return headsResult{}, false, fmt.Errorf("heads: ids tensor: %w", err) + } + defer idsT.Destroy() + maskT, err := ort.NewTensor(shape, mask) + if err != nil { + return headsResult{}, false, fmt.Errorf("heads: mask tensor: %w", err) + } + defer maskT.Destroy() + + intentT, err := ort.NewEmptyTensor[float32](ort.NewShape(1, int64(len(h.intents)))) + if err != nil { + return headsResult{}, false, fmt.Errorf("heads: intent tensor: %w", err) + } + defer intentT.Destroy() + sourceT, err := ort.NewEmptyTensor[float32](ort.NewShape(1, int64(len(h.sources)))) + if err != nil { + return headsResult{}, false, fmt.Errorf("heads: source tensor: %w", err) + } + defer sourceT.Destroy() + slotsT, err := ort.NewEmptyTensor[float32](ort.NewShape(1, headsSeq, int64(numBIOTags))) + if err != nil { + return headsResult{}, false, fmt.Errorf("heads: slots tensor: %w", err) + } + defer slotsT.Destroy() + clarifyT, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 2)) + if err != nil { + return headsResult{}, false, fmt.Errorf("heads: clarify tensor: %w", err) + } + defer clarifyT.Destroy() + + if err := h.session.Run( + []*ort.Tensor[int64]{idsT, maskT}, + []*ort.Tensor[float32]{intentT, sourceT, slotsT, clarifyT}, + ); err != nil { + return headsResult{}, false, fmt.Errorf("heads: run: %w", err) + } + + // The graph applies its own softmax, so these are probabilities and the max + // is the same number the eval calibrated the threshold against. + i, conf := argmax(intentT.GetData()) + res := headsResult{ + Intent: h.intents[i], + Confidence: conf, + } + cl := clarifyT.GetData() + res.Clarify = len(cl) == 2 && cl[1] > cl[0] + + // The destination head is trained on query rows and is meaningless on any + // other intent, the same way queryWalk is never reached by one. + if res.Intent == IntentQuery { + s, _ := argmax(sourceT.GetData()) + res.Source = h.sources[s] + } + + // The clarify head decides on its own, and it decides first. It answers a + // different question from the intent head — not which intent, but whether + // there is enough here to act on at all — so a low intent confidence is no + // reason to discard it. It is usually the same turns: "вода" reads as + // intent act at 0.23 and clarify at 0.98, and letting the intent threshold + // bury that hands the turn to the classifier, which routes it confidently + // and never asks. + if res.Clarify { + return res, true, nil + } + if conf < h.threshold { + return res, false, nil + } + return res, true, nil +} + +// numBIOTags — O plus B- and I- for each of Maven's five slots. The head is not +// read, but the graph writes it and the output tensor has to be the right size. +const numBIOTags = 11 + +func argmax(v []float32) (int, float64) { + best, bestV := 0, math.Inf(-1) + for i, x := range v { + if float64(x) > bestV { + best, bestV = i, float64(x) + } + } + return best, bestV +}