diff --git a/CLAUDE.md b/CLAUDE.md index fc4fe84..558f6e2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -308,9 +308,36 @@ the possessive agenda rules claim those cases at stage 0 and name nothing, so no label reaches the head. That is the same trade V-660 flagged and it wants the owner's call. -**Nothing of this runs in Go.** The weights are `heads.pt` and `out/body_heads/` -on workpc. Reaching the daemon needs an ONNX export and a caller. The resident -e5-small must not be replaced by the copy, because recall depends on that file. +**The heads run in Go and route every turn, since 08-08-2026** (V-664, +`docs/evals/2026-08-08-routing-heads-in-go.md`). This section used to say +nothing of it ran. `RouterHeads` in `internal/router/heads.go` loads +`router_heads.onnx` and reads intent, destination and clarify off one forward +pass. It is stage 0b: after the grammars, **before** the resident model, and the +classifier is still behind both. Through the cascade it scores intent **96.9%** +and destination **75.8%** at p50 27.9ms. That beats the gemma-4-12b cascade, +84.4% and 72.7%, at a twelfth of its 329ms. The workstation stays the better +phraser and is no longer the better router. + +Three rules around it. The **clarify head decides first**, before the intent +threshold. It answers a different question. A thin utterance scores low +intent by construction, so gating it cost 6 of 8 ambiguous cases. The +**destination head is read on `IntentQuery` only**, since no other intent +reaches `queryWalk`. And `headsThreshold` is 0.6, the measured knee: every value +to 0.85 drops right answers and keeps the same two wrong ones. + +`voice.embedder.heads_path` is the whole switch. Empty, missing or unloadable +means the heads are nil and the cascade is byte-for-byte what shipped before +them. **It must never be pointed at `model_path`.** The resident e5-small must +not be replaced by the fine-tuned copy. Recall depends on that file scoring +what it scored. + +**The hand-written tokenizer read every long word backwards** until this task +(`encodeWord`, `onnxembedder.go`). It cost recall@1 7.4 points and recall@3 11.1. +Nothing caught it because seeds and queries were mangled the same way, so cosine +survived. The heads found it. They are trained through transformers and read +through this. The embedder id now carries a tokenizer revision +(`@384/tok2`), so fixing the tokenizer triggers `ReembedAll` the way swapping the +model file does. Bump `tokenizerRev` on any change to what it emits. `Confidence: 1.0` used to be hardcoded in `llmrouter.go`, so the LLM path could never ask for clarification (6/6 refusal cases missed on the fixture) — Vikunja diff --git a/cmd/mavend/chat_degrade_test.go b/cmd/mavend/chat_degrade_test.go index c376aa0..f127ee0 100644 --- a/cmd/mavend/chat_degrade_test.go +++ b/cmd/mavend/chat_degrade_test.go @@ -19,7 +19,7 @@ func TestChatAnswersWithNoLlamaServer(t *testing.T) { dead := llm.New("http://127.0.0.1:1", 500*time.Millisecond) emb := router.NewHashEmbedder(1024) h.recall.embedder = emb - h.router = buildRouter(emb, h.matcher, 0.55, pickLLMRouter(true, dead)) + h.router = buildRouter(emb, h.matcher, 0.55, pickLLMRouter(true, dead), nil) h.replier = newLLMReplier(dead, nil) ctx := withDialogueID(context.Background(), dialogueIDFor(sourceText, "web")) diff --git a/cmd/mavend/clarify_test.go b/cmd/mavend/clarify_test.go index 946c393..ec7a47a 100644 --- a/cmd/mavend/clarify_test.go +++ b/cmd/mavend/clarify_test.go @@ -317,7 +317,7 @@ func TestClarifyExpiryIsAnnouncedAndWordsStillRoute(t *testing.T) { h, _, now := newClarifyHandler(t) emb := router.NewHashEmbedder(1024) h.recall.embedder = emb - h.router = buildRouter(emb, h.matcher, 0.55, nil) + h.router = buildRouter(emb, h.matcher, 0.55, nil, nil) if _, asked := h.askClarify(ctx, clarifyDec(router.IntentReminder, router.Slots{Text: "напомни"}, "напомни")); !asked { t.Fatal("expected a question") @@ -671,7 +671,7 @@ func TestUnresolvedActSaysItDoesNotKnowTheCommand(t *testing.T) { func newRoutingClarifyHandler(t *testing.T) (*reactiveHandler, *store.Store) { t.Helper() h, st, _ := newClarifyHandler(t) - h.router = buildRouter(router.NewHashEmbedder(1024), h.matcher, 0.55, nil) + h.router = buildRouter(router.NewHashEmbedder(1024), h.matcher, 0.55, nil, nil) h.recall = recallWiring{embedder: router.NewHashEmbedder(1024), memStore: memory.NewInMemoryStore()} return h, st } diff --git a/cmd/mavend/decisiontrace_test.go b/cmd/mavend/decisiontrace_test.go index 813d34f..eac9f2e 100644 --- a/cmd/mavend/decisiontrace_test.go +++ b/cmd/mavend/decisiontrace_test.go @@ -25,7 +25,7 @@ func traceHandler(t *testing.T, ring *decision.Ring) *reactiveHandler { return &reactiveHandler{ api: api, recall: recallWiring{embedder: emb, memStore: memory.NewInMemoryStore()}, - router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil), + router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil, nil), replier: voice.NewStubReplier(), now: func() time.Time { return now }, dataStore: st, diff --git a/cmd/mavend/dialogue_contract_test.go b/cmd/mavend/dialogue_contract_test.go index 7db4ced..b4557d9 100644 --- a/cmd/mavend/dialogue_contract_test.go +++ b/cmd/mavend/dialogue_contract_test.go @@ -171,7 +171,7 @@ func newDialogueHandler(t *testing.T) (*reactiveHandler, *store.Store, *time.Tim // and never a coincidence (V-577, V-579). checkEnd refuses any reminder // landing on it, and at 09:00 the row that answers "на 9" would trip that. *now = time.Date(2026, 7, 31, 9, 17, 0, 0, time.UTC) - h.router = buildRouter(router.NewHashEmbedder(1024), h.matcher, 0.55, nil) + h.router = buildRouter(router.NewHashEmbedder(1024), h.matcher, 0.55, nil, nil) h.recall = recallWiring{embedder: router.NewHashEmbedder(1024), memStore: memory.NewInMemoryStore()} return h, st, now } diff --git a/cmd/mavend/fact_subject_test.go b/cmd/mavend/fact_subject_test.go index 885f067..34aaed7 100644 --- a/cmd/mavend/fact_subject_test.go +++ b/cmd/mavend/fact_subject_test.go @@ -24,7 +24,7 @@ func TestApplyAction_FactCapture_QueuesEntityResolution(t *testing.T) { emb := router.NewHashEmbedder(1024) matcher := tool.NewMatcher(api) - rtr := buildRouter(emb, matcher, 0.55, nil) + rtr := buildRouter(emb, matcher, 0.55, nil, nil) h := &reactiveHandler{ api: api, diff --git a/cmd/mavend/factgate_test.go b/cmd/mavend/factgate_test.go index 842bb2c..72ecb87 100644 --- a/cmd/mavend/factgate_test.go +++ b/cmd/mavend/factgate_test.go @@ -20,7 +20,7 @@ func newFactGateHandler(t *testing.T, now time.Time) (*reactiveHandler, ipc.Core h := &reactiveHandler{ api: api, recall: recallWiring{embedder: emb, memStore: memory.NewInMemoryStore()}, - router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil), + router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil, nil), replier: voice.NewStubReplier(), now: func() time.Time { return now }, dataStore: st, diff --git a/cmd/mavend/notefragment_test.go b/cmd/mavend/notefragment_test.go index cf9becc..7e2f50d 100644 --- a/cmd/mavend/notefragment_test.go +++ b/cmd/mavend/notefragment_test.go @@ -45,7 +45,7 @@ func newNoteHandler(t *testing.T) (*reactiveHandler, *store.Store) { h := &reactiveHandler{ api: api, recall: recallWiring{embedder: emb, memStore: memory.NewInMemoryStore()}, - router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil), + router: buildRouter(emb, tool.NewMatcher(api), 0.55, nil, nil), replier: voice.NewStubReplier(), now: func() time.Time { return now }, dataStore: st, diff --git a/cmd/mavend/reactive_notes_test.go b/cmd/mavend/reactive_notes_test.go index 811e706..0449b88 100644 --- a/cmd/mavend/reactive_notes_test.go +++ b/cmd/mavend/reactive_notes_test.go @@ -22,7 +22,7 @@ func TestReactiveNotesReminders(t *testing.T) { emb := router.NewHashEmbedder(1024) matcher := tool.NewMatcher(api) - rtr := buildRouter(emb, matcher, 0.55, nil) + rtr := buildRouter(emb, matcher, 0.55, nil, nil) h := &reactiveHandler{ api: api, @@ -104,7 +104,7 @@ func TestSpokenTaskCaptureFilesATask(t *testing.T) { h := &reactiveHandler{ api: api, recall: recallWiring{embedder: emb, memStore: memory.NewInMemoryStore()}, - router: buildRouter(emb, matcher, 0.55, nil), + router: buildRouter(emb, matcher, 0.55, nil, nil), replier: voice.NewStubReplier(), now: func() time.Time { return now }, dataStore: st, diff --git a/cmd/mavend/simulator_test.go b/cmd/mavend/simulator_test.go index bf347a0..23b5957 100644 --- a/cmd/mavend/simulator_test.go +++ b/cmd/mavend/simulator_test.go @@ -474,7 +474,7 @@ func newSimWorld(t *testing.T, sc scenario) *simWorld { // used to be built on a nil API, which meant any scenario that produced an // act panicked the moment the matcher was consulted. matcher := tool.NewMatcher(api) - rtr := buildRouter(emb, matcher, config.DefaultRouterThreshold, router.NewLLMRouter(scripted)) + rtr := buildRouter(emb, matcher, config.DefaultRouterThreshold, router.NewLLMRouter(scripted), nil) w.handler = &reactiveHandler{ stt: simTranscriber{}, diff --git a/cmd/mavend/voicewire.go b/cmd/mavend/voicewire.go index 9e308e9..af7a3a4 100644 --- a/cmd/mavend/voicewire.go +++ b/cmd/mavend/voicewire.go @@ -36,7 +36,9 @@ type voiceWiring struct { sessions *voice.Sessions voiceSink delivery.Sink embedder router.Embedder - handler *reactiveHandler // the reactive handler for IPC Chat + // heads — the routing heads, nil unless embedder.heads_path is set. + heads *router.RouterHeads + handler *reactiveHandler // the reactive handler for IPC Chat // worker clients (set when configured as Remote): closed on shutdown so // mavsttd / mavttsd don't keep a stale conn into a restarting daemon. sttClient *worker.Client @@ -72,6 +74,9 @@ func (w *voiceWiring) close() { if w.embedder != nil { _ = w.embedder.Close() } + if w.heads != nil { + _ = w.heads.Close() + } if w.server != nil { _ = w.server.Close() } @@ -147,6 +152,24 @@ func wireVoice(cfg *config.Config, coreAPI ipc.CoreAPI, phr phraser.Phraser, mem emb = router.NewHashEmbedder(1024) } w.embedder = emb + + // ----- router: routing heads (only when configured, and never fatal) ----- + // A missing or broken weights file logs and leaves w.heads nil, which is + // byte-for-byte the cascade that shipped before V-664. Refusing to start + // over a routing accelerator would trade a working box for a better one. + if cfg.Voice.Embedder != nil && cfg.Voice.Embedder.HeadsPath != "" { + h, err := router.NewRouterHeads( + cfg.Voice.Embedder.HeadsPath, + cfg.Voice.Embedder.TokenizerPath, + ) + if err != nil { + log.Printf("voice: routing heads unavailable, cascade unchanged: %v", err) + } else { + log.Printf("voice: routing heads loaded from %s", cfg.Voice.Embedder.HeadsPath) + w.heads = h + } + } + repairFactVectors(dataStore, emb) checkStoredEmbedder(dataStore, emb) // Retention is enforced on write, which is not enough on its own: a box that @@ -223,7 +246,8 @@ func wireVoice(cfg *config.Config, coreAPI ipc.CoreAPI, phr phraser.Phraser, mem // against the classifier's 50.0%, at about 1s a turn instead of 30ms (see // config.VoiceConfig.LLMRouter). The classifier always stays wired as the // fallback, so a model error never breaks a turn. - rtr := buildRouter(emb, matcher, threshold, pickLLMRouter(cfg.Voice.UseLLMRouter(), hot)) + rtr := buildRouter(emb, matcher, threshold, + pickLLMRouter(cfg.Voice.UseLLMRouter(), hot), w.heads) // ----- sessions registry (shared with voicesink) ----- sessions := voice.NewSessions() @@ -390,7 +414,8 @@ func pickLLMRouter(enabled bool, c router.Completer) *router.LLMRouter { // intent from seedDir (models/seeds/.txt) — see seedClassifier // below for the current intent list and file names. // - Threshold is from voice.router_threshold config (default 0.55). -func buildRouter(emb router.Embedder, acts router.ActMatcher, threshold float64, llmR *router.LLMRouter) *router.Router { +func buildRouter(emb router.Embedder, acts router.ActMatcher, threshold float64, + llmR *router.LLMRouter, heads *router.RouterHeads) *router.Router { cls := router.NewClassifier(emb) seedClassifier(cls) grammars := router.DefaultGrammars(acts) @@ -442,6 +467,7 @@ func buildRouter(emb router.Embedder, acts router.ActMatcher, threshold float64, }, Threshold: threshold, LLM: llmR, + Heads: heads, }) } diff --git a/deploy/mavend.json b/deploy/mavend.json index a7371e3..91c06f9 100644 --- a/deploy/mavend.json +++ b/deploy/mavend.json @@ -231,7 +231,8 @@ "embedder": { "model_path": "/opt/maven/models/embedder/multilingual-e5-small/model_quantized.onnx", "tokenizer_path": "/opt/maven/models/embedder/multilingual-e5-small/tokenizer.json", - "lib_path": "/opt/maven/lib/libonnxruntime.so" + "lib_path": "/opt/maven/lib/libonnxruntime.so", + "heads_path": "/opt/maven/models/embedder/router-heads/router_heads.onnx" }, "llm_router": true, "query_min_score": 0.55, diff --git a/docs/evals/2026-08-08-routing-heads-in-go.md b/docs/evals/2026-08-08-routing-heads-in-go.md new file mode 100644 index 0000000..aea3a6a --- /dev/null +++ b/docs/evals/2026-08-08-routing-heads-in-go.md @@ -0,0 +1,146 @@ +# The routing heads, running in Go + +Date: 2026-08-08. Vikunja V-664. +Weights: `router_heads.onnx`, fp32, exported from `heads.pt` on workpc. +Fixture: `internal/router/eval/ru_routing_v1.json`, 96 cases, 33 carrying a destination. +Runner: `make t PKG=./internal/router/eval/ RUN=TestONNXRoutingHeads`. + +The four heads of V-661 ran nowhere. This is the number they score through the +Go cascade. Same fixture and same grammars as `TestONNXBaseline`, and only the +middle stage varies. + +## Headline + +| | classifier + ONNX | heads + classifier | gemma-4-12b cascade | +|---|---|---|---| +| intent | 75.0% (72/96) | **96.9% (93/96)** | 84.4% | +| destination | 33.3% (11/33) | **75.8% (25/33)** | 72.7% | +| false clarify | 0 | 1 | 2 | +| missed clarify | 8 | 1 | 1 | +| p50 | 24.5ms | 27.9ms | 329ms | + +A 118M encoder beats the 12B teacher it was distilled from. It wins on both +halves of the route, at a twelfth of the latency. The workstation stays the +better phraser and is no longer the better router. + +The p50 is not the heads. Most of it is the classifier's own embedder pass on +the turns the heads decline, plus process warm-up on the first case. The heads' +own forward pass measures 7.3ms on workpc. + +## Two defects were in the way, and the first was not in the heads + +**The tokenizer read every long word backwards.** `encodeWord` backtracks the +Viterbi path from the end of a word and prepends each piece. That puts them back +in reading order, and a second reverse after the loop undid it. So +`query: вода` tokenized to `[0 12 1294 41 12489 2]` where the reference +tokenizer gives `[0 41 1294 12 12489 2]`. + +It was found here and only here. The heads were trained through transformers and +are read through the hand-written tokenizer. So a mismatch shows up as a score +far below what Python measured on the same weights. Nothing else in the suite +compares the two. + +Measured on the recall fixture, same 27 cases either way: + +| | reversed | fixed | +|---|---|---| +| recall@1 | 70.4% (19/27) | **77.8% (21/27)** | +| recall@3 | 85.2% (23/27) | **96.3% (26/27)** | +| answered after gate | 63.0% | 66.7% | +| wrong note on top | 8 | 6 | +| false recall | 0/5 | 1/5 | + +The classifier barely moved, 76.0% to 75.0%, and destination 36.4% to 33.3%. +Both are one case on 96 and neither is a finding. Seeds and queries were mangled +the same way, so cosine survived it. Recall is where it cost, because a stored +passage and a live query are different lengths and break differently. + +The one new false recall is the honest cost and it is not being hidden. A +sharper embedder scores every candidate higher, including the ones that should +have stayed under the gate. That is the same trade `2026-08-04-recall-e5-small.md` +recorded when e5-small replaced MiniLM. + +The embedder id now carries a tokenizer revision, `model_quantized@384/tok2`. +Stored vectors were written under rev 1 and no longer sit in the same space as a +query embedded now. The model file's name never moved, so nothing would have +triggered `ReembedAll`. On the box the marker fired on start, and the re-embed +rewrote 65 notes and 19 facts in 5 seconds. + +**The clarify head was being thrown away.** It was read only when the intent head +cleared its own threshold. That cost 6 of the 8 ambiguous cases. `вода` reads as intent +`act` at 0.233 and clarify at 0.983. Burying that handed the turn to the +classifier, which routed it confidently and never asked. The clarify head answers +a different question, which is whether there is enough here to act on at all. So +it decides on its own and decides first. + +| | intent-gated | clarify decides first | +|---|---|---| +| intent | 90.6% | 96.9% | +| missed clarify | 7 | 1 | +| false clarify | 0 | 1 | + +## The threshold is measured, not chosen + +Max softmax over the intent head, on the 88 cases carrying an intent: + +| threshold | kept | accuracy kept | wrong kept | right dropped | +|---|---|---|---|---| +| 0.5 | 84 | 96.4% | 3 | 2 | +| **0.6** | **81** | **97.5%** | **2** | **4** | +| 0.7 | 75 | 97.3% | 2 | 10 | +| 0.8 | 64 | 96.9% | 2 | 21 | +| 0.9 | 46 | 100.0% | 0 | 37 | + +0.6 is the knee. Every value from 0.7 to 0.85 drops right answers and keeps the +same two wrong ones. 0.9 is the only value that clears them, and it costs 37 +correct routes to do it. + +## Quantization was measured and rejected + +| build | size | intent | destination | p50 | +|---|---|---|---|---| +| fp32 | 470MB | 83/88 (94.3%) | 28/33 (84.8%) | 7.3ms | +| int8 | 118MB | 79/88 (89.8%) | 26/33 (78.8%) | 4.0ms | +| fp16 | 235MB | will not load | — | — | + +Python numbers, on the heads alone rather than through the cascade. int8 costs +4.5 points of intent and 6 of destination to save 3ms. The cascade around it has +a p50 over a second when the resident model answers. The fp16 graph is broken: +`convert_float_to_float16` leaves a Cast node emitting float16 where the graph +expects float, and onnxruntime refuses the session. It was not worth fixing. + +The exporter also had to be told to write one file. It splits weights into a +`.onnx.data` sidecar by default. This onnxruntime resolves that path against the +process working directory rather than the model. A split graph loads from one +directory only. + +## What is still wrong + +**Four of the eight destination misses are calendar.** Training cannot move them. +The possessive agenda rules claim those cases at stage 0 and name nothing on +purpose. That caution was free while nothing downstream could name anything +either. It has now cost four points in three separate measurements. The call is +the owner's and it is still open. + +**The slot head is exported and not read.** Slots come from the stage-2 +extractor. Mapping BIO tags back to text needs character offsets the unigram +tokenizer does not keep, which is its own piece of work. + +**`поужинал` is a false clarify**, which is the same defect `thinSingleToken` +was narrowed for on 2026-08-01, arriving now from a different direction. + +## On the box + +Deployed to homesrv the same day. `voice: routing heads loaded` on start, and +`/trace` shows `routing-heads` winning or thinning every turn. The resident model +and the classifier are both marked never asked. Live probes: + +```text +что такое TCP? -> kiwix a real definition +кто такой Линус Торвальдс? -> kiwix a real answer +во сколько я лёг вчера -> personal не нашла у тебя такой записи +вода -> thinned to clarify at 0.233 / 0.983 +``` + +A missing or broken weights file logs and leaves the heads nil, which is +byte-for-byte the cascade that shipped before this. diff --git a/internal/config/voice.go b/internal/config/voice.go index 752a8f6..1e9a9e3 100644 --- a/internal/config/voice.go +++ b/internal/config/voice.go @@ -37,6 +37,16 @@ type EmbedderConfig struct { ModelPath string `json:"model_path,omitempty"` TokenizerPath string `json:"tokenizer_path,omitempty"` LibPath string `json:"lib_path,omitempty"` + + // HeadsPath — the routing heads graph, which is a fine-tuned COPY of the + // model above with four linear heads on its pooled output (V-664). Empty + // means no heads, and the cascade runs exactly as it did before they + // existed. It shares LibPath and TokenizerPath, and router_heads.json is + // read from the same directory. + // + // It must never be pointed at ModelPath. Memory recall depends on the + // resident copy scoring what it scored, and the fine-tuned one does not. + HeadsPath string `json:"heads_path,omitempty"` } // WeatherConfig configures the weather provider for voice queries. diff --git a/internal/router/decisiontrace.go b/internal/router/decisiontrace.go index f66ac2a..140ec20 100644 --- a/internal/router/decisiontrace.go +++ b/internal/router/decisiontrace.go @@ -14,12 +14,15 @@ import ( "github.com/kami/maven/internal/decision" ) -// The two routing engines, named as claimants. They are one stage and not two, -// because only one of them ever runs: the classifier is reached when the model -// is absent or errored, never alongside it. +// The three routing engines, named as claimants. The model and the classifier +// are one stage and not two, because only one of them ever runs: the classifier +// is reached when the model is absent or errored, never alongside it. The heads +// run before both and decline on low confidence, so they can appear beside +// either one in a record. const ( claimantLLM = "llm-router" claimantClassifier = "classifier" + claimantHeads = "routing-heads" ) // thinReason names which arm of gateLLMDecision cut the confidence. The gate diff --git a/internal/router/embedderid_test.go b/internal/router/embedderid_test.go index 03f1e16..6b2ecfd 100644 --- a/internal/router/embedderid_test.go +++ b/internal/router/embedderid_test.go @@ -1,10 +1,13 @@ package router -import "testing" +import ( + "strings" + "testing" +) func TestEmbedderIDFromModelPath(t *testing.T) { got := modelIDFromPath("/opt/maven/models/embedder/multilingual-e5-small.onnx") - if got != "multilingual-e5-small@384" { + if got != "multilingual-e5-small@384/tok2" { t.Fatalf("modelIDFromPath = %q", got) } // A different model file must produce a different id, even at 384 dim. @@ -12,6 +15,13 @@ func TestEmbedderIDFromModelPath(t *testing.T) { if old == got { t.Fatal("two different models share one id") } + // The tokenizer is half of what makes a vector, and it changes under a + // model file whose name never moves (V-664). An id that ignored it would + // leave stored passages in one space and every new query in another, with + // nothing to trigger the re-embed. + if !strings.Contains(got, "/tok") { + t.Fatalf("id %q does not name the tokenizer revision", got) + } } func TestEmbedderIDIncludesDim(t *testing.T) { diff --git a/internal/router/eval/heads_test.go b/internal/router/eval/heads_test.go new file mode 100644 index 0000000..c65fd36 --- /dev/null +++ b/internal/router/eval/heads_test.go @@ -0,0 +1,86 @@ +package eval + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/kami/maven/internal/config" + "github.com/kami/maven/internal/router" +) + +// TestONNXRoutingHeads — the cascade with the routing heads wired, which is +// what V-664 deploys. Opt-in via MAVEN_ONNX_LIB, same as TestONNXBaseline, and +// one TestONNX* per process. +// +// The comparison worth reading is against TestONNXBaseline, which is the same +// cascade with the same grammars and the same classifier floor and no heads. +// Only the middle arm varies. +// +// It also checks the Go unigram tokenizer against the Python one, because the +// heads were trained through transformers and are read through a hand-written +// tokenizer. A mismatch shows up here as a score below what Python measured on +// the same weights, and nowhere else. +func TestONNXRoutingHeads(t *testing.T) { + lib := os.Getenv("MAVEN_ONNX_LIB") + if lib == "" { + t.Skip("MAVEN_ONNX_LIB unset — see AGENTS.md § Embedder model for intent routing") + } + // Absolute, because onnxruntime resolves a graph's external weights file + // against the model path it was given, and a relative one lands in the + // test's working directory. + root, err := filepath.Abs("../../..") + if err != nil { + t.Fatal(err) + } + model := filepath.Join(root, "models/embedder/multilingual-e5-small/model_quantized.onnx") + tok := filepath.Join(root, "models/embedder/multilingual-e5-small/tokenizer.json") + heads := filepath.Join(root, "models/embedder/router-heads/router_heads.onnx") + for _, p := range []string{lib, model, tok, heads} { + if _, err := os.Stat(p); err != nil { + t.Skipf("missing %s: %v", p, err) + } + } + emb, err2 := router.NewONNXEmbedder(model, tok, lib) + if err2 != nil { + t.Skipf("onnx embedder unavailable: %v", err2) + } + err = nil + defer emb.Close() + + h, err := router.NewRouterHeads(heads, tok) + if err != nil { + t.Skipf("routing heads unavailable: %v", err) + } + defer h.Close() + + f, err := Load() + if err != nil { + t.Fatalf("Load: %v", err) + } + rep, err := Score(context.Background(), "heads+classifier", withHeads(t, emb, h), f) + if err != nil { + t.Fatalf("Score: %v", err) + } + t.Log("\n" + rep.String() + rep.Failures()) +} + +// withHeads mirrors newBaselineRouter and adds the one arm under test. It is a +// separate function rather than a parameter so the baseline's signature stays +// the shape every other test calls it with. +func withHeads(t *testing.T, emb router.Embedder, h *router.RouterHeads) *router.Router { + t.Helper() + acts := router.DefaultActMatcher{Fns: actFns} + return router.New(router.Config{ + Grammars: baselineGrammars(acts), + Classifier: newBaselineClassifier(t, emb), + Extractor: router.Extractor{ + Time: router.StubDateTimeParser{}, + Acts: acts, + Facts: router.DefaultFactParser{}, + }, + Threshold: config.DefaultRouterThreshold, + Heads: h, + }) +} 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 +} diff --git a/internal/router/onnxembedder.go b/internal/router/onnxembedder.go index cbbf63c..005255f 100644 --- a/internal/router/onnxembedder.go +++ b/internal/router/onnxembedder.go @@ -66,12 +66,19 @@ func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, e func (e *onnxEmbedder) Dim() int { return embedDim } // ID names the loaded model for the DB marker (Vikunja #378): the model file's -// own name plus the dimension, so pointing the config at another model changes -// the string on its own. +// 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". +// "multilingual-e5-small@384/tok2". func modelIDFromPath(modelPath string) string { name := modelPath if i := strings.LastIndexAny(name, "/\\"); i >= 0 { @@ -81,7 +88,7 @@ func modelIDFromPath(modelPath string) string { if name == "" { name = "onnx" } - return fmt.Sprintf("%s@%d", name, embedDim) + return fmt.Sprintf("%s@%d/tok%d", name, embedDim, tokenizerRev) } // Embed treats the text as a query. The classifier compares one short @@ -339,14 +346,17 @@ func (t *unigramTokenizer) encodeWord(word string) []int64 { } } + // 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...) } - // Reverse - for l, r := 0, len(result)-1; l < r; l, r = l+1, r-1 { - result[l], result[r] = result[r], result[l] - } return result } diff --git a/internal/router/router.go b/internal/router/router.go index 2a6415d..3ad9653 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -31,6 +31,12 @@ type Config struct { // error/parse failure, falls through to the classifier (never fails the // turn on the model). LLM *LLMRouter + // Heads — optional routing heads over the fine-tuned embedder copy. When + // set, Route consults them after stage 0 and before the LLM router. They + // decline below their own confidence threshold, so a low-confidence turn + // reaches the model exactly as it does today. Nil is the shipped-before + // behaviour and costs nothing. + Heads *RouterHeads } // Router — the deterministic cascade. Route never guesses: stage 0 wins @@ -42,6 +48,7 @@ type Router struct { extractor Extractor threshold float64 llm *LLMRouter + heads *RouterHeads } func New(cfg Config) *Router { @@ -51,6 +58,7 @@ func New(cfg Config) *Router { extractor: cfg.Extractor, threshold: cfg.Threshold, llm: cfg.LLM, + heads: cfg.Heads, } } @@ -98,6 +106,60 @@ func (r *Router) Route(ctx context.Context, utterance string, now time.Time) (De } r.noteGrammarOutcomes(ctx, len(r.grammars), declinedBuild, "", "") + // stage 0b — routing heads (when wired). A softmax over the label set, so + // it cannot name an intent or a destination that does not exist, and its + // max is a real confidence. It runs before the model because it is three + // orders of magnitude faster and scores better on both halves of the route. + // + // It declines below its threshold rather than clarifying. A declined turn + // carries on to the model and then the classifier, which is what a box with + // no weights file does on every turn. + if r.heads != nil { + res, ok, err := r.heads.Route(ctx, utterance) + switch { + case err != nil: + log.Printf("router: heads fell through to the rest of the cascade: %v", err) + decision.Note(ctx, decision.Claim{ + Stage: decision.StageRoute, Claimant: claimantHeads, + Outcome: decision.Declined, Reason: "error: " + err.Error(), + }) + case !ok: + decision.Note(ctx, decision.Scored(decision.StageRoute, claimantHeads, + string(res.Intent), res.Confidence, decision.Declined, + "below the heads confidence threshold")) + default: + d := Decision{ + Utterance: utterance, + Stage: 2, + Intent: res.Intent, + Confidence: res.Confidence, + Source: res.Source, + Clarify: res.Clarify, + } + r.fillSlots(ctx, &d, now) + decision.Note(ctx, decision.Claim{ + Stage: decision.StageRoute, Claimant: claimantLLM, + Outcome: decision.NeverAsked, Reason: "the routing heads answered", + }) + decision.Note(ctx, decision.Claim{ + Stage: decision.StageRoute, Claimant: claimantClassifier, + Outcome: decision.NeverAsked, Reason: "the routing heads answered", + }) + outcome, reason := decision.Won, "" + if d.Clarify { + outcome, reason = decision.Thinned, "the clarify head says there is too little here to act on" + } + decision.Note(ctx, decision.Scored(decision.StageRoute, claimantHeads, + string(d.Intent), d.Confidence, outcome, reason)) + return d, nil + } + } else { + decision.Note(ctx, decision.Claim{ + Stage: decision.StageRoute, Claimant: claimantHeads, + Outcome: decision.NeverAsked, Reason: "no routing heads are wired", + }) + } + // stage 1a — LLM router (when wired). It reasons over the utterance instead // of nearest-centroid guessing. On any error/parse-fail, fall through to the // classifier cascade (never fail the turn on the model).