scripts/mvp_aligned.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 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())