437 lines
15 KiB
Python
437 lines
15 KiB
Python
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)
|
|
class TranscriptionSegment:
|
|
start_seconds: float
|
|
end_seconds: float
|
|
text: str
|
|
speaker: Optional[str] = None
|
|
words: Optional[List["WordTimestamp"]] = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WordTimestamp:
|
|
start_seconds: float
|
|
end_seconds: float
|
|
text: str
|
|
|
|
|
|
@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.")
|
|
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"
|
|
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:
|
|
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 _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:
|
|
words = None
|
|
if segment.words:
|
|
words = [
|
|
WordTimestamp(
|
|
start_seconds=word.start_seconds + start_seconds,
|
|
end_seconds=word.end_seconds + start_seconds,
|
|
text=word.text,
|
|
)
|
|
for word in segment.words
|
|
]
|
|
segments.append(
|
|
TranscriptionSegment(
|
|
start_seconds=segment.start_seconds + start_seconds,
|
|
end_seconds=segment.end_seconds + start_seconds,
|
|
text=segment.text,
|
|
speaker=segment.speaker,
|
|
words=words,
|
|
)
|
|
)
|
|
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 _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:
|
|
word_payload = dict(data)
|
|
word_payload["timestamp_granularities[]"] = "word"
|
|
response = self._post_once(
|
|
word_payload, audio_file, filename, content_type
|
|
)
|
|
if response.status_code in (400, 422):
|
|
audio_file.seek(0)
|
|
response = self._post_once(data, audio_file, filename, content_type)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except (httpx.HTTPError, ValueError) as exc:
|
|
raise RuntimeError(f"MOSS transcription request failed: {exc}") from exc
|
|
|
|
def _post_once(self, data: dict, audio_file, filename: str, content_type: str):
|
|
return httpx.post(
|
|
self.endpoint,
|
|
data=data,
|
|
files={"file": (filename, audio_file, content_type)},
|
|
timeout=httpx.Timeout(self.timeout_seconds, connect=30),
|
|
)
|
|
|
|
|
|
def _parse_json_segments(
|
|
raw_segments, compression_limit: Optional[float] = None
|
|
) -> 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
|
|
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")
|
|
segments.append(
|
|
TranscriptionSegment(
|
|
start_seconds=start,
|
|
end_seconds=end,
|
|
text=text,
|
|
speaker=str(speaker) if speaker is not None else None,
|
|
words=_parse_word_timestamps(item.get("words")),
|
|
)
|
|
)
|
|
return segments
|
|
|
|
|
|
def _parse_word_timestamps(raw_words) -> Optional[List[WordTimestamp]]:
|
|
if not isinstance(raw_words, list):
|
|
return None
|
|
words = []
|
|
for item in raw_words:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
try:
|
|
start = float(item["start"])
|
|
end = float(item["end"])
|
|
except (KeyError, TypeError, ValueError):
|
|
continue
|
|
text = str(item.get("word") or "").strip()
|
|
if text and end > start:
|
|
words.append(WordTimestamp(start_seconds=start, end_seconds=end, text=text))
|
|
return words or None
|
|
|
|
|
|
_PERIOD_BOUNDARY = re.compile(r"\.+")
|
|
|
|
|
|
def split_segment_by_periods(
|
|
segment: TranscriptionSegment,
|
|
) -> List[TranscriptionSegment]:
|
|
"""Split one whisper segment into sentences at periods (".")."""
|
|
text = segment.text.strip()
|
|
if not text:
|
|
return []
|
|
pieces: List[str] = []
|
|
cursor = 0
|
|
for match in _PERIOD_BOUNDARY.finditer(text):
|
|
piece = text[cursor:match.end()].strip()
|
|
if piece.strip(".").strip():
|
|
pieces.append(piece)
|
|
cursor = match.end()
|
|
tail = text[cursor:].strip()
|
|
if tail:
|
|
pieces.append(tail)
|
|
if len(pieces) <= 1:
|
|
return [segment] if pieces else []
|
|
ends = _sentence_end_seconds(segment, pieces)
|
|
sub_segments: List[TranscriptionSegment] = []
|
|
start = segment.start_seconds
|
|
for piece, end in zip(pieces, ends):
|
|
if end > start:
|
|
sub_segments.append(
|
|
TranscriptionSegment(
|
|
start_seconds=start,
|
|
end_seconds=end,
|
|
text=piece,
|
|
speaker=segment.speaker,
|
|
)
|
|
)
|
|
start = end
|
|
return sub_segments
|
|
|
|
|
|
def _sentence_end_seconds(
|
|
segment: TranscriptionSegment, pieces: List[str]
|
|
) -> List[float]:
|
|
lengths = [len(piece) for piece in pieces]
|
|
total = sum(lengths)
|
|
if segment.words and len(segment.words) >= len(pieces):
|
|
word_ends = _word_boundary_ends(segment, lengths)
|
|
if _valid_boundaries(segment, word_ends):
|
|
return word_ends + [segment.end_seconds]
|
|
return _proportional_ends(segment, lengths) + [segment.end_seconds]
|
|
|
|
|
|
def _word_boundary_ends(
|
|
segment: TranscriptionSegment, lengths: List[int]
|
|
) -> List[float]:
|
|
targets = [sum(lengths[:count]) for count in range(1, len(lengths))]
|
|
buffer = ""
|
|
word_indexes: List[int] = []
|
|
target_index = 0
|
|
for word_index, word in enumerate(segment.words):
|
|
buffer += (" " if buffer else "") + word.text
|
|
while target_index < len(targets) and len(buffer) >= targets[target_index]:
|
|
word_indexes.append(word_index)
|
|
target_index += 1
|
|
if len(word_indexes) != len(targets):
|
|
return []
|
|
return [segment.words[index].end_seconds for index in word_indexes]
|
|
|
|
|
|
def _valid_boundaries(
|
|
segment: TranscriptionSegment, ends: List[float]
|
|
) -> bool:
|
|
previous = segment.start_seconds
|
|
for end in ends:
|
|
if not segment.start_seconds < end < segment.end_seconds:
|
|
return False
|
|
if end <= previous:
|
|
return False
|
|
previous = end
|
|
return True
|
|
|
|
|
|
def _proportional_ends(
|
|
segment: TranscriptionSegment, lengths: List[int]
|
|
) -> List[float]:
|
|
total = sum(lengths)
|
|
ends = []
|
|
cumulative = 0
|
|
for length in lengths[:-1]:
|
|
cumulative += length
|
|
ends.append(
|
|
segment.start_seconds
|
|
+ (cumulative / total) * (segment.end_seconds - segment.start_seconds)
|
|
)
|
|
return ends
|
|
|
|
|
|
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+)\]"
|
|
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
|