Run the routing heads in Go and route with them #206
@@ -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"))
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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{},
|
||||
|
||||
+29
-3
@@ -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/<intent>.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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
-1
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user