Maven/internal/router/heads.go

package router

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"sync"

	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]
	runtime   *ONNXRuntimeLease
	intents   []Intent
	sources   []Source
	threshold float64
	closeOnce sync.Once
	closeErr  error
}

// 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 shares the process-global ONNX environment with the embedder. The runtime
// lease is independent so shutdown order cannot unload the library while this
// graph's session is still alive.
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)
	}
	runtime, err := AcquireONNXRuntime("")
	if err != nil {
		return nil, fmt.Errorf("heads: %w", err)
	}
	session, err := ort.NewDynamicSession[int64, float32](
		modelPath,
		[]string{"input_ids", "attention_mask"},
		[]string{"intent", "source", "slots", "clarify"},
	)
	if err != nil {
		_ = runtime.Close()
		return nil, fmt.Errorf("heads: create session: %w", err)
	}

	return &RouterHeads{
		tokenizer: tok,
		session:   session,
		runtime:   runtime,
		intents:   intents,
		sources:   sources,
		threshold: headsThreshold,
	}, nil
}

func (h *RouterHeads) Close() error {
	if h == nil {
		return nil
	}
	h.closeOnce.Do(func() {
		if h.session != nil {
			h.session.Destroy()
			h.session = nil
		}
		h.closeErr = h.runtime.Close()
	})
	return h.closeErr
}

// 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
}