135 lines
4.0 KiB
Python
135 lines
4.0 KiB
Python
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
|