142 lines
6.7 KiB
Python
142 lines
6.7 KiB
Python
"""Плагины `Transcriber` на основе faster-whisper: CPU (`min`) и GPU (`medium`/`max`).
|
||
|
||
Оба плагина используют встроенный в faster-whisper Silero VAD (`vad_filter=True`)
|
||
и дополнительно отбрасывают сегменты короче `MIN_SEGMENT_DURATION_S` —
|
||
типичные галлюцинации Whisper на тишине/шуме (ТЗ §1.4). Общая логика
|
||
(ленивая загрузка модели-синглтона процесса, вызов `transcribe` с VAD,
|
||
фильтрация коротких сегментов) вынесена в `_FasterWhisperBase`; CPU/GPU-варианты
|
||
отличаются только параметрами устройства/квантизации (ADR-004,
|
||
`docs/architecture/adr/004-ai-tier-matrix.md`).
|
||
"""
|
||
|
||
from typing import TYPE_CHECKING, ClassVar
|
||
|
||
from core.plugins.factory import register_transcriber
|
||
from core.plugins.transcriber import Segment, Transcriber
|
||
|
||
if TYPE_CHECKING:
|
||
# Импорт только для проверки типов: рантайм-импорт — ленивый, см. `_get_model`,
|
||
# чтобы API-процесс, где транскрибация не используется, не тянул тяжёлую
|
||
# зависимость (ctranslate2 и т.п.) в память.
|
||
from faster_whisper import WhisperModel
|
||
|
||
MIN_SEGMENT_DURATION_S = 0.3
|
||
"""Минимальная длительность сегмента (сек); короче — отбрасывается как
|
||
вероятная галлюцинация Whisper на тишине/шуме."""
|
||
|
||
VAD_MIN_SILENCE_DURATION_MS = 500
|
||
"""Порог Silero VAD (мс) для разбиения на речевые куски внутри трека."""
|
||
|
||
|
||
class _FasterWhisperBase(Transcriber):
|
||
"""Общая логика плагинов faster-whisper: синглтон модели процесса + VAD-транскрибация.
|
||
|
||
Модель-синглтон принадлежит конкретному подклассу (`FasterWhisperCPU`,
|
||
`FasterWhisperGPU`), а не общему базовому классу: присваивание
|
||
`cls._model = ...` в `_get_model` всегда происходит через `type(self)`,
|
||
поэтому у каждого подкласса — свой атрибут класса, и CPU/GPU-плагины не
|
||
делят один кэшированный инстанс модели, даже если оба сконфигурированы в
|
||
одном процессе.
|
||
"""
|
||
|
||
MIN_SEGMENT_S: ClassVar[float] = MIN_SEGMENT_DURATION_S
|
||
_model: "ClassVar[WhisperModel | None]" = None
|
||
|
||
# Задаются наследниками в `__init__` (device — фиксированно классом,
|
||
# compute_type — либо фиксированно, либо конструкторская опция).
|
||
device: str
|
||
compute_type: str
|
||
|
||
def __init__(
|
||
self,
|
||
model: str,
|
||
language: str = "ru",
|
||
download_root: str | None = None,
|
||
) -> None:
|
||
self.model_name = model
|
||
self.language = language
|
||
self.download_root = download_root
|
||
|
||
def _get_model(self) -> "WhisperModel":
|
||
"""Лениво создать (или переиспользовать) синглтон `WhisperModel` конкретного подкласса."""
|
||
cls = type(self)
|
||
if cls._model is None:
|
||
from faster_whisper import WhisperModel # ленивый импорт тяжёлой зависимости
|
||
|
||
cls._model = WhisperModel(
|
||
self.model_name,
|
||
device=self.device,
|
||
compute_type=self.compute_type,
|
||
download_root=self.download_root,
|
||
)
|
||
return cls._model
|
||
|
||
def transcribe(self, audio_path: str, language: str = "ru") -> list[Segment]:
|
||
"""Транскрибировать аудиофайл трека, отбросив короткие сегменты-галлюцинации.
|
||
|
||
VAD (Silero, встроен в faster-whisper) включён с порогом тишины
|
||
`VAD_MIN_SILENCE_DURATION_MS`; дополнительно отбрасываются сегменты
|
||
короче `MIN_SEGMENT_DURATION_S`.
|
||
"""
|
||
model = self._get_model()
|
||
raw_segments, _info = model.transcribe(
|
||
audio_path,
|
||
language=language,
|
||
vad_filter=True,
|
||
vad_parameters={"min_silence_duration_ms": VAD_MIN_SILENCE_DURATION_MS},
|
||
)
|
||
return [
|
||
Segment(start=segment.start, end=segment.end, text=segment.text)
|
||
for segment in raw_segments
|
||
if (segment.end - segment.start) >= MIN_SEGMENT_DURATION_S
|
||
]
|
||
|
||
|
||
@register_transcriber
|
||
class FasterWhisperCPU(_FasterWhisperBase):
|
||
"""Транскрибер faster-whisper (CTranslate2) на CPU с int8-квантизацией (уровень `min`).
|
||
|
||
Модель — синглтон на процесс: создаётся лениво при первом вызове
|
||
`transcribe` и переиспользуется всеми последующими вызовами в рамках
|
||
одного процесса воркера (процесс запускается
|
||
в Celery-очереди `transcription` с `--pool=solo --concurrency=1`, поэтому
|
||
гонок за атрибут класса не возникает).
|
||
"""
|
||
|
||
provider: ClassVar[str] = "faster_whisper_cpu"
|
||
_model: "ClassVar[WhisperModel | None]" = None
|
||
|
||
def __init__(
|
||
self,
|
||
model: str | None = None,
|
||
language: str = "ru",
|
||
download_root: str | None = None,
|
||
) -> None:
|
||
super().__init__(model=model or "small", language=language, download_root=download_root)
|
||
self.device = "cpu"
|
||
self.compute_type = "int8"
|
||
|
||
|
||
@register_transcriber
|
||
class FasterWhisperGPU(_FasterWhisperBase):
|
||
"""Транскрибер faster-whisper на GPU (CUDA, уровни `medium`/`max`, ADR-004).
|
||
|
||
`compute_type` — конструкторская опция (дефолт `float16`, как в матрице
|
||
ADR-004); для экономии VRAM конфиг уровня может задать `int8_float16`
|
||
(options плагина в `TIERS`/`plugins.yaml`).
|
||
"""
|
||
|
||
provider: ClassVar[str] = "faster_whisper_gpu"
|
||
_model: "ClassVar[WhisperModel | None]" = None
|
||
|
||
def __init__(
|
||
self,
|
||
model: str | None = None,
|
||
language: str = "ru",
|
||
download_root: str | None = None,
|
||
compute_type: str = "float16",
|
||
) -> None:
|
||
super().__init__(model=model or "medium", language=language, download_root=download_root)
|
||
self.device = "cuda"
|
||
self.compute_type = compute_type
|