83 lines
3.3 KiB
Python
83 lines
3.3 KiB
Python
"""Провайдер STT на faster-whisper (CTranslate2).
|
|
|
|
Модель автоматически скачивается с HuggingFace при первом запуске
|
|
и кэшируется в ~/.cache/huggingface (small ≈ 460 МБ).
|
|
|
|
Если скачивание из HuggingFace недоступно/медленное, можно указать зеркало:
|
|
export HF_ENDPOINT=https://hf-mirror.com
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
from faster_whisper import WhisperModel
|
|
|
|
from .base import SAMPLE_RATE, AudioInput, SttProvider, TranscriptResult
|
|
|
|
|
|
class FasterWhisperProvider(SttProvider):
|
|
name = "faster-whisper"
|
|
|
|
def __init__(
|
|
self,
|
|
model_size: str = "small", # tiny | base | small | medium | large-v3
|
|
device: str = "cpu", # cpu | cuda
|
|
compute_type: str = "int8", # int8 на CPU — быстро и мало памяти
|
|
language: str = "ru", # фиксируем язык: детекция не нужна
|
|
beam_size: int = 1, # 1 = жадный поиск, заметно быстрее
|
|
vad_filter: bool = True, # Silero-VAD внутри: отрезает тишину
|
|
cpu_threads: int = 0, # 0 = по умолчанию движка
|
|
) -> None:
|
|
self._language = language
|
|
self._beam_size = beam_size
|
|
self._vad_filter = vad_filter
|
|
|
|
t0 = time.perf_counter()
|
|
self._model = WhisperModel(
|
|
model_size,
|
|
device=device,
|
|
compute_type=compute_type,
|
|
cpu_threads=cpu_threads,
|
|
)
|
|
# Время загрузки (и первого скачивания) модели — отдельно от распознавания.
|
|
self.load_sec = time.perf_counter() - t0
|
|
|
|
def transcribe(self, audio: AudioInput, sample_rate: int = SAMPLE_RATE) -> TranscriptResult:
|
|
duration: Optional[float] = None
|
|
|
|
if isinstance(audio, np.ndarray):
|
|
samples = np.asarray(audio, dtype=np.float32)
|
|
if samples.ndim > 1: # (frames, channels) → mono
|
|
samples = samples.mean(axis=1)
|
|
duration = len(samples) / sample_rate
|
|
source: AudioInput = samples
|
|
else:
|
|
path = Path(audio)
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"Аудио-файл не найден: {path}")
|
|
source = str(path)
|
|
|
|
t0 = time.perf_counter()
|
|
segments, info = self._model.transcribe(
|
|
source,
|
|
language=self._language,
|
|
beam_size=self._beam_size,
|
|
vad_filter=self._vad_filter,
|
|
)
|
|
text = "".join(segment.text for segment in segments)
|
|
text = " ".join(text.split()) # нормализуем пробелы между сегментами
|
|
processing = time.perf_counter() - t0
|
|
|
|
if duration is None: # для файлов длина известна после transcribe
|
|
duration = float(getattr(info, "duration", 0.0) or 0.0)
|
|
|
|
return TranscriptResult(
|
|
text=text,
|
|
duration_sec=duration,
|
|
processing_sec=processing,
|
|
language=info.language or self._language,
|
|
language_probability=float(info.language_probability or 0.0),
|
|
) |