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}")
    
    # Check if predictions.json or results.txt or surprisal.json exists and is non-empty
    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 to avoid duplicate downloads
    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)

    # Programmatic Hugging Face Hub Login if HF_TOKEN is in environment
    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_xpertgpt import (
        XpertGPTModel,
        XpertGPTModelConfig,
        XpertGPTConfig
    )

    # 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          
    MASK_TOKEN_ID     = 16383          
    BLOCK_SIZE        = 512
    BATCH_SIZE        = 16
    GRAD_ACCUM_STEPS  = 1  # grad_acc_step = 1
    EPOCHS            = epochs
    LEARNING_RATE     = 3e-4  # lr = 3e-4
    LR_MIN            = LEARNING_RATE * 0.05
    WARMUP_STEPS      = 800  # warmup_steps = 800
    WEIGHT_DECAY      = 0.1
    GRAD_CLIP         = 1.0

    NUM_THIN_BLOCKS   = 4             
    EC_CAPACITY_FACTOR = 2.0           
    CAUSAL_RATIO      = 1 / 1        

    MASK_PROB_START   = 0.20
    MASK_PROB_END     = 0.10

    # 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}'...")
        
        # A. Convert state dict keys to CausalLM wrapper naming
        state_dict = raw_model.state_dict()
        new_state_dict = {}
        for k, v in state_dict.items():
            name = k
            if name.startswith("_orig_mod."):
                name = name[10:]
            if name.startswith("model."):
                name = name[6:]
            
            if name == "lm_head.weight":
                new_state_dict["lm_head.weight"] = v
            else:
                new_state_dict[f"transformer.{name}"] = v
                
        torch.save(new_state_dict, os.path.join(save_dir, "pytorch_model.bin"))
        
        # B. Copy modeling.py and configuration.py
        shutil.copy("modeling_xpertgpt.py", os.path.join(save_dir, "modeling_xpertgpt.py"))
        shutil.copy("configuration_xpertgpt.py", os.path.join(save_dir, "configuration_xpertgpt.py"))
        
        # C. Create config.json
        config_dict = {
            "auto_map": {
                "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
                "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
                "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
            },
            "vocab_size": VOCAB_SIZE,
            "block_size": BLOCK_SIZE,
            "d_model": 256,        # d_model = 256
            "hidden_size": 256,    # hidden_size = 256
            "d_thin": 384,
            "num_layers": 6,
            "num_blocks": NUM_THIN_BLOCKS,
            "capacity_factor": EC_CAPACITY_FACTOR,
            "dropout": 0.1,
            "model_type": "xpertgpt"
        }
        with open(os.path.join(save_dir, "config.json"), "w") as f:
            json.dump(config_dict, f, indent=2)
            
        # D. Save tokenizer config files
        fast_tokenizer = PreTrainedTokenizerFast(
            tokenizer_object=tokenizer,
            bos_token="[CLS]",
            eos_token="[SEP]",
            unk_token="[UNK]",
            pad_token="[PAD]",
            mask_token="[MASK]"
        )
        fast_tokenizer.save_pretrained(save_dir)
        print(f"[Checkpoint] Checkpoint '{checkpoint_dir_name}' successfully saved.")

    # ─────────────────────────────────────────────────────────────
    # 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

    # ─────────────────────────────────────────────────────────────
    # Batch preparation and schedules
    # ─────────────────────────────────────────────────────────────
    def get_current_mask_prob(global_step: int, total_steps: int) -> float:
        ratio = min(1.0, global_step / total_steps)
        return MASK_PROB_START + ratio * (MASK_PROB_END - MASK_PROB_START)

    def prepare_causal_batch(x: torch.Tensor, y: torch.Tensor):
        return x, y, False

    def prepare_masked_batch(x: torch.Tensor, y: torch.Tensor, mask_prob: float, mask_token_id: int):
        B, T = x.size()
        mask = torch.rand(B, T, device=x.device) < mask_prob
        masked_x = x.clone()
        masked_x[mask] = mask_token_id

        targets = torch.full_like(y, -100)
        targets[mask] = y[mask]

        return masked_x, targets, True

    def get_hybrid_batch(train_loader: DataLoaderLite, global_step: int, total_steps: int, device: torch.device):
        x, y = train_loader.next_batch()
        x, y = x.to(device), y.to(device)

        if random.random() < CAUSAL_RATIO:
            input_ids, targets, bidir = prepare_causal_batch(x, y)
        else:
            mask_prob = get_current_mask_prob(global_step, total_steps)
            input_ids, targets, bidir = prepare_masked_batch(x, y, mask_prob, MASK_TOKEN_ID)

        return input_ids, targets, bidir

    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'])

    tokenizer = build_and_train_tokenizer(all_text)

    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, tokenizer, "train")
    val_loader   = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, val_texts, 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   = XpertGPTModelConfig()
    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 = XpertGPTModel(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")):
            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("d_model") == 256:
                        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
                    else:
                        print(f"[Training] Found checkpoint '{ckpt_name}' but it has mismatch d_model={saved_config.get('d_model')}. Starting fresh.")
                except Exception as e:
                    pass

    # Load weights if resuming
    if resume_checkpoint_dir is not None:
        print(f"[Model] Loading weights from checkpoint '{resume_checkpoint_dir}'...")
        state_dict = torch.load(os.path.join(resume_checkpoint_dir, "pytorch_model.bin"), map_location=device)
        model_state_dict = {}
        for k, v in state_dict.items():
            name = k
            if name.startswith("transformer."):
                name = name[12:]
            model_state_dict[name] = v
        model.load_state_dict(model_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")):
        print("\n[Pipeline] Final checkpoint 'main' already exists. Skipping training phase and transitioning directly to evaluations!")
    else:
        # torch.compile
        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)

        # Helper for validation loss calculation
        def evaluate_validation_loss(model_eval, val_loader_eval, dev, autocast):
            was_training = model_eval.training
            model_eval.eval()
            
            from modeling_xpertgpt import ROUTER_TRACKER
            old_tracker_enabled = ROUTER_TRACKER["enabled"]
            ROUTER_TRACKER["enabled"] = False
            
            val_loss_accum = 0.0
            val_steps = min(val_loader_eval.steps_per_epoch(), 50)
            with torch.no_grad():
                for _ in range(val_steps):
                    x, y = val_loader_eval.next_batch()
                    x, y = x.to(dev), y.to(dev)
                    with autocast:
                        _, loss = model_eval(x, y, bidirectional=False)
                    val_loss_accum += loss.item()
            if was_training:
                model_eval.train()
                
            ROUTER_TRACKER["enabled"] = old_tracker_enabled
            return val_loss_accum / val_steps

        loss_records = []
        checkpoint_perplexities = []
        
        from modeling_xpertgpt import ROUTER_TRACKER
        ROUTER_TRACKER["enabled"] = True
        ROUTER_TRACKER["expert_counts"] = []

        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 XpertGPT MoEP 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

                input_ids, targets, bidir = get_hybrid_batch(train_loader, global_step, total_steps, device)
                words_trained += input_ids.numel()

                with autocast_ctx:
                    _, loss = model(input_ids, targets, bidirectional=bidir)
                scaled_loss  = loss / GRAD_ACCUM_STEPS
                loss_accum  += scaled_loss.item()
                scaled_loss.backward()

                # Optimizer Step
                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
                    mode_tag = "MLM" if bidir else "CLM"
                    mask_p   = get_current_mask_prob(global_step, total_steps)
                    
                    current_step = (chunk_step + 1) // GRAD_ACCUM_STEPS
                    print(
                        f"[E{epoch+1:02d} {current_step:>5d}/{steps_per_epoch} G{global_step:>7d}|{mode_tag}] "
                        f"train={loss_accum:.4f}  mask={mask_p:.1%}  norm={norm:.3f}  lr={lr:.2e}  dt={dt:6.1f}ms  words={words_trained:,}"
                    )
                    
                    # Record validation and training loss every 100 steps
                    if (global_step + 1) % 100 == 0:
                        val_loss = evaluate_validation_loss(model, val_loader, device, autocast_ctx)
                        loss_records.append({
                            "step": global_step + 1,
                            "train_loss": loss_accum,
                            "val_loss": val_loss
                        })
                        with open("loss_records.json", "w") as f:
                            json.dump(loss_records, f, indent=2)
                        if os.path.exists(model_dir):
                            with open(os.path.join(model_dir, "loss_records.json"), "w") as f:
                                json.dump(loss_records, f, indent=2)
                        print(f"[Metrics G{global_step+1}] Recorded train_loss={loss_accum:.4f}, val_loss={val_loss:.4f}")
                        
                    loss_accum = 0.0
                    global_step += 1

                # Check if we passed a milestone for checkpointing
                if next_milestone_idx < len(milestones) and words_trained >= milestones[next_milestone_idx]:
                    milestone_val = milestones[next_milestone_idx]
                    if milestone_val < 10_000_000:
                        milestone_name = f"chck_{milestone_val // 1_000_000}M"
                    else:
                        milestone_name = 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)
                    
                    # Calculate validation perplexity at this milestone
                    val_loss = evaluate_validation_loss(model, val_loader, device, autocast_ctx)
                    val_ppl = math.exp(val_loss)
                    checkpoint_perplexities.append({
                        "checkpoint": milestone_name,
                        "words_trained": milestone_val,
                        "global_step": global_step,
                        "val_loss": val_loss,
                        "val_perplexity": val_ppl
                    })
                    with open("checkpoint_perplexities.json", "w") as f:
                        json.dump(checkpoint_perplexities, f, indent=2)
                    if os.path.exists(model_dir):
                        with open(os.path.join(model_dir, "checkpoint_perplexities.json"), "w") as f:
                            json.dump(checkpoint_perplexities, f, indent=2)
                    print(f"[Milestone {milestone_name}] Evaluated val_perplexity={val_ppl:.2f}")

                    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)
        
        # Calculate validation perplexity for final checkpoint
        val_loss = evaluate_validation_loss(model, val_loader, device, autocast_ctx)
        val_ppl = math.exp(val_loss)
        checkpoint_perplexities.append({
            "checkpoint": "main",
            "words_trained": words_trained,
            "global_step": global_step,
            "val_loss": val_loss,
            "val_perplexity": val_ppl
        })
        with open("checkpoint_perplexities.json", "w") as f:
            json.dump(checkpoint_perplexities, f, indent=2)
        if os.path.exists(model_dir):
            with open(os.path.join(model_dir, "checkpoint_perplexities.json"), "w") as f:
                json.dump(checkpoint_perplexities, f, indent=2)
        print(f"[Final Checkpoint main] Evaluated val_perplexity={val_ppl:.2f}")
        
        print("\n[Training] Training phase complete!")

        # ─────────────────────────────────────────────────────────────
        # Router utilization report and plot generation
        # ─────────────────────────────────────────────────────────────
        from modeling_xpertgpt import ROUTER_TRACKER
        total_expert_counts = [0, 0, 0, 0, 0]
        for cnt in ROUTER_TRACKER["expert_counts"]:
            for i in range(5):
                total_expert_counts[i] += cnt[i]
                
        total_tokens = sum(total_expert_counts)
        if total_tokens > 0:
            fractions = [total_expert_counts[i] / total_tokens for i in range(5)]
        else:
            fractions = [0.0] * 5
            
        fraction_no_expert = fractions[0]
        fraction_multiple_experts = sum(fractions[2:])
        
        report_text = f"""=== ROUTER UTILIZATION REPORT ===
Total tokens routed: {total_tokens:,}

Number of experts selected per token (0, 1, 2, 3, 4):
- 0 experts: {total_expert_counts[0]:,} tokens ({fractions[0]:.2%})
- 1 expert:  {total_expert_counts[1]:,} tokens ({fractions[1]:.2%})
- 2 experts: {total_expert_counts[2]:,} tokens ({fractions[2]:.2%})
- 3 experts: {total_expert_counts[3]:,} tokens ({fractions[3]:.2%})
- 4 experts: {total_expert_counts[4]:,} tokens ({fractions[4]:.2%})

Summary metrics:
- Fraction of tokens receiving NO expert: {fraction_no_expert:.2%}
- Fraction of tokens receiving MULTIPLE experts: {fraction_multiple_experts:.2%}
"""
        with open("router_report.txt", "w") as f:
            f.write(report_text)
        if os.path.exists(model_dir):
            with open(os.path.join(model_dir, "router_report.txt"), "w") as f:
                f.write(report_text)
        print(report_text)

        # Generate bar chart
        try:
            import matplotlib
            matplotlib.use('Agg')
            import matplotlib.pyplot as plt
            
            categories = ['0', '1', '2', '3', '4']
            plt.figure(figsize=(8, 5))
            plt.bar(categories, [f * 100 for f in fractions], color='#10b981', edgecolor='#059669', width=0.6)
            plt.xlabel('Number of Experts Selected per Token')
            plt.ylabel('Percentage of Tokens (%)')
            plt.title('Router Utilization / Expert Choices per Token')
            plt.grid(axis='y', linestyle='--', alpha=0.6)
            
            for i, f in enumerate(fractions):
                plt.text(i, f * 100 + 1, f"{f:.2%}", ha='center', fontweight='bold')
                
            plt.ylim(0, max([f * 100 for f in fractions]) + 10)
            plt.savefig('router_utilization.png', dpi=150)
            if os.path.exists(model_dir):
                plt.savefig(os.path.join(model_dir, 'router_utilization.png'), dpi=150)
            plt.close()
            print("[Plots] Router utilization plot saved successfully.")
        except Exception as e:
            print(f"[Plots] Warning: Could not generate plots: {e}")

    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()
            
        # Check if already patched
        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)
        
        # Unzip EWoK fast
        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)
            
        # Download EWoK full
        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):
        # Determine paths
        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")

        # If cached, copy it over
        cache_file = os.path.join(local_cache_path, "predictions.json")
        
        # Self-healing: invalidate old unfiltered entity_tracking caches
        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 from results/main if needed
        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")

        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
        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)

        # Relocate from results/main/main if needed
        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
        verify_eval_run(target_results_dir, f"GLUE task {task}")

        # Cache locally
        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_aoa_with_cache(cmd):
        local_cache_path = os.path.join("./results", model_name, "main", "aoa")
        target_results_dir = os.path.join(strict_dir, "results", model_name, "main", "aoa")

        if os.path.exists(os.path.join(local_cache_path, "aoa_score.json")) or os.path.exists(os.path.join(local_cache_path, "surprisal.json")):
            print(f"[Eval] AoA 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("[Eval] Running AoA task...")
        subprocess.run(cmd, cwd=strict_dir, check=True)

        # Relocate from results/main/main if needed
        possible_main_path = os.path.join(strict_dir, "results", "main", "main", "aoa")
        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
        verify_eval_run(target_results_dir, "AoA task")

        # Cache locally
        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)

    # ─────────────────────────────────────────────────────────────
    # B. FINAL MODEL 'main' FULL ZERO-SHOT EVALUATION
    # ─────────────────────────────────────────────────────────────
    main_ckpt_path = os.path.join(model_dir, "main")
    if os.path.exists(main_ckpt_path):
        print(f"[Eval] Running full zero-shot evaluation on main...")
        # blimp filtered
        run_task_with_cache(
            "main", "blimp", "blimp_filtered",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/blimp_filtered", "--save_predictions"]
        )
        # supplement filtered
        run_task_with_cache(
            "main", "blimp", "supplement_filtered",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/supplement_filtered", "--save_predictions"]
        )
        # ewok filtered
        run_task_with_cache(
            "main", "ewok", "ewok_filtered",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "ewok", "--data_path", "evaluation_data/full_eval/ewok_filtered", "--save_predictions"]
        )
        # entity tracking
        run_task_with_cache(
            "main", "entity_tracking", "entity_tracking",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "entity_tracking", "--data_path", "evaluation_data/full_eval/entity_tracking", "--save_predictions"]
        )
        # comps
        run_task_with_cache(
            "main", "comps", "comps",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "comps", "--data_path", "evaluation_data/full_eval/comps", "--save_predictions"]
        )
        # reading
        run_task_with_cache(
            "main", "reading", "reading",
            ["python", "-m", "evaluation_pipeline.reading.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--data_path", "evaluation_data/full_eval/reading/reading_data.csv"]
        )
        # global piqa parallel
        run_task_with_cache(
            "main", "global_piqa_parallel", "global_piqa_parallel",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "global_piqa_parallel", "--data_path", "evaluation_data/full_eval/global_piqa_parallel", "--save_predictions"]
        )
        # global piqa nonparallel
        run_task_with_cache(
            "main", "global_piqa_nonparallel", "global_piqa_nonparallel",
            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "global_piqa_nonparallel", "--data_path", "evaluation_data/full_eval/global_piqa_nonparallel", "--save_predictions"]
        )

    # ─────────────────────────────────────────────────────────────
    # C. FINAL MODEL GLUE AND AOA EVALUATION
    # ─────────────────────────────────────────────────────────────
    if os.path.exists(main_ckpt_path):
        # 1. GLUE fine-tuning
        if skip_glue:
            print("[Eval] Skipping GLUE fine-tuning evaluations as requested.")
        else:
            print("[Eval] Running GLUE fine-tuning evaluations on main task-by-task...")
            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)

        # 2. AoA
        if skip_aoa:
            print("[Eval] Skipping AoA evaluations as requested.")
        else:
            run_aoa_with_cache([
                "python", "-m", "evaluation_pipeline.AoA_word.run",
                "--model_name", model_dir,
                "--backend", "causal",
                "--track_name", "strict-small",
                "--word_path", "evaluation_data/full_eval/aoa/cdi_childes.json",
                "--output_dir", "results"
            ])

    # ─────────────────────────────────────────────────────────────
    # A. 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)]
    
    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", model_dir, "--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", model_dir, "--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", model_dir, "--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", model_dir, "--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", model_dir, "--backend", "causal", "--data_path", "evaluation_data/fast_eval/reading/reading_data.csv", "--revision_name", checkpoint]
        )
        # 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", model_dir, "--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", model_dir, "--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...")
    # Clean collation destination in evaluation repo
    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)
    
    # Copy from local cache results to strict results for collation
    shutil.copytree(local_results_dir, collate_results_dir)

    # Run collation
    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)
    
    # Save results to local folder
    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)
    
    # Copy final collated json to current folder
    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):
    import os
    import shutil
    import json
    import hashlib
    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}")
    
    # Create the repository if it doesn't exist
    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 verify or create repository. Error: {e}")
        return

    checkpoint_dir = os.path.abspath(f"./checkpoints/{model_name}")
    
    # Resolve the main checkpoint directory using self-healing rules
    revisions = {}
    if os.path.exists("pytorch_model.bin") and os.path.exists("config.json"):
        print("[HF] Detected weight and config files in the current working directory. Using current folder as 'main' checkpoint.")
        revisions = {"main": os.getcwd()}
    elif os.path.exists(os.path.join(checkpoint_dir, "main")):
        revisions = {"main": os.path.join(checkpoint_dir, "main")}
    elif os.path.exists(os.path.abspath("./checkpoints/msit_gptbert_fresh/main")):
        print("[HF] Using fallback checkpoint folder './checkpoints/msit_gptbert_fresh/main'...")
        revisions = {"main": os.path.abspath("./checkpoints/msit_gptbert_fresh/main")}
    elif os.path.exists(os.path.abspath("./checkpoints/main")):
        revisions = {"main": os.path.abspath("./checkpoints/main")}
    else:
        print(f"[HF] Error: Could not locate the 'main' checkpoint weights. Checked: {checkpoint_dir}/main, current directory, and fallbacks.")
        return

    # Check for intermediate checkpoints relative to the main checkpoint's parent folder
    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
        else:
            # Check if there is a .pt file in the parent folder
            pt_path = os.path.join(parent_dir, f"{ckpt_name}.pt")
            if os.path.exists(pt_path):
                print(f"[HF] Found legacy checkpoint file '{ckpt_name}.pt'. Converting to HF format for upload...")
                import torch
                from transformers import AutoTokenizer
                from tokenizers import Tokenizer
                try:
                    from modeling_xpertgpt import XpertGPTForCausalLM, XpertGPTConfig
                    cfg = XpertGPTConfig(d_model=256, d_thin=384, num_layers=6, num_blocks=4)
                    model_to_save = XpertGPTForCausalLM(cfg)
                    sd = torch.load(pt_path, map_location="cpu")
                    clean_sd = {}
                    for k, v in sd.items():
                        new_k = k.replace("module.", "")
                        clean_sd[new_k] = v
                    model_to_save.load_state_dict(clean_sd)
                    
                    vocab_path = os.path.join(parent_dir, "bpe_vocab_16k.json")
                    if not os.path.exists(vocab_path):
                        vocab_path = os.path.join(main_dir, "bpe_vocab_16k.json")
                    if os.path.exists(vocab_path):
                        tok = Tokenizer.from_file(vocab_path)
                    else:
                        tok = None
                        
                    os.makedirs(ckpt_path, exist_ok=True)
                    state_dict = model_to_save.state_dict()
                    new_state_dict = {}
                    for k, v in state_dict.items():
                        name = k
                        if name.startswith("_orig_mod."):
                            name = name[10:]
                        if name.startswith("model."):
                            name = name[6:]
                        if name == "lm_head.weight":
                            new_state_dict["lm_head.weight"] = v
                        else:
                            new_state_dict[f"transformer.{name}"] = v
                    torch.save(new_state_dict, os.path.join(ckpt_path, "pytorch_model.bin"))
                    
                    config_dict = {
                        "auto_map": {
                            "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
                            "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
                            "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
                        },
                        "vocab_size": 16384,
                        "block_size": 512,
                        "d_model": 256,
                        "hidden_size": 256,
                        "d_thin": 384,
                        "num_layers": 6,
                        "num_blocks": 4,
                        "capacity_factor": 2.0,
                        "dropout": 0.1,
                        "model_type": "xpertgpt",
                        "num_hidden_layers": 6
                    }
                    with open(os.path.join(ckpt_path, "config.json"), "w") as f:
                        json.dump(config_dict, f, indent=2)
                        
                    if tok:
                        from transformers import PreTrainedTokenizerFast
                        fast_tokenizer = PreTrainedTokenizerFast(
                            tokenizer_object=tok,
                            bos_token="[CLS]",
                            eos_token="[SEP]",
                            unk_token="[UNK]",
                            pad_token="[PAD]",
                            mask_token="[MASK]"
                        )
                        fast_tokenizer.save_pretrained(ckpt_path)
                    
                    revisions[ckpt_name] = ckpt_path
                except Exception as ex:
                    print(f"[HF] Failed to convert legacy checkpoint file '{ckpt_name}.pt': {ex}")

    # Temporary directory for staging uploads
    temp_dir = os.path.abspath("./temp_hf_upload")
    
    # LICENSE text
    license_text = """Creative Commons Attribution-NonCommercial 4.0 International Public License

By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions.

Section 1 -- Definitions.
a. Licensed Material means the artistic or literary work, database, or other material to which the Licensor applied this Public License.
b. Licensed Rights means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license.
c. NonCommercial means not primarily intended for or directed towards commercial advantage or monetary compensation.
d. Share means to provide material to the public by any means or process that requires permission under the Licensed Rights.
e. You means the individual or entity exercising the Licensed Rights under this Public License.

Section 2 -- Scope.
a. License grant.
   1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to:
      A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and
      B. Produce, reproduce, and Share Adapted Material for NonCommercial purposes only.
   2. Attribution. As a condition of the license, You must attribute the Licensor and keep intact copyright notices.
"""

    for revision_name, local_path in revisions.items():
        print(f"\n[HF] Staging files for revision '{revision_name}' from '{local_path}'...")
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
        os.makedirs(temp_dir)

        # 1. Copy weight file and save as model.safetensors if possible, otherwise pytorch_model.bin
        src_bin = os.path.join(local_path, "pytorch_model.bin")
        weight_file_dest = None
        if os.path.exists(src_bin):
            # Try to convert to safetensors
            try:
                import torch
                from safetensors.torch import save_file
                state_dict = torch.load(src_bin, map_location="cpu")
                # Clone tensors to break memory sharing (prevents shared weight memory error in safetensors)
                state_dict = {k: v.clone() for k, v in state_dict.items()}
                weight_file_dest = os.path.join(temp_dir, "model.safetensors")
                save_file(state_dict, weight_file_dest)
                print(f"[HF] Converted weights to safetensors format.")
            except Exception as e:
                print(f"[HF] Conversion to safetensors failed ({e}). Staging raw pytorch_model.bin...")
                weight_file_dest = os.path.join(temp_dir, "pytorch_model.bin")
                shutil.copy2(src_bin, weight_file_dest)
        else:
            print(f"[HF] Error: No weight file found in '{local_path}'!")
            continue

        # Calculate weight file hash
        sha256_hash = hashlib.sha256()
        with open(weight_file_dest, "rb") as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256_hash.update(byte_block)
        weight_hash = sha256_hash.hexdigest()
        print(f"[HF] Weight file SHA-256: {weight_hash}")

        # 2. Copy and patch config.json
        src_config = os.path.join(local_path, "config.json")
        if os.path.exists(src_config):
            with open(src_config, "r") as f:
                cfg_data = json.load(f)
            # Patch config
            cfg_data["auto_map"] = {
                "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
                "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
                "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
            }
            cfg_data["architectures"] = ["XpertGPTForCausalLM"]
            with open(os.path.join(temp_dir, "config.json"), "w") as f:
                json.dump(cfg_data, f, indent=2)
        else:
            # Fallback configuration
            cfg_data = {
                "auto_map": {
                    "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
                    "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
                    "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
                },
                "architectures": ["XpertGPTForCausalLM"],
                "vocab_size": 16384,
                "block_size": 512,
                "d_model": 256,
                "hidden_size": 256,
                "d_thin": 384,
                "num_layers": 6,
                "num_blocks": 4,
                "capacity_factor": 2.0,
                "dropout": 0.1,
                "model_type": "xpertgpt"
            }
            with open(os.path.join(temp_dir, "config.json"), "w") as f:
                json.dump(cfg_data, f, indent=2)

        # 3. Copy tokenizers
        for tok_file in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]:
            src_tok = os.path.join(local_path, tok_file)
            if os.path.exists(src_tok):
                shutil.copy2(src_tok, os.path.join(temp_dir, tok_file))

        # 4. Copy custom code
        shutil.copy2("modeling_xpertgpt.py", os.path.join(temp_dir, "modeling_xpertgpt.py"))
        shutil.copy2("configuration_xpertgpt.py", os.path.join(temp_dir, "configuration_xpertgpt.py"))

        # 5. Write metadata files
        with open(os.path.join(temp_dir, ".gitattributes"), "w") as f:
            f.write("*.safetensors filter=lfs diff=lfs merge=lfs -text\n")
            f.write("*.bin filter=lfs diff=lfs merge=lfs -text\n")

        with open(os.path.join(temp_dir, "LICENSE"), "w") as f:
            f.write(license_text)

        with open(os.path.join(temp_dir, "CITATION.cff"), "w") as f:
            citation_yaml = f"""cff-version: 1.2.0
message: "If you use this model or software, please cite it as below."
authors:
  - family-names: "Jain"
    given-names: "Soham"
  - family-names: "Singh"
    given-names: "Harsh"
  - family-names: "Dewan"
    given-names: "Divija"
  - family-names: "Dev"
    given-names: "Atul"
title: "XpertGPT: Mixture of Experts with Parallelized Multi-Scale Information Transmission for Data-Constrained Pretraining"
year: 2026
url: "https://huggingface.co/{repo_id}"
"""
            f.write(citation_yaml)

        with open(os.path.join(temp_dir, "PROVENANCE.md"), "w") as f:
            provenance_md = f"""# Provenance and Weight Integrity Record

This file records the provenance, cryptographic hash, and reproducibility metadata of the model weights.

## Verification Fingerprints
- **Model Weight File**: {"model.safetensors" if weight_file_dest.endswith(".safetensors") else "pytorch_model.bin"}
- **Weight SHA-256**: {weight_hash}
- **Tokenizer Vocab Size**: 16,384

## Training Run Details
- **Training Word Budget**: 10M words (BabyLM 2026 Strict-Small track)
- **Model Parameters**: ~48.4M non-embedding parameters / ~52.6M total tied parameters
- **Optimizer**: AdamW
- **Epochs**: 8
"""
            f.write(provenance_md)

        # Build README
        readme_md = f"""---
license: cc-by-nc-4.0
language:
  - en
tags:
  - babylm
  - babylm-2026
  - mixture-of-experts
  - msit
  - xpertgpt
  - custom_code
  - safetensors
library_name: transformers
pipeline_tag: text-generation
---

# XpertGPT Strict-Small

XpertGPT (Mixture of Experts with Parallelized Multi-Scale Information Transmission) is a sparse, data-efficient recurrent language model for the BabyLM 2026 challenge (Strict-Small (10M) track, 10M words). It combines sliding window attention global streams with sparse parallel expert blocks using Expert Choice routing. ~48.4M non-embedding parameters / ~52.6M total parameters. Custom code (`trust_remote_code=True`).

- **Architecture:** 6 layers of MoEP-MSIT blocks. Each block combines a lower-dimensional dense global sliding window attention layer (`dim = 256`) with 4 parallel high-dimensional sparse expert blocks (`dim = 384`) routed via Expert Choice gating.
- **Track:** BabyLM 2026 Strict-Small (10M) (10M words).
- **Tokenizer:** Custom BPE tokenizer (vocab size: 16384).
- **Revision / Checkpoint:** {revision_name}

## Usage

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("{repo_id}", revision="{revision_name}", trust_remote_code=True).eval()
tok = AutoTokenizer.from_pretrained("{repo_id}", revision="{revision_name}")
ids = tok("The quick brown fox", return_tensors="pt").input_ids
with torch.no_grad():
    logits = model(ids).logits
```

## Intermediate checkpoints

Intermediate training checkpoints are provided as git revisions named `chck_<N>M` for the BabyLM challenge fast-eval.

## License and citation

Released under CC BY-NC 4.0 (attribution required, non-commercial only). If you use this model or code, please cite (see `CITATION.cff`):

```bibtex
@misc{{jain2026xpertgpt,
  title        = {{XpertGPT: Mixture of Experts with Parallelized Multi-Scale Information Transmission for Data-Constrained Pretraining}},
  author       = {{Jain, Soham and Singh, Harsh and Dewan, Divija and Dev, Atul}},
  year         = {{2026}},
  howpublished = {{Hugging Face Repository}},
  note         = {{XpertGPT MoE language model, BabyLM 2026}}
}}
```

Provenance and integrity fingerprints are documented in `PROVENANCE.md`.
"""
        with open(os.path.join(temp_dir, "README.md"), "w") as f:
            f.write(readme_md)

        # 6. For main branch only: also upload collated predictions
        if revision_name == "main":
            for pred_file in ["all_full_preds_and_fast_scores_causal.json", "all_full_preds_and_fast_scores_causal (3).json"]:
                if os.path.exists(pred_file):
                    shutil.copy2(pred_file, os.path.join(temp_dir, "all_full_preds_and_fast_scores_causal.json"))
                    print(f"[HF] Copied predictions file '{pred_file}' to staging area.")
                    break

        # 7. Create branch if it does not exist, then upload staged files to HF under the revision
        if revision_name != "main":
            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}")

    # Cleanup temp dir
    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="xpertgpt_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="XpertGPT-BabyLM2026-Strict-Small", 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:
        run_pipeline(args.model_name, epochs=args.epochs, skip_eval=args.skip_eval, skip_aoa=args.skip_aoa, skip_glue=args.skip_glue)