kami/Maven
1
0
Fork 0
Code Issues 1 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Load the routing heads and read three of the four (V-664)

Browse Source 
The heads trained in V-661 ran nowhere. This loads the exported graph
and reads intent, destination and clarify off one forward pass. It
declines below 0.6 max softmax rather than clarifying, so a declined
turn reaches whatever is behind it.

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 tokenizer does not keep.

The clarify head decides on its own and decides first. It answers a
different question from the intent head, so a low intent confidence is
no reason to discard it. Reading it only above the intent threshold
cost 6 of the 8 ambiguous cases on the fixture: the word for water
reads as intent act at 0.23 and clarify at 0.98.

0.6 is the knee measured on the intent fixture: every higher value up
to 0.9 drops right answers and keeps the same two wrong ones.

The body is a fine-tuned COPY of the resident embedder and must never
replace it, because memory recall depends on that file scoring what it
scored.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013ptwopxyo3Z2kwFckHkLvN
This commit is contained in:
claude
2026-08-08 22:24:37 +04:00
parent feabf9f350
commit 88c086482e
1 changed files with 243 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/router/heads.go
+243
View File
@@ -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
}