Files
mediaplayer/sentence_api/transcription.py
2026-09-12 21:27:51 +08:00

648 lines
22 KiB
Python

import io
import mimetypes
import re
import wave
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Protocol, Tuple
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()
plain_text = _strip_speaker_tags(plain_text) or _strip_speaker_tags(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 = _strip_speaker_tags(
" ".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
_PUNCTUATION_BOUNDARY = re.compile(r"[.?]+|,+")
MIN_WORDS_FOR_COMMA_SPLIT = 50
_ABBREVIATION_RE = re.compile(
r"(?i)\b(?:a\.m\.|p\.m\.|mr\.|mrs\.|ms\.|dr\.|prof\.|st\.|vs\.|etc\.|"
r"e\.g\.|i\.e\.|no\.|approx\.|fig\.|inc\.|ltd\.|jr\.|sr\.|u\.s\.|u\.k\.|"
r"ph\.d\.|m\.d\.|b\.a\.|m\.a\.|d\.c\.|p\.s\.)"
r"(?![A-Za-z0-9])"
)
_TIME_MERIDIEM_RE = re.compile(
r"(?i)\b(?:[01]?\d|2[0-3])(?::[0-5]\d)?\s*(?:a\.m\.|p\.m\.|am\.|pm\.)"
r"(?![A-Za-z0-9])"
)
_TIME_CONTINUATION_RE = re.compile(
r"(?i)^(eastern|western|central|pacific|mountain|time|"
r"et|est|edt|pt|pst|pdt|ct|cst|cdt|mt|mst|mdt|gmt|utc|local)\b"
)
_SPEAKER_TAG_RE = re.compile(r"\[(S\d+)\]")
def _strip_speaker_tags(text: str) -> str:
"""Remove inline ``[S01]`` speaker tags from a text string."""
return _SPEAKER_TAG_RE.sub("", text).strip()
def _expand_speaker_segments(
segment: TranscriptionSegment,
) -> List[TranscriptionSegment]:
"""Split a segment at inline ``[S01]`` speaker tags and strip the tags.
MOSS sometimes embeds speaker tags directly in the text (rather than in
the ``speaker`` JSON field). Each tag marks the start of a new speaker
turn, so the text is split there and the tags are removed from the
visible transcription.
"""
text = segment.text
matches = list(_SPEAKER_TAG_RE.finditer(text))
if not matches:
return [segment]
sub_segments: List[TranscriptionSegment] = []
total = len(matches)
duration = segment.end_seconds - segment.start_seconds
for index, match in enumerate(matches):
start_pos = match.end()
if index + 1 < total:
end_pos = matches[index + 1].start()
else:
end_pos = len(text)
piece = text[start_pos:end_pos].strip()
if not piece:
continue
# Allocate time proportionally across sub-segments.
frac_start = index / total
frac_end = (index + 1) / total
sub_start = segment.start_seconds + duration * frac_start
sub_end = segment.start_seconds + duration * frac_end
words = None
if segment.words:
sub_words = [
w
for w in segment.words
if w.start_seconds >= sub_start and w.end_seconds <= sub_end
]
words = sub_words or None
sub_segments.append(
TranscriptionSegment(
start_seconds=sub_start,
end_seconds=sub_end,
text=piece,
speaker=match.group(1),
words=words,
)
)
return sub_segments or [segment]
def split_sentences_at_punctuation(
segments: List[TranscriptionSegment],
) -> List[TranscriptionSegment]:
"""Merge whisper segments and cut sentences at sentence punctuation.
Whisper's own segment breaks often fall mid-sentence, so text is accumulated
across segments and a sentence is closed at periods/question marks. Commas
are also eligible boundaries, but only when the text before the comma has
enough words; otherwise the comma stays with the following text.
Inline ``[S01]`` speaker tags are stripped. A new sentence starts when the
speaker actually changes; same-speaker line wraps stay merged until
sentence punctuation. Decimal points, initials, and common abbreviations
are not treated as sentence ends.
"""
sentences: List[TranscriptionSegment] = []
pieces: List[str] = []
start_seconds: Optional[float] = None
last_segment_end: Optional[float] = None
current_speaker: Optional[str] = None
def flush(end_seconds: float) -> None:
nonlocal pieces, start_seconds, current_speaker
text = " ".join(piece for piece in pieces if piece).strip()
if text and start_seconds is not None and end_seconds > start_seconds:
sentences.append(
TranscriptionSegment(
start_seconds=start_seconds,
end_seconds=end_seconds,
text=text,
speaker=current_speaker,
)
)
pieces = []
start_seconds = None
current_speaker = None
expanded: List[TranscriptionSegment] = []
for seg in segments:
expanded.extend(_expand_speaker_segments(seg))
for segment in sorted(
expanded, key=lambda item: (item.start_seconds, item.end_seconds)
):
text = segment.text.strip()
if not text:
continue
# MOSS often tags every phrase/line with [S01] even when the speaker
# did not change. Only a real speaker change starts a new sentence.
if (
pieces
and segment.speaker is not None
and current_speaker is not None
and segment.speaker != current_speaker
):
flush(segment.start_seconds)
if not pieces:
start_seconds = segment.start_seconds
current_speaker = segment.speaker
seg_pieces: List[Tuple[str, Optional[str]]] = []
cursor = 0
for match in _PUNCTUATION_BOUNDARY.finditer(text):
boundary_type: Optional[str] = None
if "," in match.group(0):
boundary_type = "comma"
elif _is_sentence_end_run(text, match.start(), match.end()):
boundary_type = "sentence"
if boundary_type is None:
continue
piece = text[cursor:match.end()].strip()
cursor = match.end()
if _has_word_token(piece):
seg_pieces.append((piece, boundary_type))
tail = text[cursor:].strip()
if tail:
seg_pieces.append((tail, False))
if not seg_pieces:
continue
lengths = [len(piece) for piece, _ in seg_pieces]
ends: List[float] = []
if segment.words and len(segment.words) >= len(seg_pieces):
ends = _word_boundary_ends(segment, lengths)
if not _valid_boundaries(segment, ends):
ends = []
if not ends:
ends = _proportional_ends(segment, lengths)
last_segment_end = segment.end_seconds
for index, (piece, boundary_type) in enumerate(seg_pieces):
pieces.append(piece)
if boundary_type == "sentence" or (
boundary_type == "comma"
and _word_count(" ".join(pieces)) >= MIN_WORDS_FOR_COMMA_SPLIT
):
end_seconds = ends[index] if index < len(ends) else segment.end_seconds
flush(end_seconds)
start_seconds = end_seconds
current_speaker = segment.speaker
if pieces:
flush(last_segment_end if last_segment_end is not None else start_seconds)
return sentences
def _word_count(text: str) -> int:
return len(text.split())
def _has_word_token(text: str) -> bool:
return any(char.isalnum() for char in text)
def _is_sentence_end_run(text: str, run_start: int, run_end: int) -> bool:
"""Decide whether a "." / "?" run really ends a sentence."""
punctuation = text[run_start:run_end]
if "?" in punctuation:
return True
if _is_decimal_point(text, run_start, run_end):
return False
if len(punctuation) > 1:
return True # ellipsis or "!?" etc.
if _is_abbreviation_run(text, run_start, run_end):
return False # e.g. "a.m.", "p.m.", "Mr."
if _is_time_meridiem_run(text, run_start, run_end):
return False # e.g. "3:00 PM.", "5 pm."
if _is_single_letter_initial(text, run_start, run_end):
return False # e.g. "Donald J. Trump"
rest = text[run_end:].lstrip()
if rest and rest[0].islower():
return False # period followed by a lowercase word is an abbreviation
return True
def _is_decimal_point(text: str, run_start: int, run_end: int) -> bool:
if run_end - run_start != 1 or text[run_start] != ".":
return False
if run_start == 0 or run_end >= len(text):
return False
return text[run_start - 1].isdigit() and text[run_end].isdigit()
def _is_single_letter_initial(text: str, run_start: int, run_end: int) -> bool:
if run_end - run_start != 1 or text[run_start] != ".":
return False
if run_start == 0 or not text[run_start - 1].isalpha():
return False
return run_start == 1 or not text[run_start - 2].isalpha()
def _meridiem_continues(text: str, after: int) -> bool:
rest = text[after:].lstrip()
if not rest:
return True
if rest[0].islower():
return True
return _TIME_CONTINUATION_RE.match(rest) is not None
def _is_time_meridiem_run(text: str, run_start: int, run_end: int) -> bool:
for match in _TIME_MERIDIEM_RE.finditer(text):
if not (match.start() < run_end <= match.end()):
continue
if run_end < match.end():
return True
return _meridiem_continues(text, match.end())
return False
def _is_abbreviation_run(text: str, run_start: int, run_end: int) -> bool:
for match in _ABBREVIATION_RE.finditer(text):
if not (match.start() < run_end <= match.end()):
continue
token = match.group(0).lower()
if token == "no.":
rest = text[match.end():].lstrip()
if not rest or not rest[0].isdigit():
return False # "No." as a reply ends the sentence
if token in {"a.m.", "p.m."}:
if run_end < match.end():
return True
return _meridiem_continues(text, match.end())
return True
return False
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