from __future__ import annotations import argparse import json import random import sys from pathlib import Path from typing import Any sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import librosa import soundfile as sf import torch from speech_bridge_gemma.ctc_gop import ( calibrate_gop_thresholds, ctc_gop_phone_scores, gop_word_events, ) PHONE_RECOGNIZER = "facebook/wav2vec2-xlsr-53-espeak-cv-ft" def load_audio_16k(path: str) -> Any: wav, sr = sf.read(path) if getattr(wav, "ndim", 1) > 1: wav = wav.mean(axis=1) wav = wav.astype("float32") if sr != 16000: wav = librosa.resample(wav, orig_sr=sr, target_sr=16000) return wav class PhoneGop: def __init__(self, device: str): from transformers import AutoModelForCTC, Wav2Vec2FeatureExtractor, Wav2Vec2PhonemeCTCTokenizer self.device = device self.extractor = Wav2Vec2FeatureExtractor.from_pretrained(PHONE_RECOGNIZER) self.tokenizer = Wav2Vec2PhonemeCTCTokenizer.from_pretrained(PHONE_RECOGNIZER) self.model = AutoModelForCTC.from_pretrained(PHONE_RECOGNIZER).to(device).eval() self.vocab = dict(self.tokenizer.get_vocab()) self.blank_id = int(self.tokenizer.pad_token_id) self.candidates = [p for p, i in self.vocab.items() if int(i) != self.blank_id and not p.startswith("<") and p not in ("|",) and not p.isdigit()] def log_probs(self, audio_path: str) -> torch.Tensor: wav = load_audio_16k(audio_path) values = self.extractor(wav, sampling_rate=16000, return_tensors="pt").input_values.to(self.device) with torch.inference_mode(): logits = self.model(values).logits[0] return logits.log_softmax(dim=-1).detach().cpu() def realized(self, audio_path: str) -> list[str]: wav = load_audio_16k(audio_path) values = self.extractor(wav, sampling_rate=16000, return_tensors="pt").input_values.to(self.device) with torch.inference_mode(): ids = self.model(values).logits.argmax(dim=-1) return [p for p in self.tokenizer.batch_decode(ids)[0].split() if p] def score(self, log_probs: torch.Tensor, expected: list[str], normalization: str | None = None) -> list[dict[str, Any]]: words = [str(index) for index in range(len(expected))] expected_word_phonemes = [[phone] for phone in expected] return ctc_gop_phone_scores(log_probs, words, expected_word_phonemes, self.vocab, normalization, self.blank_id) def perturb(expected: list[str], candidates: list[str], rate: float, seed: int) -> tuple[list[str], set[int]]: rng = random.Random(seed) out = list(expected) errors: set[int] = set() for index, phone in enumerate(expected): if rng.random() < rate: alt = rng.choice(candidates) tries = 0 while alt == phone and tries < 5: alt = rng.choice(candidates) tries += 1 out[index] = alt errors.add(index) return out, errors def read_jsonl(path: str) -> list[dict[str, Any]]: rows = [] with Path(path).open(encoding="utf-8") as handle: for line in handle: line = line.strip() if line: rows.append(json.loads(line)) return rows def calibrate(args: argparse.Namespace) -> int: gop = PhoneGop(args.device) manifest = read_jsonl(args.manifest) if args.max_rows: manifest = manifest[: args.max_rows] scores_by_id: dict[str, list[dict[str, Any]]] = {} rows: list[dict[str, Any]] = [] for index, item in enumerate(manifest): audio = str(item["audio"]) expected = gop.realized(audio) if len(expected) < 3: continue log_probs = gop.log_probs(audio) perturbed, error_positions = perturb(expected, gop.candidates, args.perturb_rate, args.seed + index) scores = gop.score(log_probs, perturbed, args.normalization) item_id = str(item.get("id") or Path(audio).stem) scores_by_id[item_id] = scores expected_errors = [{"word": str(pos), "expected": perturbed[pos]} for pos in sorted(error_positions)] rows.append({"id": item_id, "expected_errors": expected_errors}) print(json.dumps({"id": item_id, "phones": len(expected), "perturbed": len(error_positions)}, ensure_ascii=False), flush=True) thresholds = calibrate_gop_thresholds(scores_by_id, rows, args.normalization) Path(args.out).write_text(json.dumps(thresholds, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") summary = {k: thresholds[k] for k in ("positive_scores", "negative_scores")} summary["default"] = thresholds["default"] summary["per_phone_count"] = len(thresholds["per_phone"]) print(json.dumps(summary, ensure_ascii=False), flush=True) return 0 def assess(args: argparse.Namespace) -> int: gop = PhoneGop(args.device) thresholds = json.loads(Path(args.thresholds).read_text(encoding="utf-8")) expected = gop.realized(args.reference_audio) if args.reference_audio else None if expected is None: raise SystemExit("provide --reference-audio for B+ assess") log_probs = gop.log_probs(args.audio) scores = gop.score(log_probs, expected, args.normalization) words = [str(index) for index in range(len(expected))] expected_word_phonemes = [[phone] for phone in expected] events = gop_word_events(words, expected_word_phonemes, scores, thresholds, args.normalization) flagged = [{"phone": s["expected"], "heard": s["alternative"], "gop": round(s["gop"], 3)} for s in scores if s["gop"] <= float((thresholds["per_phone"].get(s["expected"]) or thresholds["default"])["threshold"])] result = { "audio": args.audio, "reference_audio": args.reference_audio, "expected": expected, "flagged": flagged, "event_count": len(events), } print(json.dumps(result, ensure_ascii=False, indent=2)) if args.out: Path(args.out).write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") return 0 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() sub = parser.add_subparsers(dest="cmd", required=True) cal = sub.add_parser("calibrate") cal.add_argument("--manifest", required=True) cal.add_argument("--out", default="job_output/omni-train-130/gop_thresholds.json") cal.add_argument("--perturb-rate", type=float, default=0.25) cal.add_argument("--normalization", default=None) cal.add_argument("--max-rows", type=int) cal.add_argument("--seed", type=int, default=1337) cal.add_argument("--device", default="cuda") cal.set_defaults(func=calibrate) asn = sub.add_parser("assess") asn.add_argument("--audio", required=True) asn.add_argument("--reference-audio", required=True) asn.add_argument("--thresholds", default="job_output/omni-train-130/gop_thresholds.json") asn.add_argument("--out") asn.add_argument("--normalization", default=None) asn.add_argument("--device", default="cuda") asn.set_defaults(func=assess) return parser.parse_args() def main() -> int: args = parse_args() return args.func(args) if __name__ == "__main__": raise SystemExit(main())