Files
pokemon_pinball_wiki/scripts/transcribe_sound_draft_labels.py
2026-07-07 00:25:38 -05:00

170 lines
6.6 KiB
Python

#!/usr/bin/env python3
"""Generate draft content labels for exported WAVs with local Whisper ASR.
These labels are not native/Radium symbols. They are content-derived draft
transcripts intended to make voice-like sounds easier to triage in the wiki.
"""
from __future__ import annotations
import argparse
import csv
import re
import sys
from pathlib import Path
def clean_text(value: str) -> str:
value = re.sub(r"\s+", " ", value).strip()
return value.strip(" \"'`")
def draft_slug(value: str, max_words: int = 8) -> str:
words = re.findall(r"[A-Za-z0-9]+", value.upper())
return "_".join(words[:max_words]) or "ASR_DRAFT"
def load_existing(path: Path) -> dict[str, dict[str, str]]:
if not path.exists():
return {}
with path.open(newline="") as f:
return {row["sound_name"]: row for row in csv.DictReader(f, delimiter="\t")}
def passes_duration(row: dict[str, str], min_duration: float, max_duration: float, channels: set[int]) -> bool:
duration = float(row["duration_seconds"])
channel_count = int(row["channels"])
return min_duration <= duration <= max_duration and channel_count in channels
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--manifest", type=Path, default=Path("analysis/radium-image/sound_export_manifest.tsv"))
parser.add_argument("--out", type=Path, default=Path("analysis/radium-image/asr_draft_sound_labels.tsv"))
parser.add_argument("--model", default="tiny.en")
parser.add_argument("--model-dir", type=Path, default=Path("/Users/jordan/.cache/whisper"))
parser.add_argument("--device", default="cpu")
parser.add_argument("--language", default="en")
parser.add_argument("--min-duration", type=float, default=0.8)
parser.add_argument("--max-duration", type=float, default=6.0)
parser.add_argument("--channels", default="1,2", help="Comma-separated channel counts to include.")
parser.add_argument("--limit", type=int, default=0)
parser.add_argument("--resume", action="store_true")
parser.add_argument("--accept-logprob", type=float, default=-0.85)
parser.add_argument("--accept-no-speech", type=float, default=0.35)
parser.add_argument("--accept-compression", type=float, default=2.0)
parser.add_argument(
"--temperature",
type=float,
default=0.0,
help="Whisper decode temperature. A single value disables the slow fallback temperature schedule.",
)
args = parser.parse_args()
try:
import whisper
except ImportError as exc:
raise SystemExit(
"Python module 'whisper' is not importable. Run with the Homebrew whisper Python, "
"for example /opt/homebrew/Cellar/openai-whisper/20250625_4/libexec/bin/python."
) from exc
channels = {int(part) for part in args.channels.split(",") if part.strip()}
manifest_rows = list(csv.DictReader(args.manifest.open(newline=""), delimiter="\t"))
candidates = [row for row in manifest_rows if passes_duration(row, args.min_duration, args.max_duration, channels)]
if args.limit:
candidates = candidates[: args.limit]
existing = load_existing(args.out) if args.resume else {}
output_rows = [existing[name] for name in sorted(existing)] if existing else []
done = set(existing)
args.out.parent.mkdir(parents=True, exist_ok=True)
fields = [
"sound_name",
"wav_path",
"section8_index",
"object_id",
"duration_seconds",
"channels",
"transcript",
"draft_label",
"accepted",
"avg_logprob",
"no_speech_prob",
"compression_ratio",
"segment_count",
"model",
]
model = whisper.load_model(args.model, download_root=str(args.model_dir), device=args.device)
with args.out.open("w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=fields, delimiter="\t")
writer.writeheader()
for row in output_rows:
writer.writerow(row)
f.flush()
processed = len(done)
for row in candidates:
if row["sound_name"] in done:
continue
result = model.transcribe(
row["wav_path"],
language=args.language,
fp16=False,
verbose=False,
condition_on_previous_text=False,
temperature=args.temperature,
)
segments = result.get("segments") or []
transcript = clean_text(str(result.get("text") or ""))
if segments:
total_duration = sum(max(0.01, float(seg["end"]) - float(seg["start"])) for seg in segments)
avg_logprob = sum(
float(seg.get("avg_logprob", -99.0)) * max(0.01, float(seg["end"]) - float(seg["start"]))
for seg in segments
) / total_duration
no_speech_prob = max(float(seg.get("no_speech_prob", 1.0)) for seg in segments)
compression_ratio = max(float(seg.get("compression_ratio", 99.0)) for seg in segments)
else:
avg_logprob = -99.0
no_speech_prob = 1.0
compression_ratio = 99.0
accepted = (
bool(transcript)
and avg_logprob >= args.accept_logprob
and no_speech_prob <= args.accept_no_speech
and compression_ratio <= args.accept_compression
and not re.fullmatch(r"[*\\W_]+", transcript)
)
out_row = {
"sound_name": row["sound_name"],
"wav_path": row["wav_path"],
"section8_index": row["section8_index"],
"object_id": row["primary_object_local_id"],
"duration_seconds": row["duration_seconds"],
"channels": row["channels"],
"transcript": transcript,
"draft_label": f"ASR_{draft_slug(transcript)}" if accepted else "",
"accepted": "yes" if accepted else "no",
"avg_logprob": f"{avg_logprob:.6f}",
"no_speech_prob": f"{no_speech_prob:.6f}",
"compression_ratio": f"{compression_ratio:.6f}",
"segment_count": len(segments),
"model": args.model,
}
writer.writerow(out_row)
f.flush()
processed += 1
if processed % 25 == 0:
print(f"processed {processed}/{len(candidates)}", file=sys.stderr)
print(f"wrote {args.out} rows={sum(1 for _ in args.out.open()) - 1} candidates={len(candidates)}")
return 0
if __name__ == "__main__":
raise SystemExit(main())