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:
@@ -3,8 +3,11 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
from local_transcriber.types import Segment, TranscribeResult
|
||||||
|
|
||||||
MODEL_ALIASES: dict[str, str] = {
|
MODEL_ALIASES: dict[str, str] = {
|
||||||
"gigaam-v3": "gigaam-v3-ctc",
|
"gigaam-v3": "gigaam-v3-ctc",
|
||||||
"parakeet-v3": "nemo-parakeet-tdt-0.6b-v3",
|
"parakeet-v3": "nemo-parakeet-tdt-0.6b-v3",
|
||||||
@@ -62,6 +65,52 @@ class OnnxAsrBackend:
|
|||||||
self._vad = vad
|
self._vad = vad
|
||||||
return model.with_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:
|
def _resolve_model(self, model_name: str) -> str:
|
||||||
"""Resolve alias to onnx-asr model name. Raw names pass through."""
|
"""Resolve alias to onnx-asr model name. Raw names pass through."""
|
||||||
if model_name in MODEL_ALIASES:
|
if model_name in MODEL_ALIASES:
|
||||||
@@ -74,3 +123,8 @@ class OnnxAsrBackend:
|
|||||||
f"Доступные алиасы: {SUPPORTED_ALIASES}. "
|
f"Доступные алиасы: {SUPPORTED_ALIASES}. "
|
||||||
f"Либо укажите полное имя модели onnx-asr."
|
f"Либо укажите полное имя модели onnx-asr."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _notify(on_status: Callable[[str], None] | None, message: str) -> None:
|
||||||
|
if on_status is not None:
|
||||||
|
on_status(message)
|
||||||
|
|||||||
@@ -1,7 +1,19 @@
|
|||||||
"""Tests for onnx-asr backend."""
|
"""Tests for onnx-asr backend."""
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
from local_transcriber.backends.onnx_asr import OnnxAsrBackend, MODEL_ALIASES
|
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:
|
class TestEnsureModelAvailable:
|
||||||
@@ -95,6 +107,113 @@ class TestCreateModel:
|
|||||||
assert calls == ["fp16"]
|
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:
|
class TestModelAliases:
|
||||||
def test_gigaam_v3_resolves(self):
|
def test_gigaam_v3_resolves(self):
|
||||||
backend = OnnxAsrBackend()
|
backend = OnnxAsrBackend()
|
||||||
|
|||||||
Reference in New Issue
Block a user