diff --git a/cmd/mavwaked/wakefeatures.go b/cmd/mavwaked/wakefeatures.go new file mode 100644 index 0000000..1ba69ae --- /dev/null +++ b/cmd/mavwaked/wakefeatures.go @@ -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 +}