fff

sentence_api/tests/test_transcription.py

@@ -0,0 +1,138 @@
+import io
+import wave
+
+import numpy as np
+import pytest
+
+from sentence_api.transcription import (
+    MossTranscriber,
+    _parse_json_segments,
+    _plan_chunks,
+    _wav_bytes,
+    _wav_duration_seconds,
+)
+
+
+SAMPLE_RATE = 16000
+
+
+def _speech(seconds: float, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
+    t = np.arange(int(seconds * sample_rate)) / sample_rate
+    return 0.25 * np.sin(2 * np.pi * 220 * t)
+
+
+def _silence(seconds: float, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
+    return np.zeros(int(seconds * sample_rate))
+
+
+def _audio(*parts: np.ndarray, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
+    return np.concatenate([np.asarray(part) for part in parts]).astype(np.float32)
+
+
+def test_plan_chunks_cuts_at_silence_gaps():
+    samples = _audio(
+        _speech(10),
+        _silence(0.8),
+        _speech(29.2),
+        _silence(1.0),
+        _speech(29.0),
+        _silence(0.6),
+        _speech(24.4),
+    )
+    chunks = _plan_chunks(samples, SAMPLE_RATE)
+    assert chunks[0][0] == pytest.approx(0.0)
+    assert chunks[-1][1] == pytest.approx(95.0)
+    for (start, end), (next_start, _) in zip(chunks, chunks[1:]):
+        assert end == pytest.approx(next_start)
+        assert end - start <= 34.0
+    assert 10.0 <= chunks[0][1] <= 10.8
+    assert 40.0 <= chunks[1][1] <= 41.0
+    assert 70.0 <= chunks[2][1] <= 70.6
+
+
+def test_plan_chunks_falls_back_to_fixed_windows():
+    chunks = _plan_chunks(_speech(95), SAMPLE_RATE)
+    assert chunks == [(0.0, 30.0), (30.0, 60.0), (60.0, 90.0), (90.0, 95.0)]
+
+
+def test_plan_chunks_merges_short_tail():
+    chunks = _plan_chunks(_speech(93), SAMPLE_RATE)
+    assert chunks == [(0.0, 30.0), (30.0, 60.0), (60.0, 93.0)]
+
+
+def test_parse_json_segments_filters_hallucinated_segments():
+    raw = [
+        {"start": 0.0, "end": 3.0, "text": "hello world", "compression_ratio": 1.4},
+        {"start": 3.5, "end": 6.0, "text": "hello world", "compression_ratio": 12.0},
+        {"start": 6.5, "end": 9.0, "text": "no ratio field"},
+    ]
+    filtered = _parse_json_segments(raw, compression_limit=2.4)
+    assert [segment.text for segment in filtered] == ["hello world", "no ratio field"]
+    assert len(_parse_json_segments(raw)) == 3
+
+
+def test_wav_bytes_roundtrip(tmp_path):
+    samples = _speech(2.0)
+    data = _wav_bytes(samples, SAMPLE_RATE)
+    assert data.startswith(b"RIFF")
+    path = tmp_path / "chunk.wav"
+    path.write_bytes(data)
+    assert _wav_duration_seconds(path) == pytest.approx(2.0, abs=0.05)
+    with wave.open(str(path), "rb") as wav:
+        assert wav.getnchannels() == 1
+        assert wav.getsampwidth() == 2
+        assert wav.getframerate() == SAMPLE_RATE
+
+
+def test_wav_duration_returns_none_for_non_wav(tmp_path):
+    path = tmp_path / "note.txt"
+    path.write_text("not audio")
+    assert _wav_duration_seconds(path) is None
+
+
+def test_chunked_transcribe_offsets_timestamps(monkeypatch):
+    requests = []
+    call_count = {"n": 0}
+
+    def fake_post(endpoint, data=None, files=None, timeout=None):
+        call_count["n"] += 1
+        requests.append(data)
+        audio = files["file"][1].read()
+        with wave.open(io.BytesIO(audio), "rb") as wav:
+            duration = wav.getnframes() / wav.getframerate()
+        if call_count["n"] == 1:
+            segments = [
+                {"start": 1.0, "end": 4.0, "text": "first part", "compression_ratio": 1.4},
+                {"start": 27.0, "end": 30.0, "text": "second part", "compression_ratio": 1.5},
+            ]
+        elif call_count["n"] == 2:
+            segments = [
+                {"start": 1.0, "end": 4.0, "text": "third part", "compression_ratio": 1.4},
+                {"start": 27.0, "end": 30.0, "text": "fourth part", "compression_ratio": 1.5},
+            ]
+        else:
+            segments = [{"start": 0.5, "end": 3.5, "text": "tail part", "compression_ratio": 1.3}]
+        assert abs(duration - 30.0) < 1.0 or abs(duration - 5.0) < 1.0
+
+        class Response:
+            def raise_for_status(self):
+                pass
+
+            def json(self):
+                return {"text": "", "segments": segments}
+
+        return Response()
+
+    monkeypatch.setattr("sentence_api.transcription.httpx.post", fake_post)
+    transcriber = MossTranscriber(endpoint="http://whisper:9000", model="whisper")
+    result = transcriber._transcribe_chunked(_speech(65), SAMPLE_RATE, None)
+    assert [(s.start_seconds, s.end_seconds, s.text) for s in result.segments] == [
+        (1.0, 4.0, "first part"),
+        (27.0, 30.0, "second part"),
+        (31.0, 34.0, "third part"),
+        (57.0, 60.0, "fourth part"),
+        (60.5, 63.5, "tail part"),
+    ]
+    assert result.text == "first part second part third part fourth part tail part"
+    assert len(requests) == 3
+    assert all(request.get("condition_on_previous_text") == "false" for request in requests)