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 run_pipeline(model_name: str, epochs: int = 10, skip_eval: bool = False, skip_aoa: bool = False, 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_msit import (
        MSITGPTBERTModel,
        MSITGPTBERTConfig,
        MSITGPTBERTHFConfig
    )

    # 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 as requested
    EPOCHS            = epochs
    LEARNING_RATE     = 1e-4
    LR_MIN            = LEARNING_RATE * 0.05
    WARMUP_STEPS      = 50           
    WEIGHT_DECAY      = 0.1
    GRAD_CLIP         = 1.0

    NUM_THIN_BLOCKS   = 6             
    EC_CAPACITY_FACTOR = 2.0           
    CAUSAL_RATIO      = 1 / 1        

    MASK_PROB_START   = 0.20
    MASK_PROB_END     = 0.10

    # Output directories locally (persistent in Studio workspace)
    model_dir = os.path.abspath(f"./checkpoints/{model_name}")
    os.makedirs(model_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
        shutil.copy("modeling_msit.py", os.path.join(save_dir, "modeling_msit.py"))
        
        # C. Create config.json
        config_dict = {
            "auto_map": {
                "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig",
                "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper",
                "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM"
            },
            "vocab_size": VOCAB_SIZE,
            "block_size": BLOCK_SIZE,
            "d_model": 384,        # 384 d_model as requested
            "hidden_size": 384,    # 384 hidden_size
            "d_thin": 192,
            "num_layers": 6,
            "num_blocks": NUM_THIN_BLOCKS,
            "capacity_factor": EC_CAPACITY_FACTOR,
            "dropout": 0.1,
            "model_type": "msit_gptbert"
        }
        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")

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

    cfg   = MSITGPTBERTConfig()  # Will use updated d_model = 384
    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 = MSITGPTBERTModel(cfg).to(device)

    # ─────────────────────────────────────────────────────────────
    # Training Resume Check
    # ─────────────────────────────────────────────────────────────
    words_trained = 0
    next_milestone_idx = 0
    global_step = 0
    
    milestones = [i * 1_000_000 for i in range(1, 10)] + [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)
                    # Verify d_model matches to avoid weight shape mismatch crashes
                    if saved_config.get("d_model") == 384:
                        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)

        # ─────────────────────────────────────────────────────────────
        # Training Loop
        # ─────────────────────────────────────────────────────────────
        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 MSIT-GPTBERT 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:,}"
                    )
                    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)
                    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)
        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")

    # Delete invalid local caches before we start harvesting/evaluating
    if not is_zero_shot_cache_valid(os.path.join(local_main_res, "zero_shot")):
        if os.path.exists(os.path.join(local_main_res, "zero_shot")):
            print("[Eval] Local zero-shot cache is incomplete or corrupt. Cleaning up...")
            shutil.rmtree(os.path.join(local_main_res, "zero_shot"))
            
    if not is_finetune_cache_valid(os.path.join(local_main_res, "finetune")):
        if os.path.exists(os.path.join(local_main_res, "finetune")):
            print("[Eval] Local finetuning cache is incomplete or corrupt. Cleaning up...")
            shutil.rmtree(os.path.join(local_main_res, "finetune"))

    # Walk the entire workspace to search for completed evaluations
    workspace_root = "/teamspace/studios/this_studio"
    print(f"[Eval] Scanning workspace {workspace_root} to harvest any completed evaluations...")
    exclude_dirs = {"miniconda3", ".git", ".cache", "hf_cache", "nltk_data", "babylm_eval_repo"}
    
    if os.path.exists(workspace_root):
        for root, dirs, files in os.walk(workspace_root):
            # Speed up walking by pruning system dirs
            dirs[:] = [d for d in dirs if d not in exclude_dirs]
            
            parts = root.split(os.sep)
            # Skip checkpoint paths
            if any(p.startswith("chck_") for p in parts):
                continue
                
            # If we find a directory ending in zero_shot
            if root.endswith(f"{os.sep}zero_shot"):
                if is_zero_shot_cache_valid(root):
                    dest = os.path.join(local_main_res, "zero_shot")
                    if not is_zero_shot_cache_valid(dest):
                        if os.path.exists(dest):
                            shutil.rmtree(dest)
                        os.makedirs(os.path.dirname(dest), exist_ok=True)
                        shutil.copytree(root, dest)
                        print(f"[Eval] Successfully harvested zero-shot results from {root}")
                        
            # If we find a directory ending in finetune
            if root.endswith(f"{os.sep}finetune"):
                if is_finetune_cache_valid(root):
                    dest = os.path.join(local_main_res, "finetune")
                    if not is_finetune_cache_valid(dest):
                        if os.path.exists(dest):
                            shutil.rmtree(dest)
                        os.makedirs(os.path.dirname(dest), exist_ok=True)
                        shutil.copytree(root, dest)
                        print(f"[Eval] Successfully harvested GLUE finetuning results from {root}")

        # Also harvest intermediate checkpoints from the old strict dir if it exists
        old_strict_dir = os.path.abspath("./babylm_eval_repo/babylm-eval/strict")
        old_model_results = os.path.join(old_strict_dir, "results", model_name)
        if os.path.exists(old_model_results):
            for item in os.listdir(old_model_results):
                src = os.path.join(old_model_results, item)
                dst = os.path.join(local_results_dir, item)
                if os.path.isdir(src) and item != "main" and (not os.path.exists(dst) or not os.listdir(dst)):
                    shutil.copytree(src, dst)
                    print(f"[Eval] Successfully harvested checkpoint '{item}' results.")

    clone_dir = os.path.abspath("./babylm_eval_repo")
    if os.path.exists(clone_dir):
        shutil.rmtree(clone_dir)
        
    print("\n[Eval] Cloning evaluation pipeline repository...")
    subprocess.run(["git", "clone", "https://github.com/atulgithub2/babylm.git", clone_dir], check=True)
    
    strict_dir = os.path.join(clone_dir, "babylm-eval", "strict")
    
    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] Skipping pipeline dependencies installation (as requested)...")
    
    print("[Eval] Downloading NLTK tokenizer resources...")
    import nltk
    nltk.download('punkt')
    nltk.download('punkt_tab')
    
    print("[Eval] Downloading zero-shot evaluation datasets...")
    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)
    
    # 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)

    # ─────────────────────────────────────────────────────────────
    # A. INTERMEDIATE CHECKPOINTS FAST EVALUATION (First - Normal Flow)
    # ─────────────────────────────────────────────────────────────
    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:
        local_ckpt_res = os.path.join(local_results_dir, checkpoint)
        
        # Caching check: if results exist locally, skip evaluations for this checkpoint!
        if os.path.exists(local_ckpt_res) and os.listdir(local_ckpt_res):
            print(f"[Eval] Checkpoint '{checkpoint}' already evaluated. Restoring results from cache...")
            target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint)
            if os.path.exists(target_results_dir):
                shutil.rmtree(target_results_dir)
            os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
            shutil.copytree(local_ckpt_res, target_results_dir)
        else:
            print(f"[Eval] Evaluating checkpoint '{checkpoint}'...")
            subprocess.run([
                "bash", "scripts/eval_zero_shot_fast.sh",
                model_dir,
                checkpoint,
                "causal",
                "evaluation_data/fast_eval"
            ], cwd=strict_dir, check=True)
            
            # Immediately cache the result
            target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint)
            if os.path.exists(target_results_dir):
                print(f"[Eval] Caching '{checkpoint}' results...")
                if os.path.exists(local_ckpt_res):
                    shutil.rmtree(local_ckpt_res)
                shutil.copytree(target_results_dir, local_ckpt_res)

    # ─────────────────────────────────────────────────────────────
    # B. FINAL MODEL 'main' EVALUATION (Second)
    # ─────────────────────────────────────────────────────────────
    main_ckpt_path = os.path.join(model_dir, "main")
    
    if os.path.exists(main_ckpt_path):
        # 1. Zero-shot check
        zero_shot_cache_file = os.path.join(local_main_res, "zero_shot")
        if is_zero_shot_cache_valid(zero_shot_cache_file):
            print("[Eval] Full zero-shot evaluation on main already completed and cached. Restoring to all expected targets...")
            # Restore to all possible directories the pipeline might seek
            targets = [
                os.path.join(strict_dir, "results", "main", "main", "zero_shot"),
                os.path.join(strict_dir, "results", "main", "zero_shot"),
                os.path.join(strict_dir, "results", model_name, "main", "zero_shot"),
            ]
            for t_zs in targets:
                if os.path.exists(t_zs):
                    shutil.rmtree(t_zs)
                os.makedirs(os.path.dirname(t_zs), exist_ok=True)
                shutil.copytree(zero_shot_cache_file, t_zs)
        else:
            print(f"[Eval] Running full zero-shot evaluation on main...")
            subprocess.run([
                "./scripts/eval_zero_shot.sh",
                main_ckpt_path,
                "causal",
                "evaluation_data/full_eval"
            ], cwd=strict_dir, check=True)
            
            # Find generated zero-shot directory and cache it
            found_zs = None
            possible_paths = [
                os.path.join(strict_dir, "results", "main", "main", "zero_shot"),
                os.path.join(strict_dir, "results", "main", "zero_shot"),
                os.path.join(strict_dir, "results", model_name, "main", "zero_shot"),
            ]
            for p in possible_paths:
                if os.path.exists(p) and os.listdir(p):
                    found_zs = p
                    break
            if found_zs:
                os.makedirs(local_main_res, exist_ok=True)
                dest = os.path.join(local_main_res, "zero_shot")
                if os.path.exists(dest):
                    shutil.rmtree(dest)
                shutil.copytree(found_zs, dest)
                print(f"[Eval] Successfully cached zero-shot results from {found_zs}")

        # 2. GLUE Fine-Tuning check
        if skip_glue:
            print("[Eval] Skipping GLUE fine-tuning evaluations as requested by --skip-glue.")
        else:
            finetune_cache_file = os.path.join(local_main_res, "finetune")
            if is_finetune_cache_valid(finetune_cache_file):
                print("[Eval] GLUE fine-tuning on main already completed and cached. Restoring to all expected targets...")
                # Restore to all possible directories
                targets = [
                    os.path.join(strict_dir, "results", "main", "main", "finetune"),
                    os.path.join(strict_dir, "results", "main", "finetune"),
                    os.path.join(strict_dir, "results", model_name, "main", "finetune"),
                ]
                for t_ft in targets:
                    if os.path.exists(t_ft):
                        shutil.rmtree(t_ft)
                    os.makedirs(os.path.dirname(t_ft), exist_ok=True)
                    shutil.copytree(finetune_cache_file, t_ft)
            else:
                print(f"[Eval] Running GLUE fine-tuning evaluations on main (enforces batch_size=32 to prevent MoE OOM)...")
                subprocess.run([
                    "./scripts/eval_finetuning.sh",
                    "--model_path", main_ckpt_path,
                    "--lr", "3e-5",
                    "--bsz", "32"
                ], cwd=strict_dir, check=True)
                
                # Find generated finetuning directory and cache it
                found_ft = None
                possible_paths = [
                    os.path.join(strict_dir, "results", "main", "main", "finetune"),
                    os.path.join(strict_dir, "results", "main", "finetune"),
                    os.path.join(strict_dir, "results", model_name, "main", "finetune"),
                ]
                for p in possible_paths:
                    if os.path.exists(p) and os.listdir(p):
                        found_ft = p
                        break
                if found_ft:
                    os.makedirs(local_main_res, exist_ok=True)
                    dest = os.path.join(local_main_res, "finetune")
                    if os.path.exists(dest):
                        shutil.rmtree(dest)
                    shutil.copytree(found_ft, dest)
                    print(f"[Eval] Successfully cached finetuning results from {found_ft}")

        # 3. AoA check
        if skip_aoa:
            print("[Eval] Skipping AoA evaluations as requested by --skip-aoa.")
        else:
            aoa_cache_file = os.path.join(local_main_res, "aoa")
            if os.path.exists(aoa_cache_file) and os.listdir(aoa_cache_file):
                print("[Eval] AoA evaluations already completed and cached. Restoring...")
                target_main_aoa = os.path.join(strict_dir, "results", model_name, "main", "aoa")
                if os.path.exists(target_main_aoa):
                    shutil.rmtree(target_main_aoa)
                os.makedirs(os.path.dirname(target_main_aoa), exist_ok=True)
                shutil.copytree(aoa_cache_file, target_main_aoa)
            else:
                print(f"[Eval] Running AoA metrics evaluations...")
                subprocess.run([
                    "./scripts/eval_aoa.sh",
                    model_dir,
                    "causal",
                    "strict-small"
                ], cwd=strict_dir, check=True)
                
                # Cache it
                target_main_aoa = os.path.join(strict_dir, "results", model_name, "main", "aoa")
                if os.path.exists(target_main_aoa):
                    os.makedirs(local_main_res, exist_ok=True)
                    dest = os.path.join(local_main_res, "aoa")
                    if os.path.exists(dest):
                        shutil.rmtree(dest)
                    shutil.copytree(target_main_aoa, dest)

    # ─────────────────────────────────────────────────────────────
    # C. COLLATE RESULTS
    # ─────────────────────────────────────────────────────────────
    print("[Eval] Collating predictions into submission file...")
    print("[Debug] Listing local cache folder:")
    if os.path.exists(local_results_dir):
        for root, dirs, files in os.walk(local_results_dir):
            rel_path = os.path.relpath(root, local_results_dir)
            print(f"  {rel_path}: {files}")
    else:
        print("  Local results dir does not exist!")

    for item in os.listdir(local_results_dir):
        src = os.path.join(local_results_dir, item)
        dst = os.path.join(strict_dir, "results", model_name, item)
        if os.path.isdir(src):
            if os.path.exists(dst):
                shutil.rmtree(dst)
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            shutil.copytree(src, dst)

    print("[Debug] Listing strict results folder before collation:")
    strict_results_dir = os.path.join(strict_dir, "results")
    if os.path.exists(strict_results_dir):
        for root, dirs, files in os.walk(strict_results_dir):
            rel_path = os.path.relpath(root, strict_results_dir)
            print(f"  {rel_path}: {files}")
    else:
        print("  Strict results dir does not exist!")

    subprocess.run([
        "bash", "scripts/collate_preds.sh",
        model_name, "causal", "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.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-name", type=str, default="msit_gptbert_model")
    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", help="Skip AoA evaluation")
    parser.add_argument("--skip-glue", action="store_true", help="Skip GLUE fine-tuning")
    args = parser.parse_args()
    run_pipeline(args.model_name, epochs=args.epochs, skip_eval=args.skip_eval, skip_aoa=args.skip_aoa, skip_glue=args.skip_glue)