111 lines
2.3 KiB
Go
111 lines
2.3 KiB
Go
package model
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"math"
|
|
"os"
|
|
"sync"
|
|
|
|
ort "github.com/yalue/onnxruntime_go"
|
|
)
|
|
|
|
type InferenceMeta struct {
|
|
Threshold float64 `json:"threshold"`
|
|
FeatureColumns []string `json:"feature_columns"`
|
|
}
|
|
|
|
type Predictor struct {
|
|
session *ort.DynamicAdvancedSession
|
|
threshold float64
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func NewPredictor(modelPath, metaPath string) (*Predictor, error) {
|
|
if lib := os.Getenv("ONNXRUNTIME_SHARED_LIBRARY"); lib != "" {
|
|
ort.SetSharedLibraryPath(lib)
|
|
}
|
|
if err := ort.InitializeEnvironment(); err != nil {
|
|
return nil, fmt.Errorf("onnx init: %w", err)
|
|
}
|
|
|
|
metaBytes, err := os.ReadFile(metaPath)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read inference meta: %w", err)
|
|
}
|
|
var meta InferenceMeta
|
|
if err := json.Unmarshal(metaBytes, &meta); err != nil {
|
|
return nil, fmt.Errorf("parse inference meta: %w", err)
|
|
}
|
|
if meta.Threshold == 0 {
|
|
meta.Threshold = 0.5
|
|
}
|
|
|
|
session, err := ort.NewDynamicAdvancedSession(
|
|
modelPath,
|
|
[]string{"float_input"},
|
|
[]string{"probabilities"},
|
|
nil,
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("onnx session: %w", err)
|
|
}
|
|
|
|
return &Predictor{session: session, threshold: meta.Threshold}, nil
|
|
}
|
|
|
|
func (p *Predictor) Threshold() float64 {
|
|
return p.threshold
|
|
}
|
|
|
|
func (p *Predictor) Close() error {
|
|
if p.session != nil {
|
|
p.session.Destroy()
|
|
}
|
|
return ort.DestroyEnvironment()
|
|
}
|
|
|
|
func (p *Predictor) PredictProba(features [][8]float32) ([]float64, error) {
|
|
if len(features) == 0 {
|
|
return nil, nil
|
|
}
|
|
|
|
n := len(features)
|
|
flat := make([]float32, 0, n*8)
|
|
for _, row := range features {
|
|
flat = append(flat, row[:]...)
|
|
}
|
|
|
|
inputShape := ort.NewShape(int64(n), 8)
|
|
inputTensor, err := ort.NewTensor(inputShape, flat)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer inputTensor.Destroy()
|
|
|
|
outputShape := ort.NewShape(int64(n), 2)
|
|
outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer outputTensor.Destroy()
|
|
|
|
p.mu.Lock()
|
|
err = p.session.Run([]ort.Value{inputTensor}, []ort.Value{outputTensor})
|
|
p.mu.Unlock()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
data := outputTensor.GetData()
|
|
probs := make([]float64, n)
|
|
for i := 0; i < n; i++ {
|
|
probs[i] = float64(data[i*2+1])
|
|
}
|
|
return probs, nil
|
|
}
|
|
|
|
func RoundProb(p float64) float64 {
|
|
return math.Round(p*10000) / 10000
|
|
}
|