change sentence cut method

sentence_api/tests/test_transcription.py

@@ -6,8 +6,12 @@ import pytest
 
 from sentence_api.transcription import (
     MossTranscriber,
+    TranscriptionSegment,
+    WordTimestamp,
     _parse_json_segments,
     _plan_chunks,
+    _parse_word_timestamps,
+    split_segment_by_periods,
     _wav_bytes,
     _wav_duration_seconds,
 )
@@ -107,7 +111,16 @@ def test_chunked_transcribe_offsets_timestamps(monkeypatch):
             ]
         elif call_count["n"] == 2:
             segments = [
-                {"start": 1.0, "end": 4.0, "text": "third part", "compression_ratio": 1.4},
+                {
+                    "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:
@@ -115,6 +128,8 @@ def test_chunked_transcribe_offsets_timestamps(monkeypatch):
         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
 
@@ -134,5 +149,162 @@ def test_chunked_transcribe_offsets_timestamps(monkeypatch):
         (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_segment_by_periods_uses_word_timestamps():
+    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_segment_by_periods(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_segment_by_periods_falls_back_to_proportional():
+    segment = TranscriptionSegment(
+        start_seconds=10.0,
+        end_seconds=20.0,
+        text="First sentence. Second sentence. Third.",
+    )
+    sentences = split_segment_by_periods(segment)
+    assert [s.text for s in sentences] == [
+        "First sentence.",
+        "Second sentence.",
+        "Third.",
+    ]
+    assert sentences[0].start_seconds == 10.0
+    assert sentences[1].start_seconds == sentences[0].end_seconds
+    assert sentences[2].end_seconds == 20.0
+    assert sentences[0].end_seconds > 10.0
+    assert sentences[1].end_seconds < 20.0
+    assert sentences[0].end_seconds < sentences[1].end_seconds
+
+
+def test_split_segment_by_periods_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_segment_by_periods(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_segment_by_periods_keeps_segment_without_period():
+    segment = TranscriptionSegment(
+        start_seconds=1.0, end_seconds=2.0, text="no period here"
+    )
+    assert split_segment_by_periods(segment) == [segment]
+
+
+def test_split_segment_by_periods_handles_ellipsis_and_dots_only():
+    segment = TranscriptionSegment(
+        start_seconds=0.0,
+        end_seconds=2.0,
+        text="Wait... What? ...",
+    )
+    sentences = split_segment_by_periods(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
+
+
+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]