record-asr/internal/models/manager.go

// Package models owns model lifecycle: it auto-downloads the SenseVoiceSmall
// int8 model and the silero VAD on first run, then builds and holds the
// sherpa-onnx recognizer and voice-activity detector for the transcriber.
package models

import (
	"archive/tar"
	"compress/bzip2"
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	sherpa "github.com/k2-fsa/sherpa-onnx-go/sherpa_onnx"

	"git.misaka.ren/m1saka/record-asr/internal/config"
)

// Source URLs for the prebuilt models. Pinning these to the k2-fsa
// "asr-models" release channel keeps the downloads stable and resumable.
const (
	senseVoiceTarballURL = "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-int8-2024-07-17.tar.bz2"
	sileroVadURL         = "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/silero_vad.onnx"
)

// Manager holds the loaded recognizer and VAD. It is safe for concurrent use;
// the underlying sherpa-onnx objects are not re-entrant per stream, but the
// transcriber creates a fresh OfflineStream per segment under its own lock.
type Manager struct {
	cfg        config.Config
	Recognizer *sherpa.OfflineRecognizer
	VAD        *sherpa.VoiceActivityDetector

	// vadMu serializes access to the stateful VAD across concurrent workers.
	vadMu sync.Mutex
	ready bool
}

// VadMu returns the mutex guarding the VAD. The transcribe pipeline locks it
// for the duration of a single file's segmentation so VAD state never
// interleaves between files.
func (m *Manager) VadMu() *sync.Mutex { return &m.vadMu }

// New ensures the model files are present (downloading if necessary) and then
// constructs the recognizer and VAD. It blocks until both are ready.
func New(ctx context.Context, cfg config.Config) (*Manager, error) {
	m := &Manager{cfg: cfg}

	senseDir := filepath.Join(cfg.ModelDir, "sensevoice")
	senseModel := filepath.Join(senseDir, "model.int8.onnx")
	senseTokens := filepath.Join(senseDir, "tokens.txt")
	sileroModel := filepath.Join(cfg.ModelDir, "silero_vad.onnx")

	if err := os.MkdirAll(senseDir, 0o755); err != nil {
		return nil, fmt.Errorf("create model dir %s: %w", senseDir, err)
	}

	// Ensure models are on disk before constructing anything. Downloading the
	// ~230MB SenseVoice tarball can take a while on first run; log progress.
	if !fileOK(senseModel) || !fileOK(senseTokens) {
		fmt.Println("model: downloading SenseVoiceSmall int8 (first run)...")
		if err := downloadSenseVoice(ctx, cfg.DownloadTimeout, senseVoiceTarballURL, senseDir); err != nil {
			return nil, fmt.Errorf("download SenseVoice: %w", err)
		}
		fmt.Println("model: SenseVoice ready")
	}
	if !fileOK(sileroModel) {
		fmt.Println("model: downloading silero VAD...")
		if err := downloadFile(ctx, cfg.DownloadTimeout, sileroVadURL, sileroModel); err != nil {
			return nil, fmt.Errorf("download silero VAD: %w", err)
		}
		fmt.Println("model: silero VAD ready")
	}
	if !fileOK(senseModel) || !fileOK(senseTokens) || !fileOK(sileroModel) {
		return nil, fmt.Errorf("model files missing after download: %s,%s,%s", senseModel, senseTokens, sileroModel)
	}

	itn := 0
	if cfg.UseITN {
		itn = 1
	}
	recCfg := &sherpa.OfflineRecognizerConfig{
		FeatConfig: sherpa.FeatureConfig{SampleRate: 16000, FeatureDim: 80},
		ModelConfig: sherpa.OfflineModelConfig{
			SenseVoice: sherpa.OfflineSenseVoiceModelConfig{
				Model:                       senseModel,
				Language:                    cfg.Language,
				UseInverseTextNormalization: itn,
			},
			Tokens:     senseTokens,
			NumThreads: cfg.NumThreads,
			Provider:   "cpu",
			Debug:      0,
		},
		DecodingMethod: "greedy_search",
	}
	rec := sherpa.NewOfflineRecognizer(recCfg)
	if rec == nil {
		return nil, fmt.Errorf("failed to create recognizer (model=%s tokens=%s)", senseModel, senseTokens)
	}

	vadCfg := &sherpa.VadModelConfig{
		SileroVad: sherpa.SileroVadModelConfig{
			Model:              sileroModel,
			Threshold:          cfg.VadThreshold,
			MinSilenceDuration: cfg.VadMinSilenceDuration,
			MinSpeechDuration:  cfg.VadMinSpeechDuration,
			WindowSize:         cfg.VadWindowSize,
			MaxSpeechDuration:  cfg.VadMaxSpeechDuration,
		},
		SampleRate: 16000,
		NumThreads: cfg.NumThreads,
		Provider:   "cpu",
	}
	// bufferSizeInSeconds: how much audio the VAD buffers internally. Sized to
	// the largest expected utterance so we never drop speech waiting on a flush.
	vad := sherpa.NewVoiceActivityDetector(vadCfg, cfg.VadBufferSeconds)
	if vad == nil {
		sherpa.DeleteOfflineRecognizer(rec)
		return nil, fmt.Errorf("failed to create VAD (model=%s)", sileroModel)
	}

	m.Recognizer = rec
	m.VAD = vad
	m.ready = true
	fmt.Printf("model: loaded (sherpa-onnx %s, onnxruntime %s)\n", sherpa.GetVersion(), sherpa.GetOnnxruntimeVersion())
	return m, nil
}

// Close releases the native resources backing the recognizer and VAD.
func (m *Manager) Close() {
	m.vadMu.Lock()
	defer m.vadMu.Unlock()
	if m.VAD != nil {
		sherpa.DeleteVoiceActivityDetector(m.VAD)
		m.VAD = nil
	}
	if m.Recognizer != nil {
		sherpa.DeleteOfflineRecognizer(m.Recognizer)
		m.Recognizer = nil
	}
	m.ready = false
}

// fileOK reports whether path exists, is a regular file, and is non-empty.
func fileOK(path string) bool {
	st, err := os.Stat(path)
	return err == nil && !st.IsDir() && st.Size() > 0
}

// downloadFile fetches a single file from url into dest via a temp file +
// atomic rename, so a partial download never replaces a good file.
func downloadFile(ctx context.Context, timeout time.Duration, url, dest string) error {
	dctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	req, err := http.NewRequestWithContext(dctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP %d for %s", resp.StatusCode, url)
	}
	tmp := dest + ".tmp"
	f, err := os.Create(tmp)
	if err != nil {
		return err
	}
	if _, err := io.Copy(f, resp.Body); err != nil {
		f.Close()
		os.Remove(tmp)
		return err
	}
	if err := f.Close(); err != nil {
		os.Remove(tmp)
		return err
	}
	return os.Rename(tmp, dest)
}

// downloadSenseVoice streams the .tar.bz2 release directly through a bzip2
// reader into a tar reader, extracting only the two files we need
// (model.int8.onnx, tokens.txt) into destDir. This avoids materializing the
// full ~230MB archive on disk.
func downloadSenseVoice(ctx context.Context, timeout time.Duration, url, destDir string) error {
	dctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	req, err := http.NewRequestWithContext(dctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	// GitHub releases 302 to a CDN; the default client follows redirects.
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP %d for %s", resp.StatusCode, url)
	}

	bz := bzip2.NewReader(resp.Body)
	tr := tar.NewReader(bz)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return fmt.Errorf("read tar: %w", err)
		}
		if hdr.Typeflag != tar.TypeReg && hdr.Typeflag != tar.TypeRegA {
			continue
		}
		base := filepath.Base(hdr.Name)
		if base != "model.int8.onnx" && base != "tokens.txt" {
			continue
		}
		dst := filepath.Join(destDir, base)
		tmp := dst + ".tmp"
		f, err := os.Create(tmp)
		if err != nil {
			return err
		}
		if _, err := io.Copy(f, tr); err != nil {
			f.Close()
			os.Remove(tmp)
			return fmt.Errorf("extract %s: %w", base, err)
		}
		if err := f.Close(); err != nil {
			os.Remove(tmp)
			return err
		}
		if err := os.Rename(tmp, dst); err != nil {
			return err
		}
	}
	// Guard against a successful walk that found nothing matching.
	if !fileOK(filepath.Join(destDir, "model.int8.onnx")) || !fileOK(filepath.Join(destDir, "tokens.txt")) {
		return fmt.Errorf("tarball did not contain model.int8.onnx/tokens.txt")
	}
	return nil
}

// ModelDir returns the configured model directory (used by the server for
// diagnostics).
func (m *Manager) ModelDir() string { return m.cfg.ModelDir }

// Ready reports whether the recognizer and VAD are loaded.
func (m *Manager) Ready() bool {
	m.vadMu.Lock()
	defer m.vadMu.Unlock()
	return m.ready
}

// VersionInfo returns sherpa-onnx + onnxruntime build versions for /healthz.
func (m *Manager) VersionInfo() string {
	return fmt.Sprintf("sherpa-onnx %s, onnxruntime %s",
		strings.TrimSpace(sherpa.GetVersion()),
		strings.TrimSpace(sherpa.GetOnnxruntimeVersion()))
}