feat: transcribe_file returns timestamped segments when with_segments=True

transcription.py

@@ -1,5 +1,6 @@
 import asyncio
 import httpx
+from typing import Union
 
 
 class TranscriptionEngine:
@@ -25,34 +26,69 @@ class TranscriptionEngine:
         model_name: str = "large-v3",
         device: str = "auto",
         base_url: str = "",
-    ) -> str:
+        with_segments: bool = False,
+    ) -> Union[str, list[dict]]:
         if base_url:
-            return await self._transcribe_remote(audio_path, language, model_name, base_url)
-        return await self._transcribe_local(audio_path, language, model_name, device)
+            return await self._transcribe_remote(
+                audio_path, language, model_name, base_url, with_segments
+            )
+        return await self._transcribe_local(
+            audio_path, language, model_name, device, with_segments
+        )
 
     async def _transcribe_remote(
-        self, audio_path: str, language: str, model_name: str, base_url: str
-    ) -> str:
+        self,
+        audio_path: str,
+        language: str,
+        model_name: str,
+        base_url: str,
+        with_segments: bool,
+    ) -> Union[str, list[dict]]:
         async with httpx.AsyncClient(timeout=300) as client:
             with open(audio_path, "rb") as f:
+                data = {"model": model_name, "language": language}
+                if with_segments:
+                    data["timestamp_granularities[]"] = "segment"
+                    data["response_format"] = "verbose_json"
                 r = await client.post(
                     f"{base_url}/v1/audio/transcriptions",
                     files={"file": ("audio.wav", f, "audio/wav")},
-                    data={"model": model_name, "language": language},
+                    data=data,
                 )
             r.raise_for_status()
-            return r.json()["text"]
+            body = r.json()
+        if not with_segments:
+            return body["text"]
+        raw_segs = body.get("segments") or []
+        if raw_segs:
+            return [
+                {"start": s["start"], "end": s["end"], "text": s["text"].strip()}
+                for s in raw_segs
+            ]
+        return [{"start": 0.0, "end": 9999.0, "text": body["text"].strip()}]
 
     async def _transcribe_local(
-        self, audio_path: str, language: str, model_name: str, device: str
-    ) -> str:
+        self,
+        audio_path: str,
+        language: str,
+        model_name: str,
+        device: str,
+        with_segments: bool,
+    ) -> Union[str, list[dict]]:
         loop = asyncio.get_running_loop()
         model = self._get_model(model_name, device)
         segments, _ = await loop.run_in_executor(
             None,
             lambda: model.transcribe(audio_path, language=language),
         )
-        return "".join(seg.text for seg in segments).strip()
+        segments = list(segments)
+        if not with_segments:
+            return "".join(seg.text for seg in segments).strip()
+        return [
+            {"start": seg.start, "end": seg.end, "text": seg.text.strip()}
+            for seg in segments
+            if seg.text.strip()
+        ]
 
 
 engine = TranscriptionEngine()