Maven/cmd/mavsttd/whisper_handler.go

package main

/*
#cgo CFLAGS: -I${SRCDIR}/../../deps/include -I${SRCDIR}/../../deps/whisper.cpp/ggml/include
#cgo LDFLAGS: -L${SRCDIR}/../../deps/lib -Wl,-rpath,${SRCDIR}/../../deps/lib -lwhisper -lggml -lggml-base -lggml-cpu -lggml-vulkan -lm -lstdc++ -fopenmp
#include <whisper.h>
#include <stdlib.h>
*/
import "C"

import (
	"context"
	"fmt"
	"log"
	"math"
	"unsafe"

	"github.com/kami/maven/internal/worker"
)

// whisperSampleRate is the rate whisper.cpp requires; the pipeline resamples
// to it before sending, so it's also the rate the silence gate assumes.
const whisperSampleRate = 16000

type whisperHandler struct {
	ctx    *C.struct_whisper_context
	minMs  int     // clips shorter than this are dropped (hallucination bait)
	minRMS float64 // clips quieter than this (normalized RMS) are treated as silence
}

func newWhisperHandler(modelPath string, minMs int, minRMS float64) (*whisperHandler, error) {
	cparams := C.whisper_context_default_params()
	cPath := C.CString(modelPath)
	defer C.free(unsafe.Pointer(cPath))

	ctx := C.whisper_init_from_file_with_params(cPath, cparams)
	if ctx == nil {
		return nil, fmt.Errorf("whisper: failed to init from %s", modelPath)
	}
	return &whisperHandler{ctx: ctx, minMs: minMs, minRMS: minRMS}, nil
}

// gateReason returns a non-empty reason when audio must NOT reach whisper:
// too short, or below the energy floor (silence / room noise). Whisper
// hallucinates subtitle-credit boilerplate ("Редактор субтитров …",
// "Субтитры сделал …") on non-speech input, so we drop it before the model
// ever sees it. Pure (no CGo) so it's unit-tested directly.
//
// ponytail: plain RMS energy + min-duration, not a real VAD. The mic floor is
// hardware-specific (both thresholds are flags on mavsttd) — upgrade to WebRTC
// VAD / whisper's no_speech_prob if energy gating proves too blunt.
func gateReason(samples []float32, rate, minMs int, minRMS float64) string {
	if len(samples) == 0 {
		return "empty"
	}
	if ms := len(samples) * 1000 / rate; ms < minMs {
		return fmt.Sprintf("too short (%dms < %dms)", ms, minMs)
	}
	var sum float64
	for _, s := range samples {
		sum += float64(s) * float64(s)
	}
	if rms := math.Sqrt(sum / float64(len(samples))); rms < minRMS {
		return fmt.Sprintf("silence (rms %.4f < floor %.4f)", rms, minRMS)
	}
	return ""
}

// pcmSamples converts canonical s16le little-endian PCM to the float32 range
// whisper wants. Shared with the golden tests: they used to carry their own
// copy, so a regression here (a /32767 divisor, a byte order slip) left the
// assertion that the fixtures clear the silence gate green.
func pcmSamples(b []byte) []float32 {
	out := make([]float32, len(b)/2)
	for i := range out {
		s := int16(b[i*2]) | int16(b[i*2+1])<<8
		out[i] = float32(s) / 32768.0
	}
	return out
}

func (h *whisperHandler) Transcribe(ctx context.Context, req worker.TranscribeReq) (worker.TranscribeResp, error) {
	if err := ctx.Err(); err != nil {
		return worker.TranscribeResp{}, fmt.Errorf("whisper: context done before transcribe: %w", err)
	}
	a := req.Audio
	if len(a.Bytes) == 0 {
		return worker.TranscribeResp{}, fmt.Errorf("whisper: empty audio")
	}

	samples := pcmSamples(a.Bytes)

	// Silence gate: drop non-speech before whisper hallucinates on it.
	if reason := gateReason(samples, whisperSampleRate, h.minMs, h.minRMS); reason != "" {
		log.Printf("mavsttd: gated audio (%s) — skipping whisper", reason)
		return worker.TranscribeResp{Text: "", Confidence: 0}, nil
	}

	params := C.whisper_full_default_params(C.WHISPER_SAMPLING_GREEDY)
	params.print_progress = false
	params.print_realtime = false
	params.print_timestamps = false
	params.print_special = false
	params.n_threads = C.int(4)
	params.single_segment = true

	lang := C.CString(req.Lang)
	defer C.free(unsafe.Pointer(lang))
	params.language = lang
	params.detect_language = false

	// CGo blocks the goroutine; whisper has no portable CGo-friendly abort
	// callback. Run in a goroutine so the caller's context cancellation at
	// least returns promptly — the CGo goroutine leaks until whisper finishes
	// but the caller doesn't hang.
	type result struct {
		code int
	}
	ch := make(chan result, 1)
	cSamples := (*C.float)(unsafe.Pointer(&samples[0]))
	go func() {
		ch <- result{code: int(C.whisper_full(h.ctx, params, cSamples, C.int(len(samples))))}
	}()
	select {
	case r := <-ch:
		if r.code != 0 {
			return worker.TranscribeResp{}, fmt.Errorf("whisper: full failed: %d", r.code)
		}
	case <-ctx.Done():
		return worker.TranscribeResp{}, fmt.Errorf("whisper: %w", ctx.Err())
	}

	nSegments := int(C.whisper_full_n_segments(h.ctx))
	if nSegments == 0 {
		return worker.TranscribeResp{Text: "", Confidence: 0}, nil
	}

	var text string
	totalLogProb := float64(0)
	totalTokens := 0

	for i := 0; i < nSegments; i++ {
		cSeg := C.whisper_full_get_segment_text(h.ctx, C.int(i))
		if cSeg != nil {
			text += C.GoString(cSeg)
		}

		nTokens := int(C.whisper_full_n_tokens(h.ctx, C.int(i)))
		for j := 0; j < nTokens; j++ {
			p := float64(C.whisper_full_get_token_p(h.ctx, C.int(i), C.int(j)))
			if p > 0 {
				totalLogProb += math.Log(p)
				totalTokens++
			}
		}
	}

	confidence := 0.0
	if totalTokens > 0 {
		avgLogProb := totalLogProb / float64(totalTokens)
		confidence = math.Exp(avgLogProb)
	}

	noSpeechProb := float64(C.whisper_full_get_segment_no_speech_prob(h.ctx, 0))
	if noSpeechProb > 0.9 {
		confidence = 0
	}

	if math.IsNaN(confidence) || math.IsInf(confidence, 0) {
		confidence = 0
	}

	return worker.TranscribeResp{
		Text:       text,
		Confidence: confidence,
	}, nil
}

func (h *whisperHandler) Close() {
	if h.ctx != nil {
		C.whisper_free(h.ctx)
		h.ctx = nil
	}
}