diff --git a/src/local_transcriber/backends/onnx_asr.py b/src/local_transcriber/backends/onnx_asr.py index 7411767..43f186c 100644 --- a/src/local_transcriber/backends/onnx_asr.py +++ b/src/local_transcriber/backends/onnx_asr.py @@ -3,8 +3,11 @@ from __future__ import annotations from collections.abc import Callable +from pathlib import Path from typing import Any +from local_transcriber.types import Segment, TranscribeResult + MODEL_ALIASES: dict[str, str] = { "gigaam-v3": "gigaam-v3-ctc", "parakeet-v3": "nemo-parakeet-tdt-0.6b-v3", @@ -62,6 +65,52 @@ class OnnxAsrBackend: self._vad = vad return model.with_vad(vad) + 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: + """Transcribes audio file using onnx-asr model with VAD. + + model: result of create_model() — a SegmentResultsAsrAdapter. + file_path: path to audio/video file (any format supported by faster-whisper decode). + language: language code (e.g. "ru", "en") — only meaningful for multilingual models. + """ + from faster_whisper import decode_audio + + _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" + + for vad_seg in model.recognize(audio_array, 16000, language=language): + seg = Segment( + start=max(0.0, vad_seg.start_ts), + end=max(0.0, vad_seg.end_ts), + text=vad_seg.text, + ) + if on_segment is not None: + on_segment(seg) + segments.append(seg) + _notify( + on_status, + f"Транскрибирую (onnx-asr)... [{len(segments)} сегм.]", + ) + + return TranscribeResult( + segments=segments, + language=detected_language, + language_probability=1.0 if language else 0.0, + duration=duration, + device_used="", # оркестратор проставит + ) + def _resolve_model(self, model_name: str) -> str: """Resolve alias to onnx-asr model name. Raw names pass through.""" if model_name in MODEL_ALIASES: @@ -74,3 +123,8 @@ class OnnxAsrBackend: f"Доступные алиасы: {SUPPORTED_ALIASES}. " f"Либо укажите полное имя модели onnx-asr." ) + + +def _notify(on_status: Callable[[str], None] | None, message: str) -> None: + if on_status is not None: + on_status(message) diff --git a/tests/test_onnx_asr.py b/tests/test_onnx_asr.py index 7a6a81b..f495519 100644 --- a/tests/test_onnx_asr.py +++ b/tests/test_onnx_asr.py @@ -1,7 +1,19 @@ """Tests for onnx-asr backend.""" import pytest +from pathlib import Path + from local_transcriber.backends.onnx_asr import OnnxAsrBackend, MODEL_ALIASES +from local_transcriber.types import Segment, TranscribeResult + + +class FakeVadSegment: + """Mimics onnx-asr SegmentResult.""" + + def __init__(self, start_ts, end_ts, text): + self.start_ts = start_ts + self.end_ts = end_ts + self.text = text class TestEnsureModelAvailable: @@ -95,6 +107,113 @@ class TestCreateModel: assert calls == ["fp16"] +class TestTranscribe: + def test_transcribe_collects_segments(self, monkeypatch, tmp_path): + """Verify transcribe maps VAD segments to project Segments.""" + wav_file = tmp_path / "test.wav" + wav_file.write_bytes(b"fake audio") + + audio_samples = [0.0] * 16000 # 1 second of silence + + def fake_decode_audio(path, sampling_rate=16000): + import numpy as np + return np.array(audio_samples, dtype=np.float32) + + class FakeModel: + def recognize(self, waveform, sample_rate, language=None): + yield FakeVadSegment(0.0, 1.0, "hello") + yield FakeVadSegment(1.0, 2.5, "world") + + monkeypatch.setattr("faster_whisper.decode_audio", fake_decode_audio) + + backend = OnnxAsrBackend() + backend.actual_compute_type = "int8" + result = backend.transcribe( + FakeModel(), wav_file, language=None, + ) + + assert isinstance(result, TranscribeResult) + assert len(result.segments) == 2 + assert result.segments[0] == Segment(start=0.0, end=1.0, text="hello") + assert result.segments[1] == Segment(start=1.0, end=2.5, text="world") + assert result.duration == 1.0 # 16000 samples / 16000 Hz + + def test_transcribe_calls_on_segment(self, monkeypatch, tmp_path): + """Verify on_segment callback is invoked per segment.""" + wav_file = tmp_path / "test.wav" + wav_file.write_bytes(b"fake audio") + + def fake_decode_audio(path, sampling_rate=16000): + import numpy as np + return np.array([0.0] * 16000, dtype=np.float32) + + segments_captured = [] + + class FakeModel: + def recognize(self, waveform, sample_rate, language=None): + yield FakeVadSegment(0.0, 2.0, "one") + yield FakeVadSegment(2.0, 4.0, "two") + + monkeypatch.setattr("faster_whisper.decode_audio", fake_decode_audio) + + backend = OnnxAsrBackend() + result = backend.transcribe( + FakeModel(), wav_file, language=None, + on_segment=lambda s: segments_captured.append(s), + ) + + assert len(segments_captured) == 2 + assert segments_captured[0].text == "one" + assert segments_captured[1].text == "two" + + def test_transcribe_passes_language(self, monkeypatch, tmp_path): + """Verify language is passed to recognize().""" + wav_file = tmp_path / "test.wav" + wav_file.write_bytes(b"fake audio") + + def fake_decode_audio(path, sampling_rate=16000): + import numpy as np + return np.array([0.0] * 16000, dtype=np.float32) + + lang_received = [] + + class FakeModel: + def recognize(self, waveform, sample_rate, language=None): + lang_received.append(language) + yield FakeVadSegment(0.0, 1.0, "text") + + monkeypatch.setattr("faster_whisper.decode_audio", fake_decode_audio) + + backend = OnnxAsrBackend() + backend.transcribe(FakeModel(), wav_file, language="ru") + + assert lang_received == ["ru"] + + def test_transcribe_empty_audio(self, monkeypatch, tmp_path): + """Verify zero segments for silent audio.""" + wav_file = tmp_path / "test.wav" + wav_file.write_bytes(b"fake audio") + + def fake_decode_audio(path, sampling_rate=16000): + import numpy as np + return np.array([0.0] * 16000, dtype=np.float32) + + class FakeModel: + def recognize(self, waveform, sample_rate, language=None): + # No segments yielded + if False: + yield + + monkeypatch.setattr("faster_whisper.decode_audio", fake_decode_audio) + + backend = OnnxAsrBackend() + result = backend.transcribe(FakeModel(), wav_file, language=None) + + assert len(result.segments) == 0 + assert result.language == "unknown" + assert result.duration == 1.0 + + class TestModelAliases: def test_gigaam_v3_resolves(self): backend = OnnxAsrBackend()