diff --git a/src/local_transcriber/backends/onnx_asr.py b/src/local_transcriber/backends/onnx_asr.py index d4469b0..ad9e0ee 100644 --- a/src/local_transcriber/backends/onnx_asr.py +++ b/src/local_transcriber/backends/onnx_asr.py @@ -2,6 +2,7 @@ from __future__ import annotations +import warnings from collections.abc import Callable from dataclasses import dataclass from pathlib import Path @@ -12,35 +13,77 @@ from local_transcriber.types import Segment, TranscribeResult @dataclass(frozen=True) class OnnxModelSpec: - """Имя onnx-asr и опубликованные варианты квантизации модели.""" + """Имя onnx-asr, варианты квантизации и поддерживаемые языки.""" model_id: str quantizations: frozenset[str | None] + supported_languages: frozenset[str] _INT8_AND_FLOAT32 = frozenset({"int8", None}) +_RUSSIAN_ONLY = frozenset({"ru"}) +_GIGAAM_MULTILINGUAL_LANGUAGES = frozenset({"ru", "en", "kk", "ky", "uz"}) +_PARAKEET_V3_LANGUAGES = frozenset( + { + "bg", + "hr", + "cs", + "da", + "nl", + "en", + "et", + "fi", + "fr", + "de", + "el", + "hu", + "it", + "lv", + "lt", + "mt", + "pl", + "pt", + "ro", + "sk", + "sl", + "es", + "sv", + "ru", + "uk", + } +) +_WHISPER_MODEL_NAMES = frozenset( + {"tiny", "base", "small", "medium", "large-v3", "large-v3-turbo"} +) MODEL_CATALOG: dict[str, OnnxModelSpec] = { - "gigaam-v3": OnnxModelSpec("gigaam-v3-ctc", _INT8_AND_FLOAT32), + "gigaam-v3": OnnxModelSpec( + "gigaam-v3-ctc", _INT8_AND_FLOAT32, _RUSSIAN_ONLY + ), "parakeet-v3": OnnxModelSpec( "nemo-parakeet-tdt-0.6b-v3", _INT8_AND_FLOAT32, + _PARAKEET_V3_LANGUAGES, ), "gigaam-multilingual-ctc": OnnxModelSpec( "gigaam-multilingual-ctc", _INT8_AND_FLOAT32, + _GIGAAM_MULTILINGUAL_LANGUAGES, ), "gigaam-multilingual-large-ctc": OnnxModelSpec( "gigaam-multilingual-large-ctc", _INT8_AND_FLOAT32, + _GIGAAM_MULTILINGUAL_LANGUAGES, ), "gigaam-v3-e2e-ctc": OnnxModelSpec( "gigaam-v3-e2e-ctc", _INT8_AND_FLOAT32, + _RUSSIAN_ONLY, ), "gigaam-v3-e2e-rnnt": OnnxModelSpec( "gigaam-v3-e2e-rnnt", _INT8_AND_FLOAT32, + _RUSSIAN_ONLY, ), } @@ -85,6 +128,8 @@ class OnnxAsrBackend: self._compute_type_explicit = compute_type_explicit self.actual_compute_type: str | None = None self._resolved_model_id: str | None = None + self._model_name: str | None = None + self._model_spec: OnnxModelSpec | None = None self._vad: Any = None def ensure_model_available( @@ -119,6 +164,8 @@ class OnnxAsrBackend: self.actual_compute_type = resolved_compute_type self._resolved_model_id = self._resolve_model(model_name) + self._model_name = model_name + self._model_spec = spec return self._resolved_model_id def create_model( @@ -162,13 +209,14 @@ class OnnxAsrBackend: """ from faster_whisper import decode_audio + self._warn_if_language_unsupported(language) _notify(on_status, "Загружаю аудио...") audio_array = decode_audio(str(file_path), sampling_rate=16000) duration = len(audio_array) / 16000.0 _notify(on_status, "Транскрибирую (onnx-asr)...") segments: list[Segment] = [] - detected_language = language or "unknown" + result_language = language or _model_language(self._model_spec) or "unknown" for vad_seg in model.recognize( audio_array, sample_rate=16000, language=language @@ -192,7 +240,7 @@ class OnnxAsrBackend: return TranscribeResult( segments=segments, - language=detected_language, + language=result_language, language_probability=1.0 if language else 0.0, duration=duration, device_used="", # оркестратор проставит @@ -202,6 +250,13 @@ class OnnxAsrBackend: """Resolve alias to onnx-asr model name. Raw names pass through.""" if model_name in MODEL_ALIASES: return MODEL_ALIASES[model_name] + if model_name in _WHISPER_MODEL_NAMES: + raise ValueError( + f"Модель '{model_name}' относится к Whisper и не поддерживается " + "ONNX-бэкендом. Без CUDA --device auto выбирает ONNX; " + f"укажите --device openvino-cpu --model {model_name} " + "или --device cuda --model medium." + ) if "/" in model_name or model_name.count("-") >= 2: # Looks like a raw onnx-asr name — allow passthrough return model_name @@ -211,6 +266,25 @@ class OnnxAsrBackend: f"Либо укажите полное имя модели onnx-asr." ) + def _warn_if_language_unsupported(self, language: str | None) -> None: + if ( + language is None + or self._model_spec is None + or language in self._model_spec.supported_languages + ): + return + + supported = ", ".join(sorted(self._model_spec.supported_languages)) + warnings.warn( + f"Язык '{language}' не поддерживается моделью '{self._model_name}' " + f"(поддерживаются: {supported}). Результат может быть некорректным. " + "Для других языков используйте " + "--device openvino-cpu --model medium " + "или --device cuda --model medium.", + UserWarning, + stacklevel=2, + ) + def _format_compute_types(quantizations: frozenset[str | None]) -> str: values = [_compute_type_for_quantization(value) for value in quantizations] @@ -228,6 +302,12 @@ def _preferred_compute_type(quantizations: frozenset[str | None]) -> str: raise ValueError("Для ONNX-модели не указаны доступные квантизации") +def _model_language(spec: OnnxModelSpec | None) -> str | None: + if spec is not None and len(spec.supported_languages) == 1: + return next(iter(spec.supported_languages)) + return None + + def _notify(on_status: Callable[[str], None] | None, message: str) -> None: if on_status is not None: on_status(message) diff --git a/src/local_transcriber/cli.py b/src/local_transcriber/cli.py index 65c041a..3adebf1 100644 --- a/src/local_transcriber/cli.py +++ b/src/local_transcriber/cli.py @@ -57,6 +57,19 @@ def _format_device_info(device_used: str) -> str: return "CPU" +def _format_language_mode( + requested_language: str, result: TranscribeResult +) -> str: + """Описывает источник языка, не выдавая профиль модели за детектор.""" + if requested_language != "auto": + return "forced" + if result.language_probability > 0: + return "detected" + if result.language not in {"", "auto", "unknown"}: + return "из профиля модели" + return "не определён" + + def _format_repetition_blocks( blocks: list[RepetitionBlock], use_hours: bool, @@ -315,7 +328,7 @@ def _run_single( ) device_info = _format_device_info(result.device_used) - language_mode = "detected" if defaults["language"] == "auto" else "forced" + language_mode = _format_language_mode(defaults["language"], result) content = format_transcript( result=result, @@ -401,8 +414,6 @@ def _run_batch( # Phase 3: Transcribe processed = 0 failed = 0 - language_mode = "detected" if defaults["language"] == "auto" else "forced" - batch_start = time.monotonic() for i, file in enumerate(to_process, 1): @@ -442,6 +453,7 @@ def _run_batch( model_path = tfr.model_path result = tfr.result + language_mode = _format_language_mode(defaults["language"], result) if len(result.segments) == 0: console.print( diff --git a/src/local_transcriber/formatter.py b/src/local_transcriber/formatter.py index 9e0240a..55b0295 100644 --- a/src/local_transcriber/formatter.py +++ b/src/local_transcriber/formatter.py @@ -80,7 +80,7 @@ def format_transcript( source_filename: str, model_name: str, device_info: str, - language_mode: str, # "detected" | "forced" + language_mode: str, # detected | forced | из профиля модели | не определён transcription_date: datetime | None = None, # None -> datetime.now() ) -> str: """Собирает markdown-транскрипт: шапка с метаданными + абзацы с таймкодами.""" diff --git a/tests/test_backend_openvino.py b/tests/test_backend_openvino.py index 79c0799..5d7ad4a 100644 --- a/tests/test_backend_openvino.py +++ b/tests/test_backend_openvino.py @@ -17,15 +17,10 @@ from local_transcriber.types import Segment # === _resolve_repo === -def test_model_catalog_contains_supported_profiles(): - assert MODEL_REPOS == { - ("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", +def test_model_catalog_contains_large_v3_turbo_profiles(): + assert { + pair: repo for pair, repo in MODEL_REPOS.items() if pair[0] == "large-v3-turbo" + } == { ("large-v3-turbo", "int8"): "OpenVINO/whisper-large-v3-turbo-int8-ov", ("large-v3-turbo", "fp16"): "OpenVINO/whisper-large-v3-turbo-fp16-ov", } @@ -61,11 +56,20 @@ def test_resolve_repo_implicit_fallback(): assert backend._resolve_repo("base", "int8") == ("OpenVINO/whisper-base-fp16-ov", "fp16") -def test_resolve_repo_implicit_large_v3_prefers_fp16(): - """Неявный compute_type: large-v3 автоматически получает fp16.""" +@pytest.mark.parametrize( + ("model_name", "expected_compute_type"), + [("large-v3", "fp16"), ("large-v3-turbo", "int8")], +) +def test_resolve_repo_implicit_large_v3_profiles( + model_name, expected_compute_type +): + """Неявный compute_type различает обычную и turbo-модель.""" backend = OpenVINOBackend(compute_type_explicit=False) - # Дефолт int8, но для large-v3 override на fp16 - assert backend._resolve_repo("large-v3", "int8") == ("OpenVINO/whisper-large-v3-fp16-ov", "fp16") + + assert backend._resolve_repo(model_name, "int8") == ( + f"OpenVINO/whisper-{model_name}-{expected_compute_type}-ov", + expected_compute_type, + ) def test_resolve_repo_explicit_large_v3_int8_respected(): diff --git a/tests/test_cli.py b/tests/test_cli.py index 040c06c..1e5f3cd 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -5,7 +5,7 @@ import pytest from rich.console import Console from typer.testing import CliRunner -from local_transcriber.cli import _format_device_info, app +from local_transcriber.cli import _format_device_info, _format_language_mode, app from local_transcriber.transcriber import Segment, TranscribeFileResult, TranscribeResult runner = CliRunner() @@ -42,6 +42,24 @@ def _make_tfr(result=None, model=None, actual_device="cpu", backend=None, model_ ) +@pytest.mark.parametrize( + ("requested_language", "language", "probability", "expected"), + [ + ("ru", "ru", 1.0, "forced"), + ("auto", "ru", 0.95, "detected"), + ("auto", "ru", 0.0, "из профиля модели"), + ("auto", "unknown", 0.0, "не определён"), + ], +) +def test_format_language_mode( + requested_language, language, probability, expected +): + result = _make_result(language=language) + result.language_probability = probability + + assert _format_language_mode(requested_language, result) == expected + + def _single_patches(result=None, tmp_file=None, actual_device="cpu"): """Patches for a standard single-file CLI happy path.""" if result is None: diff --git a/tests/test_onnx_asr.py b/tests/test_onnx_asr.py index dc64b45..2d1e611 100644 --- a/tests/test_onnx_asr.py +++ b/tests/test_onnx_asr.py @@ -1,5 +1,7 @@ """Tests for onnx-asr backend.""" +import warnings + import pytest from local_transcriber.backends.onnx_asr import OnnxAsrBackend @@ -224,6 +226,68 @@ class TestCreateModel: class TestTranscribe: + @pytest.mark.parametrize( + ("model_name", "language", "expects_warning"), + [ + ("gigaam-v3-e2e-rnnt", "en", True), + ("gigaam-v3-e2e-rnnt", "ru", False), + ("gigaam-multilingual-ctc", "en", False), + ], + ) + def test_warns_when_language_is_not_supported( + self, monkeypatch, tmp_path, model_name, language, expects_warning + ): + wav_file = tmp_path / "test.wav" + wav_file.write_bytes(b"fake audio") + + monkeypatch.setattr( + "faster_whisper.decode_audio", + lambda path, sampling_rate=16000: [0.0] * 16000, + ) + + class FakeModel: + def recognize(self, waveform, sample_rate, language=None): + return iter(()) + + backend = OnnxAsrBackend() + backend.ensure_model_available(model_name, "int8") + + if expects_warning: + with pytest.warns( + UserWarning, + match=( + r"Язык 'en'.*--device openvino-cpu --model medium.*" + r"--device cuda --model medium" + ), + ): + backend.transcribe(FakeModel(), wav_file, language=language) + else: + with warnings.catch_warnings(record=True) as caught: + backend.transcribe(FakeModel(), wav_file, language=language) + assert caught == [] + + def test_auto_language_uses_single_supported_model_language( + self, monkeypatch, tmp_path + ): + wav_file = tmp_path / "test.wav" + wav_file.write_bytes(b"fake audio") + monkeypatch.setattr( + "faster_whisper.decode_audio", + lambda path, sampling_rate=16000: [0.0] * 16000, + ) + + class FakeModel: + def recognize(self, waveform, sample_rate, language=None): + return iter(()) + + backend = OnnxAsrBackend() + backend.ensure_model_available("gigaam-v3-e2e-rnnt", "int8") + + result = backend.transcribe(FakeModel(), wav_file, language=None) + + assert result.language == "ru" + assert result.language_probability == 0.0 + def test_transcribe_collects_segments(self, monkeypatch, tmp_path): """Verify transcribe maps VAD segments to project Segments.""" wav_file = tmp_path / "test.wav" @@ -386,3 +450,23 @@ class TestModelAliases: backend = OnnxAsrBackend() with pytest.raises(ValueError, match="Неподдерживаемая модель"): backend._resolve_model("nonexistent-model") + + def test_whisper_alias_error_suggests_explicit_backend(self): + backend = OnnxAsrBackend() + + with pytest.raises( + ValueError, + match=r"Whisper.*--device openvino-cpu.*--device cuda", + ): + backend._resolve_model("medium") + + def test_turbo_whisper_error_suggests_models_supported_by_backends(self): + backend = OnnxAsrBackend() + + with pytest.raises(ValueError) as exc_info: + backend._resolve_model("large-v3-turbo") + + message = str(exc_info.value) + assert "--device openvino-cpu --model large-v3-turbo" in message + assert "--device cuda --model medium" in message + assert "--device cuda --model large-v3-turbo" not in message diff --git a/tests/test_utils.py b/tests/test_utils.py index f95f089..0d919ea 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -66,41 +66,17 @@ def test_detect_device_explicit_passthrough(): assert detect_device("openvino-cpu") == "openvino-cpu" -def test_detect_device_auto_onnx_even_with_openvino_gpu(): - """auto + нет nvidia-smi → onnx, даже если доступен OpenVINO GPU.""" - with ( - patch("local_transcriber.utils.shutil.which", return_value=None), - patch("local_transcriber.utils._is_openvino_gpu_available", return_value=True), - ): - assert detect_device("auto") == "onnx" - - -def test_detect_device_auto_onnx_even_with_openvino_cpu(): - """auto + нет nvidia-smi → onnx, даже если доступен OpenVINO CPU.""" - with ( - patch("local_transcriber.utils.shutil.which", return_value=None), - patch("local_transcriber.utils._is_openvino_gpu_available", return_value=False), - patch("local_transcriber.utils._is_openvino_available", return_value=True), - ): - assert detect_device("auto") == "onnx" - - -def test_detect_device_cuda_over_openvino(): - """nvidia-smi доступен и openvino тоже → cuda побеждает.""" - with ( - patch("local_transcriber.utils.shutil.which", return_value="/usr/bin/nvidia-smi"), - patch("local_transcriber.utils._is_openvino_gpu_available", return_value=True), +def test_detect_device_auto_cuda_when_nvidia_smi_available(): + """При доступном nvidia-smi auto выбирает CUDA.""" + with patch( + "local_transcriber.utils.shutil.which", return_value="/usr/bin/nvidia-smi" ): assert detect_device("auto") == "cuda" -def test_detect_device_auto_onnx_without_accelerators(): +def test_detect_device_auto_onnx_without_cuda(): """Без CUDA auto выбирает ONNX CPU.""" - with ( - patch("local_transcriber.utils.shutil.which", return_value=None), - patch("local_transcriber.utils._is_openvino_gpu_available", return_value=False), - patch("local_transcriber.utils._is_openvino_available", return_value=False), - ): + with patch("local_transcriber.utils.shutil.which", return_value=None): assert detect_device("auto") == "onnx"