add full source: transcribe pipeline, HTTP API, OpenAI endpoint, web UI, Docker
This commit is contained in:
@@ -0,0 +1,272 @@
|
||||
// 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()))
|
||||
}
|
||||
Reference in New Issue
Block a user