fff

sentence_api/transcription.py

@@ -1,10 +1,28 @@
+import io
 import mimetypes
 import re
+import wave
 from dataclasses import dataclass
 from pathlib import Path
 from typing import List, Optional, Protocol
 
 import httpx
+import numpy as np
+
+from .audio_metrics import AudioAnalysisError, decode_audio_mono, speech_frame_mask
+
+
+# Whisper was trained on 30-second clips. Sending longer audio in a single request
+# makes the model enter a repetition loop (every later segment becomes the same
+# sentence while timestamps keep advancing), so long audio is split into short
+# chunks at silence gaps and each chunk is transcribed independently.
+CHUNK_SECONDS = 30.0
+MIN_CHUNK_SECONDS = 5.0
+MIN_SILENCE_SECONDS = 0.5
+MIN_TAIL_SECONDS = 4.0
+# Whisper's own repetition-loop detector: segments that compress this well are
+# almost certainly hallucinated repeated text (healthy speech is ~1.3-1.8).
+COMPRESSION_RATIO_LIMIT = 2.4
 
 
 @dataclass(frozen=True)
@@ -50,29 +68,25 @@ class MossTranscriber:
     def transcribe(self, audio_path: Path, language: Optional[str] = None) -> Transcript:
         if not self.available:
             raise RuntimeError("MOSS transcription is not configured on this server.")
+        duration = _wav_duration_seconds(audio_path)
+        if duration is None or duration > CHUNK_SECONDS + 2.0:
+            try:
+                samples, sample_rate = decode_audio_mono(audio_path)
+            except AudioAnalysisError:
+                return self._transcribe_single(audio_path, language)
+            if samples.size / sample_rate > CHUNK_SECONDS + 2.0:
+                try:
+                    return self._transcribe_chunked(samples, sample_rate, language)
+                except AudioAnalysisError:
+                    return self._transcribe_single(audio_path, language)
+        return self._transcribe_single(audio_path, language)
+
+    def _transcribe_single(
+        self, audio_path: Path, language: Optional[str]
+    ) -> Transcript:
         content_type = mimetypes.guess_type(audio_path.name)[0] or "application/octet-stream"
-        data = {
-            "model": self.model,
-            "response_format": "verbose_json",
-            "temperature": "0",
-            "max_new_tokens": str(self.max_new_tokens),
-        }
-        if language:
-            data["language"] = language
-
-        try:
-            with audio_path.open("rb") as audio_file:
-                response = httpx.post(
-                    self.endpoint,
-                    data=data,
-                    files={"file": (audio_path.name, audio_file, content_type)},
-                    timeout=httpx.Timeout(self.timeout_seconds, connect=30),
-                )
-            response.raise_for_status()
-            payload = response.json()
-        except (httpx.HTTPError, ValueError) as exc:
-            raise RuntimeError(f"MOSS transcription request failed: {exc}") from exc
-
+        with audio_path.open("rb") as audio_file:
+            payload = self._post_audio(audio_file, audio_path.name, content_type, language)
         raw_text = str(payload.get("text") or "").strip()
         segments = _parse_json_segments(payload.get("segments"))
         if not segments:
@@ -80,8 +94,71 @@ class MossTranscriber:
         plain_text = " ".join(segment.text for segment in segments).strip() or raw_text
         return Transcript(text=plain_text, segments=segments)
 
+    def _transcribe_chunked(
+        self, samples: np.ndarray, sample_rate: int, language: Optional[str]
+    ) -> Transcript:
+        segments: List[TranscriptionSegment] = []
+        for start_seconds, end_seconds in _plan_chunks(samples, sample_rate):
+            chunk = samples[
+                int(round(start_seconds * sample_rate)) : int(round(end_seconds * sample_rate))
+            ]
+            payload = self._post_audio(
+                io.BytesIO(_wav_bytes(chunk, sample_rate)),
+                "chunk.wav",
+                "audio/wav",
+                language,
+            )
+            chunk_segments = _parse_json_segments(
+                payload.get("segments"), compression_limit=COMPRESSION_RATIO_LIMIT
+            )
+            if not chunk_segments:
+                chunk_segments = _parse_compact_segments(str(payload.get("text") or "").strip())
+            for segment in chunk_segments:
+                segments.append(
+                    TranscriptionSegment(
+                        start_seconds=segment.start_seconds + start_seconds,
+                        end_seconds=segment.end_seconds + start_seconds,
+                        text=segment.text,
+                        speaker=segment.speaker,
+                    )
+                )
+        segments.sort(key=lambda item: (item.start_seconds, item.end_seconds))
+        plain_text = " ".join(segment.text for segment in segments).strip()
+        return Transcript(text=plain_text, segments=segments)
 
-def _parse_json_segments(raw_segments) -> List[TranscriptionSegment]:
+    def _post_audio(
+        self,
+        audio_file,
+        filename: str,
+        content_type: str,
+        language: Optional[str],
+    ) -> dict:
+        data = {
+            "model": self.model,
+            "response_format": "verbose_json",
+            "temperature": "0",
+            "condition_on_previous_text": "false",
+            "max_new_tokens": str(self.max_new_tokens),
+        }
+        if language:
+            data["language"] = language
+
+        try:
+            response = httpx.post(
+                self.endpoint,
+                data=data,
+                files={"file": (filename, audio_file, content_type)},
+                timeout=httpx.Timeout(self.timeout_seconds, connect=30),
+            )
+            response.raise_for_status()
+            return response.json()
+        except (httpx.HTTPError, ValueError) as exc:
+            raise RuntimeError(f"MOSS transcription request failed: {exc}") from exc
+
+
+def _parse_json_segments(
+    raw_segments, compression_limit: Optional[float] = None
+) -> List[TranscriptionSegment]:
     if not isinstance(raw_segments, list):
         return []
     segments = []
@@ -93,6 +170,13 @@ def _parse_json_segments(raw_segments) -> List[TranscriptionSegment]:
             end = float(item["end"])
         except (KeyError, TypeError, ValueError):
             continue
+        compression_ratio = item.get("compression_ratio")
+        if (
+            compression_limit is not None
+            and isinstance(compression_ratio, (int, float))
+            and compression_ratio > compression_limit
+        ):
+            continue
         text = str(item.get("text") or "").strip()
         if text and end > start >= 0:
             speaker = item.get("speaker")
@@ -107,6 +191,80 @@ def _parse_json_segments(raw_segments) -> List[TranscriptionSegment]:
     return segments
 
 
+def _wav_duration_seconds(path: Path) -> Optional[float]:
+    try:
+        with wave.open(str(path), "rb") as wav:
+            if wav.getframerate() > 0:
+                return wav.getnframes() / wav.getframerate()
+    except (wave.Error, OSError):
+        pass
+    return None
+
+
+def _wav_bytes(samples: np.ndarray, sample_rate: int) -> bytes:
+    pcm = np.clip(np.round(samples * 32767), -32768, 32767).astype("<i2")
+    buffer = io.BytesIO()
+    with wave.open(buffer, "wb") as wav:
+        wav.setnchannels(1)
+        wav.setsampwidth(2)
+        wav.setframerate(sample_rate)
+        wav.writeframes(pcm.tobytes())
+    return buffer.getvalue()
+
+
+def _plan_chunks(samples: np.ndarray, sample_rate: int) -> List[tuple]:
+    speech = speech_frame_mask(samples, sample_rate)
+    frame_seconds = 0.03
+    gap_frames = max(1, int(round(MIN_SILENCE_SECONDS / frame_seconds)))
+
+    gaps: List[tuple] = []
+    run_start: Optional[int] = None
+    for index, is_speech in enumerate(speech):
+        if not is_speech and run_start is None:
+            run_start = index
+        elif is_speech and run_start is not None:
+            if index - run_start >= gap_frames:
+                gaps.append((run_start * frame_seconds, index * frame_seconds))
+            run_start = None
+    if run_start is not None and len(speech) - run_start >= gap_frames:
+        gaps.append((run_start * frame_seconds, len(speech) * frame_seconds))
+
+    duration = samples.size / sample_rate
+    cuts: List[float] = []
+    cursor = 0.0
+    while duration - cursor > CHUNK_SECONDS + 1e-6:
+        target = cursor + CHUNK_SECONDS
+        best_gap = None
+        for gap_start, gap_end in gaps:
+            if gap_start < cursor + MIN_CHUNK_SECONDS:
+                continue
+            if gap_start > target:
+                break
+            best_gap = (gap_start, gap_end)
+        if best_gap is None:
+            cut = min(target, duration)
+        else:
+            gap_start, gap_end = best_gap
+            cut = min(target, (gap_start + gap_end) / 2.0)
+        cuts.append(cut)
+        cursor = cut
+
+    boundaries = [0.0] + cuts + [duration]
+    chunks = [
+        (boundaries[index], boundaries[index + 1])
+        for index in range(len(boundaries) - 1)
+        if boundaries[index + 1] > boundaries[index] + 0.05
+    ]
+    if (
+        len(chunks) > 1
+        and chunks[-1][1] - chunks[-1][0] < MIN_TAIL_SECONDS
+        and chunks[-1][1] - chunks[-2][0] <= CHUNK_SECONDS + MIN_TAIL_SECONDS
+    ):
+        chunks[-2] = (chunks[-2][0], chunks[-1][1])
+        chunks.pop()
+    return chunks
+
+
 _COMPACT_SEGMENT = re.compile(
     r"\[(?P<start>\d+(?:\.\d+)?)\]"
     r"\[(?P<speaker>S\d+)\]"