run_batch2.py

"""
TinyThinker — Batch 2 experiments
1. Power → Normal curriculum (3-phase manual)
2. Think + Power inside
3. Think + Mixed inside (self-modulation in thinking space)
4. Think + Reflect (retrospective verification)
5. Power → Think/speak curriculum (Vygotsky arc)
"""

import sys
import json
import time
import shutil
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from experiments import run_experiment, calculate_token_budget, EXPERIMENTS_DIR
import train as train_module
from prepare import (prepare_data, collate_fn, evaluate_model, generate,
                      extract_answer, generate_dataset, WordTokenizer,
                      ReasoningDataset)
from train import TinyThinker, Config, compute_loss, get_lr

import torch
from torch.utils.data import DataLoader

# Common config
COMMON = dict(
    exception_prob=0.4,
    varied_vocab=True,
    math_ratio=0.4,
)


def run_curriculum(name, phases, num_train=50000, num_val=1000, num_test=1000):
    """Run a multi-phase curriculum experiment.

    Args:
        phases: list of (trace_mode, num_steps) tuples
    """
    exp_dir = EXPERIMENTS_DIR / name
    exp_dir.mkdir(parents=True, exist_ok=True)

    print(f"\n{'='*60}")
    print(f"CURRICULUM: {name}")
    print(f"Phases: {phases}")
    print(f"{'='*60}")

    # We need a tokenizer that covers ALL phases' vocabularies
    all_texts = []
    datasets_by_mode = {}
    for mode, _ in phases:
        if mode not in datasets_by_mode:
            exs = generate_dataset(num_train, seed=42, trace_mode=mode, **COMMON)
            val = generate_dataset(num_val, seed=1337, trace_mode=mode, **COMMON)
            tst = generate_dataset(num_test, seed=7, trace_mode=mode, **COMMON)
            all_texts.extend([ex["text"] for ex in exs + val + tst])
            datasets_by_mode[mode] = (exs, val, tst)

    tokenizer = WordTokenizer()
    tokenizer.build(all_texts)
    print(f"Combined vocabulary: {tokenizer.vocab_size} tokens")

    # Calculate max seq_len across all phases
    max_seq = 256
    for mode, _ in phases:
        budget = calculate_token_budget(mode, **COMMON)
        # Account for generous gen_len
        gen_len = int(budget["recommended_max_gen_len"] * 1.5)
        prompt_tokens = budget["p99_total_tokens"] - budget["p99_trace_tokens"]
        needed = ((prompt_tokens + gen_len + 31) // 32) * 32
        max_seq = max(max_seq, needed, budget["recommended_max_seq_len"])

    print(f"Max seq_len across phases: {max_seq}")

    # Build model
    cfg = Config()
    cfg.n_layer = 6
    cfg.n_head = 8
    cfg.n_embd = 256
    cfg.max_seq_len = max_seq
    cfg.compile_model = False  # avoid recompile between phases

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = TinyThinker(cfg, tokenizer.vocab_size).to(device)
    n_params = model.count_parameters()
    print(f"Model: {n_params:,} parameters")

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=3e-4, weight_decay=0.01, betas=(0.9, 0.95)
    )

    total_steps = sum(steps for _, steps in phases)
    global_step = 0
    best_accuracy = 0.0
    t0 = time.time()

    for phase_idx, (mode, phase_steps) in enumerate(phases):
        print(f"\n--- Phase {phase_idx+1}: {mode} for {phase_steps} steps ---")

        exs, val, tst = datasets_by_mode[mode]
        train_ds = ReasoningDataset(exs, tokenizer, max_seq)
        loader = DataLoader(train_ds, batch_size=64, shuffle=True,
                           collate_fn=collate_fn, num_workers=2, pin_memory=True)
        data_iter = iter(loader)

        budget = calculate_token_budget(mode, **COMMON)
        gen_len = max(int(budget["recommended_max_gen_len"] * 1.5), 150)

        phase_start = global_step
        while global_step < phase_start + phase_steps:
            model.train()
            try:
                batch = next(data_iter)
            except StopIteration:
                data_iter = iter(loader)
                batch = next(data_iter)

            batch = {k: v.to(device) for k, v in batch.items()}

            lr = get_lr(global_step, 400, total_steps, 3e-4)
            for pg in optimizer.param_groups:
                pg['lr'] = lr

            loss = compute_loss(model, batch, tokenizer.pad_id)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)

            if global_step % 100 == 0:
                elapsed = time.time() - t0
                print(f"  step {global_step:5d} | loss {loss.item():.4f} | "
                      f"lr {lr:.2e} | phase={mode} | {elapsed:.0f}s", flush=True)

            if global_step > 0 and global_step % 500 == 0:
                # Eval using the CURRENT phase's test data
                results = evaluate_model(model, tokenizer, tst[:300],
                                        device, max_gen_len=gen_len)
                print(f"  EVAL step {global_step}: accuracy={results['accuracy']:.3f} "
                      f"({results['correct']}/{results['total']}) "
                      f"parse_fail={results['parse_failures']} "
                      f"[phase={mode}]", flush=True)

                if results['accuracy'] > best_accuracy:
                    best_accuracy = results['accuracy']
                    torch.save(model.state_dict(), exp_dir / "best.pt")

            global_step += 1

    # Final eval on the LAST phase's test data
    _, _, final_test = datasets_by_mode[phases[-1][0]]
    final_budget = calculate_token_budget(phases[-1][0], **COMMON)
    final_gen = max(int(final_budget["recommended_max_gen_len"] * 1.5), 150)
    results = evaluate_model(model, tokenizer, final_test,
                            device, max_gen_len=final_gen)

    elapsed = time.time() - t0
    print(f"\nCurriculum complete in {elapsed:.0f}s ({elapsed/60:.1f}m)")
    print(f"Final accuracy ({phases[-1][0]}): {results['accuracy']:.3f}")
    print(f"Best accuracy during training: {best_accuracy:.3f}")

    # Save results
    with open(exp_dir / "results.json", "w") as f:
        json.dump({
            "name": name,
            "phases": phases,
            "final_accuracy": results["accuracy"],
            "best_accuracy": best_accuracy,
            "final_mode": phases[-1][0],
        }, f, indent=2)

    print(f"\n>>> {name} DONE: {results['accuracy']:.4f}")
    return results["accuracy"]


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

    # Experiment 1: Power → Normal curriculum
    acc = run_curriculum(
        "curriculum_power_to_normal",
        phases=[
            ("power", 4000),    # Phase 1: learn structure
            ("mixed", 3000),    # Phase 2: bridge
            ("normal", 6000),   # Phase 3: full NL
        ],
    )
    results["curriculum_power_normal"] = acc

    # Experiment 2: Think + Power inside
    acc = run_experiment(
        "think_power_crossmodal",
        trace_mode="think_power",
        override_gen_len=200,
        **{**COMMON, "num_train": 50000, "num_val": 1000, "num_test": 1000,
           "max_steps": 40000, "early_stop_threshold": 0.005,
           "early_stop_min_evals": 20, "early_stop_window": 10,
           "eval_interval": 500, "eval_samples": 300},
    )
    results["think_power"] = acc

    # Experiment 3: Think + Mixed inside
    acc = run_experiment(
        "think_mixed_crossmodal",
        trace_mode="think_mixed",
        override_gen_len=200,
        **{**COMMON, "num_train": 50000, "num_val": 1000, "num_test": 1000,
           "max_steps": 40000, "early_stop_threshold": 0.005,
           "early_stop_min_evals": 20, "early_stop_window": 10,
           "eval_interval": 500, "eval_samples": 300},
    )
    results["think_mixed"] = acc

    # Experiment 4: Think + Reflect
    acc = run_experiment(
        "think_reflect_crossmodal",
        trace_mode="think_reflect",
        override_gen_len=250,
        **{**COMMON, "num_train": 50000, "num_val": 1000, "num_test": 1000,
           "max_steps": 40000, "early_stop_threshold": 0.005,
           "early_stop_min_evals": 20, "early_stop_window": 10,
           "eval_interval": 500, "eval_samples": 300},
    )
    results["think_reflect"] = acc

    # Experiment 5: Power → Think/speak curriculum (Vygotsky arc)
    acc = run_curriculum(
        "curriculum_power_to_think",
        phases=[
            ("power", 4000),        # Phase 1: learn structure
            ("think_power", 3000),  # Phase 2: internalize into thinking
            ("think", 6000),        # Phase 3: free thinking
        ],
    )
    results["curriculum_power_think"] = acc

    # Summary
    print(f"\n{'='*60}")
    print("BATCH 2 COMPLETE")
    print(f"{'='*60}")
    for name, acc in sorted(results.items(), key=lambda x: -x[1]):
        print(f"  {name:<30s}: {acc:.4f}")

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