be9f14ce34
A small helpdesk system: an office Pixel running a patched GrapheneOS Dialer answers technician calls, records both call legs as separate channels, and a Ruby backend transcribes them through Whisper and files an AI summary against the caller. Squashed to a single commit for sharing. No credentials are included; secrets live outside the repo in /etc/helpdesk/env on the server or a gitignored .claude/env.local locally. See .claude/env.local.example for the shape. Start at README.md, then docs/architecture.md.
244 lines
8.6 KiB
Python
Executable file
244 lines
8.6 KiB
Python
Executable file
#!/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: <input-stem>.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()
|