#!/usr/bin/env python3 """Local Czech audio transcription with speaker diarization (WhisperX).""" from __future__ import annotations import argparse import os import shutil import sys from pathlib import Path os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3") os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") def _patch_torch_load_weights_only() -> None: """PyTorch ≥2.6 defaults torch.load(weights_only=True), which rejects pyannote 4.x's checkpoints. Pyannote's load path doesn't expose a way to override this, and allowlisting individual globals is whack-a-mole (omegaconf, list, dict, defaultdict, ...). We monkey-patch torch.load to default weights_only=False, matching the behavior every pyannote install relied on prior to torch 2.6. Trust justification: weights come only from HF gated repos the user has explicitly accepted.""" import torch _orig_load = torch.load def _patched(*args, **kwargs): kwargs["weights_only"] = False return _orig_load(*args, **kwargs) torch.load = _patched _patch_torch_load_weights_only() HF_TOKEN_URL = "https://huggingface.co/settings/tokens" HF_GATED_URLS = ( "https://huggingface.co/pyannote/speaker-diarization-community-1", "https://huggingface.co/pyannote/segmentation-3.0", ) def eprint(*a, **kw): print(*a, file=sys.stderr, flush=True, **kw) def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser( prog="transcribe.py", description="Local Czech audio transcription with speaker diarization " "(WhisperX + pyannote community-1).", ) p.add_argument("input", type=Path, help="audio file (any format ffmpeg can read: " "mp3, wav, m4a, ogg, flac, mp4, …)") p.add_argument("-o", "--output", type=Path, help="output .txt path (default: .txt)") p.add_argument("--min-speakers", type=int, default=2) p.add_argument("--max-speakers", type=int, default=3) p.add_argument("--model", choices=["large-v2", "large-v3"], default="large-v2", help="Whisper model. Default large-v2 (less aggressive " "filtering, more accurate on Czech in practice). " "Switch to large-v3 for the newer model.") p.add_argument("--timestamps", action="store_true", help="prepend [HH:MM:SS] to each turn") p.add_argument("--hf-token", default=os.environ.get("HF_TOKEN"), help="Hugging Face read token (else $HF_TOKEN)") p.add_argument("--language", default="cs", help="ISO 639-1 language code (default cs)") args = p.parse_args() if args.min_speakers < 1 or args.max_speakers < args.min_speakers: p.error("require 1 <= --min-speakers <= --max-speakers") if not args.input.exists(): p.error(f"input file not found: {args.input}") if args.output is None: args.output = args.input.with_suffix(".txt") return args def require_token(token: str | None) -> str: if token: return token eprint("ERROR: missing Hugging Face token.") eprint(" Pass --hf-token TOKEN or set HF_TOKEN environment variable.") eprint(f" Create a read token at: {HF_TOKEN_URL}") eprint(" Then accept terms (while logged in) on:") for u in HF_GATED_URLS: eprint(f" - {u}") sys.exit(2) def require_ffmpeg() -> None: if shutil.which("ffmpeg") is None: eprint("ERROR: ffmpeg not found in PATH.") eprint(" Install: sudo apt install ffmpeg (or brew install ffmpeg)") sys.exit(3) def detect_device() -> tuple[str, str, int]: import torch if torch.cuda.is_available(): return "cuda", "float16", 16 eprint("WARNING: no CUDA GPU detected — running on CPU with int8.") eprint(" Expect 5–20× slower than realtime. Be patient.") return "cpu", "int8", 8 def fmt_hms(seconds: float | None) -> str: if seconds is None or seconds != seconds: # None or NaN return "00:00:00" s = max(0, int(seconds)) return f"{s // 3600:02d}:{(s % 3600) // 60:02d}:{s % 60:02d}" def group_turns(segments) -> list[dict]: """Collapse contiguous same-speaker segments and remap pyannote labels (SPEAKER_00, SPEAKER_01, …) to consecutive integers in first-seen order.""" label_map: dict[str, int] = {} next_idx = 1 turns: list[dict] = [] for seg in segments: text = (seg.get("text") or "").strip() if not text: continue spk_raw = seg.get("speaker") if spk_raw is None: label = "?" else: if spk_raw not in label_map: label_map[spk_raw] = next_idx next_idx += 1 label = str(label_map[spk_raw]) start = seg.get("start") end = seg.get("end", start) if turns and turns[-1]["label"] == label: turns[-1]["text"] += " " + text turns[-1]["end"] = end else: turns.append({"label": label, "text": text, "start": start, "end": end}) return turns def render(turns: list[dict], *, timestamps: bool) -> str: out: list[str] = [] for t in turns: head = f"[Speaker {t['label']}]" if timestamps: head = f"[{fmt_hms(t['start'])}] {head}" out.append(f"{head}: {t['text']}") return "\n\n".join(out) + ("\n" if out else "") def is_gated_error(exc: BaseException) -> bool: msg = str(exc).lower() return any(s in msg for s in ( "401", "403", "gated", "access to this resource", "you are not in the authorized list", "unauthorized", )) def transcribe(args: argparse.Namespace) -> None: require_ffmpeg() token = require_token(args.hf_token) device, compute_type, batch_size = detect_device() eprint(f"[1/5] loading audio: {args.input}") import whisperx audio = whisperx.load_audio(str(args.input)) eprint(f"[2/5] transcribing with whisper-{args.model} " f"({device}/{compute_type})…") asr = whisperx.load_model( args.model, device, compute_type=compute_type, language=args.language, threads=os.cpu_count() or 4, ) result = asr.transcribe(audio, batch_size=batch_size, language=args.language, print_progress=True) if not result.get("segments"): eprint("WARNING: no speech detected. Writing empty output.") args.output.write_text("", encoding="utf-8") return del asr # free memory before next stage eprint("[3/5] aligning words…") try: align_model, align_meta = whisperx.load_align_model( language_code=args.language, device=device, ) except Exception as e: eprint(f"ERROR: failed to load alignment model for " f"language='{args.language}': {e}") sys.exit(6) aligned = whisperx.align( result["segments"], align_model, align_meta, audio, device, return_char_alignments=False, print_progress=True, ) del align_model eprint(f"[4/5] diarizing speakers " f"(min={args.min_speakers}, max={args.max_speakers})…") from whisperx.diarize import DiarizationPipeline, assign_word_speakers try: diarizer = DiarizationPipeline(token=token, device=device) diarize_df = diarizer(audio, min_speakers=args.min_speakers, max_speakers=args.max_speakers) except Exception as e: if is_gated_error(e): eprint("ERROR: Hugging Face access denied. " "Accept terms (while logged in) on:") for u in HF_GATED_URLS: eprint(f" - {u}") eprint(f" Underlying error: {e}") sys.exit(5) eprint(f"ERROR: diarization failed: {e}") sys.exit(7) eprint("[5/5] assigning speakers and rendering…") final = assign_word_speakers(diarize_df, aligned) turns = group_turns(final.get("segments", [])) if not turns: eprint("WARNING: speakers could not be assigned. Writing empty output.") args.output.write_text("", encoding="utf-8") return text = render(turns, timestamps=args.timestamps) args.output.write_text(text, encoding="utf-8") n_speakers = len({t["label"] for t in turns if t["label"] != "?"}) eprint(f"OK: {len(turns)} turns, {n_speakers} speakers " f"→ {args.output}") def main() -> None: args = parse_args() try: transcribe(args) except KeyboardInterrupt: eprint("\ninterrupted") sys.exit(130) if __name__ == "__main__": main()