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 torch from scripts.overfit_aligned import PHONE_VOCAB, AlignedSesameAR, frame_phone_tensor, load_rows from speech_bridge_gemma.mimi_overfit import compute_accuracy, compute_loss, selected_decoder_layers from speech_bridge_gemma.qwen3_tts_tokenizer_smoke import ensure_audio_frame_token, synthesis_condition_text, tokenize_qwen3_batch def pad_batch(rows: list[dict[str, Any]], audio_slots: int, num_quantizers: int, device: str) -> tuple[torch.Tensor, torch.Tensor]: b = len(rows) targets = torch.full((b, audio_slots, num_quantizers), -100, dtype=torch.long) phones = torch.zeros((b, audio_slots), dtype=torch.long) for i, r in enumerate(rows): f = r["frames"] targets[i, :f] = r["codes"].transpose(0, 1)[:f] phones[i] = frame_phone_tensor(r["align"], audio_slots) return targets.to(device), phones.to(device) def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--codes-pt", required=True) parser.add_argument("--align-pt", required=True) parser.add_argument("--out-dir", required=True) parser.add_argument("--llm-model", default="Qwen/Qwen3.5-0.8B") parser.add_argument("--max-rows", type=int, default=3600) parser.add_argument("--heldout", type=int, default=200) parser.add_argument("--steps", type=int, default=4000) parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--lr", type=float, default=4e-4) parser.add_argument("--backbone-lr", type=float, default=2e-5) parser.add_argument("--train-last-n-layers", type=int, default=8) parser.add_argument("--num-quantizers", type=int, default=16) parser.add_argument("--codebook-size", type=int, default=2048) parser.add_argument("--speech-conditioning", default="prompt_answer") parser.add_argument("--save-steps", type=int, default=1000) parser.add_argument("--device", default="cuda") args = parser.parse_args() import torch.nn as nn from transformers import AutoModelForCausalLM, AutoTokenizer out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) rows, max_frames = load_rows(args.codes_pt, args.align_pt, args.max_rows) audio_slots = max_frames + 2 held = rows[-args.heldout:] if args.heldout else [] train = rows[: len(rows) - len(held)] print(json.dumps({"event": "data", "train": len(train), "held": len(held), "audio_slots": audio_slots}), flush=True) tokenizer = AutoTokenizer.from_pretrained(args.llm_model, trust_remote_code=True) llm = AutoModelForCausalLM.from_pretrained(args.llm_model, torch_dtype=torch.bfloat16, trust_remote_code=True, low_cpu_mem_usage=True).to(args.device) llm.config.use_cache = False frame_token_id = ensure_audio_frame_token(tokenizer, llm) for p in llm.parameters(): p.requires_grad_(False) backbone_params: list[nn.Parameter] = [] if args.train_last_n_layers: for layer in selected_decoder_layers(llm, args.train_last_n_layers): for p in layer.parameters(): p.requires_grad_(True) backbone_params.append(p) model = AlignedSesameAR(llm=llm, frame_token_id=frame_token_id, audio_slots=audio_slots, num_quantizers=args.num_quantizers, codebook_size=args.codebook_size, depth_dim=512, depth_layers=4, depth_heads=8, depth_ff_mult=4).to(args.device) model.train() groups = [{"params": model.audio_parameters(), "lr": args.lr}] if backbone_params: groups.append({"params": backbone_params, "lr": args.backbone_lr}) optimizer = torch.optim.AdamW(groups) order = list(range(len(train))) step = 0 running = 0.0 cursor = 0 while step < args.steps: step += 1 if cursor + args.batch_size > len(order): cursor = 0 order = order[::-1] batch_rows = [train[order[cursor + k]] for k in range(args.batch_size)] cursor += args.batch_size conds = [synthesis_condition_text(r["question"], r["answer"], args.speech_conditioning) for r in batch_rows] tok = tokenize_qwen3_batch(tokenizer, conds, audio_slots, args.device, True) targets, phones = pad_batch(batch_rows, audio_slots, args.num_quantizers, args.device) model.set_align(phones) logits = model(tok["input_ids"], tok["attention_mask"], targets=targets) loss = compute_loss(logits=logits, targets=targets, codebook_size=args.codebook_size, first_codebook_weight=1.5, depth_weight=1.0) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_([p for g in groups for p in g["params"]], 1.0) optimizer.step() running += float(loss.detach().cpu()) if step % 100 == 0 or step == 1: with torch.no_grad(): acc, per_q = compute_accuracy(logits.detach(), targets) print(json.dumps({"step": step, "loss": round(running / (100 if step % 100 == 0 else 1), 4), "acc": round(acc, 4), "cb0": round(per_q[0], 4)}), flush=True) running = 0.0 if args.save_steps and step % args.save_steps == 0: torch.save({"audio_head": model.audio_head.state_dict(), "phone_embed": model.phone_embed.state_dict()}, out_dir / "aligned_head.pt") torch.save({"audio_head": model.audio_head.state_dict(), "phone_embed": model.phone_embed.state_dict()}, out_dir / "aligned_head.pt") model.eval() pred_dir = out_dir / "heldout_pred" pred_dir.mkdir(exist_ok=True) for r in held: phones = frame_phone_tensor(r["align"], audio_slots).unsqueeze(0).to(args.device) cond = synthesis_condition_text(r["question"], r["answer"], args.speech_conditioning) tok = tokenize_qwen3_batch(tokenizer, [cond], audio_slots, args.device, True) model.set_align(phones) with torch.inference_mode(): logits = model(tok["input_ids"], tok["attention_mask"], frames=r["frames"]) codes = logits[0, : r["frames"]].argmax(dim=-1).transpose(0, 1).contiguous().cpu() torch.save(codes, pred_dir / f"{r['id']}_codes.pt") (pred_dir / f"{r['id']}.json").write_text(json.dumps({"id": r["id"], "expected": r["answer"]}, ensure_ascii=False), encoding="utf-8") print(json.dumps({"event": "heldout_predicted", "n": len(held), "dir": str(pred_dir)}), flush=True) return 0 if __name__ == "__main__": raise SystemExit(main())