Files
Maven/internal/router/eval/heads_test.go
T
claude 28c2ffb84f Make aggregate ONNX gates execute for real
Reference-count the process-global ONNX Runtime across embedder and routing-head sessions, make close idempotent, and require named proof that both aggregate routing gates executed rather than self-skipped (V-716). Owner explicitly requested direct commits to master.
2026-08-13 03:03:25 +04:00

88 lines
2.8 KiB
Go

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()
recordONNXGateExecuted(t)
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,
})
}