add test module

sentence_api/transcription.py

@@ -0,0 +1,134 @@
+import mimetypes
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Protocol
+
+import httpx
+
+
+@dataclass(frozen=True)
+class TranscriptionSegment:
+    start_seconds: float
+    end_seconds: float
+    text: str
+    speaker: Optional[str] = None
+
+
+@dataclass(frozen=True)
+class Transcript:
+    text: str
+    segments: List[TranscriptionSegment]
+
+
+class Transcriber(Protocol):
+    @property
+    def available(self) -> bool:
+        ...
+
+    def transcribe(self, audio_path: Path, language: Optional[str] = None) -> Transcript:
+        ...
+
+
+class MossTranscriber:
+    def __init__(
+        self,
+        endpoint: str,
+        model: str,
+        timeout_seconds: float = 1800,
+        max_new_tokens: int = 65536,
+    ):
+        self.endpoint = endpoint.strip()
+        self.model = model
+        self.timeout_seconds = timeout_seconds
+        self.max_new_tokens = max_new_tokens
+
+    @property
+    def available(self) -> bool:
+        return bool(self.endpoint)
+
+    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.")
+        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
+
+        raw_text = str(payload.get("text") or "").strip()
+        segments = _parse_json_segments(payload.get("segments"))
+        if not segments:
+            segments = _parse_compact_segments(raw_text)
+        plain_text = " ".join(segment.text for segment in segments).strip() or raw_text
+        return Transcript(text=plain_text, segments=segments)
+
+
+def _parse_json_segments(raw_segments) -> List[TranscriptionSegment]:
+    if not isinstance(raw_segments, list):
+        return []
+    segments = []
+    for item in raw_segments:
+        if not isinstance(item, dict):
+            continue
+        try:
+            start = float(item["start"])
+            end = float(item["end"])
+        except (KeyError, TypeError, ValueError):
+            continue
+        text = str(item.get("text") or "").strip()
+        if text and end > start >= 0:
+            speaker = item.get("speaker")
+            segments.append(
+                TranscriptionSegment(
+                    start_seconds=start,
+                    end_seconds=end,
+                    text=text,
+                    speaker=str(speaker) if speaker is not None else None,
+                )
+            )
+    return segments
+
+
+_COMPACT_SEGMENT = re.compile(
+    r"\[(?P<start>\d+(?:\.\d+)?)\]"
+    r"\[(?P<speaker>S\d+)\]"
+    r"(?P<text>.*?)"
+    r"\[(?P<end>\d+(?:\.\d+)?)\]",
+    flags=re.DOTALL,
+)
+
+
+def _parse_compact_segments(text: str) -> List[TranscriptionSegment]:
+    segments = []
+    for match in _COMPACT_SEGMENT.finditer(text):
+        start = float(match.group("start"))
+        end = float(match.group("end"))
+        segment_text = match.group("text").strip()
+        if segment_text and end > start:
+            segments.append(
+                TranscriptionSegment(
+                    start_seconds=start,
+                    end_seconds=end,
+                    text=segment_text,
+                    speaker=match.group("speaker"),
+                )
+            )
+    return segments