82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
|
"""STTProvider interface — faster-whisper implementation + test fake (§4).
|
||
|
|
|
||
|
|
Language hint 'sr' covers spoken Bosnian/Serbian/Croatian (§6.1). The model
|
||
|
|
runs on GPU in the founder's DC (large-v3); CPU with a small model for dev.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
from typing import Protocol
|
||
|
|
|
||
|
|
from gogo.voice.audio import SAMPLE_RATE, resample
|
||
|
|
|
||
|
|
log = logging.getLogger("gogo.voice.stt")
|
||
|
|
|
||
|
|
|
||
|
|
class STTProvider(Protocol):
|
||
|
|
async def transcribe(self, pcm_8k: bytes) -> str:
|
||
|
|
"""SLIN 8 kHz PCM → text (empty string when nothing intelligible)."""
|
||
|
|
...
|
||
|
|
|
||
|
|
|
||
|
|
class WhisperSTT:
|
||
|
|
"""faster-whisper; model size/device via env (GOGO_WHISPER_MODEL, GOGO_WHISPER_DEVICE)."""
|
||
|
|
|
||
|
|
def __init__(self, model_name: str | None = None, device: str | None = None):
|
||
|
|
from faster_whisper import WhisperModel # heavy import, voice extra
|
||
|
|
|
||
|
|
model_name = model_name or os.environ.get("GOGO_WHISPER_MODEL", "small")
|
||
|
|
device = device or os.environ.get("GOGO_WHISPER_DEVICE", "auto")
|
||
|
|
compute = "float16" if device == "cuda" else "int8"
|
||
|
|
log.info("loading whisper model=%s device=%s", model_name, device)
|
||
|
|
self._model = WhisperModel(model_name, device=device, compute_type=compute)
|
||
|
|
self._lock = asyncio.Lock()
|
||
|
|
|
||
|
|
async def transcribe(self, pcm_8k: bytes) -> str:
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
pcm_16k = resample(pcm_8k, SAMPLE_RATE, 16000)
|
||
|
|
audio = np.frombuffer(pcm_16k, dtype=np.int16).astype(np.float32) / 32768.0
|
||
|
|
|
||
|
|
def run() -> str:
|
||
|
|
segments, _info = self._model.transcribe(
|
||
|
|
audio,
|
||
|
|
language="sr", # covers bs/sr/hr as one spoken language (§6.1)
|
||
|
|
beam_size=5,
|
||
|
|
vad_filter=True,
|
||
|
|
)
|
||
|
|
return " ".join(s.text.strip() for s in segments).strip()
|
||
|
|
|
||
|
|
async with self._lock: # one decode at a time per worker process
|
||
|
|
return await asyncio.get_running_loop().run_in_executor(None, run)
|
||
|
|
|
||
|
|
|
||
|
|
class FakeSTT:
|
||
|
|
"""Test fake: returns queued texts, one per transcribe() call."""
|
||
|
|
|
||
|
|
def __init__(self, texts: list[str] | None = None):
|
||
|
|
self.texts = list(texts or [])
|
||
|
|
self.received: list[bytes] = []
|
||
|
|
|
||
|
|
async def transcribe(self, pcm_8k: bytes) -> str:
|
||
|
|
self.received.append(pcm_8k)
|
||
|
|
return self.texts.pop(0) if self.texts else ""
|
||
|
|
|
||
|
|
|
||
|
|
_stt: STTProvider | None = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_stt() -> STTProvider:
|
||
|
|
global _stt
|
||
|
|
if _stt is None:
|
||
|
|
_stt = WhisperSTT()
|
||
|
|
return _stt
|
||
|
|
|
||
|
|
|
||
|
|
def set_stt(provider: STTProvider | None) -> None:
|
||
|
|
global _stt
|
||
|
|
_stt = provider
|