experiments.py

"""
TinyThinker — experiments.py
Non-destructive experiment runner. Each experiment gets its own results
directory. Previous results are never overwritten.
"""

import sys
import json
import time
import os
from pathlib import Path
from datetime import datetime

sys.path.insert(0, str(Path(__file__).parent))

import torch
from prepare import generate_dataset, WordTokenizer, ReasoningDataset, collate_fn
from train import TinyThinker, Config, train, compute_loss, get_lr


EXPERIMENTS_DIR = Path(__file__).parent / "experiments"


def calculate_token_budget(trace_mode, exception_prob=0.4, varied_vocab=True,
                           math_ratio=0.4, sample_size=1000):
    """Calculate required max_seq_len and max_gen_len for a trace mode.
    Samples the data and returns generous limits."""
    exs = generate_dataset(sample_size, seed=99, trace_mode=trace_mode,
                           exception_prob=exception_prob,
                           varied_vocab=varied_vocab,
                           math_ratio=math_ratio)

    total_lens = []
    trace_lens = []
    for ex in exs:
        tokens = ex["text"].split()
        total_lens.append(len(tokens))
        parts = ex["text"].split(" A: ")
        if len(parts) == 2:
            trace_lens.append(len(parts[1].split()))

    max_total = max(total_lens)
    max_trace = max(trace_lens) if trace_lens else 0
    p99_total = sorted(total_lens)[int(len(total_lens) * 0.99)]
    p99_trace = sorted(trace_lens)[int(len(trace_lens) * 0.99)] if trace_lens else 0
    mean_total = sum(total_lens) / len(total_lens)

    # Use p99 + 20% margin, rounded up to nearest 32
    seq_len = ((int(p99_total * 1.2) + 31) // 32) * 32
    gen_len = int(p99_trace * 1.3)

    report = {
        "trace_mode": trace_mode,
        "mean_total_tokens": round(mean_total, 1),
        "max_total_tokens": max_total,
        "p99_total_tokens": p99_total,
        "max_trace_tokens": max_trace,
        "p99_trace_tokens": p99_trace,
        "recommended_max_seq_len": seq_len,
        "recommended_max_gen_len": gen_len,
    }
    return report


def run_experiment(name, trace_mode, num_train=50000, num_val=1000, num_test=1000,
                   max_steps=40000, exception_prob=0.4, varied_vocab=True,
                   math_ratio=0.4, n_layer=6, n_head=8, n_embd=256,
                   learning_rate=3e-4, batch_size=64,
                   early_stop_threshold=0.005, early_stop_min_evals=20,
                   early_stop_window=10, eval_interval=500, eval_samples=300,
                   checkpoint_path=None, override_seq_len=None,
                   override_gen_len=None):
    """Run a single experiment with full logging.

    Args:
        name: experiment name (used for directory)
        trace_mode: trace format to use
        checkpoint_path: optional path to load weights from (for curriculum)
        override_seq_len: override calculated max_seq_len
        override_gen_len: override calculated max_gen_len
    """
    # Create experiment directory
    exp_dir = EXPERIMENTS_DIR / name
    exp_dir.mkdir(parents=True, exist_ok=True)

    # Calculate token budget
    budget = calculate_token_budget(trace_mode, exception_prob, varied_vocab,
                                    math_ratio)
    max_seq_len = override_seq_len or budget["recommended_max_seq_len"]
    max_gen_len = override_gen_len or budget["recommended_max_gen_len"]
    # Ensure max_seq_len can accommodate prompt + full generation
    # Prompt is roughly (total - trace) tokens, gen needs to cover trace + margin
    min_seq_for_gen = budget["p99_total_tokens"] - budget["p99_trace_tokens"] + max_gen_len + 16
    if max_seq_len < min_seq_for_gen:
        max_seq_len = ((min_seq_for_gen + 31) // 32) * 32

    print(f"\n{'='*60}")
    print(f"EXPERIMENT: {name}")
    print(f"{'='*60}")
    print(f"Token budget: seq_len={max_seq_len}, gen_len={max_gen_len}")
    print(json.dumps(budget, indent=2))

    # Save config
    config = {
        "name": name,
        "trace_mode": trace_mode,
        "num_train": num_train,
        "max_steps": max_steps,
        "exception_prob": exception_prob,
        "varied_vocab": varied_vocab,
        "math_ratio": math_ratio,
        "n_layer": n_layer,
        "n_head": n_head,
        "n_embd": n_embd,
        "max_seq_len": max_seq_len,
        "max_gen_len": max_gen_len,
        "learning_rate": learning_rate,
        "checkpoint_path": checkpoint_path,
        "token_budget": budget,
        "timestamp": datetime.now().isoformat(),
    }
    with open(exp_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    # Configure training
    import train as train_module
    cfg = train_module.cfg
    cfg.trace_mode = trace_mode
    cfg.exception_prob = exception_prob
    cfg.varied_vocab = varied_vocab
    cfg.math_ratio = math_ratio
    cfg.num_train = num_train
    cfg.num_val = num_val
    cfg.num_test = num_test
    cfg.max_steps = max_steps
    cfg.warmup_steps = min(400, max_steps // 10)
    cfg.eval_interval = eval_interval
    cfg.eval_samples = eval_samples
    cfg.max_seq_len = max_seq_len
    cfg.n_layer = n_layer
    cfg.n_head = n_head
    cfg.n_embd = n_embd
    cfg.learning_rate = learning_rate
    cfg.batch_size = batch_size
    cfg.compile_model = True
    cfg.early_stop_window = early_stop_window
    cfg.early_stop_min_evals = early_stop_min_evals
    cfg.early_stop_threshold = early_stop_threshold
    cfg.max_gen_len = max_gen_len

    # Run training
    acc = train_module.train()

    # Save results
    results = {
        "final_accuracy": acc,
        "name": name,
        "trace_mode": trace_mode,
    }
    with open(exp_dir / "results.json", "w") as f:
        json.dump(results, f, indent=2)

    # Copy best checkpoint to experiment directory
    src = Path(__file__).parent / "checkpoints" / f"best_{trace_mode}.pt"
    if src.exists():
        import shutil
        shutil.copy(src, exp_dir / "best.pt")

    print(f"\n>>> {name} DONE: {acc:.4f}")
    print(f"Results saved to {exp_dir}/")

    return acc


if __name__ == "__main__":
    # Show token budgets for all modes
    for mode in ["normal", "verbose", "power", "mixed"]:
        budget = calculate_token_budget(mode)
        print(f"{mode:>10s}: seq_len={budget['recommended_max_seq_len']:>4d}  "
              f"gen_len={budget['recommended_max_gen_len']:>4d}  "
              f"mean_tokens={budget['mean_total_tokens']:.0f}  "
              f"p99={budget['p99_total_tokens']}")