feat(onnx-asr): реализовать transcribe с сегментацией через VAD

- Зачем:
  - Основной метод бэкенда — транскрипция аудиофайла в сегменты с временными метками.
- Что:
  - метод transcribe декодирует аудио через faster_whisper.decode_audio, затем вызывает model.recognize() с VAD-сегментацией.
  - каждый VAD-сегмент преобразуется в проектную структуру Segment.
  - поддержка колбэков on_segment, on_status.
  - написаны 4 теста: сбор сегментов, вызов on_segment, передача языка, обработка пустого аудио.
- Проверка:
  - uv run pytest tests/test_onnx_asr.py -v (14 passed)
This commit is contained in:
2026-04-25 21:24:31 +03:00
parent d9c9aefdb3
commit 9e04dc8c25
2 changed files with 173 additions and 0 deletions
@@ -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)
+119
View File
@@ -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()