scripts/gop_logit_bplus.py

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())