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"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
"golang.org/x/text/unicode/norm"
|
||||
@@ -33,17 +34,21 @@ const (
|
||||
type onnxEmbedder struct {
|
||||
tokenizer *unigramTokenizer
|
||||
session *ort.DynamicSession[int64, float32]
|
||||
runtime *ONNXRuntimeLease
|
||||
id string
|
||||
closeOnce sync.Once
|
||||
closeErr error
|
||||
}
|
||||
|
||||
func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, error) {
|
||||
ort.SetSharedLibraryPath(libPath)
|
||||
if err := ort.InitializeEnvironment(); err != nil {
|
||||
return nil, fmt.Errorf("onnx: init environment: %w", err)
|
||||
runtime, err := AcquireONNXRuntime(libPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
tok, err := newUnigramTokenizer(tokenizerPath)
|
||||
if err != nil {
|
||||
_ = runtime.Close()
|
||||
return nil, fmt.Errorf("tokenizer: %w", err)
|
||||
}
|
||||
|
||||
@@ -53,12 +58,14 @@ func NewONNXEmbedder(modelPath, tokenizerPath, libPath string) (*onnxEmbedder, e
|
||||
[]string{"last_hidden_state"},
|
||||
)
|
||||
if err != nil {
|
||||
_ = runtime.Close()
|
||||
return nil, fmt.Errorf("onnx: create session: %w", err)
|
||||
}
|
||||
|
||||
return &onnxEmbedder{
|
||||
tokenizer: tok,
|
||||
session: session,
|
||||
runtime: runtime,
|
||||
id: modelIDFromPath(modelPath),
|
||||
}, nil
|
||||
}
|
||||
@@ -148,8 +155,17 @@ func (e *onnxEmbedder) embed(ctx context.Context, text string) ([]float32, error
|
||||
}
|
||||
|
||||
func (e *onnxEmbedder) Close() error {
|
||||
e.session.Destroy()
|
||||
return nil
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
e.closeOnce.Do(func() {
|
||||
if e.session != nil {
|
||||
e.session.Destroy()
|
||||
e.session = nil
|
||||
}
|
||||
e.closeErr = e.runtime.Close()
|
||||
})
|
||||
return e.closeErr
|
||||
}
|
||||
|
||||
func meanPool(hidden []float32, mask []int64, seqLen, dim int) []float32 {
|
||||
|
||||
Reference in New Issue
Block a user