scripts/build_frame_alignment.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import librosa
import numpy as np
import torch

from speech_bridge_gemma.ctc_gop import token_id_for_phone
from speech_bridge_gemma.qwen3_tts_tokenizer_smoke import decode_qwen3_codes, load_qwen3_codec, qwen3_codes_to_qt

PHONE_RECOGNIZER = "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
CODEC_HZ = 12.5


def forced_ctc_align(log_probs: torch.Tensor, target_ids: torch.Tensor, blank_id: int) -> torch.Tensor:
    device = log_probs.device
    target = [int(v) for v in target_ids.tolist()]
    frames = int(log_probs.shape[0])
    if frames <= 0 or not target:
        return torch.zeros((max(0, frames),), dtype=torch.long)
    expanded = [blank_id]
    for token in target:
        expanded.append(int(token))
        expanded.append(blank_id)
    ext = torch.tensor(expanded, device=device, dtype=torch.long)
    states = int(ext.shape[0])
    emit = log_probs[:, ext]
    neg = torch.tensor(-1.0e30, device=device)
    skip_allowed = torch.zeros(states, dtype=torch.bool, device=device)
    if states > 2:
        skip_allowed[2:] = (ext[2:] != blank_id) & (ext[2:] != ext[:-2])
    dp = torch.full((states,), -1.0e30, device=device)
    dp[0] = emit[0, 0]
    if states > 1:
        dp[1] = emit[0, 1]
    back = torch.zeros((frames, states), dtype=torch.long, device=device)
    idx_states = torch.arange(states, device=device)
    for frame in range(1, frames):
        step = torch.cat([neg.view(1), dp[:-1]])
        skip = torch.cat([neg.view(1), neg.view(1), dp[:-2]]) if states > 2 else torch.full((states,), -1.0e30, device=device)
        skip = torch.where(skip_allowed, skip, neg)
        best_val, best_choice = torch.stack([dp, step, skip], dim=0).max(dim=0)
        dp = best_val + emit[frame]
        back[frame] = idx_states - best_choice
    last_state = states - 1 if states == 1 else (states - 1 if float(dp[states - 1]) >= float(dp[states - 2]) else states - 2)
    back_cpu = back.cpu()
    ext_cpu = ext.cpu()
    path = [last_state]
    state = last_state
    for frame in range(frames - 1, 0, -1):
        state = int(back_cpu[frame, state])
        path.append(state)
    path.reverse()
    aligned = torch.full((frames,), -1, dtype=torch.long)
    for frame, idx in enumerate(path):
        token = int(ext_cpu[idx])
        if token != blank_id:
            aligned[frame] = token
    known = torch.nonzero(aligned >= 0).flatten()
    if known.numel() == 0:
        return torch.full((frames,), int(target[0]), dtype=torch.long)
    first = int(known[0])
    aligned[:first] = aligned[first]
    last = int(aligned[first])
    for frame in range(first, frames):
        if int(aligned[frame]) >= 0:
            last = int(aligned[frame])
        else:
            aligned[frame] = last
    return aligned.long()


def resample_to_codec(aligned_ctc: torch.Tensor, codec_frames: int) -> torch.Tensor:
    n = int(aligned_ctc.shape[0])
    if n == 0 or codec_frames <= 0:
        return torch.zeros((max(0, codec_frames),), dtype=torch.long)
    idx = (torch.arange(codec_frames).float() * (n / codec_frames)).long().clamp(max=n - 1)
    return aligned_ctc[idx]


def run_lengths(ids: torch.Tensor, id_to_phone: dict[int, str]) -> list[tuple[str, int]]:
    out: list[tuple[str, int]] = []
    for v in ids.tolist():
        ph = id_to_phone.get(int(v), str(v))
        if out and out[-1][0] == ph:
            out[-1] = (ph, out[-1][1] + 1)
        else:
            out.append((ph, 1))
    return out


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--codes-pt", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--max-rows", type=int, default=12)
    parser.add_argument("--espeak-lang", default="pt-br")
    parser.add_argument("--device", default="cuda")
    args = parser.parse_args()

    from phonemizer.backend import EspeakBackend
    from phonemizer.separator import Separator
    from transformers import AutoModelForCTC, Wav2Vec2FeatureExtractor, Wav2Vec2PhonemeCTCTokenizer

    extractor = Wav2Vec2FeatureExtractor.from_pretrained(PHONE_RECOGNIZER)
    tokenizer = Wav2Vec2PhonemeCTCTokenizer.from_pretrained(PHONE_RECOGNIZER)
    model = AutoModelForCTC.from_pretrained(PHONE_RECOGNIZER).to(args.device).eval()
    vocab = tokenizer.get_vocab()
    id_to_phone = {int(i): p for p, i in vocab.items()}
    blank_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
    g2p = EspeakBackend(args.espeak_lang, preserve_punctuation=False, with_stress=False)
    sep = Separator(phone=" ", word="", syllable="")
    codec = load_qwen3_codec("Qwen/Qwen3-TTS-Tokenizer-12Hz", args.device)

    blob = torch.load(args.codes_pt, map_location="cpu")
    samples, codes = blob["samples"], blob["codes"]
    alignments: dict[str, torch.Tensor] = {}
    rows_out = []
    for i in range(min(args.max_rows, len(samples))):
        meta = samples[i]
        sid = str(meta.get("id") or i)
        answer = str(meta.get("answer") or "")
        qt = qwen3_codes_to_qt(codes[i])
        codec_frames = int(qt.shape[1])
        wav, sr = decode_qwen3_codes(codec, qt)
        wav16 = librosa.resample(np.asarray(wav, dtype="float32"), orig_sr=sr, target_sr=16000) if sr != 16000 else np.asarray(wav, dtype="float32")
        values = extractor(wav16, sampling_rate=16000, return_tensors="pt").input_values.to(args.device)
        with torch.inference_mode():
            log_probs = torch.log_softmax(model(values).logits[0].float(), dim=-1).cpu()
        phones = [p for p in g2p.phonemize([answer], separator=sep, strip=True)[0].split() if p]
        target_ids = [token_id_for_phone(p, vocab, None) for p in phones]
        target_ids = [t for t in target_ids if t is not None]
        if not target_ids:
            continue
        aligned_ctc = forced_ctc_align(log_probs, torch.tensor(target_ids), blank_id)
        aligned_codec = resample_to_codec(aligned_ctc, codec_frames)
        alignments[sid] = aligned_codec
        rows_out.append({"id": sid, "answer": answer, "codec_frames": codec_frames, "ctc_frames": int(aligned_ctc.shape[0]), "n_target_phones": len(target_ids)})
        if i == 0:
            print("SANITY", sid, "| frames", codec_frames, "| runs:", run_lengths(aligned_codec, id_to_phone)[:20], flush=True)

    torch.save({"version": 1, "alignments": alignments, "id_to_phone": id_to_phone}, args.out)
    print(json.dumps({"aligned": len(alignments), "rows": rows_out[:3]}, ensure_ascii=False), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())