import io import wave import numpy as np import pytest from sentence_api.transcription import ( MossTranscriber, TranscriptionSegment, WordTimestamp, _parse_json_segments, _plan_chunks, _parse_word_timestamps, split_sentences_at_punctuation, _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, "words": [ {"word": "third", "start": 1.0, "end": 2.5}, {"word": " part", "start": 2.6, "end": 4.0}, ], }, {"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: status_code = 200 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 [(w.start_seconds, w.end_seconds, w.text) for w in result.segments[2].words] == [ (31.0, 32.5, "third"), (32.6, 34.0, "part"), ] assert result.segments[0].words is None assert len(requests) == 3 assert all(request.get("condition_on_previous_text") == "false" for request in requests) assert all(request.get("timestamp_granularities[]") == "word" for request in requests) def test_parse_json_segments_parses_word_timestamps(): raw = [ { "start": 0.0, "end": 4.0, "text": "Hello world. Good day.", "words": [ {"word": "Hello", "start": 0.0, "end": 0.6}, {"word": " world.", "start": 0.7, "end": 1.5}, {"word": " Good", "start": 1.8, "end": 2.4}, {"word": " day.", "start": 2.5, "end": 3.2}, ], }, {"start": 4.5, "end": 6.0, "text": "no words"}, ] segments = _parse_json_segments(raw) assert [word.text for word in segments[0].words] == ["Hello", "world.", "Good", "day."] assert segments[1].words is None def test_parse_word_timestamps_ignores_invalid_entries(): words = _parse_word_timestamps( [ {"word": "ok", "start": 0.0, "end": 0.5}, {"word": "bad"}, {"word": "", "start": 1.0, "end": 1.5}, {"word": "flat", "start": 2.0, "end": 2.0}, ] ) assert [(word.text, word.start_seconds, word.end_seconds) for word in words] == [ ("ok", 0.0, 0.5) ] assert _parse_word_timestamps(None) is None assert _parse_word_timestamps("nope") is None def test_split_sentences_merges_segments_without_punctuation(): segments = [ TranscriptionSegment( start_seconds=0.0, end_seconds=8.0, text=( "So I've been single for about four years now and I find it hard " "to meet a guy especially in" ), ), TranscriptionSegment( start_seconds=8.0, end_seconds=9.5, text="California.", words=[WordTimestamp(8.0, 9.5, "California.")], ), TranscriptionSegment( start_seconds=9.5, end_seconds=20.0, text=( "I've tried dating apps, I've met friends through friends, " "I've done online dating" ), ), TranscriptionSegment( start_seconds=20.0, end_seconds=24.0, text=( "and I just feel like most of the guys that I meet aren't serious " "about a relationship." ), words=[WordTimestamp(20.0, 24.0, "relationship.")], ), ] sentences = split_sentences_at_punctuation(segments) assert [(s.text, s.start_seconds, s.end_seconds) for s in sentences] == [ ( "So I've been single for about four years now and I find it hard to meet " "a guy especially in California.", 0.0, 9.5, ), ( "I've tried dating apps, I've met friends through friends, I've done " "online dating and I just feel like most of the guys that I meet aren't " "serious about a relationship.", 9.5, 24.0, ), ] def test_split_sentences_uses_word_timestamps_within_segment(): segment = TranscriptionSegment( start_seconds=0.0, end_seconds=4.0, text="Hello world. Good day. Nice to meet you.", words=[ WordTimestamp(0.0, 0.6, "Hello"), WordTimestamp(0.7, 1.5, "world."), WordTimestamp(1.8, 2.4, "Good"), WordTimestamp(2.5, 3.2, "day."), WordTimestamp(3.3, 3.8, "Nice"), WordTimestamp(3.9, 4.0, "you."), ], ) sentences = split_sentences_at_punctuation([segment]) assert [(s.text, s.start_seconds, s.end_seconds) for s in sentences] == [ ("Hello world.", 0.0, 1.5), ("Good day.", 1.5, 3.2), ("Nice to meet you.", 3.2, 4.0), ] def test_split_sentences_splits_at_question_marks(): segment = TranscriptionSegment( start_seconds=0.0, end_seconds=10.0, text="Really? Wait a minute. What about now?", ) sentences = split_sentences_at_punctuation([segment]) assert [s.text for s in sentences] == [ "Really?", "Wait a minute.", "What about now?", ] assert sentences[0].start_seconds == 0.0 assert sentences[0].end_seconds == sentences[1].start_seconds assert sentences[1].end_seconds == sentences[2].start_seconds assert sentences[2].end_seconds == 10.0 def test_split_sentences_falls_back_to_proportional(): segments = [ TranscriptionSegment( start_seconds=10.0, end_seconds=14.0, text="This is the first" ), TranscriptionSegment( start_seconds=14.0, end_seconds=18.0, text="sentence. And this is the second one.", ), ] sentences = split_sentences_at_punctuation(segments) assert [s.text for s in sentences] == [ "This is the first sentence.", "And this is the second one.", ] assert sentences[0].start_seconds == 10.0 assert sentences[0].end_seconds == sentences[1].start_seconds assert sentences[1].end_seconds == 18.0 def test_split_sentences_falls_back_when_word_ends_are_invalid(): segment = TranscriptionSegment( start_seconds=0.0, end_seconds=2.0, text="One. Two.", words=[ WordTimestamp(0.0, 2.5, "One."), WordTimestamp(2.6, 3.0, "Two."), ], ) sentences = split_sentences_at_punctuation([segment]) assert [s.text for s in sentences] == ["One.", "Two."] assert sentences[0].end_seconds == sentences[1].start_seconds assert sentences[1].end_seconds == 2.0 def test_split_sentences_keeps_incomplete_tail_as_sentence(): segment = TranscriptionSegment( start_seconds=1.0, end_seconds=2.0, text="no period here" ) sentences = split_sentences_at_punctuation([segment]) assert [(s.text, s.start_seconds, s.end_seconds) for s in sentences] == [ ("no period here", 1.0, 2.0) ] assert split_sentences_at_punctuation([]) == [] def test_split_sentences_handles_ellipsis_and_dots_only(): segment = TranscriptionSegment( start_seconds=0.0, end_seconds=2.0, text="Wait... What? ...", ) sentences = split_sentences_at_punctuation([segment]) assert [s.text for s in sentences] == ["Wait...", "What?"] assert sentences[0].start_seconds == 0.0 assert sentences[0].end_seconds == sentences[1].start_seconds assert sentences[1].end_seconds == 2.0 assert split_sentences_at_punctuation( [TranscriptionSegment(0.0, 1.0, "...")] ) == [] def test_split_sentences_ignores_decimal_points(): segment = TranscriptionSegment( start_seconds=0.0, end_seconds=6.0, text="The rate rose by 0.4%. That is a big jump.", ) sentences = split_sentences_at_punctuation([segment]) assert [s.text for s in sentences] == [ "The rate rose by 0.4%.", "That is a big jump.", ] def test_split_sentences_ignores_am_pm_abbreviations(): segment = TranscriptionSegment( start_seconds=0.0, end_seconds=8.0, text=( "It was at 9 o'clock p.m. Eastern, President Trump was speaking. " "The room was quiet." ), ) sentences = split_sentences_at_punctuation([segment]) assert [s.text for s in sentences] == [ "It was at 9 o'clock p.m. Eastern, President Trump was speaking.", "The room was quiet.", ] morning = TranscriptionSegment( start_seconds=0.0, end_seconds=4.0, text="I woke up at 7 a.m. and had coffee. Then I left.", ) sentences = split_sentences_at_punctuation([morning]) assert [s.text for s in sentences] == [ "I woke up at 7 a.m. and had coffee.", "Then I left.", ] def test_split_sentences_merges_segment_ending_with_abbreviation(): segments = [ TranscriptionSegment(0.0, 5.0, "The meeting ends at 5 p.m."), TranscriptionSegment(5.0, 9.0, "Eastern time. Then we go home."), ] sentences = split_sentences_at_punctuation(segments) assert [s.text for s in sentences] == [ "The meeting ends at 5 p.m. Eastern time.", "Then we go home.", ] def test_split_sentences_distinguishes_no_reply_from_number(): reply = TranscriptionSegment(0.0, 4.0, "No. I don't think so.") assert [s.text for s in split_sentences_at_punctuation([reply])] == [ "No.", "I don't think so.", ] numbered = TranscriptionSegment(0.0, 4.0, "See no. 5 on the list.") assert [s.text for s in split_sentences_at_punctuation([numbered])] == [ "See no. 5 on the list." ] def test_split_sentences_ignores_common_abbreviations(): segment = TranscriptionSegment( start_seconds=0.0, end_seconds=4.0, text="Dr. Smith said the U.S. economy is growing. He was right.", ) sentences = split_sentences_at_punctuation([segment]) assert [s.text for s in sentences] == [ "Dr. Smith said the U.S. economy is growing.", "He was right.", ] def test_post_audio_requests_word_timestamps_and_falls_back(monkeypatch): calls = [] class RejectedResponse: status_code = 400 def raise_for_status(self): pass def json(self): return {"error": "unknown parameter"} class OkResponse: status_code = 200 def raise_for_status(self): pass def json(self): return {"text": "hi.", "segments": []} def fake_post(endpoint, data=None, files=None, timeout=None): calls.append(dict(data)) if "timestamp_granularities[]" in data: return RejectedResponse() return OkResponse() monkeypatch.setattr("sentence_api.transcription.httpx.post", fake_post) transcriber = MossTranscriber(endpoint="http://whisper:9000", model="whisper") result = transcriber._post_audio( io.BytesIO(b"fake-audio"), "clip.wav", "audio/wav", "en" ) assert result["text"] == "hi." assert calls[0]["timestamp_granularities[]"] == "word" assert "timestamp_granularities[]" not in calls[1]