86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
import sys
|
|
import threading
|
|
import warnings
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Protocol
|
|
|
|
|
|
class SpeechRecognitionError(RuntimeError):
|
|
"""Raised when a voice message cannot be transcribed."""
|
|
|
|
|
|
class SpeechTranscriber(Protocol):
|
|
def transcribe(self, audio_path: str | Path) -> str:
|
|
...
|
|
|
|
|
|
def load_whisper_model_class() -> Any:
|
|
# CTranslate2 imports torch for optional model converters. Speech inference
|
|
# does not need it, so do not let an unrelated torch installation prevent
|
|
# faster-whisper from loading.
|
|
missing = object()
|
|
loaded_torch = sys.modules.get("torch", missing)
|
|
if loaded_torch is missing:
|
|
sys.modules["torch"] = None
|
|
try:
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings(
|
|
"ignore",
|
|
message="pkg_resources is deprecated as an API.*",
|
|
category=UserWarning,
|
|
)
|
|
from faster_whisper import WhisperModel
|
|
finally:
|
|
if loaded_torch is missing:
|
|
sys.modules.pop("torch", None)
|
|
return WhisperModel
|
|
|
|
|
|
class SpeechRecognizer:
|
|
def __init__(
|
|
self,
|
|
model_name: str,
|
|
device: str,
|
|
compute_type: str,
|
|
language: str | None,
|
|
model_factory: Callable[..., Any] | None = None,
|
|
) -> None:
|
|
self.model_name = model_name
|
|
self.device = device
|
|
self.compute_type = compute_type
|
|
self.language = language
|
|
self._model_factory = model_factory
|
|
self._model: Any | None = None
|
|
self._lock = threading.Lock()
|
|
|
|
def _create_model(self) -> Any:
|
|
factory = self._model_factory
|
|
if factory is None:
|
|
factory = load_whisper_model_class()
|
|
return factory(
|
|
self.model_name,
|
|
device=self.device,
|
|
compute_type=self.compute_type,
|
|
)
|
|
|
|
def transcribe(self, audio_path: str | Path) -> str:
|
|
try:
|
|
with self._lock:
|
|
if self._model is None:
|
|
self._model = self._create_model()
|
|
segments, _info = self._model.transcribe(
|
|
str(audio_path),
|
|
language=self.language,
|
|
beam_size=5,
|
|
vad_filter=True,
|
|
)
|
|
parts = [
|
|
segment.text.strip()
|
|
for segment in segments
|
|
if segment.text.strip()
|
|
]
|
|
except Exception as exc:
|
|
raise SpeechRecognitionError("Voice transcription failed") from exc
|
|
|
|
return " ".join(parts)
|