refactor(transcriber): введена pluggable-архитектура бэкендов транскрипции
- Зачем: - подготовка к добавлению OpenVINO бэкенда для ускорения на x86 CPU без CUDA. - архитектура должна позволять добавлять новые бэкенды (CoreML, AMD XDNA) без переписывания кода. - Что: - создан types.py с общими типами (Segment, TranscribeResult, TranscribeFileResult). - создан backends/base.py с Backend Protocol (3 метода: ensure_model_available, create_model, transcribe). - создан backends/faster_whisper.py — текущий код вынесен из transcriber.py в FasterWhisperBackend. - transcriber.py переделан в оркестратор: load_model() владеет полным пайплайном (ensure + create), CLI больше не вызывает ensure_model_available() отдельно. - TranscribeFileResult расширен полями backend и model_path для корректного cross-backend fallback в батч-режиме. - device_used проставляется оркестратором, а не бэкендом. - cli.py: вынесен _format_device_info(), подготовлен к openvino. - тесты обновлены: mock-точки перенесены с WhisperModel на get_backend/бэкенд-объекты. - Проверка: - uv run pytest -v — 98 passed. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,28 @@
|
||||
"""Реестр бэкендов транскрипции и выбор бэкенда по устройству."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .base import Backend
|
||||
|
||||
|
||||
def get_backend(device: str) -> Backend:
|
||||
"""Возвращает экземпляр бэкенда для указанного устройства.
|
||||
|
||||
Импорты ленивые — бэкенд загружается только при запросе.
|
||||
"""
|
||||
if device == "openvino":
|
||||
try:
|
||||
from .openvino import OpenVINOBackend
|
||||
except ImportError:
|
||||
raise ValueError(
|
||||
"OpenVINO бэкенд недоступен. Установите: pip install openvino-genai"
|
||||
) from None
|
||||
return OpenVINOBackend()
|
||||
|
||||
# cuda, cpu и всё остальное → faster-whisper
|
||||
from .faster_whisper import FasterWhisperBackend
|
||||
|
||||
return FasterWhisperBackend()
|
||||
@@ -0,0 +1,46 @@
|
||||
"""Протокол бэкенда транскрипции."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol
|
||||
|
||||
from local_transcriber.types import Segment, TranscribeResult
|
||||
|
||||
|
||||
class Backend(Protocol):
|
||||
"""Минимальный интерфейс бэкенда транскрипции.
|
||||
|
||||
Бэкенды реализуют этот протокол (structural typing) —
|
||||
наследование не требуется.
|
||||
"""
|
||||
|
||||
def ensure_model_available(
|
||||
self,
|
||||
model_name: str,
|
||||
compute_type: str,
|
||||
on_status: Callable[[str], None] | None = None,
|
||||
) -> str:
|
||||
"""Гарантирует наличие модели, возвращает путь к файлам."""
|
||||
...
|
||||
|
||||
def create_model(
|
||||
self,
|
||||
model_path: str,
|
||||
device: str,
|
||||
compute_type: str,
|
||||
) -> Any:
|
||||
"""Создаёт модель. Возвращает backend-специфичный объект."""
|
||||
...
|
||||
|
||||
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:
|
||||
"""Транскрибирует файл, возвращает результат."""
|
||||
...
|
||||
@@ -0,0 +1,181 @@
|
||||
"""Бэкенд транскрипции на основе faster-whisper (CTranslate2)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import io
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
# CUDA bootstrap — должен быть ДО импорта faster_whisper / ctranslate2
|
||||
from local_transcriber._cuda_bootstrap import ensure_cublas_loadable
|
||||
|
||||
ensure_cublas_loadable()
|
||||
|
||||
from faster_whisper import WhisperModel # noqa: E402
|
||||
from huggingface_hub import snapshot_download # noqa: E402
|
||||
from huggingface_hub.errors import LocalEntryNotFoundError # noqa: E402
|
||||
|
||||
from local_transcriber.types import Segment, TranscribeResult # noqa: E402
|
||||
|
||||
MODEL_REPOS = {
|
||||
"tiny": "Systran/faster-whisper-tiny",
|
||||
"base": "Systran/faster-whisper-base",
|
||||
"small": "Systran/faster-whisper-small",
|
||||
"medium": "Systran/faster-whisper-medium",
|
||||
"large-v3": "Systran/faster-whisper-large-v3",
|
||||
}
|
||||
|
||||
MODEL_ALLOW_PATTERNS = [
|
||||
"config.json",
|
||||
"preprocessor_config.json",
|
||||
"model.bin",
|
||||
"tokenizer.json",
|
||||
"vocabulary.*",
|
||||
]
|
||||
|
||||
MODEL_REQUIRED_FILES = [
|
||||
"config.json",
|
||||
"model.bin",
|
||||
"tokenizer.json",
|
||||
]
|
||||
|
||||
|
||||
class FasterWhisperBackend:
|
||||
"""Бэкенд транскрипции через faster-whisper (CTranslate2)."""
|
||||
|
||||
def ensure_model_available(
|
||||
self,
|
||||
model_name: str,
|
||||
compute_type: str,
|
||||
on_status: Callable[[str], None] | None = None,
|
||||
) -> str:
|
||||
"""Резолвит alias модели в repo_id и гарантирует наличие файлов."""
|
||||
local_path = Path(model_name).expanduser()
|
||||
if local_path.is_dir():
|
||||
_validate_model_dir(local_path)
|
||||
return str(local_path)
|
||||
|
||||
repo_id = _resolve_model_repo(model_name)
|
||||
|
||||
try:
|
||||
_notify(on_status, f"Проверяю кэш модели {model_name}...")
|
||||
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} из Hugging Face...")
|
||||
downloaded_path = Path(_snapshot_download(repo_id, local_files_only=False))
|
||||
_validate_model_dir(downloaded_path)
|
||||
return str(downloaded_path)
|
||||
|
||||
def create_model(
|
||||
self,
|
||||
model_path: str,
|
||||
device: str,
|
||||
compute_type: str,
|
||||
) -> Any:
|
||||
"""Создаёт WhisperModel."""
|
||||
try:
|
||||
return WhisperModel(model_path, device=device, compute_type=compute_type)
|
||||
except ImportError as exc:
|
||||
if _is_missing_socksio_error(exc):
|
||||
raise RuntimeError(
|
||||
"Обнаружен SOCKS proxy, но не установлена зависимость `socksio`, "
|
||||
"нужная для загрузки модели из Hugging Face через proxy. "
|
||||
"Обновите окружение: `uv sync`."
|
||||
) from exc
|
||||
raise
|
||||
|
||||
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:
|
||||
"""Транскрибирует файл через faster-whisper."""
|
||||
segment_generator, info = model.transcribe(
|
||||
str(file_path), language=language,
|
||||
)
|
||||
total_duration = info.duration
|
||||
segments: list[Segment] = []
|
||||
for raw_seg in segment_generator:
|
||||
seg = Segment(start=raw_seg.start, end=raw_seg.end, text=raw_seg.text)
|
||||
if on_segment is not None:
|
||||
on_segment(seg)
|
||||
segments.append(seg)
|
||||
_notify(
|
||||
on_status,
|
||||
f"Транскрибирую... {_fmt_time(seg.end)} / {_fmt_time(total_duration)}"
|
||||
f" [{len(segments)} сегм.]",
|
||||
)
|
||||
|
||||
return TranscribeResult(
|
||||
segments=segments,
|
||||
language=info.language,
|
||||
language_probability=info.language_probability,
|
||||
duration=info.duration,
|
||||
device_used="", # оркестратор проставит actual_device
|
||||
)
|
||||
|
||||
|
||||
def _notify(on_status: Callable[[str], None] | None, message: str) -> None:
|
||||
if on_status is not None:
|
||||
on_status(message)
|
||||
|
||||
|
||||
def _fmt_time(seconds: float) -> str:
|
||||
m, s = divmod(int(seconds), 60)
|
||||
h, m = divmod(m, 60)
|
||||
return f"{h}:{m:02d}:{s:02d}" if h else f"{m:02d}:{s:02d}"
|
||||
|
||||
|
||||
def _resolve_model_repo(model_name: str) -> str:
|
||||
if "/" in model_name:
|
||||
return model_name
|
||||
repo_id = MODEL_REPOS.get(model_name)
|
||||
if repo_id is None:
|
||||
expected = ", ".join(MODEL_REPOS)
|
||||
raise ValueError(f"Неподдерживаемая модель '{model_name}'. Ожидалось одно из: {expected}")
|
||||
return repo_id
|
||||
|
||||
|
||||
def _snapshot_download(repo_id: str, local_files_only: bool) -> str:
|
||||
try:
|
||||
return snapshot_download(
|
||||
repo_id,
|
||||
local_files_only=local_files_only,
|
||||
allow_patterns=MODEL_ALLOW_PATTERNS,
|
||||
)
|
||||
except ImportError as exc:
|
||||
if _is_missing_socksio_error(exc):
|
||||
raise RuntimeError(
|
||||
"Обнаружен SOCKS proxy, но не установлена зависимость `socksio`, "
|
||||
"нужная для загрузки модели из Hugging Face через proxy. "
|
||||
"Обновите окружение: `uv sync`."
|
||||
) from exc
|
||||
raise
|
||||
|
||||
|
||||
def _validate_model_dir(model_dir: Path) -> None:
|
||||
missing = [
|
||||
filename for filename in MODEL_REQUIRED_FILES if not (model_dir / filename).exists()
|
||||
]
|
||||
if not any(model_dir.glob("vocabulary.*")):
|
||||
missing.append("vocabulary.*")
|
||||
if missing:
|
||||
missing_str = ", ".join(missing)
|
||||
raise ValueError(f"Неполная локальная модель в '{model_dir}': отсутствуют {missing_str}")
|
||||
|
||||
|
||||
def _is_missing_socksio_error(exc: BaseException) -> bool:
|
||||
msg = str(exc).lower()
|
||||
return "socks proxy" in msg and "socksio" in msg
|
||||
@@ -14,9 +14,7 @@ from .transcriber import (
|
||||
Segment,
|
||||
_is_cuda_error,
|
||||
_transcribe_file,
|
||||
ensure_model_available,
|
||||
load_model,
|
||||
transcribe,
|
||||
)
|
||||
from .utils import (
|
||||
build_output_path,
|
||||
@@ -31,6 +29,16 @@ app = typer.Typer()
|
||||
console = Console(stderr=True)
|
||||
|
||||
|
||||
def _format_device_info(device_used: str) -> str:
|
||||
"""Формирует строку устройства для шапки транскрипта."""
|
||||
if device_used == "cuda":
|
||||
gpu_name = get_gpu_name()
|
||||
return f"CUDA ({gpu_name or 'Unknown GPU'})"
|
||||
if device_used == "openvino":
|
||||
return "OpenVINO (CPU)"
|
||||
return "CPU"
|
||||
|
||||
|
||||
@app.command()
|
||||
def main(
|
||||
files: list[Path] = typer.Argument(..., help="Пути к аудио/видеофайлам"),
|
||||
@@ -42,7 +50,8 @@ def main(
|
||||
),
|
||||
output: Path | None = typer.Option(None, "--output", "-o", help="Путь к выходному файлу"),
|
||||
device: str | None = typer.Option(
|
||||
None, "--device", "-d", show_default=False, help="Устройство (auto|cpu|cuda) [по умолч.: auto]"
|
||||
None, "--device", "-d", show_default=False,
|
||||
help="Устройство (auto|cpu|cuda) [по умолч.: auto]"
|
||||
),
|
||||
compute_type: str | None = typer.Option(
|
||||
None, "--compute-type", show_default=False,
|
||||
@@ -120,7 +129,6 @@ def _run_single(
|
||||
validated_file = validate_input_file(file)
|
||||
requested_device = defaults["device"]
|
||||
resolved_device = detect_device(requested_device)
|
||||
# Если пользователь явно указал устройство — запрещаем fallback на CPU
|
||||
strict = requested_device != "auto"
|
||||
output_path = build_output_path(validated_file, output)
|
||||
|
||||
@@ -131,15 +139,11 @@ def _run_single(
|
||||
f"Compute: [bold]{defaults['compute_type']}[/bold]"
|
||||
)
|
||||
|
||||
model_path = ensure_model_available(
|
||||
defaults["model"], on_status=lambda message: console.print(message)
|
||||
)
|
||||
|
||||
def on_segment(seg: Segment) -> None:
|
||||
console.print(f" [{seg.start:.2f}s] {seg.text.strip()}")
|
||||
|
||||
model_obj, actual_device = load_model(
|
||||
model_path, resolved_device, defaults["compute_type"],
|
||||
model_obj, actual_device, backend, model_path = load_model(
|
||||
defaults["model"], resolved_device, defaults["compute_type"],
|
||||
on_status=lambda msg: console.print(msg), strict_device=strict,
|
||||
)
|
||||
|
||||
@@ -147,8 +151,10 @@ def _run_single(
|
||||
tfr = _transcribe_file(
|
||||
model=model_obj,
|
||||
actual_device=actual_device,
|
||||
backend=backend,
|
||||
model_path=model_path,
|
||||
file_path=validated_file,
|
||||
model_name=model_path,
|
||||
model_name=defaults["model"],
|
||||
compute_type=defaults["compute_type"],
|
||||
language=defaults["language"] if defaults["language"] != "auto" else None,
|
||||
on_segment=on_segment if verbose else None,
|
||||
@@ -176,12 +182,7 @@ def _run_single(
|
||||
f"Речь не обнаружена в файле {validated_file.name}", style="yellow"
|
||||
)
|
||||
|
||||
if result.device_used == "cuda":
|
||||
gpu_name = get_gpu_name()
|
||||
device_info = f"CUDA ({gpu_name or 'Unknown GPU'})"
|
||||
else:
|
||||
device_info = "CPU"
|
||||
|
||||
device_info = _format_device_info(result.device_used)
|
||||
language_mode = "detected" if defaults["language"] == "auto" else "forced"
|
||||
|
||||
content = format_transcript(
|
||||
@@ -232,15 +233,12 @@ def _run_batch(
|
||||
raise SystemExit(1)
|
||||
return
|
||||
|
||||
# Phase 2: Load model
|
||||
# Phase 2: Load model (ensure + create в одном вызове)
|
||||
requested_device = defaults["device"]
|
||||
resolved_device = detect_device(requested_device)
|
||||
strict = requested_device != "auto"
|
||||
model_path = ensure_model_available(
|
||||
defaults["model"], on_status=lambda msg: console.print(msg)
|
||||
)
|
||||
model_obj, actual_device = load_model(
|
||||
model_path, resolved_device, defaults["compute_type"],
|
||||
model_obj, actual_device, backend, model_path = load_model(
|
||||
defaults["model"], resolved_device, defaults["compute_type"],
|
||||
on_status=lambda msg: console.print(msg), strict_device=strict,
|
||||
)
|
||||
|
||||
@@ -277,8 +275,10 @@ def _run_batch(
|
||||
tfr = _transcribe_file(
|
||||
model=model_obj,
|
||||
actual_device=actual_device,
|
||||
backend=backend,
|
||||
model_path=model_path,
|
||||
file_path=file,
|
||||
model_name=model_path,
|
||||
model_name=defaults["model"],
|
||||
compute_type=defaults["compute_type"],
|
||||
language=defaults["language"] if defaults["language"] != "auto" else None,
|
||||
on_segment=on_segment if verbose else None,
|
||||
@@ -291,8 +291,11 @@ def _run_batch(
|
||||
f" {file.name}: fallback на {tfr.actual_device} при транскрипции",
|
||||
style="yellow",
|
||||
)
|
||||
# Обновляем после возможного mid-stream fallback на CPU
|
||||
model_obj, actual_device = tfr.model, tfr.actual_device
|
||||
# Обновляем после возможного mid-stream fallback
|
||||
model_obj = tfr.model
|
||||
actual_device = tfr.actual_device
|
||||
backend = tfr.backend
|
||||
model_path = tfr.model_path
|
||||
|
||||
result = tfr.result
|
||||
|
||||
@@ -301,11 +304,7 @@ def _run_batch(
|
||||
f" Речь не обнаружена: {file.name}", style="yellow"
|
||||
)
|
||||
|
||||
if result.device_used == "cuda":
|
||||
gpu_name = get_gpu_name()
|
||||
device_info = f"CUDA ({gpu_name or 'Unknown GPU'})"
|
||||
else:
|
||||
device_info = "CPU"
|
||||
device_info = _format_device_info(result.device_used)
|
||||
|
||||
content = format_transcript(
|
||||
result=result,
|
||||
|
||||
@@ -4,7 +4,7 @@ from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from .transcriber import Segment, TranscribeResult
|
||||
from .types import Segment, TranscribeResult
|
||||
|
||||
_PAUSE_THRESHOLD_S = 2.0 # пауза между сегментами для разбиения на абзацы
|
||||
_MAX_PARAGRAPH_S = 60.0 # максимальная длительность абзаца
|
||||
|
||||
@@ -1,65 +1,18 @@
|
||||
"""Обёртка над faster-whisper: загрузка моделей, транскрипция, CUDA fallback."""
|
||||
"""Оркестрация транскрипции: выбор бэкенда, загрузка модели, fallback."""
|
||||
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
# Должен быть ДО импорта faster_whisper / ctranslate2
|
||||
from local_transcriber._cuda_bootstrap import ensure_cublas_loadable
|
||||
from local_transcriber.backends import get_backend
|
||||
|
||||
ensure_cublas_loadable()
|
||||
|
||||
from faster_whisper import WhisperModel # noqa: E402
|
||||
from huggingface_hub import snapshot_download
|
||||
from huggingface_hub.errors import LocalEntryNotFoundError
|
||||
|
||||
MODEL_REPOS = {
|
||||
"tiny": "Systran/faster-whisper-tiny",
|
||||
"base": "Systran/faster-whisper-base",
|
||||
"small": "Systran/faster-whisper-small",
|
||||
"medium": "Systran/faster-whisper-medium",
|
||||
"large-v3": "Systran/faster-whisper-large-v3",
|
||||
}
|
||||
|
||||
# allow — фильтр для snapshot_download (какие файлы скачивать из репозитория);
|
||||
# required — для валидации (что обязано быть после скачивания/в локальной модели)
|
||||
MODEL_ALLOW_PATTERNS = [
|
||||
"config.json",
|
||||
"preprocessor_config.json",
|
||||
"model.bin",
|
||||
"tokenizer.json",
|
||||
"vocabulary.*",
|
||||
]
|
||||
|
||||
MODEL_REQUIRED_FILES = [
|
||||
"config.json",
|
||||
"model.bin",
|
||||
"tokenizer.json",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Segment:
|
||||
start: float # seconds
|
||||
end: float # seconds
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class TranscribeResult:
|
||||
segments: list[Segment]
|
||||
language: str
|
||||
language_probability: float
|
||||
duration: float # seconds
|
||||
device_used: str # "cpu" / "cuda"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TranscribeFileResult:
|
||||
result: TranscribeResult
|
||||
model: WhisperModel
|
||||
actual_device: str
|
||||
# Re-export из types.py для обратной совместимости
|
||||
from local_transcriber.types import ( # noqa: F401
|
||||
Segment,
|
||||
TranscribeFileResult,
|
||||
TranscribeResult,
|
||||
)
|
||||
|
||||
|
||||
def load_model(
|
||||
@@ -68,15 +21,21 @@ def load_model(
|
||||
compute_type: str,
|
||||
on_status: Callable[[str], None] | None = None,
|
||||
strict_device: bool = False,
|
||||
) -> tuple[WhisperModel, str]:
|
||||
"""Загружает модель с CUDA-фолбеком. Возвращает (model, actual_device)."""
|
||||
) -> tuple[Any, str, Any, str]:
|
||||
"""Загружает модель: ensure + create с fallback.
|
||||
|
||||
Возвращает (model, actual_device, backend, model_path).
|
||||
"""
|
||||
backend = get_backend(device)
|
||||
actual_device = device
|
||||
|
||||
model_path = backend.ensure_model_available(model_name, compute_type, on_status)
|
||||
|
||||
try:
|
||||
_notify_status(on_status, f"Инициализирую модель на {device}...")
|
||||
model = _create_model(model_name, device, compute_type)
|
||||
model = backend.create_model(model_path, device, compute_type)
|
||||
except (RuntimeError, ValueError) as exc:
|
||||
# strict — пользователь явно указал устройство, fallback запрещён
|
||||
if device != "cpu" and _is_cuda_error(exc):
|
||||
if device != "cpu" and _is_backend_error(exc, device):
|
||||
if strict_device:
|
||||
raise
|
||||
warnings.warn(
|
||||
@@ -85,16 +44,21 @@ def load_model(
|
||||
stacklevel=2,
|
||||
)
|
||||
actual_device = "cpu"
|
||||
backend = get_backend("cpu")
|
||||
model_path = backend.ensure_model_available(model_name, compute_type, on_status)
|
||||
_notify_status(on_status, "Инициализирую модель на cpu...")
|
||||
model = _create_model(model_name, "cpu", compute_type)
|
||||
model = backend.create_model(model_path, "cpu", compute_type)
|
||||
else:
|
||||
raise
|
||||
return model, actual_device
|
||||
|
||||
return model, actual_device, backend, model_path
|
||||
|
||||
|
||||
def _transcribe_file(
|
||||
model: WhisperModel,
|
||||
model: Any,
|
||||
actual_device: str,
|
||||
backend: Any,
|
||||
model_path: str,
|
||||
file_path: Path,
|
||||
model_name: str,
|
||||
compute_type: str,
|
||||
@@ -103,39 +67,40 @@ def _transcribe_file(
|
||||
on_status: Callable[[str], None] | None = None,
|
||||
strict_device: bool = False,
|
||||
) -> TranscribeFileResult:
|
||||
"""Транскрибирует один файл. При mid-stream CUDA fallback перезагружает модель."""
|
||||
"""Транскрибирует один файл. При mid-stream fallback перезагружает модель."""
|
||||
lang_arg = language if language and language != "auto" else None
|
||||
|
||||
try:
|
||||
_notify_status(on_status, "Транскрибирую...")
|
||||
segments, info = _run_transcription(model, file_path, lang_arg, on_segment, on_status)
|
||||
result = backend.transcribe(model, file_path, lang_arg, on_segment, on_status)
|
||||
result.device_used = actual_device
|
||||
except (RuntimeError, ValueError) as exc:
|
||||
# Mid-stream fallback: GPU может упасть с OOM уже во время транскрипции,
|
||||
# поэтому перезагружаем модель на CPU и начинаем сначала
|
||||
if actual_device != "cpu" and _is_cuda_error(exc):
|
||||
if actual_device != "cpu" and _is_backend_error(exc, actual_device):
|
||||
if strict_device:
|
||||
raise
|
||||
warnings.warn(
|
||||
f"CUDA ошибка при транскрипции: {exc}. "
|
||||
f"Ошибка при транскрипции на {actual_device}: {exc}. "
|
||||
"Переключение на CPU и повтор.",
|
||||
stacklevel=2,
|
||||
)
|
||||
actual_device = "cpu"
|
||||
backend = get_backend("cpu")
|
||||
model_path = backend.ensure_model_available(model_name, compute_type, on_status)
|
||||
_notify_status(on_status, "Инициализирую модель на cpu...")
|
||||
model = _create_model(model_name, "cpu", compute_type)
|
||||
model = backend.create_model(model_path, "cpu", compute_type)
|
||||
_notify_status(on_status, "Транскрибирую...")
|
||||
segments, info = _run_transcription(model, file_path, lang_arg, on_segment, on_status)
|
||||
result = backend.transcribe(model, file_path, lang_arg, on_segment, on_status)
|
||||
result.device_used = actual_device
|
||||
else:
|
||||
raise
|
||||
|
||||
result = TranscribeResult(
|
||||
segments=segments,
|
||||
language=info.language,
|
||||
language_probability=info.language_probability,
|
||||
duration=info.duration,
|
||||
device_used=actual_device,
|
||||
return TranscribeFileResult(
|
||||
result=result,
|
||||
model=model,
|
||||
actual_device=actual_device,
|
||||
backend=backend,
|
||||
model_path=model_path,
|
||||
)
|
||||
return TranscribeFileResult(result=result, model=model, actual_device=actual_device)
|
||||
|
||||
|
||||
def transcribe(
|
||||
@@ -149,9 +114,12 @@ def transcribe(
|
||||
strict_device: bool = False,
|
||||
) -> TranscribeResult:
|
||||
"""High-level API: загрузка модели + транскрипция за один вызов."""
|
||||
model, actual_device = load_model(model_name, device, compute_type, on_status, strict_device)
|
||||
model, actual_device, backend, model_path = load_model(
|
||||
model_name, device, compute_type, on_status, strict_device,
|
||||
)
|
||||
tfr = _transcribe_file(
|
||||
model, actual_device, file_path, model_name, compute_type,
|
||||
model, actual_device, backend, model_path,
|
||||
file_path, model_name, compute_type,
|
||||
language, on_segment, on_status, strict_device,
|
||||
)
|
||||
return tfr.result
|
||||
@@ -159,127 +127,30 @@ def transcribe(
|
||||
|
||||
def ensure_model_available(
|
||||
model_name: str,
|
||||
device: str = "cpu",
|
||||
compute_type: str = "float32",
|
||||
on_status: Callable[[str], None] | None = None,
|
||||
) -> str:
|
||||
"""Резолвит alias модели в repo_id и гарантирует наличие файлов.
|
||||
|
||||
Стратегия: cache-first (``local_files_only=True``), затем download.
|
||||
Два вызова ``snapshot_download`` — чтобы не лезть в сеть, если модель уже в кэше.
|
||||
"""
|
||||
local_path = Path(model_name).expanduser()
|
||||
if local_path.is_dir():
|
||||
_validate_model_dir(local_path)
|
||||
return str(local_path)
|
||||
|
||||
repo_id = _resolve_model_repo(model_name)
|
||||
|
||||
try:
|
||||
_notify_status(on_status, f"Проверяю кэш модели {model_name}...")
|
||||
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_status(on_status, f"Кэш модели {model_name} неполный, докачиваю...")
|
||||
|
||||
_notify_status(on_status, f"Скачиваю модель {model_name} из Hugging Face...")
|
||||
downloaded_path = Path(_snapshot_download(repo_id, local_files_only=False))
|
||||
_validate_model_dir(downloaded_path)
|
||||
return str(downloaded_path)
|
||||
|
||||
|
||||
def _run_transcription(model, file_path, lang_arg, on_segment, on_status=None):
|
||||
"""Run model.transcribe and iterate segments. Returns (segments, info)."""
|
||||
segment_generator, info = model.transcribe(str(file_path), language=lang_arg)
|
||||
total_duration = info.duration
|
||||
segments: list[Segment] = []
|
||||
for raw_seg in segment_generator:
|
||||
seg = Segment(start=raw_seg.start, end=raw_seg.end, text=raw_seg.text)
|
||||
if on_segment is not None:
|
||||
on_segment(seg)
|
||||
segments.append(seg)
|
||||
_notify_status(
|
||||
on_status,
|
||||
f"Транскрибирую... {_fmt_time(seg.end)} / {_fmt_time(total_duration)}"
|
||||
f" [{len(segments)} сегм.]",
|
||||
)
|
||||
return segments, info
|
||||
|
||||
|
||||
def _create_model(model_name: str, device: str, compute_type: str):
|
||||
try:
|
||||
return WhisperModel(model_name, device=device, compute_type=compute_type)
|
||||
except ImportError as exc:
|
||||
# WhisperModel при инициализации может загружать файлы через HF Hub;
|
||||
# если в системе настроен SOCKS proxy, но socksio не установлен,
|
||||
# HF Hub бросает ImportError — оборачиваем в понятное сообщение
|
||||
if _is_missing_socksio_error(exc):
|
||||
raise RuntimeError(
|
||||
"Обнаружен SOCKS proxy, но не установлена зависимость `socksio`, "
|
||||
"нужная для загрузки модели из Hugging Face через proxy. "
|
||||
"Обновите окружение: `uv sync`."
|
||||
) from exc
|
||||
raise
|
||||
"""Публичный helper: гарантирует наличие модели для указанного бэкенда."""
|
||||
backend = get_backend(device)
|
||||
return backend.ensure_model_available(model_name, compute_type, on_status)
|
||||
|
||||
|
||||
def _is_cuda_error(exc: BaseException) -> bool:
|
||||
"""Проверка CUDA ошибок — используется в cli.py для Windows-диагностики."""
|
||||
msg = str(exc).lower()
|
||||
return any(k in msg for k in ("cuda", "cublas", "cudnn", "out of memory"))
|
||||
|
||||
|
||||
def _is_missing_socksio_error(exc: BaseException) -> bool:
|
||||
msg = str(exc).lower()
|
||||
return "socks proxy" in msg and "socksio" in msg
|
||||
|
||||
|
||||
def _fmt_time(seconds: float) -> str:
|
||||
m, s = divmod(int(seconds), 60)
|
||||
h, m = divmod(m, 60)
|
||||
return f"{h}:{m:02d}:{s:02d}" if h else f"{m:02d}:{s:02d}"
|
||||
def _is_backend_error(exc: BaseException, device: str) -> bool:
|
||||
"""Определяет, связана ли ошибка с конкретным бэкендом (а не с пользовательскими данными)."""
|
||||
if device in ("cuda", "cpu"):
|
||||
return _is_cuda_error(exc)
|
||||
# openvino и другие бэкенды: конкретные паттерны ошибок добавим
|
||||
# при реализации бэкенда; пока — не маскируем ошибки
|
||||
return False
|
||||
|
||||
|
||||
def _notify_status(on_status: Callable[[str], None] | None, message: str) -> None:
|
||||
if on_status is not None:
|
||||
on_status(message)
|
||||
|
||||
|
||||
def _resolve_model_repo(model_name: str) -> str:
|
||||
if "/" in model_name:
|
||||
return model_name
|
||||
|
||||
repo_id = MODEL_REPOS.get(model_name)
|
||||
if repo_id is None:
|
||||
expected = ", ".join(MODEL_REPOS)
|
||||
raise ValueError(f"Неподдерживаемая модель '{model_name}'. Ожидалось одно из: {expected}")
|
||||
|
||||
return repo_id
|
||||
|
||||
|
||||
def _snapshot_download(repo_id: str, local_files_only: bool) -> str:
|
||||
try:
|
||||
return snapshot_download(
|
||||
repo_id,
|
||||
local_files_only=local_files_only,
|
||||
allow_patterns=MODEL_ALLOW_PATTERNS,
|
||||
)
|
||||
except ImportError as exc:
|
||||
if _is_missing_socksio_error(exc):
|
||||
raise RuntimeError(
|
||||
"Обнаружен SOCKS proxy, но не установлена зависимость `socksio`, "
|
||||
"нужная для загрузки модели из Hugging Face через proxy. "
|
||||
"Обновите окружение: `uv sync`."
|
||||
) from exc
|
||||
raise
|
||||
|
||||
|
||||
def _validate_model_dir(model_dir: Path) -> None:
|
||||
missing = [
|
||||
filename for filename in MODEL_REQUIRED_FILES if not (model_dir / filename).exists()
|
||||
]
|
||||
if not any(model_dir.glob("vocabulary.*")):
|
||||
missing.append("vocabulary.*")
|
||||
|
||||
if missing:
|
||||
missing_str = ", ".join(missing)
|
||||
raise ValueError(f"Неполная локальная модель в '{model_dir}': отсутствуют {missing_str}")
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Общие типы данных для всех бэкендов транскрипции."""
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class Segment:
|
||||
start: float # seconds
|
||||
end: float # seconds
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class TranscribeResult:
|
||||
segments: list[Segment]
|
||||
language: str
|
||||
language_probability: float
|
||||
duration: float # seconds
|
||||
device_used: str # "cpu" / "cuda" / "openvino"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TranscribeFileResult:
|
||||
result: TranscribeResult
|
||||
model: Any # backend-specific model handle
|
||||
actual_device: str
|
||||
backend: Any = None # backend instance (для переиспользования в батче)
|
||||
model_path: str = "" # путь к модели (меняется при cross-backend fallback)
|
||||
|
||||
|
||||
StatusCallback = Callable[[str], None] | None
|
||||
SegmentCallback = Callable[[Segment], None] | None
|
||||
Reference in New Issue
Block a user