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.
This commit is contained in:
@@ -7,6 +7,7 @@ import (
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
)
|
||||
@@ -54,9 +55,12 @@ const (
|
||||
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.
|
||||
@@ -70,8 +74,9 @@ type headsMeta struct {
|
||||
// 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.
|
||||
// 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)
|
||||
@@ -106,18 +111,24 @@ func NewRouterHeads(modelPath, tokenizerPath string) (*RouterHeads, error) {
|
||||
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,
|
||||
@@ -128,8 +139,14 @@ func (h *RouterHeads) Close() error {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
h.session.Destroy()
|
||||
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.
|
||||
|
||||
Reference in New Issue
Block a user