Dmitry Dementiev
a65e5039f5
- Зачем: - нужен быстрый качественный профиль Whisper для локального распознавания. - Что: - добавлены официальные варианты large-v3-turbo INT8 и FP16. - покрыты каталог, выбор квантизации, ошибка и создание модели. - Проверка: - uv run pytest.
254 lines
9.4 KiB
Python
254 lines
9.4 KiB
Python
"""Бэкенд транскрипции на основе OpenVINO GenAI."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import threading
|
||
import time
|
||
import warnings
|
||
from collections.abc import Callable
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from huggingface_hub import snapshot_download
|
||
from huggingface_hub.errors import LocalEntryNotFoundError
|
||
|
||
from local_transcriber.types import Segment, TranscribeResult
|
||
|
||
# (model_alias, compute_type) → HF repo
|
||
MODEL_REPOS: dict[tuple[str, str], str] = {
|
||
("tiny", "int8"): "OpenVINO/whisper-tiny-int8-ov",
|
||
("base", "fp16"): "OpenVINO/whisper-base-fp16-ov",
|
||
("small", "int8"): "OpenVINO/whisper-small-int8-ov",
|
||
("medium", "int8"): "OpenVINO/whisper-medium-int8-ov",
|
||
("medium", "fp16"): "OpenVINO/whisper-medium-fp16-ov",
|
||
("large-v3", "int8"): "OpenVINO/whisper-large-v3-int8-ov",
|
||
("large-v3", "fp16"): "OpenVINO/whisper-large-v3-fp16-ov",
|
||
("large-v3-turbo", "int8"): "OpenVINO/whisper-large-v3-turbo-int8-ov",
|
||
("large-v3-turbo", "fp16"): "OpenVINO/whisper-large-v3-turbo-fp16-ov",
|
||
}
|
||
|
||
# Fallback: если точная пара не найдена, пробуем альтернативный compute_type
|
||
_COMPUTE_TYPE_FALLBACKS: dict[str, list[str]] = {
|
||
"float32": ["fp16", "int8"],
|
||
"float16": ["fp16", "int8"],
|
||
"fp16": ["fp16", "int8"],
|
||
"int8": ["int8", "fp16"],
|
||
}
|
||
|
||
# large-v3: при неявном compute_type предпочитаем fp16 (стабильнее по качеству)
|
||
_IMPLICIT_COMPUTE_TYPE_OVERRIDES: dict[str, str] = {
|
||
"large-v3": "fp16",
|
||
}
|
||
|
||
MODEL_REQUIRED_FILES = [
|
||
"openvino_encoder_model.xml",
|
||
"openvino_decoder_model.xml",
|
||
]
|
||
|
||
|
||
class OpenVINOBackend:
|
||
"""Бэкенд транскрипции через openvino-genai WhisperPipeline."""
|
||
|
||
def __init__(
|
||
self,
|
||
ov_device: str = "openvino-cpu",
|
||
compute_type_explicit: bool = True,
|
||
):
|
||
"""ov_device: "openvino", "openvino-gpu" или "openvino-cpu".
|
||
compute_type_explicit=False означает, что compute_type пришёл из дефолтов.
|
||
"""
|
||
self._ov_device = ov_device
|
||
self._compute_type_explicit = compute_type_explicit
|
||
self.actual_compute_type: str | None = None
|
||
self.actual_ov_device: str | None = None
|
||
|
||
def ensure_model_available(
|
||
self,
|
||
model_name: str,
|
||
compute_type: str,
|
||
on_status: Callable[[str], None] | None = None,
|
||
) -> str:
|
||
"""Скачивает/находит OpenVINO модель нужной квантизации."""
|
||
repo_id, resolved_ct = self._resolve_repo(model_name, compute_type)
|
||
self.actual_compute_type = resolved_ct
|
||
|
||
try:
|
||
_notify(on_status, f"Проверяю кэш модели {model_name} (OpenVINO)...")
|
||
cached_path = Path(snapshot_download(repo_id, local_files_only=True))
|
||
_validate_model_dir(cached_path)
|
||
return str(cached_path)
|
||
except LocalEntryNotFoundError:
|
||
pass
|
||
except ValueError:
|
||
_notify(on_status, f"Кэш модели {model_name} неполный, докачиваю...")
|
||
|
||
_notify(on_status, f"Скачиваю модель {model_name} (OpenVINO) из Hugging Face...")
|
||
downloaded_path = Path(snapshot_download(repo_id, local_files_only=False))
|
||
_validate_model_dir(downloaded_path)
|
||
return str(downloaded_path)
|
||
|
||
def _resolve_ov_device(self) -> str:
|
||
"""Резолвит device string в OpenVINO device: ``"GPU"`` или ``"CPU"``."""
|
||
if self._ov_device == "openvino-gpu":
|
||
return "GPU"
|
||
if self._ov_device == "openvino-cpu":
|
||
return "CPU"
|
||
# "openvino" → авто-детект
|
||
try:
|
||
from openvino import Core
|
||
|
||
if "GPU" in Core().available_devices:
|
||
return "GPU"
|
||
except Exception:
|
||
pass
|
||
return "CPU"
|
||
|
||
def create_model(
|
||
self,
|
||
model_path: str,
|
||
device: str,
|
||
compute_type: str,
|
||
cpu_threads: int = 0,
|
||
) -> Any:
|
||
"""Создаёт WhisperPipeline. cpu_threads не используется (OpenVINO управляет сам)."""
|
||
import openvino_genai as ov_genai
|
||
|
||
ov_dev = self._resolve_ov_device()
|
||
self.actual_ov_device = ov_dev
|
||
return ov_genai.WhisperPipeline(model_path, ov_dev)
|
||
|
||
def transcribe(
|
||
self,
|
||
model: Any,
|
||
file_path: Path,
|
||
language: str | None,
|
||
on_segment: Callable[[Segment], None] | None = None,
|
||
on_status: Callable[[str], None] | None = None,
|
||
) -> TranscribeResult:
|
||
"""Транскрибирует файл через OpenVINO GenAI."""
|
||
from faster_whisper import decode_audio
|
||
|
||
_notify(on_status, "Загружаю аудио...")
|
||
raw_speech = decode_audio(str(file_path), sampling_rate=16000)
|
||
duration = len(raw_speech) / 16000.0
|
||
|
||
kwargs: dict[str, Any] = {"return_timestamps": True}
|
||
if language:
|
||
kwargs["language"] = f"<|{language}|>"
|
||
|
||
dur_min = int(duration // 60)
|
||
duration_str = f"{dur_min} мин" if dur_min > 0 else f"{int(duration)} сек"
|
||
pcm_list = raw_speech.tolist()
|
||
result = _generate_with_progress(model, pcm_list, kwargs, duration_str, on_status)
|
||
|
||
segments: list[Segment] = []
|
||
if hasattr(result, "chunks") and result.chunks:
|
||
for chunk in result.chunks:
|
||
start = max(0.0, chunk.start_ts)
|
||
end = max(start, chunk.end_ts)
|
||
seg = Segment(
|
||
start=start,
|
||
end=end,
|
||
text=chunk.text,
|
||
)
|
||
if on_segment is not None:
|
||
on_segment(seg)
|
||
segments.append(seg)
|
||
_notify(
|
||
on_status,
|
||
f"Транскрибирую (OpenVINO)... [{len(segments)} сегм.]",
|
||
)
|
||
|
||
detected_language = language or "auto"
|
||
language_probability = 1.0 if language else 0.0
|
||
|
||
return TranscribeResult(
|
||
segments=segments,
|
||
language=detected_language,
|
||
language_probability=language_probability,
|
||
duration=duration,
|
||
device_used="", # оркестратор проставит
|
||
)
|
||
|
||
def _resolve_repo(self, model_name: str, compute_type: str) -> tuple[str, str]:
|
||
"""Находит HF repo для пары (model, compute_type) с fallback.
|
||
|
||
Возвращает (repo_id, actual_compute_type).
|
||
"""
|
||
# Для неявного compute_type: override для конкретных моделей
|
||
if not self._compute_type_explicit and model_name in _IMPLICIT_COMPUTE_TYPE_OVERRIDES:
|
||
compute_type = _IMPLICIT_COMPUTE_TYPE_OVERRIDES[model_name]
|
||
|
||
# Точное совпадение
|
||
repo = MODEL_REPOS.get((model_name, compute_type))
|
||
if repo:
|
||
return repo, compute_type
|
||
|
||
# Fallback только для неявного compute_type
|
||
if not self._compute_type_explicit:
|
||
fallbacks = _COMPUTE_TYPE_FALLBACKS.get(compute_type, [])
|
||
for fallback_ct in fallbacks:
|
||
repo = MODEL_REPOS.get((model_name, fallback_ct))
|
||
if repo:
|
||
return repo, fallback_ct
|
||
|
||
# Явный --compute-type с несуществующей парой → ошибка
|
||
available = [ct for (m, ct) in MODEL_REPOS if m == model_name]
|
||
if available:
|
||
raise ValueError(
|
||
f"Модель '{model_name}' недоступна с compute_type='{compute_type}' для OpenVINO. "
|
||
f"Доступные варианты: {', '.join(sorted(set(available)))}"
|
||
)
|
||
|
||
all_models = sorted({m for m, _ in MODEL_REPOS})
|
||
raise ValueError(
|
||
f"Модель '{model_name}' не найдена для OpenVINO. "
|
||
f"Доступные модели: {', '.join(all_models)}"
|
||
)
|
||
|
||
|
||
def _generate_with_progress(
|
||
model: Any,
|
||
pcm_list: list[float],
|
||
kwargs: dict[str, Any],
|
||
duration_str: str,
|
||
on_status: Callable[[str], None] | None,
|
||
) -> Any:
|
||
"""Запускает model.generate() в потоке, обновляя статус с elapsed time."""
|
||
result_box: list[Any] = [None]
|
||
error_box: list[BaseException | None] = [None]
|
||
|
||
def run() -> None:
|
||
try:
|
||
result_box[0] = model.generate(pcm_list, **kwargs)
|
||
except BaseException as exc:
|
||
error_box[0] = exc
|
||
|
||
thread = threading.Thread(target=run)
|
||
start = time.monotonic()
|
||
thread.start()
|
||
|
||
while thread.is_alive():
|
||
elapsed = int(time.monotonic() - start)
|
||
elapsed_str = f"{elapsed // 60:02d}:{elapsed % 60:02d}"
|
||
_notify(on_status, f"Транскрибирую {duration_str} аудио (OpenVINO)... прошло {elapsed_str}")
|
||
thread.join(timeout=1.0)
|
||
|
||
if error_box[0] is not None:
|
||
raise error_box[0]
|
||
|
||
return result_box[0]
|
||
|
||
|
||
def _notify(on_status: Callable[[str], None] | None, message: str) -> None:
|
||
if on_status is not None:
|
||
on_status(message)
|
||
|
||
|
||
def _validate_model_dir(model_dir: Path) -> None:
|
||
missing = [f for f in MODEL_REQUIRED_FILES if not (model_dir / f).exists()]
|
||
if missing:
|
||
raise ValueError(
|
||
f"Неполная OpenVINO модель в '{model_dir}': отсутствуют {', '.join(missing)}"
|
||
)
|