v3.0.0: Scripts

scripts/benchmark.py

@@ -0,0 +1,186 @@
+"""
+Comprehensive benchmark script for Q-TensorFormer v3.
+
+Runs multi-model comparison against all baselines and produces
+a full evaluation report with Pareto frontier analysis.
+
+Usage:
+    python scripts/benchmark.py --preset small --epochs 5 --output results/
+"""
+
+import sys
+import os
+import argparse
+import json
+from pathlib import Path
+
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.config import ExperimentConfig, ModelConfig, TrainingConfig, PRESETS
+from src.models import create_model
+from src.baselines import StandardTransformer, DistilledTransformer, PrunedTransformer
+from src.data import load_wikitext2, load_synthetic_data
+from src.training import Trainer
+from src.metrics import (
+    evaluate_model, compare_models, compute_pareto_frontier,
+    compute_efficiency_score, print_comparison_table,
+    rank_trajectory_analysis,
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Q-TensorFormer Benchmark")
+    parser.add_argument("--preset", type=str, default="small",
+                        choices=["tiny", "small", "medium"],
+                        help="Configuration preset")
+    parser.add_argument("--epochs", type=int, default=5,
+                        help="Training epochs")
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--seq-len", type=int, default=128)
+    parser.add_argument("--output", type=str, default="./outputs/benchmark/",
+                        help="Output directory")
+    parser.add_argument("--device", type=str, default="cpu",
+                        help="Device (cpu, cuda)")
+    parser.add_argument("--synthetic", action="store_true",
+                        help="Use synthetic data (faster)")
+    parser.add_argument("--seed", type=int, default=42)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    torch.manual_seed(args.seed)
+
+    # Load config
+    config = PRESETS[args.preset]()
+    config.training.max_epochs = args.epochs
+    config.training.batch_size = args.batch_size
+    config.model.max_seq_len = args.seq_len
+
+    print(f"Config: {config.experiment_name}")
+    print(f"Model: d_model={config.model.d_model}, "
+          f"n_layers={config.model.n_layers}, "
+          f"tt_rank={config.model.tt_rank}")
+
+    # Load data
+    print("\nLoading data...")
+    if args.synthetic:
+        train_loader = load_synthetic_data(
+            vocab_size=config.model.vocab_size,
+            seq_len=args.seq_len,
+            n_samples=2000,
+            batch_size=args.batch_size,
+        )
+        val_loader = None
+        test_loader = train_loader  # Same for synthetic
+        tokenizer = None
+    else:
+        train_loader, val_loader, test_loader, tokenizer = load_wikitext2(
+            seq_len=args.seq_len,
+            batch_size=args.batch_size,
+        )
+        config.model.vocab_size = tokenizer.vocab_size
+
+    # Create models
+    print("\nCreating models...")
+    models = {}
+
+    # Q-TensorFormer (hybrid)
+    models["QTensorFormer"] = create_model(config, "qtensor")
+    print(f"  QTensorFormer: {models['QTensorFormer'].total_params:,} params")
+
+    # TT-Only (no quantum)
+    models["TensorOnly"] = create_model(config, "tensor_only")
+    print(f"  TensorOnly: {models['TensorOnly'].total_params:,} params")
+
+    # Standard transformer (dense)
+    models["StandardTransformer"] = StandardTransformer(
+        vocab_size=config.model.vocab_size,
+        d_model=config.model.d_model,
+        n_heads=config.model.n_heads,
+        n_layers=config.model.n_layers,
+        max_seq_len=config.model.max_seq_len,
+    )
+    print(f"  StandardTransformer: {models['StandardTransformer'].total_params:,} params")
+
+    # Distilled (smaller dense)
+    models["Distilled"] = DistilledTransformer(
+        vocab_size=config.model.vocab_size,
+        d_model=max(64, config.model.d_model // 2),
+        n_heads=config.model.n_heads,
+        n_layers=config.model.n_layers,
+        max_seq_len=config.model.max_seq_len,
+    )
+    print(f"  Distilled: {models['Distilled'].total_params:,} params")
+
+    # Train all models
+    print(f"\n{'='*60}")
+    print("Training models...")
+    print(f"{'='*60}")
+
+    trained_models = {}
+    for name, model in models.items():
+        print(f"\n--- Training {name} ---")
+        trainer = Trainer(
+            model, config,
+            train_loader=train_loader,
+            val_loader=val_loader,
+            test_loader=test_loader,
+            device=args.device,
+            output_dir=f"{args.output}/{name}",
+        )
+        trainer.train()
+        trained_models[name] = model
+
+    # Evaluate
+    print(f"\n{'='*60}")
+    print("Evaluating models...")
+    print(f"{'='*60}")
+
+    results = {}
+    for name, model in trained_models.items():
+        results[name] = evaluate_model(model, test_loader, args.device)
+
+    # Print comparison
+    print_comparison_table(results)
+
+    # Pareto frontier
+    pareto = compute_pareto_frontier(results)
+    print(f"\nPareto-optimal models: {pareto}")
+
+    # Efficiency ranking
+    efficiency = {name: compute_efficiency_score(r) for name, r in results.items()}
+    best = max(efficiency, key=efficiency.get)
+    print(f"Most efficient: {best} (score={efficiency[best]:.1f})")
+
+    # Save results
+    os.makedirs(args.output, exist_ok=True)
+    with open(f"{args.output}/results.json", "w") as f:
+        # Convert float32 to native float
+        clean = {}
+        for name, r in results.items():
+            clean[name] = {k: (float(v) if hasattr(v, 'item') else v) for k, v in r.items()}
+        json.dump({
+            "config": config.experiment_name,
+            "results": clean,
+            "pareto": pareto,
+            "efficiency": {k: float(v) for k, v in efficiency.items()},
+            "best": best,
+        }, f, indent=2)
+
+    print(f"\nResults saved to {args.output}/results.json")
+
+    # Summary
+    print(f"\n{'='*60}")
+    print("SUMMARY")
+    print(f"{'='*60}")
+    for name in results:
+        ppl = results[name]["test_ppl"]
+        params = results[name]["total_params"]
+        lat = results[name].get("latency_ms_mean", 0)
+        print(f"  {name:<25} PPL={ppl:.2f}  Params={params:,}  Lat={lat:.1f}ms")
+
+
+if __name__ == "__main__":
+    main()

scripts/distill.py

@@ -0,0 +1,107 @@
+"""
+Knowledge distillation training script.
+
+Trains a compressed Q-TensorFormer student using a dense teacher model.
+Matches the student's parameter budget to ~50% of the teacher.
+
+Usage:
+    python scripts/distill.py --teacher_config small --student_rank 4
+"""
+
+import sys
+import os
+import argparse
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+import torch
+from src.config import ExperimentConfig, PRESETS
+from src.models import create_model
+from src.baselines import StandardTransformer
+from src.data import load_wikitext2, load_synthetic_data
+from src.training import DistillationTrainer
+from src.metrics import evaluate_model
+
+
+def main():
+    parser = argparse.ArgumentParser(description="KD for Q-TensorFormer")
+    parser.add_argument("--teacher_config", type=str, default="small")
+    parser.add_argument("--student_rank", type=int, default=4)
+    parser.add_argument("--alpha", type=float, default=0.5,
+                        help="Distillation loss weight")
+    parser.add_argument("--temperature", type=float, default=3.0)
+    parser.add_argument("--epochs", type=int, default=8)
+    parser.add_argument("--batch_size", type=int, default=16)
+    parser.add_argument("--device", type=str, default="cpu")
+    parser.add_argument("--output", type=str, default="./outputs/distill/")
+    parser.add_argument("--synthetic", action="store_true")
+    args = parser.parse_args()
+
+    torch.manual_seed(42)
+
+    # Teacher: dense baseline
+    teacher_config = PRESETS[args.teacher_config]()
+    print(f"Teacher config: {teacher_config.experiment_name}")
+
+    # Load data
+    if args.synthetic:
+        train_loader = load_synthetic_data(batch_size=args.batch_size)
+        test_loader = train_loader
+    else:
+        train_loader, val_loader, test_loader, tokenizer = load_wikitext2(
+            batch_size=args.batch_size
+        )
+        teacher_config.model.vocab_size = tokenizer.vocab_size
+
+    # Create teacher (dense)
+    teacher = StandardTransformer(
+        vocab_size=teacher_config.model.vocab_size,
+        d_model=teacher_config.model.d_model,
+        n_heads=teacher_config.model.n_heads,
+        n_layers=teacher_config.model.n_layers,
+    )
+    print(f"Teacher params: {teacher.total_params:,}")
+
+    # Student: compressed Q-TensorFormer
+    student_config = ExperimentConfig(
+        model=type(teacher_config.model)(
+            **{k: v for k, v in teacher_config.model.__dict__.items()}
+        ),
+        training=type(teacher_config.training)(
+            **{k: v for k, v in teacher_config.training.__dict__.items()}
+        ),
+    )
+    student_config.model.tt_rank = args.student_rank
+    student_config.model.use_quantum = True
+    student_config.training.max_epochs = args.epochs
+
+    student = create_model(student_config, "qtensor")
+    print(f"Student params: {student.total_params:,}")
+    print(f"Compression: {teacher.total_params / student.total_params:.1f}x")
+
+    # Train with distillation
+    trainer = DistillationTrainer(
+        student=student,
+        teacher=teacher,
+        config=student_config,
+        train_loader=train_loader,
+        val_loader=val_loader if not args.synthetic else None,
+        test_loader=test_loader,
+        device=args.device,
+        output_dir=args.output,
+        alpha=args.alpha,
+        temperature=args.temperature,
+    )
+    trainer.train()
+
+    # Evaluate
+    print("\nEvaluating knowledge-distilled model...")
+    results = evaluate_model(student, test_loader, args.device)
+    print(f"Student PPL: {results['test_ppl']:.2f}")
+    print(f"Student params: {results['total_params']:,}")
+    print(f"Compression vs teacher: {teacher.total_params / results['total_params']:.1f}x")
+
+
+if __name__ == "__main__":
+    main()

scripts/sweep.py

@@ -0,0 +1,167 @@
+"""
+Hyperparameter sweep script for Q-TensorFormer v3.
+
+Runs a grid/search over key hyperparameters and produces
+comparative evaluation results.
+
+Usage:
+    python scripts/sweep.py --preset sweep --output results/
+"""
+
+import sys
+import os
+import json
+import itertools
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+import torch
+from src.config import ExperimentConfig, ModelConfig, TrainingConfig
+from src.models import create_model
+from src.baselines import StandardTransformer
+from src.data import load_wikitext2, load_synthetic_data
+from src.training import Trainer
+from src.metrics import evaluate_model, print_comparison_table, compute_pareto_frontier
+
+
+def run_sweep(base_config, sweep_params, train_loader, val_loader, test_loader,
+              device="cpu", output_dir="./outputs/sweep/"):
+    """
+    Run a hyperparameter sweep.
+
+    Args:
+        base_config: Base ExperimentConfig.
+        sweep_params: Dict of param_name → [values].
+    """
+    keys = list(sweep_params.keys())
+    values = list(sweep_params.values())
+    os.makedirs(output_dir, exist_ok=True)
+
+    results = {}
+    configs = []
+
+    for combo in itertools.product(*values):
+        config = ExperimentConfig(
+            model=ModelConfig(**base_config.model.__dict__),
+            training=TrainingConfig(**base_config.training.__dict__),
+        )
+
+        # Apply sweep params
+        param_dict = dict(zip(keys, combo))
+        for k, v in param_dict.items():
+            if "." in k:
+                section, key = k.split(".")
+                getattr(getattr(config, section), key).__class__.__dict__
+                setattr(getattr(config, section), key, v)
+            else:
+                if hasattr(config.model, k):
+                    setattr(config.model, k, v)
+                elif hasattr(config.training, k):
+                    setattr(config.training, k, v)
+
+        name = "_".join(f"{k}={v}" for k, v in param_dict.items())
+        config.experiment_name = name
+        configs.append((name, config))
+
+    print(f"Running {len(configs)} configurations...")
+    for i, (name, config) in enumerate(configs):
+        print(f"\n[{i+1}/{len(configs)}] {name}")
+
+        # Create model
+        model = create_model(config, "qtensor")
+
+        # Train
+        trainer = Trainer(
+            model, config,
+            train_loader=train_loader,
+            val_loader=val_loader,
+            test_loader=test_loader,
+            device=device,
+            output_dir=f"{output_dir}/{name}",
+        )
+        trainer.train()
+
+        # Evaluate
+        results[name] = evaluate_model(model, test_loader, device)
+
+    # Save sweep results
+    with open(f"{output_dir}/sweep_results.json", "w") as f:
+        clean = {}
+        for name, r in results.items():
+            clean[name] = {k: (float(v) if hasattr(v, "item") else v) for k, v in r.items()}
+        json.dump(clean, f, indent=2)
+
+    # Print summary
+    print("\n" + "=" * 70)
+    print("SWEEP RESULTS")
+    print("=" * 70)
+    print_comparison_table(results)
+
+    pareto = compute_pareto_frontier(results)
+    print(f"\nPareto-optimal: {pareto}")
+
+    # Best by metric
+    best_ppl = min(results.items(), key=lambda x: x[1]["test_ppl"])
+    best_params = min(results.items(), key=lambda x: x[1]["total_params"])
+    print(f"\nBest PPL: {best_ppl[0]} ({best_ppl[1]['test_ppl']:.2f})")
+    print(f"Fewest params: {best_params[0]} ({best_params[1]['total_params']:,})")
+
+    return results
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--epochs", type=int, default=5)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--device", type=str, default="cpu")
+    parser.add_argument("--output", type=str, default="./outputs/sweep/")
+    parser.add_argument("--synthetic", action="store_true")
+    args = parser.parse_args()
+
+    torch.manual_seed(42)
+
+    # Base config
+    config = ExperimentConfig(
+        model=ModelConfig(d_model=128, n_layers=2, n_heads=4, tt_rank=8,
+                          vocab_size=10000, max_seq_len=128),
+        training=TrainingConfig(max_epochs=args.epochs, batch_size=args.batch_size),
+    )
+
+    # Load data
+    if args.synthetic:
+        train_loader = load_synthetic_data(batch_size=args.batch_size)
+        val_loader = None
+        test_loader = train_loader
+    else:
+        train_loader, val_loader, test_loader, tokenizer = load_wikitext2(
+            seq_len=128, batch_size=args.batch_size
+        )
+        config.model.vocab_size = tokenizer.vocab_size
+
+    # Sweep parameters
+    sweep = {
+        "tt_rank": [2, 4, 8, 16],
+        "use_quantum": [True, False],
+        "quantum_sparsity": [0.5, 0.7, 0.9],
+        "rank_alpha": [1.0, 2.0, 3.0],
+    }
+
+    # Limit combinations for manageable runtime
+    # Full sweep: 4 * 2 * 3 * 3 = 72 combos
+    # Reduced: tt_rank vs quantum vs alpha
+    sweep = {
+        "tt_rank": [2, 4, 8, 16],
+        "use_quantum": [True, False],
+        "quantum_sparsity": [0.7],   # Fixed for now
+        "rank_alpha": [2.0],         # Fixed for now
+    }
+    # 4 * 2 = 8 combos
+
+    run_sweep(config, sweep, train_loader, val_loader, test_loader,
+              args.device, args.output)
+
+
+if __name__ == "__main__":
+    main()