local-transcriber/src/local_transcriber/transcriber.py

import warnings
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path

from faster_whisper import WhisperModel
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",
}

MODEL_ALLOW_PATTERNS = [
    "config.json",
    "preprocessor_config.json",
    "model.bin",
    "tokenizer.json",
    "vocabulary.*",
]

MODEL_REQUIRED_FILES = [
    "config.json",
    "preprocessor_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"


def transcribe(
    file_path: Path,
    model_name: str = "large-v3",
    device: str = "auto",
    compute_type: str = "int8",
    language: str | None = None,
    on_segment: Callable[[Segment], None] | None = None,
    on_status: Callable[[str], None] | None = None,
) -> TranscribeResult:
    actual_device = device
    lang_arg = language if language and language != "auto" else None

    try:
        _notify_status(on_status, f"Загружаю модель на {device}...")
        model = _create_model(model_name, device, compute_type)
    except (RuntimeError, ValueError) as exc:
        if device != "cpu" and _is_cuda_error(exc):
            warnings.warn(
                f"Не удалось загрузить модель на {device}: {exc}. "
                "Переключение на CPU.",
                stacklevel=2,
            )
            actual_device = "cpu"
            _notify_status(on_status, "Загружаю модель на cpu...")
            model = _create_model(model_name, "cpu", compute_type)
        else:
            raise

    try:
        _notify_status(on_status, "Транскрибирую...")
        segments, info = _run_transcription(model, file_path, lang_arg, on_segment)
    except (RuntimeError, ValueError) as exc:
        if actual_device != "cpu" and _is_cuda_error(exc):
            warnings.warn(
                f"CUDA ошибка при транскрипции: {exc}. "
                "Переключение на CPU и повтор.",
                stacklevel=2,
            )
            actual_device = "cpu"
            _notify_status(on_status, "Загружаю модель на cpu...")
            model = _create_model(model_name, "cpu", compute_type)
            _notify_status(on_status, "Транскрибирую...")
            segments, info = _run_transcription(model, file_path, lang_arg, on_segment)
        else:
            raise

    return TranscribeResult(
        segments=segments,
        language=info.language,
        language_probability=info.language_probability,
        duration=info.duration,
        device_used=actual_device,
    )


def ensure_model_available(
    model_name: str,
    on_status: Callable[[str], None] | None = None,
) -> str:
    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):
    """Run model.transcribe and iterate segments. Returns (segments, info)."""
    segment_generator, info = model.transcribe(str(file_path), language=lang_arg)
    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)
    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:
        if _is_missing_socksio_error(exc):
            raise RuntimeError(
                "Обнаружен SOCKS proxy, но не установлена зависимость `socksio`, "
                "нужная для загрузки модели из Hugging Face через proxy. "
                "Обновите окружение: `uv sync`."
            ) from exc
        raise


def _is_cuda_error(exc: BaseException) -> bool:
    msg = str(exc).lower()
    return "cuda" in msg or "out of memory" in msg


def _is_missing_socksio_error(exc: BaseException) -> bool:
    msg = str(exc).lower()
    return "socks proxy" in msg and "socksio" in msg


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}")