119 lines
3.3 KiB
Python
119 lines
3.3 KiB
Python
import argparse
|
|
import json
|
|
import re
|
|
from dataclasses import dataclass
|
|
from itertools import groupby
|
|
from pathlib import Path
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Turn:
|
|
start: float
|
|
end: float
|
|
speaker: str
|
|
|
|
|
|
@dataclass
|
|
class Token:
|
|
start: float
|
|
end: float
|
|
text: str
|
|
|
|
|
|
def load_tokens(path: Path) -> list[Token]:
|
|
with path.open(encoding="utf-8") as file:
|
|
transcript = json.load(file)
|
|
|
|
tokens: list[Token] = []
|
|
pending = ""
|
|
# Whisper reserves all IDs from end-of-text onward for control tokens.
|
|
end_of_text = 50256 + int(transcript["model"]["multilingual"])
|
|
|
|
for segment in transcript["transcription"]:
|
|
for raw_token in segment["tokens"]:
|
|
if raw_token["id"] >= end_of_text:
|
|
continue
|
|
text = raw_token["text"]
|
|
offsets = raw_token.get("offsets")
|
|
|
|
if not text:
|
|
continue
|
|
if offsets is None:
|
|
if tokens:
|
|
tokens[-1].text += text
|
|
else:
|
|
pending += text
|
|
continue
|
|
|
|
tokens.append(
|
|
Token(offsets["from"] / 1000, offsets["to"] / 1000, pending + text)
|
|
)
|
|
pending = ""
|
|
|
|
if not tokens:
|
|
raise ValueError("Whisper produced no timed text")
|
|
|
|
return tokens
|
|
|
|
|
|
def load_turns(audio: Path, model: Path) -> list[Turn]:
|
|
from pyannote.audio import Pipeline
|
|
|
|
pipeline = Pipeline.from_pretrained(model)
|
|
result = pipeline(str(audio))
|
|
turns = [
|
|
Turn(segment.start, segment.end, speaker)
|
|
for segment, _, speaker in result.exclusive_speaker_diarization.itertracks(
|
|
yield_label=True
|
|
)
|
|
]
|
|
if not turns:
|
|
raise RuntimeError("Pyannote detected no speech")
|
|
|
|
return turns
|
|
|
|
|
|
def token_speaker(token: Token, turns: list[Turn]) -> str:
|
|
midpoint = (token.start + token.end) / 2
|
|
|
|
def rank(turn: Turn) -> tuple[float, float]:
|
|
overlap = max(0.0, min(token.end, turn.end) - max(token.start, turn.start))
|
|
distance = max(turn.start - midpoint, midpoint - turn.end, 0.0)
|
|
return overlap, -distance
|
|
|
|
return max(turns, key=rank).speaker
|
|
|
|
|
|
def render(tokens: list[Token], turns: list[Turn]) -> str:
|
|
speaker_numbers: dict[str, int] = {}
|
|
rendered = []
|
|
for speaker, group in groupby(
|
|
tokens, key=lambda token: token_speaker(token, turns)
|
|
):
|
|
number = speaker_numbers.setdefault(speaker, len(speaker_numbers) + 1)
|
|
text = "".join(token.text for token in group)
|
|
text = re.sub(r"\s+", " ", text).strip()
|
|
if text:
|
|
rendered.append(f"**Speaker {number}:**\n\n{text}")
|
|
|
|
if not rendered:
|
|
raise ValueError("Whisper produced no transcript text")
|
|
return "\n\n".join(rendered) + "\n"
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--audio", required=True, type=Path)
|
|
parser.add_argument("--model", required=True, type=Path)
|
|
parser.add_argument("--transcript", required=True, type=Path)
|
|
parser.add_argument("--output", required=True, type=Path)
|
|
args = parser.parse_args()
|
|
|
|
tokens = load_tokens(args.transcript)
|
|
turns = load_turns(args.audio, args.model)
|
|
args.output.write_text(render(tokens, turns), encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|