diff --git a/src/local_transcriber/backends/onnx_asr.py b/src/local_transcriber/backends/onnx_asr.py index 43f186c..664cba0 100644 --- a/src/local_transcriber/backends/onnx_asr.py +++ b/src/local_transcriber/backends/onnx_asr.py @@ -59,7 +59,6 @@ class OnnxAsrBackend: model = onnx_asr.load_model( model=model_path, quantization=ct, - cpu_preprocessing=True, ) vad = onnx_asr.load_vad("silero") self._vad = vad @@ -89,10 +88,10 @@ class OnnxAsrBackend: segments: list[Segment] = [] detected_language = language or "unknown" - for vad_seg in model.recognize(audio_array, 16000, language=language): + for vad_seg in model.recognize(audio_array, sample_rate=16000, language=language): seg = Segment( - start=max(0.0, vad_seg.start_ts), - end=max(0.0, vad_seg.end_ts), + start=max(0.0, vad_seg.start), + end=max(0.0, vad_seg.end), text=vad_seg.text, ) if on_segment is not None: diff --git a/src/local_transcriber/config.py b/src/local_transcriber/config.py index 5d7ab3f..8b85234 100644 --- a/src/local_transcriber/config.py +++ b/src/local_transcriber/config.py @@ -22,11 +22,12 @@ DEVICE_DEFAULTS: dict[str, dict[str, str]] = { "openvino": {"model": "medium", "compute_type": "int8"}, "openvino-gpu": {"model": "medium", "compute_type": "int8"}, "openvino-cpu": {"model": "medium", "compute_type": "int8"}, + "onnx": {"model": "gigaam-v3", "compute_type": "int8"}, } # Одно место правды для допустимых ключей конфига _VALID_KEYS = set(HARDCODED_DEFAULTS) -_VALID_DEVICES = {"auto", "cpu", "cuda", "openvino", "openvino-gpu", "openvino-cpu"} +_VALID_DEVICES = {"auto", "cpu", "cuda", "openvino", "openvino-gpu", "openvino-cpu", "onnx"} def find_config_file() -> Path | None: diff --git a/tests/test_onnx_asr.py b/tests/test_onnx_asr.py index 31e31f3..207e60f 100644 --- a/tests/test_onnx_asr.py +++ b/tests/test_onnx_asr.py @@ -10,9 +10,9 @@ 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 + def __init__(self, start, end, text): + self.start = start + self.end = end self.text = text @@ -40,10 +40,9 @@ class TestCreateModel: calls = [] def fake_load_model(model=None, path=None, quantization=None, - cpu_preprocessing=None, **kwargs): + **kwargs): calls.append({ "model": model, "path": path, "quantization": quantization, - "cpu_preprocessing": cpu_preprocessing, }) return FakeAsrAdapter() @@ -59,7 +58,6 @@ class TestCreateModel: assert len(calls) == 1 assert calls[0]["quantization"] == "int8" - assert calls[0]["cpu_preprocessing"] is True assert model is not None def test_loads_silero_vad(self, monkeypatch):