training_scripts/precompute_pseudo_labels.py

"""Stage 1: precompute (audio_embeds, teacher_token_ids) per clip on CPU.

For each audio clip in the training corpus:
  1. Load audio, run Voxtral audio_tower + multi_modal_projector to get audio_embeds.
  2. Run full Voxtral end-to-end on CPU (greedy) to get teacher's transcript
     token IDs. This is the "pseudo-label" for distillation.
  3. Also save the full input_ids sequence and the audio_token_id mask so
     downstream code can splice audio_embeds into the LLM input.
  4. Save each clip as `<clip>.pt`.

CPU-only. This is slow (Voxtral is 3B params on CPU) but only runs once per
clip and can be run offline in parallel while we develop the training loop.
Only rerun if the target model or prompt template changes.

Usage:
    python precompute_pseudo_labels.py \
        --model mistralai/Voxtral-Mini-3B-2507 \
        --manifest /mnt/data/librispeech_manifest.csv \
        --dataset-root / \
        --output-dir /mnt/data/pseudo_labels \
        --language en \
        --max-new-tokens 256

Output per clip: /mnt/data/pseudo_labels/<clip_name>.pt with keys:
    - input_ids: [1, prefix_len] LongTensor (Voxtral prompt + audio placeholder tokens)
    - audio_embeds: [n_audio_tokens, hidden_size] BFloat16Tensor
    - audio_token_id: int
    - teacher_token_ids: [n_generated] LongTensor (target's greedy transcript)
    - reference: str (ground-truth transcript, informational only)
    - duration_sec: float
    - metadata: {model_id, language, prompt_variant}
"""

from __future__ import annotations

import argparse
import csv
import gc
import json
import time
from pathlib import Path

import torch
from transformers import VoxtralForConditionalGeneration, AutoProcessor


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="mistralai/Voxtral-Mini-3B-2507")
    p.add_argument("--manifest", required=True, type=Path)
    p.add_argument("--dataset-root", required=True, type=Path)
    p.add_argument("--output-dir", required=True, type=Path)
    p.add_argument("--language", default="en")
    p.add_argument("--max-new-tokens", type=int, default=256)
    p.add_argument("--limit", type=int, default=None, help="Limit to first N clips")
    p.add_argument("--skip-existing", action="store_true", default=True)
    return p.parse_args()


def read_manifest(path: Path, root: Path, limit: int | None):
    rows = []
    with open(path) as f:
        r = csv.DictReader(f)
        for row in r:
            rel = row["audio_path"]
            rows.append({
                "audio_path": root / rel,
                "clip_name": Path(rel).stem,
                "duration_sec": float(row["duration_sec"]),
                "reference": row["transcript"],
            })
            if limit is not None and len(rows) >= limit:
                break
    return rows


def main() -> int:
    args = parse_args()
    torch.set_grad_enabled(False)

    print(f"[precompute] Loading {args.model} on CPU (bfloat16)")
    dtype = torch.bfloat16
    t0 = time.perf_counter()
    model = VoxtralForConditionalGeneration.from_pretrained(
        args.model, torch_dtype=dtype, low_cpu_mem_usage=True,
    ).eval()
    processor = AutoProcessor.from_pretrained(args.model)
    print(f"[precompute] Loaded in {time.perf_counter()-t0:.1f}s")

    audio_token_id = model.config.audio_token_id
    rows = read_manifest(args.manifest, args.dataset_root, args.limit)
    print(f"[precompute] {len(rows)} clips to process")

    args.output_dir.mkdir(parents=True, exist_ok=True)

    metadata = {
        "model_id": args.model,
        "language": args.language,
        "max_new_tokens": args.max_new_tokens,
        "audio_token_id": audio_token_id,
        "hidden_size": model.config.text_config.hidden_size,
        "vocab_size": model.config.text_config.vocab_size,
        "num_hidden_layers": model.config.text_config.num_hidden_layers,
    }
    with open(args.output_dir / "_metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    processed = 0
    for i, row in enumerate(rows):
        out_path = args.output_dir / f"{row['clip_name']}.pt"
        if args.skip_existing and out_path.exists():
            print(f"[precompute] {i+1}/{len(rows)}: {row['clip_name']}.pt exists, skipping")
            processed += 1
            continue

        print(f"[precompute] {i+1}/{len(rows)}: {row['clip_name']} ({row['duration_sec']:.2f}s)")
        t_clip = time.perf_counter()

        # -- Build transcription prompt
        inputs = processor.apply_transcription_request(
            language=args.language,
            audio=str(row["audio_path"]),
            model_id=args.model,
        )
        input_ids = inputs["input_ids"]  # [1, prefix_len]
        input_features = inputs["input_features"].to(dtype)  # [1, 128, 3000]

        # -- Compute audio embeds
        with torch.no_grad():
            audio_out = model.model.get_audio_features(input_features)
            audio_embeds = audio_out.pooler_output  # [N_audio_tokens, hidden]

        # -- Run greedy generation to get teacher_token_ids
        t_gen = time.perf_counter()
        with torch.no_grad():
            out = model.generate(
                input_ids=input_ids,
                input_features=input_features,
                max_new_tokens=args.max_new_tokens,
                do_sample=False,
                temperature=None,
                top_p=None,
                pad_token_id=processor.tokenizer.pad_token_id if processor.tokenizer.pad_token_id else processor.tokenizer.eos_token_id,
            )
        teacher_token_ids = out[0, input_ids.shape[1]:].clone()  # [n_generated]
        gen_t = time.perf_counter() - t_gen

        teacher_text = processor.tokenizer.decode(teacher_token_ids, skip_special_tokens=True).strip()
        print(f"[precompute]   audio_embeds={tuple(audio_embeds.shape)}, "
              f"prefix_len={input_ids.shape[1]}, "
              f"teacher_tokens={teacher_token_ids.shape[0]}, "
              f"gen_time={gen_t:.1f}s")
        print(f"[precompute]   teacher: {teacher_text[:100]!r}")

        # -- Save
        torch.save({
            "input_ids": input_ids.cpu(),  # [1, prefix_len]
            "audio_embeds": audio_embeds.cpu().to(dtype),  # [n_audio, hidden]
            "audio_token_id": audio_token_id,
            "teacher_token_ids": teacher_token_ids.cpu(),  # [n_generated]
            "teacher_text": teacher_text,
            "reference": row["reference"],
            "clip_name": row["clip_name"],
            "duration_sec": row["duration_sec"],
        }, out_path)
        processed += 1
        print(f"[precompute]   saved {out_path} (clip time: {time.perf_counter()-t_clip:.1f}s)")

    print(f"\n[precompute] Done. {processed}/{len(rows)} clips written to {args.output_dir}")
    return 0


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