File size: 5,110 Bytes
f9e3832 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """Validate a trained draft checkpoint: does it produce coherent output?
Runs the draft standalone (greedy generation) on a single audio clip and compares
to the target's output. Also runs one prefill of both on the same audio to
measure top-1 argmax agreement (a proxy for acceptance rate).
Everything runs on CPU (audio encoder can't run on trn2, and for validation
CPU is fine).
"""
from __future__ import annotations
import argparse
import time
import torch
from pathlib import Path
from transformers import VoxtralForConditionalGeneration, AutoProcessor
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--target", default="mistralai/Voxtral-Mini-3B-2507")
p.add_argument("--draft", required=True, help="Path to trained draft checkpoint dir")
p.add_argument("--audio", default="/mnt/data/LibriSpeech/dev-clean/2035/147960/2035-147960-0015.flac")
p.add_argument("--max-new-tokens", type=int, default=64)
p.add_argument("--language", default="en")
return p.parse_args()
def main() -> int:
args = parse_args()
torch.set_grad_enabled(False)
dtype = torch.bfloat16
print(f"[validate] Loading target {args.target}")
target = VoxtralForConditionalGeneration.from_pretrained(
args.target, torch_dtype=dtype, low_cpu_mem_usage=True,
).eval()
proc = AutoProcessor.from_pretrained(args.target)
print(f"[validate] Loading draft {args.draft}")
draft = VoxtralForConditionalGeneration.from_pretrained(
args.draft, torch_dtype=dtype, low_cpu_mem_usage=True,
).eval()
print(f"[validate] Target layers: {target.config.text_config.num_hidden_layers}")
print(f"[validate] Draft layers: {draft.config.text_config.num_hidden_layers}")
inputs = proc.apply_transcription_request(
language=args.language, audio=args.audio, model_id=args.target,
)
# -- Target greedy
print("\n[validate] Target greedy generation")
t0 = time.perf_counter()
out_t = target.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False,
temperature=None, top_p=None)
gen_t = out_t[0, inputs["input_ids"].shape[1]:]
target_text = proc.tokenizer.decode(gen_t, skip_special_tokens=True).strip()
print(f"[validate] Target ({time.perf_counter()-t0:.1f}s, {gen_t.shape[0]} tok):")
print(f" {target_text!r}")
# -- Draft greedy (standalone)
print("\n[validate] Draft greedy generation (standalone)")
t0 = time.perf_counter()
out_d = draft.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False,
temperature=None, top_p=None)
gen_d = out_d[0, inputs["input_ids"].shape[1]:]
draft_text = proc.tokenizer.decode(gen_d, skip_special_tokens=True).strip()
print(f"[validate] Draft ({time.perf_counter()-t0:.1f}s, {gen_d.shape[0]} tok):")
print(f" {draft_text!r}")
print(f"[validate] First 20 raw tokens: {gen_d[:20].tolist()}")
# -- Top-1 agreement: how often does draft's argmax match target's argmax
# given TEACHER-FORCED context = target's greedy sequence?
print("\n[validate] Top-1 agreement (teacher-forced)")
# Build full sequence: prefix + target's generated tokens
full_ids = out_t.clone() # [1, prefix_len + n_gen_target]
prefix_len = inputs["input_ids"].shape[1]
# Forward through target (get logits at every position) and same for draft
with torch.no_grad():
target_out = target(input_ids=full_ids, input_features=inputs["input_features"])
draft_out = draft(input_ids=full_ids, input_features=inputs["input_features"])
target_argmax = target_out.logits.argmax(-1) # [1, seq_len]
draft_argmax = draft_out.logits.argmax(-1)
# Agreement over positions [prefix_len - 1 ... end - 1] (positions that predict target tokens)
agree_positions = target_argmax[0, prefix_len-1:-1] == draft_argmax[0, prefix_len-1:-1]
total_positions = agree_positions.numel()
n_agree = agree_positions.sum().item()
print(f"[validate] Top-1 agreement: {n_agree}/{total_positions} = {n_agree/total_positions*100:.1f}%")
# -- Also print divergences
if total_positions > 0:
print(f"[validate] First 20 positions (target vs draft argmax):")
for i in range(min(20, total_positions)):
tgt_id = target_argmax[0, prefix_len-1+i].item()
dft_id = draft_argmax[0, prefix_len-1+i].item()
match = "✓" if tgt_id == dft_id else "✗"
tgt_tok = proc.tokenizer.decode([tgt_id])
dft_tok = proc.tokenizer.decode([dft_id])
print(f" pos {i:3d} {match} tgt={tgt_id} ({tgt_tok!r}) dft={dft_id} ({dft_tok!r})")
print(f"\n[validate] Summary:")
print(f" draft coherent: {'yes' if not draft_text.startswith(('\",\"', '.,.', ',,,')) else 'no (gibberish)'}")
print(f" byte-identical to target: {draft_text == target_text}")
print(f" top-1 agreement: {n_agree/total_positions*100:.1f}%")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|