train.py

import os
import math
import time
import json
import random
import inspect
import shutil
import subprocess
import argparse

# ─────────────────────────────────────────────────────────────
# 1. CORE PIPELINE FUNCTION AND HELPERS
# ─────────────────────────────────────────────────────────────

def is_zero_shot_cache_valid(cache_dir):
    if not os.path.exists(cache_dir):
        return False
    pred_count = 0
    for root, dirs, files in os.walk(cache_dir):
        if "predictions.json" in files:
            pred_count += 1
    return pred_count >= 3

def is_finetune_cache_valid(cache_dir):
    if not os.path.exists(cache_dir):
        return False
    pred_count = 0
    for root, dirs, files in os.walk(cache_dir):
        if "predictions.json" in files:
            pred_count += 1
    return pred_count >= 3

def verify_eval_run(path_to_check, description):
    print(f"[Eval] Verifying {description} path: {path_to_check}...")
    if not os.path.exists(path_to_check):
        raise RuntimeError(f"CRITICAL ERROR: {description} directory was NOT created at: {path_to_check}")
    
    found_valid = False
    for root, dirs, files in os.walk(path_to_check):
        for f in files:
            if f in ["predictions.json", "results.txt", "surprisal.json"]:
                file_path = os.path.join(root, f)
                if os.path.getsize(file_path) > 0:
                    found_valid = True
                    break
        if found_valid:
            break
            
    if not found_valid:
        raise RuntimeError(f"CRITICAL ERROR: {description} completed but no valid prediction/result files were written in: {path_to_check}")
    print(f"[Eval] Success! Verified {description} results are stored correctly.")

def run_pipeline(model_name: str, epochs: int = 10, skip_eval: bool = False, skip_aoa: bool = True, skip_glue: bool = False):
    # Configure persistent cache paths locally
    os.environ["HF_HOME"] = os.path.abspath("./hf_cache")
    os.environ["NLTK_DATA"] = os.path.abspath("./nltk_data")
    os.makedirs("./hf_cache", exist_ok=True)
    os.makedirs("./nltk_data", exist_ok=True)

    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        try:
            from huggingface_hub import login
            login(token=hf_token)
            print("[HF] Programmatic login successful using HF_TOKEN.")
        except Exception as e:
            print(f"[HF] Warning: Programmatic login failed: {e}")

    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from datasets import load_dataset
    from tokenizers import Tokenizer
    from transformers import PreTrainedTokenizerFast
    
    # Import model architecture
    from modeling_gpt2 import GPT2CustomLMHeadModel
    from configuration_gpt2 import GPT2CustomConfig

    # Print GPU details
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        print(f"\n[GPU] CUDA is available! Using GPU: {gpu_name}\n")
    else:
        print("\n[GPU] Warning: CUDA is NOT available! Running on CPU.\n")

    # ─────────────────────────────────────────────────────────────
    # GLOBAL HYPERPARAMETERS
    # ─────────────────────────────────────────────────────────────
    VOCAB_SIZE        = 16384          
    BLOCK_SIZE        = 512
    BATCH_SIZE        = 16
    GRAD_ACCUM_STEPS  = 1  
    EPOCHS            = epochs
    LEARNING_RATE     = 3e-4  
    LR_MIN            = LEARNING_RATE * 0.05
    WARMUP_STEPS      = 800  
    WEIGHT_DECAY      = 0.1
    GRAD_CLIP         = 1.0

    # Output directories locally
    model_dir = os.path.abspath(f"./checkpoints/{model_name}")
    os.makedirs(model_dir, exist_ok=True)
    local_results_dir = os.path.abspath(f"./results/{model_name}")
    os.makedirs(local_results_dir, exist_ok=True)

    # ─────────────────────────────────────────────────────────────
    # Helper: Save Hugging Face Compliant Checkpoint
    # ─────────────────────────────────────────────────────────────
    def save_hf_checkpoint(raw_model, checkpoint_dir_name, tokenizer):
        save_dir = os.path.join(model_dir, checkpoint_dir_name)
        os.makedirs(save_dir, exist_ok=True)
        print(f"\n[Checkpoint] Saving Hugging Face format checkpoint to '{save_dir}'...")
        
        # Save standard HF model format (bin & configs) for Hub upload
        raw_model.save_pretrained(save_dir)
        tokenizer.save_pretrained(save_dir)
        
        # Copy modeling.py and configuration.py
        shutil.copy("modeling_gpt2.py", os.path.join(save_dir, "modeling_gpt2.py"))
        shutil.copy("configuration_gpt2.py", os.path.join(save_dir, "configuration_gpt2.py"))

    # ─────────────────────────────────────────────────────────────
    # Tokenizer Training
    # ─────────────────────────────────────────────────────────────
    def build_and_train_tokenizer(texts: list) -> Tokenizer:
        from tokenizers.models import BPE
        from tokenizers.trainers import BpeTrainer
        from tokenizers.pre_tokenizers import Whitespace

        vocab_path = os.path.join(model_dir, "bpe_vocab_16k.json")
        if os.path.exists(vocab_path):
            print(f"[Tokenizer] Loading trained BPE model layout from '{vocab_path}'...")
            return Tokenizer.from_file(vocab_path)

        print(f"[Tokenizer] Generating fresh HuggingFace BPE Tokenizer model with {VOCAB_SIZE} slots...")
        tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
        tokenizer.pre_tokenizer = Whitespace()

        trainer = BpeTrainer(
            vocab_size=VOCAB_SIZE,
            special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
        )
        tokenizer.train_from_iterator(texts, trainer)
        tokenizer.save(vocab_path)
        print(f"[Tokenizer] Tokenizer training completed and saved to '{vocab_path}'.")
        return tokenizer

    # ─────────────────────────────────────────────────────────────
    # Data Loader Setup
    # ─────────────────────────────────────────────────────────────
    class DataLoaderLite:
        def __init__(self, B: int, T: int, texts: list, tokenizer: Tokenizer, name: str):
            self.B = B
            self.T = T

            print(f"[DataLoader:{name}] Tokenising dataset sequences...")
            all_ids = []
            for t in texts:
                if t.strip():
                    encoded = tokenizer.encode(t).ids
                    all_ids.extend(encoded)

            self.tokens = torch.tensor(all_ids, dtype=torch.long)
            self.chunk_size = B * T
            self.n_chunks = (len(self.tokens) - 1) // self.chunk_size
            self.indices = list(range(self.n_chunks))
            self.pos = 0
            self._shuffle()

            print(f"[DataLoader:{name}] Total tokens: {len(self.tokens):,} | Epoch steps: {self.n_chunks:,}")

        def _shuffle(self):
            random.shuffle(self.indices)
            self.pos = 0

        def steps_per_epoch(self) -> int:
            return self.n_chunks

        def next_batch(self):
            B, T = self.B, self.T
            if self.pos >= len(self.indices):
                self._shuffle()

            chunk_idx = self.indices[self.pos]
            self.pos += 1

            start_pos = chunk_idx * self.chunk_size
            temp = self.tokens[start_pos : start_pos + self.chunk_size + 1]

            x = temp[:-1].view(B, T)
            y = temp[1:].view(B, T)
            return x, y

    # Learning Rate Cosine Schedule
    def get_lr(it: int, total_steps: int) -> float:
        if it < WARMUP_STEPS:
            return LEARNING_RATE * (it + 1) / WARMUP_STEPS
        if it >= total_steps:
            return LR_MIN
        decay_ratio = (it - WARMUP_STEPS) / (total_steps - WARMUP_STEPS)
        coeff        = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
        return LR_MIN + coeff * (LEARNING_RATE - LR_MIN)

    # ─────────────────────────────────────────────────────────────
    # Dataset Preparation
    # ─────────────────────────────────────────────────────────────
    print("\n[Data] Loading BabyLM-2026-Strict-Small ...")
    ds = load_dataset("BabyLM-community/BabyLM-2026-Strict-Small")
    all_text = list(ds['train']['text'])

    raw_tokenizer = build_and_train_tokenizer(all_text)
    tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=raw_tokenizer,
        bos_token="[CLS]",
        eos_token="[SEP]",
        unk_token="[UNK]",
        pad_token="[PAD]",
        mask_token="[MASK]"
    )
    BOS_ID = 2
    EOS_ID = 3

    split       = int(len(all_text) * 0.95)
    train_texts = all_text[:split]
    val_texts   = all_text[split:]

    train_loader = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, train_texts, raw_tokenizer, "train")
    val_loader   = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, val_texts, raw_tokenizer, "val")

    chunks_per_epoch = train_loader.steps_per_epoch()
    steps_per_epoch  = chunks_per_epoch // GRAD_ACCUM_STEPS 
    total_steps      = steps_per_epoch * EPOCHS

    cfg = GPT2CustomConfig(
        vocab_size=VOCAB_SIZE,
        n_positions=BLOCK_SIZE,
        n_embd=512,
        n_layer=12,
        n_head=8,
        n_inner=1360,
        bos_token_id=BOS_ID,
        eos_token_id=EOS_ID,
    )
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch.manual_seed(42)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(42)
    random.seed(42)
    if hasattr(torch, 'set_float32_matmul_precision'):
        torch.set_float32_matmul_precision('high')

    model = GPT2CustomLMHeadModel(cfg).to(device)

    # ─────────────────────────────────────────────────────────────
    # Training Resume Check
    # ─────────────────────────────────────────────────────────────
    words_trained = 0
    next_milestone_idx = 0
    global_step = 0
    
    milestones = sorted(list(set([i * 1_000_000 for i in range(1, 11)] + [i * 10_000_000 for i in range(1, 11)])))
    
    resume_checkpoint_dir = None
    for idx in range(len(milestones) - 1, -1, -1):
        m = milestones[idx]
        ckpt_name = f"chck_{m // 1_000_000}M"
        ckpt_path = os.path.join(model_dir, ckpt_name)
        if os.path.exists(os.path.join(ckpt_path, "pytorch_model.bin")) or os.path.exists(os.path.join(ckpt_path, "model.safetensors")):
            config_json_path = os.path.join(ckpt_path, "config.json")
            if os.path.exists(config_json_path):
                try:
                    with open(config_json_path, "r") as f:
                        saved_config = json.load(f)
                    if saved_config.get("n_embd") == 512:
                        resume_checkpoint_dir = ckpt_path
                        next_milestone_idx = idx + 1
                        words_trained = m
                        global_step = words_trained // (BATCH_SIZE * BLOCK_SIZE)
                        print(f"[Training] Found existing milestone checkpoint '{ckpt_name}'. Resuming from step {global_step:,} ({words_trained:,} tokens trained)...")
                        break
                except Exception:
                    pass

    # Load weights if resuming
    if resume_checkpoint_dir is not None:
        print(f"[Model] Loading weights from checkpoint '{resume_checkpoint_dir}'...")
        if os.path.exists(os.path.join(resume_checkpoint_dir, "pytorch_model.bin")):
            state_dict = torch.load(os.path.join(resume_checkpoint_dir, "pytorch_model.bin"), map_location=device)
        else:
            from safetensors.torch import load_file
            state_dict = load_file(os.path.join(resume_checkpoint_dir, "model.safetensors"), device=device)
        model.load_state_dict(state_dict)

    # Check if final main model exists
    main_ckpt_path = os.path.join(model_dir, "main")
    if os.path.exists(os.path.join(main_ckpt_path, "pytorch_model.bin")) or os.path.exists(os.path.join(main_ckpt_path, "model.safetensors")):
        print("\n[Pipeline] Final checkpoint 'main' already exists. Skipping training phase and transitioning directly to evaluations!")
    else:
        try:
            model = torch.compile(model)
            print("[Model] torch.compile() successfully verified graph optimizations")
        except Exception as e:
            print(f"[Model] torch.compile() skipped ({e})")

        # Optimizer
        param_dict     = {n: p for n, p in model.named_parameters() if p.requires_grad}
        decay_params   = [p for p in param_dict.values() if p.dim() >= 2]
        nodecay_params = [p for p in param_dict.values() if p.dim() < 2]
        groups = [
            {'params': decay_params,   'weight_decay': WEIGHT_DECAY},
            {'params': nodecay_params, 'weight_decay': 0.0},
        ]
        fused_ok  = 'fused' in inspect.signature(torch.optim.AdamW).parameters
        use_fused = fused_ok and ('cuda' in device)
        optimizer = torch.optim.AdamW(groups, lr=LEARNING_RATE, betas=(0.9, 0.95), eps=1e-8, fused=use_fused)

        # ─────────────────────────────────────────────────────────────
        # Training Loop
        # ─────────────────────────────────────────────────────────────
        def evaluate_validation_loss(model_eval, val_loader_eval, dev, autocast):
            was_training = model_eval.training
            model_eval.eval()
            
            val_loss_accum = 0.0
            val_steps = val_loader_eval.steps_per_epoch()
            max_val_steps = min(val_steps, 100)
            
            with torch.no_grad():
                for _ in range(max_val_steps):
                    x, y = val_loader_eval.next_batch()
                    x, y = x.to(dev), y.to(dev)
                    with autocast:
                        out = model_eval(x, labels=y)
                        loss = out.loss
                    val_loss_accum += loss.item()
            
            avg_val_loss = val_loss_accum / max_val_steps
            val_perplexity = math.exp(avg_val_loss)
            
            if was_training:
                model_eval.train()
            return avg_val_loss, val_perplexity

        validation_logs = []
        training_logs = []

        model.train()
        autocast_ctx = torch.autocast(device_type="cuda" if "cuda" in device else "cpu", dtype=torch.bfloat16, enabled=True)

        start_epoch = global_step // steps_per_epoch
        start_chunk = (global_step % steps_per_epoch) * GRAD_ACCUM_STEPS

        print(f"\n[Training] Starting GPT-2 Dense Baseline training for {EPOCHS} epochs...")
        for epoch in range(start_epoch, EPOCHS):
            train_loader._shuffle()
            if epoch == start_epoch and start_chunk > 0:
                print(f"[Training] Fast-forwarding dataloader to chunk index {start_chunk}...")
                train_loader.pos = start_chunk

            optimizer.zero_grad(set_to_none=True)
            loss_accum = 0.0
            
            start_chunk_idx = start_chunk if epoch == start_epoch else 0
            for chunk_step in range(start_chunk_idx, chunks_per_epoch):
                t0 = time.perf_counter()

                lr = get_lr(global_step, total_steps)
                for pg in optimizer.param_groups:
                    pg['lr'] = lr

                x, y = train_loader.next_batch()
                input_ids, targets = x.to(device), y.to(device)
                words_trained += input_ids.numel()

                with autocast_ctx:
                    out = model(input_ids, labels=targets)
                    loss = out.loss
                scaled_loss  = loss / GRAD_ACCUM_STEPS
                loss_accum  += scaled_loss.item()
                scaled_loss.backward()

                if (chunk_step + 1) % GRAD_ACCUM_STEPS == 0:
                    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
                    optimizer.step()
                    optimizer.zero_grad(set_to_none=True)
                    if "cuda" in device:
                        torch.cuda.synchronize()

                    dt = (time.perf_counter() - t0) * 1000
                    
                    current_step = (chunk_step + 1) // GRAD_ACCUM_STEPS
                    print(
                        f"[E{epoch+1:02d} {current_step:>5d}/{steps_per_epoch} G{global_step:>7d}] "
                        f"train={loss_accum:.4f}  norm={norm:.3f}  lr={lr:.2e}  dt={dt:6.1f}ms  words={words_trained:,}"
                    )

                    training_logs.append({
                        "step": global_step,
                        "train_loss": loss_accum,
                        "lr": lr,
                        "words": words_trained
                    })
                    loss_accum = 0.0
                    global_step += 1

                # Validation Loss every 100 steps
                if global_step > 0 and global_step % 100 == 0:
                    v_loss, v_ppl = evaluate_validation_loss(model, val_loader, device, autocast_ctx)
                    validation_logs.append({
                        "step": global_step,
                        "val_loss": v_loss,
                        "val_perplexity": v_ppl
                    })
                    with open(os.path.join(model_dir, "validation_metrics.json"), "w") as f:
                        json.dump(validation_logs, f, indent=2)
                    with open("validation_metrics.json", "w") as f:
                        json.dump(validation_logs, f, indent=2)
                    print(f"\n[Validation Step {global_step}] Loss: {v_loss:.4f} | Perplexity: {v_ppl:.2f}\n")

                # Checkpoint Milestones
                if next_milestone_idx < len(milestones) and words_trained >= milestones[next_milestone_idx]:
                    milestone_val = milestones[next_milestone_idx]
                    milestone_name = f"chck_{milestone_val // 1_000_000}M" if milestone_val < 10_000_000 else f"chck_{(milestone_val // 10_000_000) * 10}M"
                    raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
                    save_hf_checkpoint(raw_model, milestone_name, tokenizer)
                    next_milestone_idx += 1

        # Save final model as 'main'
        raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
        save_hf_checkpoint(raw_model, "main", tokenizer)
        
        # Save training logs
        with open(os.path.join(model_dir, "training_metrics.json"), "w") as f:
            json.dump(training_logs, f, indent=2)
        with open("training_metrics.json", "w") as f:
            json.dump(training_logs, f, indent=2)

        # Generate plots
        try:
            import matplotlib
            matplotlib.use('Agg')
            import matplotlib.pyplot as plt

            if len(validation_logs) > 0:
                steps = [log["step"] for log in validation_logs]
                perplexities = [log["val_perplexity"] for log in validation_logs]

                plt.figure(figsize=(8, 5))
                plt.plot(steps, perplexities, marker='o', color='#3b82f6', linewidth=2)
                plt.xlabel('Global Step')
                plt.ylabel('Validation Perplexity')
                plt.title('Validation Perplexity Learning Curve')
                plt.grid(True, linestyle='--', alpha=0.6)
                plt.savefig(os.path.join(model_dir, 'validation_perplexity.png'), dpi=150)
                plt.savefig('validation_perplexity.png', dpi=150)
                plt.close()
                print("[Plots] Validation perplexity curve saved successfully.")
        except Exception as e:
            print(f"[Plots] Warning: Could not generate plots: {e}")

        print("\n[Training] Training phase complete!")

    if skip_eval:
        print("[Pipeline] Skipping evaluations phase as requested.")
        return

    # ─────────────────────────────────────────────────────────────
    # 2. RUN EVALUATION PIPELINE
    # ─────────────────────────────────────────────────────────────
    local_results_dir = os.path.abspath(f"./results/{model_name}")
    os.makedirs(local_results_dir, exist_ok=True)
    local_main_res = os.path.join(local_results_dir, "main")

    # Ensure clone_dir exists and has the global_piqa files
    clone_parent = os.path.abspath("./babylm_eval_repo")
    
    # Self-healing check for global_piqa presence
    has_global_piqa = False
    for potential_strict in [os.path.join(clone_parent, "babylm-eval", "strict"), os.path.join(clone_parent, "strict")]:
        if os.path.exists(os.path.join(potential_strict, "evaluation_pipeline", "global_piqa")):
            has_global_piqa = True
            break
            
    if not has_global_piqa:
        print("[Eval] Cloned repository does not contain global_piqa tasks.")
        print("[Eval] Deleting and cloning official main branch...")
        if os.path.exists(clone_parent):
            shutil.rmtree(clone_parent)
        subprocess.run([
            "git", "clone", "-b", "main",
            "https://github.com/babylm-org/babylm-eval.git",
            clone_parent
        ], check=True)

    # Determine strict_dir path dynamically
    if os.path.exists(os.path.join(clone_parent, "strict")):
        strict_dir = os.path.join(clone_parent, "strict")
    else:
        strict_dir = os.path.join(clone_parent, "babylm-eval", "strict")
        
    print(f"[Eval] Using strict directory: {strict_dir}")
    os.environ["PYTHONPATH"] = strict_dir

    def patch_evaluation_run_script(strict_dir):
        import pathlib
        run_file = os.path.join(strict_dir, "evaluation_pipeline", "sentence_zero_shot", "run.py")
        if not os.path.exists(run_file):
            print(f"[GlobalPIQA] Warning: {run_file} not found. Cannot patch.")
            return
            
        print(f"[GlobalPIQA] Patching local checkpoint loader in {run_file}...")
        with open(run_file, "r") as f:
            content = f.read()
            
        if "Local checkpoint directory patch" in content:
            print("[GlobalPIQA] Script already patched.")
            return
            
        target_str = """def main():
    args = _parse_arguments()
    if args.images_path is not None:
        assert args.batch_size == 1, "Multimodal only works in batch size 1!"
    dataset = args.data_path.stem
    args.model_name = pathlib.Path(args.model_path_or_name).stem
    if args.revision_name is None:
        revision_name = "main"
    else:
        revision_name = args.revision_name"""
            
        patch_str = """def main():
    args = _parse_arguments()
    if args.images_path is not None:
        assert args.batch_size == 1, "Multimodal only works in batch size 1!"
    dataset = args.data_path.stem
    
    # Local checkpoint directory patch
    import os
    model_path = args.model_path_or_name
    args.model_name = pathlib.Path(model_path).stem
    revision_name = args.revision_name if args.revision_name else "main"
    
    if os.path.isdir(model_path):
        target_revision = args.revision_name if args.revision_name else "main"
        if os.path.exists(os.path.join(model_path, target_revision)):
            args.model_path_or_name = os.path.join(model_path, target_revision)
            args.revision_name = None"""
                
        if target_str in content:
            new_content = content.replace(target_str, patch_str)
            with open(run_file, "w") as f:
                f.write(new_content)
            print("[GlobalPIQA] Successfully patched run.py")
        else:
            print("[GlobalPIQA] Warning: Could not find target pattern in run.py. Manual patch may be needed.")

    # Patch sentence zero shot loader inside cloned repo
    patch_evaluation_run_script(strict_dir)

    print("[Eval] Stripping Windows-specific packages from requirements.txt...")
    req_file_path = os.path.join(strict_dir, "requirements.txt")
    if os.path.exists(req_file_path):
        with open(req_file_path, "r") as f:
            lines = f.readlines()
        with open(req_file_path, "w") as f:
            for line in lines:
                if "pywin" not in line.lower() and "wintypes" not in line.lower():
                    f.write(line)

    print("[Eval] Verifying and installing evaluation dependencies programmatically...")
    required_packages = {
        "nltk": "nltk",
        "pandas": "pandas",
        "statsmodels": "statsmodels",
        "sklearn": "scikit-learn",
        "scipy": "scipy"
    }
    for pkg_import, pkg_install in required_packages.items():
        try:
            __import__(pkg_import)
        except ImportError:
            print(f"[Eval] Package '{pkg_install}' not found. Installing it programmatically...")
            import sys
            subprocess.run([sys.executable, "-m", "pip", "install", pkg_install], check=True)
    
    print("[Eval] Downloading NLTK tokenizer resources...")
    import nltk
    nltk.download('punkt', download_dir=os.environ["NLTK_DATA"])
    nltk.download('punkt_tab', download_dir=os.environ["NLTK_DATA"])

    # Ensure standard zero-shot datasets are downloaded
    blimp_fast_dir = os.path.join(strict_dir, "evaluation_data", "fast_eval", "blimp_fast")
    if not os.path.exists(blimp_fast_dir) or not os.listdir(blimp_fast_dir):
        print("[Eval] Standard zero-shot datasets not found. Downloading...")
        subprocess.run(["python", "-m", "scripts.download_evals"], cwd=strict_dir, check=True)
        
        ewok_zip = os.path.join(strict_dir, "evaluation_data/fast_eval/ewok_fast.zip")
        if os.path.exists(ewok_zip):
            print("[Eval] Unzipping EWoK fast data...")
            bad_nested_dir = os.path.join(strict_dir, "evaluation_data/fast_eval/evaluation_data")
            if os.path.exists(bad_nested_dir):
                shutil.rmtree(bad_nested_dir)
            subprocess.run(["unzip", "-o", "-P", "BabyLM2025", "evaluation_data/fast_eval/ewok_fast.zip", "-d", "."], cwd=strict_dir, check=True)
            
        print("[Eval] Downloading and filtering full EWoK dataset...")
        subprocess.run(["python", "-m", "evaluation_pipeline.ewok.dl_and_filter"], cwd=strict_dir, check=True)

    # Download GlobalPIQA dataset
    global_piqa_parallel_dir = os.path.join(strict_dir, "evaluation_data", "fast_eval", "global_piqa_parallel")
    if not os.path.exists(global_piqa_parallel_dir) or not os.listdir(global_piqa_parallel_dir):
        print("[Eval] GlobalPIQA dataset not found. Downloading...")
        subprocess.run(["python", "evaluation_pipeline/global_piqa/dl.py"], cwd=strict_dir, check=True)

    # Ensure all scripts are executable
    print("[Eval] Making evaluation shell scripts executable...")
    subprocess.run("chmod +x scripts/*.sh", shell=True, cwd=strict_dir, check=True)

    def run_task_with_cache(checkpoint, task, output_subpath, cmd):
        local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", task, output_subpath)
        if task == "reading":
            local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", "reading")
        elif task == "comps":
            local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", "comps", "comps")
            
        target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", task, output_subpath)
        if task == "reading":
            target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", "reading")
        elif task == "comps":
            target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", "comps", "comps")

        cache_file = os.path.join(local_cache_path, "predictions.json")
        
        # Self-healing for entity_tracking cache
        if task == "entity_tracking" and os.path.exists(cache_file):
            try:
                import json
                with open(cache_file, "r") as f:
                    preds = json.load(f)
                is_valid_cache = True
                for k, v in preds.items():
                    if len(v.get("predictions", [])) in [605, 606, 607, 615, 529, 156, 187, 159]:
                        is_valid_cache = False
                        break
                if not is_valid_cache:
                    print(f"[Eval] Cached entity_tracking for '{checkpoint}' has incorrect old sizes. Invalidate and re-run fresh...")
                    shutil.rmtree(local_cache_path, ignore_errors=True)
            except Exception:
                pass

        if os.path.exists(cache_file):
            print(f"[Eval] Task '{task}' ({output_subpath}) for checkpoint '{checkpoint}' is cached. Restoring...")
            if os.path.exists(target_results_dir):
                shutil.rmtree(target_results_dir)
            os.makedirs(target_results_dir, exist_ok=True)
            for item in os.listdir(local_cache_path):
                s = os.path.join(local_cache_path, item)
                d = os.path.join(target_results_dir, item)
                if os.path.isdir(s):
                    shutil.copytree(s, d)
                else:
                    shutil.copy2(s, d)
            return

        print(f"[Eval] Running task '{task}' ({output_subpath}) for checkpoint '{checkpoint}'...")
        subprocess.run(cmd, cwd=strict_dir, check=True)

        # Relocate results
        actual_model_basename = os.path.basename(model_dir.rstrip("/"))
        possible_actual_path = os.path.join(strict_dir, "results", actual_model_basename, checkpoint, "zero_shot", "causal", task, output_subpath)
        if task == "reading":
            possible_actual_path = os.path.join(strict_dir, "results", actual_model_basename, checkpoint, "zero_shot", "causal", "reading")
        elif task == "comps":
            possible_actual_path = os.path.join(strict_dir, "results", actual_model_basename, checkpoint, "zero_shot", "causal", "comps", "comps")

        possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", task, output_subpath)
        if task == "reading":
            possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", "reading")
        elif task == "comps":
            possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", "comps", "comps")

        possible_local_reading_path = os.path.join(strict_dir, "results", checkpoint, "main", "zero_shot", "causal", "reading")

        for p_path in [possible_actual_path, possible_main_path, possible_local_reading_path]:
            if os.path.exists(p_path) and p_path != target_results_dir:
                print(f"[Eval] Relocating results from {p_path} to {target_results_dir}...")
                if os.path.exists(target_results_dir):
                    shutil.rmtree(target_results_dir)
                os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
                shutil.move(p_path, target_results_dir)
                break

        verify_eval_run(target_results_dir, f"{checkpoint} {task} ({output_subpath})")

        # Save to local cache
        if os.path.exists(local_cache_path):
            shutil.rmtree(local_cache_path)
        os.makedirs(local_cache_path, exist_ok=True)
        for item in os.listdir(target_results_dir):
            s = os.path.join(target_results_dir, item)
            d = os.path.join(local_cache_path, item)
            if os.path.isdir(s):
                shutil.copytree(s, d)
            else:
                shutil.copy2(s, d)

    def run_finetune_task_with_cache(task, cmd):
        local_cache_path = os.path.join("./results", model_name, "main", "finetune", task)
        target_results_dir = os.path.join(strict_dir, "results", model_name, "main", "finetune", task)

        if os.path.exists(os.path.join(local_cache_path, "predictions.json")):
            print(f"[Eval] GLUE task '{task}' is cached. Restoring...")
            if os.path.exists(target_results_dir):
                shutil.rmtree(target_results_dir)
            os.makedirs(target_results_dir, exist_ok=True)
            for item in os.listdir(local_cache_path):
                s = os.path.join(local_cache_path, item)
                d = os.path.join(target_results_dir, item)
                if os.path.isdir(s):
                    shutil.copytree(s, d)
                else:
                    shutil.copy2(s, d)
            return

        print(f"[Eval] Running GLUE task '{task}'...")
        subprocess.run(cmd, cwd=strict_dir, check=True)

        possible_main_path = os.path.join(strict_dir, "results", "main", "main", "finetune", task)
        if os.path.exists(possible_main_path) and possible_main_path != target_results_dir:
            print(f"[Eval] Relocating results from {possible_main_path} to {target_results_dir}...")
            if os.path.exists(target_results_dir):
                shutil.rmtree(target_results_dir)
            os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
            shutil.move(possible_main_path, target_results_dir)

        verify_eval_run(target_results_dir, f"GLUE fine-tuning {task}")

        if os.path.exists(local_cache_path):
            shutil.rmtree(local_cache_path)
        os.makedirs(local_cache_path, exist_ok=True)
        for item in os.listdir(target_results_dir):
            s = os.path.join(target_results_dir, item)
            d = os.path.join(local_cache_path, item)
            if os.path.isdir(s):
                shutil.copytree(s, d)
            else:
                shutil.copy2(s, d)

    # ─────────────────────────────────────────────────────────────
    # A. ZERO-SHOT EVALUATIONS
    # ─────────────────────────────────────────────────────────────
    eval_checkpoints = ["main"] + [f"chck_{m}M" for m in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 30, 35, 40, 45, 50, 60, 70, 80, 90, 100]]
    existing_eval_checkpoints = []
    for checkpoint in eval_checkpoints:
        checkpoint_path = os.path.join(model_dir, checkpoint)
        if os.path.exists(os.path.join(checkpoint_path, "pytorch_model.bin")) or os.path.exists(os.path.join(checkpoint_path, "model.safetensors")):
            existing_eval_checkpoints.append(checkpoint)

    print(f"\n[Eval] Found {len(existing_eval_checkpoints)} checkpoints ready for zero-shot evaluation: {existing_eval_checkpoints}")

    # Self-healing: Patch model_type and auto_map in existing checkpoints' config.json
    # and copy code files if missing.
    for checkpoint in existing_eval_checkpoints:
        checkpoint_path = os.path.join(model_dir, checkpoint)
        
        # A. Overwrite modeling_gpt2.py and configuration_gpt2.py unconditionally
        for code_file in ["modeling_gpt2.py", "configuration_gpt2.py"]:
            dest_code_file = os.path.join(checkpoint_path, code_file)
            if os.path.exists(code_file):
                try:
                    shutil.copy(code_file, dest_code_file)
                    print(f"[Self-Healing] Copied and updated {code_file} in checkpoint '{checkpoint}' directory.")
                except Exception as e:
                    print(f"[Self-Healing] Warning: Failed to copy {code_file} to {checkpoint}: {e}")
                    
        # B. Patch config.json
        config_path = os.path.join(checkpoint_path, "config.json")
        if os.path.exists(config_path):
            try:
                with open(config_path, "r") as f:
                    config_dict = json.load(f)
                
                # Check if it needs patching
                if config_dict.get("model_type") != "gpt2_custom" or "auto_map" not in config_dict:
                    config_dict["model_type"] = "gpt2_custom"
                    config_dict["auto_map"] = {
                        "AutoConfig": "configuration_gpt2.GPT2CustomConfig",
                        "AutoModelForCausalLM": "modeling_gpt2.GPT2CustomLMHeadModel"
                    }
                    with open(config_path, "w") as f:
                        json.dump(config_dict, f, indent=2)
                    print(f"[Self-Healing] Successfully patched 'auto_map' and 'model_type' in checkpoint '{checkpoint}' config.json.")
            except Exception as e:
                print(f"[Self-Healing] Warning: Failed to patch {config_path}: {e}")

    for checkpoint in existing_eval_checkpoints:
        eval_model_path = os.path.join(model_dir, checkpoint)
        print(f"\n{'='*65}\n  RUNNING EVALUATIONS FOR CHECKPOINT: {checkpoint}\n{'='*65}")

        # blimp filtered
        run_task_with_cache(
            checkpoint, "blimp", "blimp_filtered",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/blimp_filtered", "--save_predictions", "--revision_name", checkpoint]
        )
        # supplement filtered
        run_task_with_cache(
            checkpoint, "blimp", "supplement_filtered",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/supplement_filtered", "--save_predictions", "--revision_name", checkpoint]
        )
        # ewok filtered
        run_task_with_cache(
            checkpoint, "ewok", "ewok_filtered",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "ewok", "--data_path", "evaluation_data/full_eval/ewok_filtered", "--save_predictions", "--revision_name", checkpoint]
        )
        # entity tracking
        run_task_with_cache(
            checkpoint, "entity_tracking", "entity_tracking",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "entity_tracking", "--data_path", "evaluation_data/full_eval/entity_tracking", "--save_predictions", "--revision_name", checkpoint]
        )
        # comps
        run_task_with_cache(
            checkpoint, "comps", "comps",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "comps", "--data_path", "evaluation_data/full_eval/comps", "--save_predictions", "--revision_name", checkpoint]
        )
        # reading
        run_task_with_cache(
            checkpoint, "reading", "reading",
            ["python", "-m", "evaluation_pipeline.reading.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--data_path", "evaluation_data/full_eval/reading/reading_data.csv"]
        )
        # global piqa parallel
        run_task_with_cache(
            checkpoint, "global_piqa_parallel", "global_piqa_parallel",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "global_piqa_parallel", "--data_path", "evaluation_data/full_eval/global_piqa_parallel", "--save_predictions", "--revision_name", checkpoint]
        )
        # global piqa nonparallel
        run_task_with_cache(
            checkpoint, "global_piqa_nonparallel", "global_piqa_nonparallel",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "global_piqa_nonparallel", "--data_path", "evaluation_data/full_eval/global_piqa_nonparallel", "--save_predictions", "--revision_name", checkpoint]
        )

    # ─────────────────────────────────────────────────────────────
    # B. GLUE FINE-TUNING EVALUATIONS
    # ─────────────────────────────────────────────────────────────
    if not skip_glue:
        main_ckpt_path = os.path.join(model_dir, "main")
        if os.path.exists(main_ckpt_path):
            print(f"\n{'='*65}\n  RUNNING GLUE FINE-TUNING FOR CHECKPOINT: main\n{'='*65}")
            
            glue_tasks = {
                "boolq": ["boolq", "16", "10"],
                "multirc": ["multirc", "16", "10"],
                "rte": ["rte", "32", "10"],
                "wsc": ["wsc", "32", "30"],
                "mrpc": ["mrpc", "32", "10"],
                "qqp": ["qqp", "32", "10"],
                "mnli": ["mnli", "32", "10"]
            }
            for task_name, (task, bsz, max_epochs) in glue_tasks.items():
                num_labels = "3" if task == "mnli" else "2"
                metric_for_valid = "accuracy"
                if task in ["mrpc", "qqp"]:
                    metric_for_valid = "f1"
                metrics = ["accuracy"]
                if task != "mnli":
                    metrics = ["accuracy", "f1", "mcc"]

                cmd = [
                    "python", "-m", "evaluation_pipeline.finetune.run",
                    "--model_name_or_path", main_ckpt_path,
                    "--train_data", f"evaluation_data/full_eval/glue_filtered/{task}.train.jsonl",
                    "--valid_data", f"evaluation_data/full_eval/glue_filtered/{task}.valid.jsonl",
                    "--predict_data", f"evaluation_data/full_eval/glue_filtered/{task}.valid.jsonl",
                    "--task", task,
                    "--num_labels", num_labels,
                    "--batch_size", bsz,
                    "--learning_rate", "3e-5",
                    "--num_epochs", max_epochs,
                    "--sequence_length", "512",
                    "--results_dir", "results",
                    "--save",
                    "--save_dir", "models",
                    "--metric_for_valid", metric_for_valid,
                    "--seed", "42",
                    "--verbose",
                    "--padding_side", "left",
                    "--take_final"
                ]
                cmd.append("--metrics")
                cmd.extend(metrics)

                run_finetune_task_with_cache(task_name, cmd)

    # ─────────────────────────────────────────────────────────────
    # C. INTERMEDIATE CHECKPOINTS FAST EVALUATION
    # ─────────────────────────────────────────────────────────────
    print(f"[Eval] Running zero-shot fast evaluations on intermediate checkpoints...")
    checkpoints = [f"chck_{i}M" for i in range(1, 10)] + [f"chck_{i}M" for i in range(10, 110, 10)]

    eval_model_path = model_dir

    for checkpoint in checkpoints:
        ckpt_full_path = os.path.join(model_dir, checkpoint)
        if not os.path.exists(ckpt_full_path):
            continue

        # blimp fast
        run_task_with_cache(
            checkpoint, "blimp", "blimp_fast",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/fast_eval/blimp_fast", "--save_predictions", "--revision_name", checkpoint]
        )
        # supplement fast
        run_task_with_cache(
            checkpoint, "blimp", "supplement_fast",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/fast_eval/supplement_fast", "--save_predictions", "--revision_name", checkpoint]
        )
        # ewok fast
        run_task_with_cache(
            checkpoint, "ewok", "ewok_fast",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "ewok", "--data_path", "evaluation_data/fast_eval/ewok_fast", "--save_predictions", "--revision_name", checkpoint]
        )
        # entity tracking fast
        run_task_with_cache(
            checkpoint, "entity_tracking", "entity_tracking_fast",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "entity_tracking", "--data_path", "evaluation_data/fast_eval/entity_tracking_fast", "--save_predictions", "--revision_name", checkpoint]
        )
        # reading fast
        run_task_with_cache(
            checkpoint, "reading", "reading",
            ["python", "-m", "evaluation_pipeline.reading.run", "--model_path_or_name", ckpt_full_path, "--backend", "causal", "--data_path", "evaluation_data/fast_eval/reading/reading_data.csv"]
        )
        # global piqa parallel fast
        run_task_with_cache(
            checkpoint, "global_piqa_parallel", "global_piqa_parallel",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "global_piqa_parallel", "--data_path", "evaluation_data/fast_eval/global_piqa_parallel", "--save_predictions", "--revision_name", checkpoint]
        )
        # global piqa nonparallel fast
        run_task_with_cache(
            checkpoint, "global_piqa_nonparallel", "global_piqa_nonparallel",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "global_piqa_nonparallel", "--data_path", "evaluation_data/fast_eval/global_piqa_nonparallel", "--save_predictions", "--revision_name", checkpoint]
        )

    # ─────────────────────────────────────────────────────────────
    # D. COLLATE RESULTS AND CLEANUP
    # ─────────────────────────────────────────────────────────────
    print("[Eval] Collating predictions into submission file...")
    collate_results_dir = os.path.join(strict_dir, "results", model_name)
    if os.path.exists(collate_results_dir):
        shutil.rmtree(collate_results_dir)
    os.makedirs(os.path.dirname(collate_results_dir), exist_ok=True)

    shutil.copytree(local_results_dir, collate_results_dir)

    subprocess.run([
        "python", "-m", "evaluation_pipeline.collate_preds",
        "--model_path_or_name", model_name,
        "--backend", "causal",
        "--track", "strict-small",
        "--fast"
    ], cwd=strict_dir, check=True)

    results_src = os.path.join(strict_dir, "results")
    results_dest = os.path.abspath("./results")
    if os.path.exists(results_dest):
        shutil.rmtree(results_dest)
    shutil.copytree(results_src, results_dest)

    collated_json = os.path.join(strict_dir, "all_full_preds_and_fast_scores_causal.json")
    if os.path.exists(collated_json):
        shutil.copy(collated_json, "./all_full_preds_and_fast_scores_causal.json")
        print("\n[Eval] Success! Collation completed! Final file is at './all_full_preds_and_fast_scores_causal.json'")

    print("\n[Eval] Pipeline evaluation run finished.")

def upload_pipeline(model_name, repo_name, token=None):
    from huggingface_hub import HfApi, create_repo
    if not token:
        token = os.environ.get("HF_TOKEN")

    api = HfApi(token=token)
    try:
        user_info = api.whoami()
        username = user_info["name"]
        print(f"[HF] Authenticated successfully as user: {username}")
    except Exception as e:
        print(f"[HF] Authentication failed. Error: {e}")
        return

    repo_id = f"{username}/{repo_name}"
    print(f"[HF] Target Repository ID: {repo_id}")

    try:
        create_repo(repo_id=repo_id, repo_type="model", token=token, exist_ok=True)
        print(f"[HF] Repository '{repo_id}' is ready.")
    except Exception as e:
        print(f"[HF] Failed to create repository: {e}")
        return

    checkpoint_dir = os.path.abspath(f"./checkpoints/{model_name}")

    revisions = {}
    if os.path.exists(os.path.join(checkpoint_dir, "main")):
        revisions = {"main": os.path.join(checkpoint_dir, "main")}
    else:
        print(f"[HF] Error: Could not locate 'main' checkpoint weights under {checkpoint_dir}/main")
        return

    main_dir = revisions["main"]
    parent_dir = os.path.dirname(main_dir)

    for m in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 30, 35, 40, 45, 50, 60, 70, 80, 90, 100]:
        ckpt_name = f"chck_{m}M"
        ckpt_path = os.path.join(parent_dir, ckpt_name)
        if os.path.exists(ckpt_path):
            revisions[ckpt_name] = ckpt_path

    temp_dir = os.path.abspath("./hf_upload_temp")
    
    for revision_name, local_path in revisions.items():
        print(f"\n[HF] Preparing revision '{revision_name}' from folder: {local_path}")
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
        os.makedirs(temp_dir, exist_ok=True)
        
        for item in os.listdir(local_path):
            shutil.copy2(os.path.join(local_path, item), os.path.join(temp_dir, item))
            
        try:
            api.create_branch(
                repo_id=repo_id,
                repo_type="model",
                branch=revision_name,
                exist_ok=True
            )
            print(f"[HF] Created branch/revision '{revision_name}' on repository.")
        except Exception as branch_err:
            print(f"[HF] Info: Branch creation failed or exists: {branch_err}")

        print(f"[HF] Uploading staged folder to '{repo_id}' revision '{revision_name}'...")
        try:
            api.upload_folder(
                folder_path=temp_dir,
                repo_id=repo_id,
                repo_type="model",
                revision=revision_name
            )
            print(f"[HF] Successfully uploaded revision '{revision_name}' to repository.")
        except Exception as e:
            print(f"[HF] Failed to upload revision '{revision_name}': {e}")

    if os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
    print(f"\n[HF] All uploads finished! View your repository at https://huggingface.co/{repo_id}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-name", type=str, default="gpt2_dense_fresh")
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--skip-eval", action="store_true", help="Skip evaluation phase after training")
    parser.add_argument("--skip-aoa", action="store_true", default=True, help="Skip AoA evaluation")
    parser.add_argument("--skip-glue", action="store_true", default=False, help="Skip GLUE fine-tuning")
    parser.add_argument("--upload", action="store_true", help="Upload model repository to Hugging Face")
    parser.add_argument("--upload-repo", type=str, default="gpt2_dense_50M_same_total_param", help="Hugging Face repository name")
    parser.add_argument("--upload-token", type=str, default=None, help="Hugging Face API token")
    args = parser.parse_args()

    if args.upload:
        upload_pipeline(args.model_name, args.upload_repo, args.upload_token)
    else:
        class Tee:
            def __init__(self, filepath, original_stream):
                self.file = open(filepath, "a", encoding="utf-8", buffering=1)
                self.original_stream = original_stream

            def write(self, data):
                self.original_stream.write(data)
                self.file.write(data)

            def flush(self):
                self.original_stream.flush()
                self.file.flush()

        import sys
        sys.stdout = Tee("training_terminal.log", sys.stdout)
        sys.stderr = Tee("training_terminal.log", sys.stderr)

        print("\n=== STARTING NEW TRAINING RUN LOGGING TO training_terminal.log ===")
        run_pipeline(args.model_name, epochs=args.epochs, skip_eval=args.skip_eval, skip_aoa=args.skip_aoa, skip_glue=args.skip_glue)