Load openWakeWord's three models and run their tensors (V-487)

The two feature models are frozen and pretrained; only the 100KB head was
trained here. The shapes were measured rather than assumed: 2.0s of 16kHz
audio gives 197 mel frames, and 76-frame windows at stride 8 give exactly the
16 embeddings the head was fitted on.

This file knows tensors and nothing about the 80ms cadence, which is why the
scaling openWakeWord applies between the two feature models lives here.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013ptwopxyo3Z2kwFckHkLvN
This commit is contained in:
2026-08-09 14:18:46 +04:00
parent b8279f6a22
commit 21a42cb3e6
+187
View File
@@ -0,0 +1,187 @@
package main
// The three models behind the wake word (V-487 stage two).
//
// openWakeWord's pipeline, run in a row:
//
// audio -> melspectrogram.onnx -> 32-bin mel frames, one per 10ms
// 76 frames -> embedding_model.onnx -> one 96-dim embedding per 80ms
// 16 embeds -> maven_wakeword.onnx -> one score
//
// The first two are frozen and pretrained. Only the last was trained here,
// which is why it is 100KB and the other two are megabytes. The shapes are
// not guesses: 2.0s of 16kHz audio measures 197 mel frames, and 76-frame
// windows at stride 8 give exactly the 16 embeddings the head was fitted on.
//
// This file knows ONNX and nothing about the 80ms cadence. wakeword.go knows
// the cadence and nothing about tensors.
import (
"fmt"
ort "github.com/yalue/onnxruntime_go"
)
const (
// melHop — samples per mel frame. 10ms at 16kHz.
melHop = 160
// melBins — mel bins per frame, fixed by melspectrogram.onnx.
melBins = 32
// embedFrames — mel frames one embedding is computed over, 760ms.
embedFrames = 76
// embedStride — mel frames between embeddings, 80ms.
embedStride = 8
// embedDim — the embedding width.
embedDim = 96
// headWindow — embeddings the head scores at once, 1.28s of audio.
headWindow = 16
// melContext — samples of history prepended to each incremental mel
// call, chosen so the eight frames this call yields continue exactly
// where the previous call's eight stopped.
//
// melspectrogram.onnx returns N/160-3 frames for N samples, and frame i
// covers [i*160, i*160+400). With 480 samples of history the buffer is
// 1760 samples, which is 8 frames, and the oldest of them starts one hop
// after the newest of the previous call. Less history leaves a gap: the
// first frames of a bare chunk would be computed against silence.
melContext = 480
// chunkSamples — audio per embedding step, 80ms.
chunkSamples = embedStride * melHop
)
// wakeModels holds the three ONNX sessions. It runs on CPU threads beside
// silero and never touches the GPU. That is a rule, not a result: a wake word
// that waits on card admission is not a wake word.
type wakeModels struct {
mel *ort.DynamicAdvancedSession
emb *ort.DynamicAdvancedSession
head *ort.DynamicAdvancedSession
}
// newWakeModels loads all three. melPath and embedPath are openWakeWord's
// frozen feature models; headPath is the keyword head trained for "Мэйвен".
func newWakeModels(melPath, embedPath, headPath, libPath string) (*wakeModels, error) {
if !ort.IsInitialized() {
if libPath != "" {
ort.SetSharedLibraryPath(libPath)
}
if err := ort.InitializeEnvironment(); err != nil {
return nil, fmt.Errorf("wake word: onnx runtime: %w", err)
}
}
open := func(p string, in, out []string) (*ort.DynamicAdvancedSession, error) {
s, err := ort.NewDynamicAdvancedSession(p, in, out, nil)
if err != nil {
return nil, fmt.Errorf("wake word: load %s: %w", p, err)
}
return s, nil
}
m := &wakeModels{}
var err error
if m.mel, err = open(melPath, []string{"input"}, []string{"output"}); err != nil {
return nil, err
}
if m.emb, err = open(embedPath, []string{"input_1"}, []string{"conv2d_19"}); err != nil {
m.Close()
return nil, err
}
if m.head, err = open(headPath, []string{"embeddings"}, []string{"score"}); err != nil {
m.Close()
return nil, err
}
return m, nil
}
// Close releases the three sessions.
func (m *wakeModels) Close() {
if m == nil {
return
}
for _, s := range []*ort.DynamicAdvancedSession{m.mel, m.emb, m.head} {
if s != nil {
s.Destroy()
}
}
m.mel, m.emb, m.head = nil, nil, nil
}
// melFrames runs one buffer of samples and returns the mel frames it yielded.
func (m *wakeModels) melFrames(buf []float32) ([][melBins]float32, error) {
in, err := ort.NewTensor(ort.NewShape(1, int64(len(buf))), buf)
if err != nil {
return nil, err
}
defer in.Destroy()
n := int64(len(buf)/melHop - 3)
if n < 1 {
return nil, fmt.Errorf("wake word: %d samples yield no mel frames", len(buf))
}
out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1, n, melBins))
if err != nil {
return nil, err
}
defer out.Destroy()
if err := m.mel.Run([]ort.Value{in}, []ort.Value{out}); err != nil {
return nil, err
}
data := out.GetData()
frames := make([][melBins]float32, n)
for i := range frames {
for j := 0; j < melBins; j++ {
// The scaling openWakeWord applies between the two feature
// models, and the head was fitted on its output.
frames[i][j] = data[i*melBins+j]/10.0 + 2.0
}
}
return frames, nil
}
// embedding runs embedFrames mel frames through the frozen embedder.
func (m *wakeModels) embedding(mels [][melBins]float32) ([embedDim]float32, error) {
var e [embedDim]float32
flat := make([]float32, 0, embedFrames*melBins)
for _, f := range mels {
flat = append(flat, f[:]...)
}
in, err := ort.NewTensor(ort.NewShape(1, embedFrames, melBins, 1), flat)
if err != nil {
return e, err
}
defer in.Destroy()
out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1, 1, embedDim))
if err != nil {
return e, err
}
defer out.Destroy()
if err := m.emb.Run([]ort.Value{in}, []ort.Value{out}); err != nil {
return e, err
}
copy(e[:], out.GetData())
return e, nil
}
// score runs the trained head over headWindow embeddings.
func (m *wakeModels) score(embeds [][embedDim]float32) (float64, error) {
flat := make([]float32, 0, headWindow*embedDim)
for _, e := range embeds {
flat = append(flat, e[:]...)
}
in, err := ort.NewTensor(ort.NewShape(1, headWindow, embedDim), flat)
if err != nil {
return 0, err
}
defer in.Destroy()
out, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 1))
if err != nil {
return 0, err
}
defer out.Destroy()
if err := m.head.Run([]ort.Value{in}, []ort.Value{out}); err != nil {
return 0, err
}
return float64(out.GetData()[0]), nil
}