| from __future__ import annotations | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| from typing import Any | |
| import torch | |
| from conveyer.checkpoint import load_checkpoint | |
| from conveyer.config import load_config | |
| from conveyer.data import build_dataset_and_tokenizer | |
| from conveyer.eval import evaluate_closed_loop_generation | |
| from conveyer.train import train_model | |
| def main() -> None: | |
| parser = argparse.ArgumentParser( | |
| description="Head-to-Head Benchmark: 1-Pass AR Baseline vs. 64-Pass Conveyor Diffusion LLM & Ablations" | |
| ) | |
| parser.add_argument("--config", type=str, default="configs/synthetic.yaml", help="Base config file") | |
| parser.add_argument("--steps", type=int, default=400, help="Training steps per variant") | |
| parser.add_argument("--episodes", type=int, default=64, help="Closed-loop evaluation episodes") | |
| parser.add_argument("--out-dir", type=str, default="runs/sweep", help="Output directory for sweep") | |
| args = parser.parse_args() | |
| variants = [ | |
| { | |
| "name": "1_ar_baseline", | |
| "label": "1-Pass Causal AR Baseline (B=1)", | |
| "overrides": ["train.ar_baseline=true", "train.unroll_steps=8"], | |
| "bake_steps": 1, | |
| "overlap": 0, | |
| }, | |
| { | |
| "name": "2_conveyer_soft_only", | |
| "label": "64-Slot Conveyor (Soft-Vocab Only, alpha_k=0)", | |
| "overrides": ["train.no_thought_highway=true"], | |
| "bake_steps": 64, | |
| "overlap": 16, | |
| }, | |
| { | |
| "name": "3_conveyer_nobptt", | |
| "label": "64-Slot Conveyor (Detached W=1, No Belt BPTT)", | |
| "overrides": ["train.no_belt_bptt=true"], | |
| "bake_steps": 64, | |
| "overlap": 16, | |
| }, | |
| { | |
| "name": "4_conveyer_dual_full", | |
| "label": "64-Slot Dual-Channel Conveyor (Soft-Vocab + Gated Thought + BPTT)", | |
| "overrides": [], | |
| "bake_steps": 64, | |
| "overlap": 16, | |
| }, | |
| ] | |
| summary: list[dict[str, Any]] = [] | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| for var in variants: | |
| run_dir = Path(args.out_dir) / var["name"] | |
| overrides = [ | |
| f"train.max_steps={args.steps}", | |
| f"train.eval_every={max(50, args.steps // 2)}", | |
| f"train.out_dir={run_dir.as_posix()}", | |
| ] + var["overrides"] | |
| print("\n" + "=" * 90) | |
| print(f"Training Variant: {var['label']}") | |
| print("=" * 90) | |
| cfg = load_config(args.config, overrides) | |
| val_m = train_model(cfg) | |
| ckpt_path = run_dir / "best.pt" | |
| model, loaded_cfg, _ = load_checkpoint(ckpt_path, device=device) | |
| _, val_ds, tokenizer = build_dataset_and_tokenizer(loaded_cfg) | |
| cl_metrics = evaluate_closed_loop_generation( | |
| model=model, | |
| val_ds=val_ds, | |
| tokenizer=tokenizer, | |
| device=device, | |
| num_episodes=args.episodes, | |
| prompt_overlap=var["overlap"], | |
| bake_steps=var["bake_steps"], | |
| ) | |
| row = { | |
| "name": var["name"], | |
| "label": var["label"], | |
| "val_slot0_acc": val_m.get("val_slot0_acc", 0.0), | |
| "val_slot0_ppl": val_m.get("val_slot0_ppl", 0.0), | |
| **cl_metrics, | |
| } | |
| summary.append(row) | |
| print("\n" + "=" * 100) | |
| print(f"{'Model / Ablation Variant':<56} | {'Pos-0 Acc':<10} | {'Token Acc':<10} | {'Exact Seq':<10}") | |
| print("-" * 100) | |
| for r in summary: | |
| print( | |
| f"{r['label']:<56} | {r['first_token_acc']*100:8.2f}% | " | |
| f"{r['token_acc']*100:8.2f}% | {r['exact_seq_acc']*100:8.2f}%" | |
| ) | |
| print("=" * 100) | |
| out_path = Path(args.out_dir) / "sweep_summary.json" | |
| out_path.parent.mkdir(parents=True, exist_ok=True) | |
| out_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") | |
| if __name__ == "__main__": | |
| main() | |