run_experiments.py

"""
TinyThinker — run_experiments.py
Queue and run all experiments sequentially.

Usage: uv run python run_experiments.py
"""

import sys
import json
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from experiments import run_experiment, calculate_token_budget

# Common config for all experiments
COMMON = dict(
    num_train=50000,
    num_val=1000,
    num_test=1000,
    max_steps=40000,
    exception_prob=0.4,
    varied_vocab=True,
    math_ratio=0.4,
    early_stop_threshold=0.005,
    early_stop_min_evals=20,
    early_stop_window=10,
    eval_interval=500,
    eval_samples=300,
)

EXPERIMENTS = [
    # 1. Verbose baseline — properly sized, generous gen_len
    {
        "name": "verbose_crossmodal",
        "trace_mode": "verbose",
        "override_gen_len": 400,
        **COMMON,
    },
    # 2. Self-check: VERIFY after each step
    {
        "name": "selfcheck_crossmodal",
        "trace_mode": "selfcheck",
        "override_gen_len": 300,
        **COMMON,
    },
    # 3. NASA redundant: maximally paranoid
    {
        "name": "redundant_crossmodal",
        "trace_mode": "redundant",
        "override_gen_len": 500,
        **COMMON,
    },
    # 4. Think/speak: inner speech vs output
    {
        "name": "think_crossmodal",
        "trace_mode": "think",
        "override_gen_len": 200,
        **COMMON,
    },
]

if __name__ == "__main__":
    results = {}

    for exp in EXPERIMENTS:
        name = exp["name"]
        mode = exp["trace_mode"]

        # Calculate token budget
        budget = calculate_token_budget(mode, exp.get("exception_prob", 0.4),
                                        exp.get("varied_vocab", True),
                                        exp.get("math_ratio", 0.4))
        print(f"\n{'='*60}")
        print(f"Queued: {name} (mode={mode})")
        print(f"Token budget: seq={budget['recommended_max_seq_len']} "
              f"gen={budget['recommended_max_gen_len']}")
        print(f"{'='*60}")

        acc = run_experiment(**exp)
        results[name] = {
            "accuracy": acc,
            "trace_mode": mode,
            "seq_len": budget["recommended_max_seq_len"],
            "gen_len": budget["recommended_max_gen_len"],
            "mean_tokens": budget["mean_total_tokens"],
        }

    # Summary
    print(f"\n{'='*60}")
    print("ALL EXPERIMENTS COMPLETE")
    print(f"{'='*60}")
    print(f"{'Name':<30s} {'Mode':<12s} {'Acc':>6s} {'Tokens':>8s}")
    print("-" * 60)
    for name, r in sorted(results.items(), key=lambda x: -x[1]["accuracy"]):
        print(f"{name:<30s} {r['trace_mode']:<12s} {r['accuracy']:>6.1%} "
              f"{r['mean_tokens']:>8.0f}")

    with open("experiments/summary.json", "w") as f:
        json.dump(results, f, indent=2)