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