Initial commit of GPT-2 46M SwiGLU Dense baseline model architecture and configs

configuration_gpt2.py

@@ -0,0 +1,40 @@
+from transformers import PretrainedConfig
+
+class GPT2CustomConfig(PretrainedConfig):
+    model_type = "gpt2_custom"
+    auto_map = {
+        "AutoConfig": "configuration_gpt2.GPT2CustomConfig",
+        "AutoModelForCausalLM": "modeling_gpt2.GPT2CustomLMHeadModel"
+    }
+    
+    def __init__(
+        self,
+        vocab_size=16384,
+        n_positions=512,
+        n_embd=384,
+        n_layer=12,
+        n_head=6,
+        n_inner=1280,
+        activation_function="swiglu",
+        resid_pdrop=0.1,
+        embd_pdrop=0.1,
+        attn_pdrop=0.1,
+        layer_norm_epsilon=1e-5,
+        initializer_range=0.02,
+        bos_token_id=2,
+        eos_token_id=3,
+        **kwargs
+    ):
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+        self.vocab_size = vocab_size
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_inner = n_inner
+        self.activation_function = activation_function
+        self.resid_pdrop = resid_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.attn_pdrop = attn_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range

generate_plots_and_report.py

@@ -0,0 +1,208 @@
+import os
+import json
+import math
+import torch
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate training reports, learning curves, and router plots.")
+    parser.add_argument("--model-dir", type=str, default="./checkpoints/gpt2_dense_31M", help="Path to checkpoints directory")
+    args = parser.parse_args()
+
+    print("=== 1. PARAMETER COUNT & MODEL METRICS REPORT ===")
+    model_path = os.path.join(args.model_dir, "main")
+    weight_file = os.path.join(model_path, "pytorch_model.bin")
+    
+    if os.path.exists(weight_file):
+        try:
+            state_dict = torch.load(weight_file, map_location="cpu")
+            total_params = 0
+            non_emb_params = 0
+            for name, param in state_dict.items():
+                size = param.numel()
+                total_params += size
+                # Exclude embedding and head tie
+                if "wte.weight" not in name and "lm_head.weight" not in name:
+                    non_emb_params += size
+            print(f"Model weights loaded successfully from: {weight_file}")
+            print(f"- Total Parameters (with tied embeddings): {total_params:,}")
+            print(f"- Non-Embedding Parameters: {non_emb_params:,}")
+        except Exception as e:
+            print(f"Warning: Could not read parameter counts from weight file: {e}")
+            # Fallback to config math
+            fallback_param_math()
+    else:
+        print(f"Info: Checked {weight_file} but weights not found. Using fallback config computation:")
+        fallback_param_math()
+
+    # Load validation metrics
+    val_file = "validation_metrics.json"
+    if not os.path.exists(val_file):
+        val_file = os.path.join(args.model_dir, "validation_metrics.json")
+        
+    if os.path.exists(val_file):
+        with open(val_file, "r") as f:
+            val_logs = json.load(f)
+        if len(val_logs) > 0:
+            final_val = val_logs[-1]
+            print(f"- Final Validation Step: {final_val['step']}")
+            print(f"- Final Validation Loss: {final_val['val_loss']:.4f}")
+            print(f"- Final Validation Perplexity: {final_val['val_perplexity']:.2f}")
+    else:
+        print(f"Warning: {val_file} not found! Cannot report validation perplexity.")
+
+    print("\n=== 2. GENERATING LEARNING CURVES ===")
+    train_file = "training_metrics.json"
+    if not os.path.exists(train_file):
+        train_file = os.path.join(args.model_dir, "training_metrics.json")
+
+    try:
+        import matplotlib
+        matplotlib.use('Agg')
+        import matplotlib.pyplot as plt
+        
+        has_plots = False
+        
+        # Setup plot grid
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
+        
+        if os.path.exists(train_file):
+            with open(train_file, "r") as f:
+                train_logs = json.load(f)
+            steps = [log["step"] for log in train_logs]
+            losses = [log["train_loss"] for log in train_logs]
+            ax1.plot(steps, losses, color="#ef4444", linewidth=1.5, label="Train Loss")
+            ax1.set_xlabel("Global Step")
+            ax1.set_ylabel("Cross Entropy Loss")
+            ax1.set_title("Training Loss Curve")
+            ax1.grid(True, linestyle="--", alpha=0.6)
+            ax1.legend()
+            has_plots = True
+        else:
+            ax1.text(0.5, 0.5, "training_metrics.json not found", ha="center")
+            
+        if os.path.exists(val_file):
+            with open(val_file, "r") as f:
+                val_logs = json.load(f)
+            steps = [log["step"] for log in val_logs]
+            ppl = [log["val_perplexity"] for log in val_logs]
+            ax2.plot(steps, ppl, marker="o", color="#3b82f6", linewidth=2, label="Val Perplexity")
+            ax2.set_xlabel("Global Step")
+            ax2.set_ylabel("Perplexity")
+            ax2.set_title("Validation Perplexity Curve")
+            ax2.grid(True, linestyle="--", alpha=0.6)
+            ax2.legend()
+            has_plots = True
+        else:
+            ax2.text(0.5, 0.5, "validation_metrics.json not found", ha="center")
+
+        if has_plots:
+            plt.tight_layout()
+            plt.savefig("learning_curves.png", dpi=150)
+            plt.close()
+            print("Successfully saved learning curve diagrams to 'learning_curves.png'.")
+    except Exception as e:
+        print(f"Warning: Could not plot learning curves: {e}")
+
+    print("\n=== 3. ROUTER UTILIZATION ANALYSIS ===")
+    router_details_file = "router_details.json"
+    if not os.path.exists(router_details_file):
+        router_details_file = os.path.join(args.model_dir, "router_details.json")
+        
+    if os.path.exists(router_details_file):
+        try:
+            with open(router_details_file, "r") as f:
+                records = json.load(f)
+                
+            print(f"Loaded router mapping records containing {len(records)} logged steps.")
+            
+            # Reconstruct selection count mapping
+            # records format: [{"step": s, "layer": l, "expert_assignments": [[tokens for exp1], [exp2], ...]}]
+            # Find the total number of unique token indices in each step
+            total_expert_counts = [0, 0, 0, 0, 0] # indices representing 0, 1, 2, 3, 4 experts
+            
+            for record in records:
+                assignments = record["expert_assignments"]
+                # Number of experts is 4
+                num_blocks = len(assignments)
+                # Find maximum token index present in assignments to know the token range
+                max_token_idx = 0
+                for exp_list in assignments:
+                    if len(exp_list) > 0:
+                        max_token_idx = max(max_token_idx, max(exp_list))
+                
+                n_tokens = max_token_idx + 1
+                token_counts = [0] * n_tokens
+                for exp_list in assignments:
+                    for token_idx in exp_list:
+                        token_counts[token_idx] += 1
+                        
+                for count in token_counts:
+                    if count < 5:
+                        total_expert_counts[count] += 1
+                        
+            total_tokens_sum = sum(total_expert_counts)
+            if total_tokens_sum > 0:
+                fractions = [total_expert_counts[i] / total_tokens_sum 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 tracked: {total_tokens_sum:,}
+
+Number of experts selected per token:
+- 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 (2, 3, or 4): {fraction_multiple_experts:.2%}
+"""
+            with open("router_report.txt", "w") as f:
+                f.write(report_text)
+            print(report_text)
+            
+            # Plot utilization 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)
+                plt.close()
+                print("Successfully saved router utilization plot to 'router_utilization.png'.")
+            except Exception as plot_err:
+                print(f"Warning: Could not plot router utilization: {plot_err}")
+                
+        except Exception as e:
+            print(f"Warning: Could not analyze router assignments file: {e}")
+    else:
+        print(f"Warning: {router_details_file} not found! Cannot analyze router utilization.")
+
+def fallback_param_math():
+    # Fallback mathematical projection computation based on model config
+    print("Fallback configuration architecture math:")
+    print("- Config: 12 layers, d_model=384, d_inner=1280, 6 heads, tied embeddings, SwiGLU")
+    print("- Non-Embedding parameter count is approximately ~24.8M parameters.")
+    print("- Total parameter count is approximately ~31.1M parameters.")
+
+if __name__ == "__main__":
+    main()

modeling_gpt2.py

@@ -0,0 +1,225 @@
+import math
+from typing import Optional, Tuple, Dict, Any
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, GenerationMixin, AutoConfig, AutoModelForCausalLM
+from transformers import PretrainedConfig
+
+class GPT2CustomConfig(PretrainedConfig):
+    model_type = "gpt2_custom"
+    auto_map = {
+        "AutoConfig": "configuration_gpt2.GPT2CustomConfig",
+        "AutoModelForCausalLM": "modeling_gpt2.GPT2CustomLMHeadModel"
+    }
+    
+    def __init__(
+        self,
+        vocab_size=16384,
+        n_positions=512,
+        n_embd=384,
+        n_layer=12,
+        n_head=6,
+        n_inner=1280,
+        activation_function="gelu_new",
+        resid_pdrop=0.1,
+        embd_pdrop=0.1,
+        attn_pdrop=0.1,
+        layer_norm_epsilon=1e-5,
+        initializer_range=0.02,
+        bos_token_id=2,
+        eos_token_id=3,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_inner = n_inner
+        self.activation_function = activation_function
+        self.resid_pdrop = resid_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.attn_pdrop = attn_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+        self.hidden_size = n_embd
+        self.num_attention_heads = n_head
+        self.num_hidden_layers = n_layer
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1)
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1)
+
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+
+    return (x * cos) + (rotated_x * sin)
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        self.num_heads = config.n_head
+        self.head_dim = config.n_embd // config.n_head
+        
+        self.q_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+        self.k_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+        self.v_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+        self.o_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+        
+        self.attn_dropout = nn.Dropout(config.attn_pdrop)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+
+    def forward(self, x, past_kv=None, use_cache: bool = False):
+        B, L, D = x.size()
+        
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        # Standard full context causal attention mask (no sliding window)
+        past_len = L_kv - L
+        pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+        pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+        causal_mask = (pos_i - pos_j) >= 0
+
+        scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+
+        attn = torch.softmax(scores, dim=-1)
+        attn = self.attn_dropout(attn)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.resid_dropout(self.o_proj(out))
+
+        present_kv = (k, v) if use_cache else None
+        return out, present_kv
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        self.fc1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc2 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc3 = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, x):
+        return self.fc3(F.silu(self.fc1(x)) * self.fc2(x))
+
+class GPT2Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.mlp = SwiGLU(config.n_embd, config.n_inner)
+
+    def forward(self, x, past_kv=None, use_cache: bool = False):
+        attn_out, present_kv = self.attn(self.ln_1(x), past_kv=past_kv, use_cache=use_cache)
+        x = x + attn_out
+        x = x + self.mlp(self.ln_2(x))
+        return x, present_kv
+
+class GPT2CustomLMHeadModel(PreTrainedModel, GenerationMixin):
+    config_class = GPT2CustomConfig
+    base_model_prefix = "transformer"
+    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            drop = nn.Dropout(config.embd_pdrop),
+            h = nn.ModuleList([GPT2Block(config) for _ in range(config.n_layer)]),
+            ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+
+    def forward(self, input_ids, labels=None, past_key_values=None, use_cache: bool = False, **kwargs):
+        device = input_ids.device
+        
+        x = self.transformer.wte(input_ids)
+        x = self.transformer.drop(x)
+        
+        new_kvs = []
+        for i, block in enumerate(self.transformer.h):
+            pkv = past_key_values[i] if past_key_values is not None else None
+            x, nkv = block(x, past_kv=pkv, use_cache=use_cache)
+            if use_cache:
+                new_kvs.append(nkv)
+            
+        x = self.transformer.ln_f(x)
+        logits = self.lm_head(x)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100)
+            
+        present_kvs = tuple(new_kvs) if use_cache else None
+        from transformers.modeling_outputs import CausalLMOutputWithPast
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=present_kvs
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        return {"input_ids": input_ids, "past_key_values": past_key_values}
+
+# Register configuration and model for auto mapping
+AutoConfig.register("gpt2_custom", GPT2CustomConfig)
+AutoModelForCausalLM.register(GPT2CustomConfig, GPT2CustomLMHeadModel)

train.py

@@ -0,0 +1,1086 @@
+import os
+import math
+import time
+import json
+import random
+import inspect
+import shutil
+import subprocess
+import argparse
+
+# ─────────────────────────────────────────────────────────────
+# 1. CORE PIPELINE FUNCTION AND HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def is_zero_shot_cache_valid(cache_dir):
+    if not os.path.exists(cache_dir):
+        return False
+    pred_count = 0
+    for root, dirs, files in os.walk(cache_dir):
+        if "predictions.json" in files:
+            pred_count += 1
+    return pred_count >= 3
+
+def is_finetune_cache_valid(cache_dir):
+    if not os.path.exists(cache_dir):
+        return False
+    pred_count = 0
+    for root, dirs, files in os.walk(cache_dir):
+        if "predictions.json" in files:
+            pred_count += 1
+    return pred_count >= 3
+
+def verify_eval_run(path_to_check, description):
+    print(f"[Eval] Verifying {description} path: {path_to_check}...")
+    if not os.path.exists(path_to_check):
+        raise RuntimeError(f"CRITICAL ERROR: {description} directory was NOT created at: {path_to_check}")
+    
+    found_valid = False
+    for root, dirs, files in os.walk(path_to_check):
+        for f in files:
+            if f in ["predictions.json", "results.txt", "surprisal.json"]:
+                file_path = os.path.join(root, f)
+                if os.path.getsize(file_path) > 0:
+                    found_valid = True
+                    break
+        if found_valid:
+            break
+            
+    if not found_valid:
+        raise RuntimeError(f"CRITICAL ERROR: {description} completed but no valid prediction/result files were written in: {path_to_check}")
+    print(f"[Eval] Success! Verified {description} results are stored correctly.")
+
+def run_pipeline(model_name: str, epochs: int = 10, skip_eval: bool = False, skip_aoa: bool = True, skip_glue: bool = False):
+    # Configure persistent cache paths locally
+    os.environ["HF_HOME"] = os.path.abspath("./hf_cache")
+    os.environ["NLTK_DATA"] = os.path.abspath("./nltk_data")
+    os.makedirs("./hf_cache", exist_ok=True)
+    os.makedirs("./nltk_data", exist_ok=True)
+
+    hf_token = os.environ.get("HF_TOKEN")
+    if hf_token:
+        try:
+            from huggingface_hub import login
+            login(token=hf_token)
+            print("[HF] Programmatic login successful using HF_TOKEN.")
+        except Exception as e:
+            print(f"[HF] Warning: Programmatic login failed: {e}")
+
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    from datasets import load_dataset
+    from tokenizers import Tokenizer
+    from transformers import PreTrainedTokenizerFast
+    
+    # Import model architecture
+    from modeling_gpt2 import GPT2CustomLMHeadModel
+    from configuration_gpt2 import GPT2CustomConfig
+
+    # Print GPU details
+    if torch.cuda.is_available():
+        gpu_name = torch.cuda.get_device_name(0)
+        print(f"\n[GPU] CUDA is available! Using GPU: {gpu_name}\n")
+    else:
+        print("\n[GPU] Warning: CUDA is NOT available! Running on CPU.\n")
+
+    # ─────────────────────────────────────────────────────────────
+    # GLOBAL HYPERPARAMETERS
+    # ─────────────────────────────────────────────────────────────
+    VOCAB_SIZE        = 16384          
+    BLOCK_SIZE        = 512
+    BATCH_SIZE        = 16
+    GRAD_ACCUM_STEPS  = 1  
+    EPOCHS            = epochs
+    LEARNING_RATE     = 3e-4  
+    LR_MIN            = LEARNING_RATE * 0.05
+    WARMUP_STEPS      = 800  
+    WEIGHT_DECAY      = 0.1
+    GRAD_CLIP         = 1.0
+
+    # Output directories locally
+    model_dir = os.path.abspath(f"./checkpoints/{model_name}")
+    os.makedirs(model_dir, exist_ok=True)
+    local_results_dir = os.path.abspath(f"./results/{model_name}")
+    os.makedirs(local_results_dir, exist_ok=True)
+
+    # ─────────────────────────────────────────────────────────────
+    # Helper: Save Hugging Face Compliant Checkpoint
+    # ─────────────────────────────────────────────────────────────
+    def save_hf_checkpoint(raw_model, checkpoint_dir_name, tokenizer):
+        save_dir = os.path.join(model_dir, checkpoint_dir_name)
+        os.makedirs(save_dir, exist_ok=True)
+        print(f"\n[Checkpoint] Saving Hugging Face format checkpoint to '{save_dir}'...")
+        
+        # Save standard HF model format (bin & configs) for Hub upload
+        raw_model.save_pretrained(save_dir)
+        tokenizer.save_pretrained(save_dir)
+        
+        # Copy modeling.py and configuration.py
+        shutil.copy("modeling_gpt2.py", os.path.join(save_dir, "modeling_gpt2.py"))
+        shutil.copy("configuration_gpt2.py", os.path.join(save_dir, "configuration_gpt2.py"))
+
+    # ─────────────────────────────────────────────────────────────
+    # Tokenizer Training
+    # ─────────────────────────────────────────────────────────────
+    def build_and_train_tokenizer(texts: list) -> Tokenizer:
+        from tokenizers.models import BPE
+        from tokenizers.trainers import BpeTrainer
+        from tokenizers.pre_tokenizers import Whitespace
+
+        vocab_path = os.path.join(model_dir, "bpe_vocab_16k.json")
+        if os.path.exists(vocab_path):
+            print(f"[Tokenizer] Loading trained BPE model layout from '{vocab_path}'...")
+            return Tokenizer.from_file(vocab_path)
+
+        print(f"[Tokenizer] Generating fresh HuggingFace BPE Tokenizer model with {VOCAB_SIZE} slots...")
+        tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
+        tokenizer.pre_tokenizer = Whitespace()
+
+        trainer = BpeTrainer(
+            vocab_size=VOCAB_SIZE,
+            special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
+        )
+        tokenizer.train_from_iterator(texts, trainer)
+        tokenizer.save(vocab_path)
+        print(f"[Tokenizer] Tokenizer training completed and saved to '{vocab_path}'.")
+        return tokenizer
+
+    # ─────────────────────────────────────────────────────────────
+    # Data Loader Setup
+    # ─────────────────────────────────────────────────────────────
+    class DataLoaderLite:
+        def __init__(self, B: int, T: int, texts: list, tokenizer: Tokenizer, name: str):
+            self.B = B
+            self.T = T
+
+            print(f"[DataLoader:{name}] Tokenising dataset sequences...")
+            all_ids = []
+            for t in texts:
+                if t.strip():
+                    encoded = tokenizer.encode(t).ids
+                    all_ids.extend(encoded)
+
+            self.tokens = torch.tensor(all_ids, dtype=torch.long)
+            self.chunk_size = B * T
+            self.n_chunks = (len(self.tokens) - 1) // self.chunk_size
+            self.indices = list(range(self.n_chunks))
+            self.pos = 0
+            self._shuffle()
+
+            print(f"[DataLoader:{name}] Total tokens: {len(self.tokens):,} | Epoch steps: {self.n_chunks:,}")
+
+        def _shuffle(self):
+            random.shuffle(self.indices)
+            self.pos = 0
+
+        def steps_per_epoch(self) -> int:
+            return self.n_chunks
+
+        def next_batch(self):
+            B, T = self.B, self.T
+            if self.pos >= len(self.indices):
+                self._shuffle()
+
+            chunk_idx = self.indices[self.pos]
+            self.pos += 1
+
+            start_pos = chunk_idx * self.chunk_size
+            temp = self.tokens[start_pos : start_pos + self.chunk_size + 1]
+
+            x = temp[:-1].view(B, T)
+            y = temp[1:].view(B, T)
+            return x, y
+
+    # Learning Rate Cosine Schedule
+    def get_lr(it: int, total_steps: int) -> float:
+        if it < WARMUP_STEPS:
+            return LEARNING_RATE * (it + 1) / WARMUP_STEPS
+        if it >= total_steps:
+            return LR_MIN
+        decay_ratio = (it - WARMUP_STEPS) / (total_steps - WARMUP_STEPS)
+        coeff        = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+        return LR_MIN + coeff * (LEARNING_RATE - LR_MIN)
+
+    # ─────────────────────────────────────────────────────────────
+    # Dataset Preparation
+    # ─────────────────────────────────────────────────────────────
+    print("\n[Data] Loading BabyLM-2026-Strict-Small ...")
+    ds = load_dataset("BabyLM-community/BabyLM-2026-Strict-Small")
+    all_text = list(ds['train']['text'])
+
+    raw_tokenizer = build_and_train_tokenizer(all_text)
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_object=raw_tokenizer,
+        bos_token="[CLS]",
+        eos_token="[SEP]",
+        unk_token="[UNK]",
+        pad_token="[PAD]",
+        mask_token="[MASK]"
+    )
+    BOS_ID = 2
+    EOS_ID = 3
+
+    split       = int(len(all_text) * 0.95)
+    train_texts = all_text[:split]
+    val_texts   = all_text[split:]
+
+    train_loader = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, train_texts, raw_tokenizer, "train")
+    val_loader   = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, val_texts, raw_tokenizer, "val")
+
+    chunks_per_epoch = train_loader.steps_per_epoch()
+    steps_per_epoch  = chunks_per_epoch // GRAD_ACCUM_STEPS 
+    total_steps      = steps_per_epoch * EPOCHS
+
+    cfg = GPT2CustomConfig(
+        vocab_size=VOCAB_SIZE,
+        n_positions=BLOCK_SIZE,
+        n_embd=384,
+        n_layer=12,
+        n_head=6,
+        n_inner=1280,
+        bos_token_id=BOS_ID,
+        eos_token_id=EOS_ID,
+    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.manual_seed(42)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    random.seed(42)
+    if hasattr(torch, 'set_float32_matmul_precision'):
+        torch.set_float32_matmul_precision('high')
+
+    model = GPT2CustomLMHeadModel(cfg).to(device)
+
+    # ─────────────────────────────────────────────────────────────
+    # Training Resume Check
+    # ─────────────────────────────────────────────────────────────
+    words_trained = 0
+    next_milestone_idx = 0
+    global_step = 0
+    
+    milestones = sorted(list(set([i * 1_000_000 for i in range(1, 11)] + [i * 10_000_000 for i in range(1, 11)])))
+    
+    resume_checkpoint_dir = None
+    for idx in range(len(milestones) - 1, -1, -1):
+        m = milestones[idx]
+        ckpt_name = f"chck_{m // 1_000_000}M"
+        ckpt_path = os.path.join(model_dir, ckpt_name)
+        if os.path.exists(os.path.join(ckpt_path, "pytorch_model.bin")) or os.path.exists(os.path.join(ckpt_path, "model.safetensors")):
+            config_json_path = os.path.join(ckpt_path, "config.json")
+            if os.path.exists(config_json_path):
+                try:
+                    with open(config_json_path, "r") as f:
+                        saved_config = json.load(f)
+                    if saved_config.get("n_embd") == 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
+                except Exception:
+                    pass
+
+    # Load weights if resuming
+    if resume_checkpoint_dir is not None:
+        print(f"[Model] Loading weights from checkpoint '{resume_checkpoint_dir}'...")
+        if os.path.exists(os.path.join(resume_checkpoint_dir, "pytorch_model.bin")):
+            state_dict = torch.load(os.path.join(resume_checkpoint_dir, "pytorch_model.bin"), map_location=device)
+        else:
+            from safetensors.torch import load_file
+            state_dict = load_file(os.path.join(resume_checkpoint_dir, "model.safetensors"), device=device)
+        model.load_state_dict(state_dict)
+
+    # Check if final main model exists
+    main_ckpt_path = os.path.join(model_dir, "main")
+    if os.path.exists(os.path.join(main_ckpt_path, "pytorch_model.bin")) or os.path.exists(os.path.join(main_ckpt_path, "model.safetensors")):
+        print("\n[Pipeline] Final checkpoint 'main' already exists. Skipping training phase and transitioning directly to evaluations!")
+    else:
+        try:
+            model = torch.compile(model)
+            print("[Model] torch.compile() successfully verified graph optimizations")
+        except Exception as e:
+            print(f"[Model] torch.compile() skipped ({e})")
+
+        # Optimizer
+        param_dict     = {n: p for n, p in model.named_parameters() if p.requires_grad}
+        decay_params   = [p for p in param_dict.values() if p.dim() >= 2]
+        nodecay_params = [p for p in param_dict.values() if p.dim() < 2]
+        groups = [
+            {'params': decay_params,   'weight_decay': WEIGHT_DECAY},
+            {'params': nodecay_params, 'weight_decay': 0.0},
+        ]
+        fused_ok  = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_ok and ('cuda' in device)
+        optimizer = torch.optim.AdamW(groups, lr=LEARNING_RATE, betas=(0.9, 0.95), eps=1e-8, fused=use_fused)
+
+        # ─────────────────────────────────────────────────────────────
+        # Training Loop
+        # ─────────────────────────────────────────────────────────────
+        def evaluate_validation_loss(model_eval, val_loader_eval, dev, autocast):
+            was_training = model_eval.training
+            model_eval.eval()
+            
+            val_loss_accum = 0.0
+            val_steps = val_loader_eval.steps_per_epoch()
+            max_val_steps = min(val_steps, 100)
+            
+            with torch.no_grad():
+                for _ in range(max_val_steps):
+                    x, y = val_loader_eval.next_batch()
+                    x, y = x.to(dev), y.to(dev)
+                    with autocast:
+                        out = model_eval(x, labels=y)
+                        loss = out.loss
+                    val_loss_accum += loss.item()
+            
+            avg_val_loss = val_loss_accum / max_val_steps
+            val_perplexity = math.exp(avg_val_loss)
+            
+            if was_training:
+                model_eval.train()
+            return avg_val_loss, val_perplexity
+
+        validation_logs = []
+        training_logs = []
+
+        model.train()
+        autocast_ctx = torch.autocast(device_type="cuda" if "cuda" in device else "cpu", dtype=torch.bfloat16, enabled=True)
+
+        start_epoch = global_step // steps_per_epoch
+        start_chunk = (global_step % steps_per_epoch) * GRAD_ACCUM_STEPS
+
+        print(f"\n[Training] Starting GPT-2 Dense Baseline training for {EPOCHS} epochs...")
+        for epoch in range(start_epoch, EPOCHS):
+            train_loader._shuffle()
+            if epoch == start_epoch and start_chunk > 0:
+                print(f"[Training] Fast-forwarding dataloader to chunk index {start_chunk}...")
+                train_loader.pos = start_chunk
+
+            optimizer.zero_grad(set_to_none=True)
+            loss_accum = 0.0
+            
+            start_chunk_idx = start_chunk if epoch == start_epoch else 0
+            for chunk_step in range(start_chunk_idx, chunks_per_epoch):
+                t0 = time.perf_counter()
+
+                lr = get_lr(global_step, total_steps)
+                for pg in optimizer.param_groups:
+                    pg['lr'] = lr
+
+                x, y = train_loader.next_batch()
+                input_ids, targets = x.to(device), y.to(device)
+                words_trained += input_ids.numel()
+
+                with autocast_ctx:
+                    out = model(input_ids, labels=targets)
+                    loss = out.loss
+                scaled_loss  = loss / GRAD_ACCUM_STEPS
+                loss_accum  += scaled_loss.item()
+                scaled_loss.backward()
+
+                if (chunk_step + 1) % GRAD_ACCUM_STEPS == 0:
+                    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
+                    optimizer.step()
+                    optimizer.zero_grad(set_to_none=True)
+                    if "cuda" in device:
+                        torch.cuda.synchronize()
+
+                    dt = (time.perf_counter() - t0) * 1000
+                    
+                    current_step = (chunk_step + 1) // GRAD_ACCUM_STEPS
+                    print(
+                        f"[E{epoch+1:02d} {current_step:>5d}/{steps_per_epoch} G{global_step:>7d}] "
+                        f"train={loss_accum:.4f}  norm={norm:.3f}  lr={lr:.2e}  dt={dt:6.1f}ms  words={words_trained:,}"
+                    )
+
+                    training_logs.append({
+                        "step": global_step,
+                        "train_loss": loss_accum,
+                        "lr": lr,
+                        "words": words_trained
+                    })
+                    loss_accum = 0.0
+                    global_step += 1
+
+                # Validation Loss every 100 steps
+                if global_step > 0 and global_step % 100 == 0:
+                    v_loss, v_ppl = evaluate_validation_loss(model, val_loader, device, autocast_ctx)
+                    validation_logs.append({
+                        "step": global_step,
+                        "val_loss": v_loss,
+                        "val_perplexity": v_ppl
+                    })
+                    with open(os.path.join(model_dir, "validation_metrics.json"), "w") as f:
+                        json.dump(validation_logs, f, indent=2)
+                    with open("validation_metrics.json", "w") as f:
+                        json.dump(validation_logs, f, indent=2)
+                    print(f"\n[Validation Step {global_step}] Loss: {v_loss:.4f} | Perplexity: {v_ppl:.2f}\n")
+
+                # Checkpoint Milestones
+                if next_milestone_idx < len(milestones) and words_trained >= milestones[next_milestone_idx]:
+                    milestone_val = milestones[next_milestone_idx]
+                    milestone_name = f"chck_{milestone_val // 1_000_000}M" if milestone_val < 10_000_000 else f"chck_{(milestone_val // 10_000_000) * 10}M"
+                    raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
+                    save_hf_checkpoint(raw_model, milestone_name, tokenizer)
+                    next_milestone_idx += 1
+
+        # Save final model as 'main'
+        raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
+        save_hf_checkpoint(raw_model, "main", tokenizer)
+        
+        # Save training logs
+        with open(os.path.join(model_dir, "training_metrics.json"), "w") as f:
+            json.dump(training_logs, f, indent=2)
+        with open("training_metrics.json", "w") as f:
+            json.dump(training_logs, f, indent=2)
+
+        # Generate plots
+        try:
+            import matplotlib
+            matplotlib.use('Agg')
+            import matplotlib.pyplot as plt
+
+            if len(validation_logs) > 0:
+                steps = [log["step"] for log in validation_logs]
+                perplexities = [log["val_perplexity"] for log in validation_logs]
+
+                plt.figure(figsize=(8, 5))
+                plt.plot(steps, perplexities, marker='o', color='#3b82f6', linewidth=2)
+                plt.xlabel('Global Step')
+                plt.ylabel('Validation Perplexity')
+                plt.title('Validation Perplexity Learning Curve')
+                plt.grid(True, linestyle='--', alpha=0.6)
+                plt.savefig(os.path.join(model_dir, 'validation_perplexity.png'), dpi=150)
+                plt.savefig('validation_perplexity.png', dpi=150)
+                plt.close()
+                print("[Plots] Validation perplexity curve saved successfully.")
+        except Exception as e:
+            print(f"[Plots] Warning: Could not generate plots: {e}")
+
+        print("\n[Training] Training phase complete!")
+
+    if skip_eval:
+        print("[Pipeline] Skipping evaluations phase as requested.")
+        return
+
+    # ─────────────────────────────────────────────────────────────
+    # 2. RUN EVALUATION PIPELINE
+    # ─────────────────────────────────────────────────────────────
+    local_results_dir = os.path.abspath(f"./results/{model_name}")
+    os.makedirs(local_results_dir, exist_ok=True)
+    local_main_res = os.path.join(local_results_dir, "main")
+
+    # Ensure clone_dir exists and has the global_piqa files
+    clone_parent = os.path.abspath("./babylm_eval_repo")
+    
+    # Self-healing check for global_piqa presence
+    has_global_piqa = False
+    for potential_strict in [os.path.join(clone_parent, "babylm-eval", "strict"), os.path.join(clone_parent, "strict")]:
+        if os.path.exists(os.path.join(potential_strict, "evaluation_pipeline", "global_piqa")):
+            has_global_piqa = True
+            break
+            
+    if not has_global_piqa:
+        print("[Eval] Cloned repository does not contain global_piqa tasks.")
+        print("[Eval] Deleting and cloning official main branch...")
+        if os.path.exists(clone_parent):
+            shutil.rmtree(clone_parent)
+        subprocess.run([
+            "git", "clone", "-b", "main",
+            "https://github.com/babylm-org/babylm-eval.git",
+            clone_parent
+        ], check=True)
+
+    # Determine strict_dir path dynamically
+    if os.path.exists(os.path.join(clone_parent, "strict")):
+        strict_dir = os.path.join(clone_parent, "strict")
+    else:
+        strict_dir = os.path.join(clone_parent, "babylm-eval", "strict")
+        
+    print(f"[Eval] Using strict directory: {strict_dir}")
+    os.environ["PYTHONPATH"] = strict_dir
+
+    def patch_evaluation_run_script(strict_dir):
+        import pathlib
+        run_file = os.path.join(strict_dir, "evaluation_pipeline", "sentence_zero_shot", "run.py")
+        if not os.path.exists(run_file):
+            print(f"[GlobalPIQA] Warning: {run_file} not found. Cannot patch.")
+            return
+            
+        print(f"[GlobalPIQA] Patching local checkpoint loader in {run_file}...")
+        with open(run_file, "r") as f:
+            content = f.read()
+            
+        if "Local checkpoint directory patch" in content:
+            print("[GlobalPIQA] Script already patched.")
+            return
+            
+        target_str = """def main():
+    args = _parse_arguments()
+    if args.images_path is not None:
+        assert args.batch_size == 1, "Multimodal only works in batch size 1!"
+    dataset = args.data_path.stem
+    args.model_name = pathlib.Path(args.model_path_or_name).stem
+    if args.revision_name is None:
+        revision_name = "main"
+    else:
+        revision_name = args.revision_name"""
+            
+        patch_str = """def main():
+    args = _parse_arguments()
+    if args.images_path is not None:
+        assert args.batch_size == 1, "Multimodal only works in batch size 1!"
+    dataset = args.data_path.stem
+    
+    # Local checkpoint directory patch
+    import os
+    model_path = args.model_path_or_name
+    args.model_name = pathlib.Path(model_path).stem
+    revision_name = args.revision_name if args.revision_name else "main"
+    
+    if os.path.isdir(model_path):
+        target_revision = args.revision_name if args.revision_name else "main"
+        if os.path.exists(os.path.join(model_path, target_revision)):
+            args.model_path_or_name = os.path.join(model_path, target_revision)
+            args.revision_name = None"""
+                
+        if target_str in content:
+            new_content = content.replace(target_str, patch_str)
+            with open(run_file, "w") as f:
+                f.write(new_content)
+            print("[GlobalPIQA] Successfully patched run.py")
+        else:
+            print("[GlobalPIQA] Warning: Could not find target pattern in run.py. Manual patch may be needed.")
+
+    # Patch sentence zero shot loader inside cloned repo
+    patch_evaluation_run_script(strict_dir)
+
+    print("[Eval] Stripping Windows-specific packages from requirements.txt...")
+    req_file_path = os.path.join(strict_dir, "requirements.txt")
+    if os.path.exists(req_file_path):
+        with open(req_file_path, "r") as f:
+            lines = f.readlines()
+        with open(req_file_path, "w") as f:
+            for line in lines:
+                if "pywin" not in line.lower() and "wintypes" not in line.lower():
+                    f.write(line)
+
+    print("[Eval] Verifying and installing evaluation dependencies programmatically...")
+    required_packages = {
+        "nltk": "nltk",
+        "pandas": "pandas",
+        "statsmodels": "statsmodels",
+        "sklearn": "scikit-learn",
+        "scipy": "scipy"
+    }
+    for pkg_import, pkg_install in required_packages.items():
+        try:
+            __import__(pkg_import)
+        except ImportError:
+            print(f"[Eval] Package '{pkg_install}' not found. Installing it programmatically...")
+            import sys
+            subprocess.run([sys.executable, "-m", "pip", "install", pkg_install], check=True)
+    
+    print("[Eval] Downloading NLTK tokenizer resources...")
+    import nltk
+    nltk.download('punkt', download_dir=os.environ["NLTK_DATA"])
+    nltk.download('punkt_tab', download_dir=os.environ["NLTK_DATA"])
+
+    # Ensure standard zero-shot datasets are downloaded
+    blimp_fast_dir = os.path.join(strict_dir, "evaluation_data", "fast_eval", "blimp_fast")
+    if not os.path.exists(blimp_fast_dir) or not os.listdir(blimp_fast_dir):
+        print("[Eval] Standard zero-shot datasets not found. Downloading...")
+        subprocess.run(["python", "-m", "scripts.download_evals"], cwd=strict_dir, check=True)
+        
+        ewok_zip = os.path.join(strict_dir, "evaluation_data/fast_eval/ewok_fast.zip")
+        if os.path.exists(ewok_zip):
+            print("[Eval] Unzipping EWoK fast data...")
+            bad_nested_dir = os.path.join(strict_dir, "evaluation_data/fast_eval/evaluation_data")
+            if os.path.exists(bad_nested_dir):
+                shutil.rmtree(bad_nested_dir)
+            subprocess.run(["unzip", "-o", "-P", "BabyLM2025", "evaluation_data/fast_eval/ewok_fast.zip", "-d", "."], cwd=strict_dir, check=True)
+            
+        print("[Eval] Downloading and filtering full EWoK dataset...")
+        subprocess.run(["python", "-m", "evaluation_pipeline.ewok.dl_and_filter"], cwd=strict_dir, check=True)
+
+    # Download GlobalPIQA dataset
+    global_piqa_parallel_dir = os.path.join(strict_dir, "evaluation_data", "fast_eval", "global_piqa_parallel")
+    if not os.path.exists(global_piqa_parallel_dir) or not os.listdir(global_piqa_parallel_dir):
+        print("[Eval] GlobalPIQA dataset not found. Downloading...")
+        subprocess.run(["python", "evaluation_pipeline/global_piqa/dl.py"], cwd=strict_dir, check=True)
+
+    # Ensure all scripts are executable
+    print("[Eval] Making evaluation shell scripts executable...")
+    subprocess.run("chmod +x scripts/*.sh", shell=True, cwd=strict_dir, check=True)
+
+    def run_task_with_cache(checkpoint, task, output_subpath, cmd):
+        local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", task, output_subpath)
+        if task == "reading":
+            local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", "reading")
+        elif task == "comps":
+            local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", "comps", "comps")
+            
+        target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", task, output_subpath)
+        if task == "reading":
+            target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", "reading")
+        elif task == "comps":
+            target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", "comps", "comps")
+
+        cache_file = os.path.join(local_cache_path, "predictions.json")
+        
+        # Self-healing for entity_tracking cache
+        if task == "entity_tracking" and os.path.exists(cache_file):
+            try:
+                import json
+                with open(cache_file, "r") as f:
+                    preds = json.load(f)
+                is_valid_cache = True
+                for k, v in preds.items():
+                    if len(v.get("predictions", [])) in [605, 606, 607, 615, 529, 156, 187, 159]:
+                        is_valid_cache = False
+                        break
+                if not is_valid_cache:
+                    print(f"[Eval] Cached entity_tracking for '{checkpoint}' has incorrect old sizes. Invalidate and re-run fresh...")
+                    shutil.rmtree(local_cache_path, ignore_errors=True)
+            except Exception:
+                pass
+
+        if os.path.exists(cache_file):
+            print(f"[Eval] Task '{task}' ({output_subpath}) for checkpoint '{checkpoint}' is cached. Restoring...")
+            if os.path.exists(target_results_dir):
+                shutil.rmtree(target_results_dir)
+            os.makedirs(target_results_dir, exist_ok=True)
+            for item in os.listdir(local_cache_path):
+                s = os.path.join(local_cache_path, item)
+                d = os.path.join(target_results_dir, item)
+                if os.path.isdir(s):
+                    shutil.copytree(s, d)
+                else:
+                    shutil.copy2(s, d)
+            return
+
+        print(f"[Eval] Running task '{task}' ({output_subpath}) for checkpoint '{checkpoint}'...")
+        subprocess.run(cmd, cwd=strict_dir, check=True)
+
+        # Relocate results
+        actual_model_basename = os.path.basename(model_dir.rstrip("/"))
+        possible_actual_path = os.path.join(strict_dir, "results", actual_model_basename, checkpoint, "zero_shot", "causal", task, output_subpath)
+        if task == "reading":
+            possible_actual_path = os.path.join(strict_dir, "results", actual_model_basename, checkpoint, "zero_shot", "causal", "reading")
+        elif task == "comps":
+            possible_actual_path = os.path.join(strict_dir, "results", actual_model_basename, checkpoint, "zero_shot", "causal", "comps", "comps")
+
+        possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", task, output_subpath)
+        if task == "reading":
+            possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", "reading")
+        elif task == "comps":
+            possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", "comps", "comps")
+
+        possible_local_reading_path = os.path.join(strict_dir, "results", checkpoint, "main", "zero_shot", "causal", "reading")
+
+        for p_path in [possible_actual_path, possible_main_path, possible_local_reading_path]:
+            if os.path.exists(p_path) and p_path != target_results_dir:
+                print(f"[Eval] Relocating results from {p_path} to {target_results_dir}...")
+                if os.path.exists(target_results_dir):
+                    shutil.rmtree(target_results_dir)
+                os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
+                shutil.move(p_path, target_results_dir)
+                break
+
+        verify_eval_run(target_results_dir, f"{checkpoint} {task} ({output_subpath})")
+
+        # Save to local cache
+        if os.path.exists(local_cache_path):
+            shutil.rmtree(local_cache_path)
+        os.makedirs(local_cache_path, exist_ok=True)
+        for item in os.listdir(target_results_dir):
+            s = os.path.join(target_results_dir, item)
+            d = os.path.join(local_cache_path, item)
+            if os.path.isdir(s):
+                shutil.copytree(s, d)
+            else:
+                shutil.copy2(s, d)
+
+    def run_finetune_task_with_cache(task, cmd):
+        local_cache_path = os.path.join("./results", model_name, "main", "finetune", task)
+        target_results_dir = os.path.join(strict_dir, "results", model_name, "main", "finetune", task)
+
+        if os.path.exists(os.path.join(local_cache_path, "predictions.json")):
+            print(f"[Eval] GLUE task '{task}' is cached. Restoring...")
+            if os.path.exists(target_results_dir):
+                shutil.rmtree(target_results_dir)
+            os.makedirs(target_results_dir, exist_ok=True)
+            for item in os.listdir(local_cache_path):
+                s = os.path.join(local_cache_path, item)
+                d = os.path.join(target_results_dir, item)
+                if os.path.isdir(s):
+                    shutil.copytree(s, d)
+                else:
+                    shutil.copy2(s, d)
+            return
+
+        print(f"[Eval] Running GLUE task '{task}'...")
+        subprocess.run(cmd, cwd=strict_dir, check=True)
+
+        possible_main_path = os.path.join(strict_dir, "results", "main", "main", "finetune", task)
+        if os.path.exists(possible_main_path) and possible_main_path != target_results_dir:
+            print(f"[Eval] Relocating results from {possible_main_path} to {target_results_dir}...")
+            if os.path.exists(target_results_dir):
+                shutil.rmtree(target_results_dir)
+            os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
+            shutil.move(possible_main_path, target_results_dir)
+
+        verify_eval_run(target_results_dir, f"GLUE fine-tuning {task}")
+
+        if os.path.exists(local_cache_path):
+            shutil.rmtree(local_cache_path)
+        os.makedirs(local_cache_path, exist_ok=True)
+        for item in os.listdir(target_results_dir):
+            s = os.path.join(target_results_dir, item)
+            d = os.path.join(local_cache_path, item)
+            if os.path.isdir(s):
+                shutil.copytree(s, d)
+            else:
+                shutil.copy2(s, d)
+
+    # ─────────────────────────────────────────────────────────────
+    # A. ZERO-SHOT EVALUATIONS
+    # ─────────────────────────────────────────────────────────────
+    eval_checkpoints = ["main"] + [f"chck_{m}M" for m in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 30, 35, 40, 45, 50, 60, 70, 80, 90, 100]]
+    existing_eval_checkpoints = []
+    for checkpoint in eval_checkpoints:
+        checkpoint_path = os.path.join(model_dir, checkpoint)
+        if os.path.exists(os.path.join(checkpoint_path, "pytorch_model.bin")) or os.path.exists(os.path.join(checkpoint_path, "model.safetensors")):
+            existing_eval_checkpoints.append(checkpoint)
+
+    print(f"\n[Eval] Found {len(existing_eval_checkpoints)} checkpoints ready for zero-shot evaluation: {existing_eval_checkpoints}")
+
+    # Self-healing: Patch model_type and auto_map in existing checkpoints' config.json
+    # and copy code files if missing.
+    for checkpoint in existing_eval_checkpoints:
+        checkpoint_path = os.path.join(model_dir, checkpoint)
+        
+        # A. Overwrite modeling_gpt2.py and configuration_gpt2.py unconditionally
+        for code_file in ["modeling_gpt2.py", "configuration_gpt2.py"]:
+            dest_code_file = os.path.join(checkpoint_path, code_file)
+            if os.path.exists(code_file):
+                try:
+                    shutil.copy(code_file, dest_code_file)
+                    print(f"[Self-Healing] Copied and updated {code_file} in checkpoint '{checkpoint}' directory.")
+                except Exception as e:
+                    print(f"[Self-Healing] Warning: Failed to copy {code_file} to {checkpoint}: {e}")
+                    
+        # B. Patch config.json
+        config_path = os.path.join(checkpoint_path, "config.json")
+        if os.path.exists(config_path):
+            try:
+                with open(config_path, "r") as f:
+                    config_dict = json.load(f)
+                
+                # Check if it needs patching
+                if config_dict.get("model_type") != "gpt2_custom" or "auto_map" not in config_dict:
+                    config_dict["model_type"] = "gpt2_custom"
+                    config_dict["auto_map"] = {
+                        "AutoConfig": "configuration_gpt2.GPT2CustomConfig",
+                        "AutoModelForCausalLM": "modeling_gpt2.GPT2CustomLMHeadModel"
+                    }
+                    with open(config_path, "w") as f:
+                        json.dump(config_dict, f, indent=2)
+                    print(f"[Self-Healing] Successfully patched 'auto_map' and 'model_type' in checkpoint '{checkpoint}' config.json.")
+            except Exception as e:
+                print(f"[Self-Healing] Warning: Failed to patch {config_path}: {e}")
+
+    for checkpoint in existing_eval_checkpoints:
+        eval_model_path = os.path.join(model_dir, checkpoint)
+        print(f"\n{'='*65}\n  RUNNING EVALUATIONS FOR CHECKPOINT: {checkpoint}\n{'='*65}")
+
+        # blimp filtered
+        run_task_with_cache(
+            checkpoint, "blimp", "blimp_filtered",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/blimp_filtered", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # supplement filtered
+        run_task_with_cache(
+            checkpoint, "blimp", "supplement_filtered",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/supplement_filtered", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # ewok filtered
+        run_task_with_cache(
+            checkpoint, "ewok", "ewok_filtered",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "ewok", "--data_path", "evaluation_data/full_eval/ewok_filtered", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # entity tracking
+        run_task_with_cache(
+            checkpoint, "entity_tracking", "entity_tracking",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "entity_tracking", "--data_path", "evaluation_data/full_eval/entity_tracking", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # comps
+        run_task_with_cache(
+            checkpoint, "comps", "comps",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "comps", "--data_path", "evaluation_data/full_eval/comps", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # reading
+        run_task_with_cache(
+            checkpoint, "reading", "reading",
+            ["python", "-m", "evaluation_pipeline.reading.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--data_path", "evaluation_data/full_eval/reading/reading_data.csv"]
+        )
+        # global piqa parallel
+        run_task_with_cache(
+            checkpoint, "global_piqa_parallel", "global_piqa_parallel",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "global_piqa_parallel", "--data_path", "evaluation_data/full_eval/global_piqa_parallel", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # global piqa nonparallel
+        run_task_with_cache(
+            checkpoint, "global_piqa_nonparallel", "global_piqa_nonparallel",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "global_piqa_nonparallel", "--data_path", "evaluation_data/full_eval/global_piqa_nonparallel", "--save_predictions", "--revision_name", checkpoint]
+        )
+
+    # ─────────────────────────────────────────────────────────────
+    # B. GLUE FINE-TUNING EVALUATIONS
+    # ─────────────────────────────────────────────────────────────
+    if not skip_glue:
+        main_ckpt_path = os.path.join(model_dir, "main")
+        if os.path.exists(main_ckpt_path):
+            print(f"\n{'='*65}\n  RUNNING GLUE FINE-TUNING FOR CHECKPOINT: main\n{'='*65}")
+            
+            glue_tasks = {
+                "boolq": ["boolq", "16", "10"],
+                "multirc": ["multirc", "16", "10"],
+                "rte": ["rte", "32", "10"],
+                "wsc": ["wsc", "32", "30"],
+                "mrpc": ["mrpc", "32", "10"],
+                "qqp": ["qqp", "32", "10"],
+                "mnli": ["mnli", "32", "10"]
+            }
+            for task_name, (task, bsz, max_epochs) in glue_tasks.items():
+                num_labels = "3" if task == "mnli" else "2"
+                metric_for_valid = "accuracy"
+                if task in ["mrpc", "qqp"]:
+                    metric_for_valid = "f1"
+                metrics = ["accuracy"]
+                if task != "mnli":
+                    metrics = ["accuracy", "f1", "mcc"]
+
+                cmd = [
+                    "python", "-m", "evaluation_pipeline.finetune.run",
+                    "--model_name_or_path", main_ckpt_path,
+                    "--train_data", f"evaluation_data/full_eval/glue_filtered/{task}.train.jsonl",
+                    "--valid_data", f"evaluation_data/full_eval/glue_filtered/{task}.valid.jsonl",
+                    "--predict_data", f"evaluation_data/full_eval/glue_filtered/{task}.valid.jsonl",
+                    "--task", task,
+                    "--num_labels", num_labels,
+                    "--batch_size", bsz,
+                    "--learning_rate", "3e-5",
+                    "--num_epochs", max_epochs,
+                    "--sequence_length", "512",
+                    "--results_dir", "results",
+                    "--save",
+                    "--save_dir", "models",
+                    "--metric_for_valid", metric_for_valid,
+                    "--seed", "42",
+                    "--verbose",
+                    "--padding_side", "left",
+                    "--take_final"
+                ]
+                cmd.append("--metrics")
+                cmd.extend(metrics)
+
+                run_finetune_task_with_cache(task_name, cmd)
+
+    # ─────────────────────────────────────────────────────────────
+    # C. INTERMEDIATE CHECKPOINTS FAST EVALUATION
+    # ─────────────────────────────────────────────────────────────
+    print(f"[Eval] Running zero-shot fast evaluations on intermediate checkpoints...")
+    checkpoints = [f"chck_{i}M" for i in range(1, 10)] + [f"chck_{i}M" for i in range(10, 110, 10)]
+
+    eval_model_path = model_dir
+
+    for checkpoint in checkpoints:
+        ckpt_full_path = os.path.join(model_dir, checkpoint)
+        if not os.path.exists(ckpt_full_path):
+            continue
+
+        # blimp fast
+        run_task_with_cache(
+            checkpoint, "blimp", "blimp_fast",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/fast_eval/blimp_fast", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # supplement fast
+        run_task_with_cache(
+            checkpoint, "blimp", "supplement_fast",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/fast_eval/supplement_fast", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # ewok fast
+        run_task_with_cache(
+            checkpoint, "ewok", "ewok_fast",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "ewok", "--data_path", "evaluation_data/fast_eval/ewok_fast", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # entity tracking fast
+        run_task_with_cache(
+            checkpoint, "entity_tracking", "entity_tracking_fast",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "entity_tracking", "--data_path", "evaluation_data/fast_eval/entity_tracking_fast", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # reading fast
+        run_task_with_cache(
+            checkpoint, "reading", "reading",
+            ["python", "-m", "evaluation_pipeline.reading.run", "--model_path_or_name", ckpt_full_path, "--backend", "causal", "--data_path", "evaluation_data/fast_eval/reading/reading_data.csv"]
+        )
+        # global piqa parallel fast
+        run_task_with_cache(
+            checkpoint, "global_piqa_parallel", "global_piqa_parallel",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "global_piqa_parallel", "--data_path", "evaluation_data/fast_eval/global_piqa_parallel", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # global piqa nonparallel fast
+        run_task_with_cache(
+            checkpoint, "global_piqa_nonparallel", "global_piqa_nonparallel",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", eval_model_path, "--backend", "causal", "--task", "global_piqa_nonparallel", "--data_path", "evaluation_data/fast_eval/global_piqa_nonparallel", "--save_predictions", "--revision_name", checkpoint]
+        )
+
+    # ─────────────────────────────────────────────────────────────
+    # D. COLLATE RESULTS AND CLEANUP
+    # ─────────────────────────────────────────────────────────────
+    print("[Eval] Collating predictions into submission file...")
+    collate_results_dir = os.path.join(strict_dir, "results", model_name)
+    if os.path.exists(collate_results_dir):
+        shutil.rmtree(collate_results_dir)
+    os.makedirs(os.path.dirname(collate_results_dir), exist_ok=True)
+
+    shutil.copytree(local_results_dir, collate_results_dir)
+
+    subprocess.run([
+        "python", "-m", "evaluation_pipeline.collate_preds",
+        "--model_path_or_name", model_name,
+        "--backend", "causal",
+        "--track", "strict-small",
+        "--fast"
+    ], cwd=strict_dir, check=True)
+
+    results_src = os.path.join(strict_dir, "results")
+    results_dest = os.path.abspath("./results")
+    if os.path.exists(results_dest):
+        shutil.rmtree(results_dest)
+    shutil.copytree(results_src, results_dest)
+
+    collated_json = os.path.join(strict_dir, "all_full_preds_and_fast_scores_causal.json")
+    if os.path.exists(collated_json):
+        shutil.copy(collated_json, "./all_full_preds_and_fast_scores_causal.json")
+        print("\n[Eval] Success! Collation completed! Final file is at './all_full_preds_and_fast_scores_causal.json'")
+
+    print("\n[Eval] Pipeline evaluation run finished.")
+
+def upload_pipeline(model_name, repo_name, token=None):
+    from huggingface_hub import HfApi, create_repo
+    if not token:
+        token = os.environ.get("HF_TOKEN")
+
+    api = HfApi(token=token)
+    try:
+        user_info = api.whoami()
+        username = user_info["name"]
+        print(f"[HF] Authenticated successfully as user: {username}")
+    except Exception as e:
+        print(f"[HF] Authentication failed. Error: {e}")
+        return
+
+    repo_id = f"{username}/{repo_name}"
+    print(f"[HF] Target Repository ID: {repo_id}")
+
+    try:
+        create_repo(repo_id=repo_id, repo_type="model", token=token, exist_ok=True)
+        print(f"[HF] Repository '{repo_id}' is ready.")
+    except Exception as e:
+        print(f"[HF] Failed to create repository: {e}")
+        return
+
+    checkpoint_dir = os.path.abspath(f"./checkpoints/{model_name}")
+
+    revisions = {}
+    if os.path.exists(os.path.join(checkpoint_dir, "main")):
+        revisions = {"main": os.path.join(checkpoint_dir, "main")}
+    else:
+        print(f"[HF] Error: Could not locate 'main' checkpoint weights under {checkpoint_dir}/main")
+        return
+
+    main_dir = revisions["main"]
+    parent_dir = os.path.dirname(main_dir)
+
+    for m in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 30, 35, 40, 45, 50, 60, 70, 80, 90, 100]:
+        ckpt_name = f"chck_{m}M"
+        ckpt_path = os.path.join(parent_dir, ckpt_name)
+        if os.path.exists(ckpt_path):
+            revisions[ckpt_name] = ckpt_path
+
+    temp_dir = os.path.abspath("./hf_upload_temp")
+    
+    for revision_name, local_path in revisions.items():
+        print(f"\n[HF] Preparing revision '{revision_name}' from folder: {local_path}")
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+        os.makedirs(temp_dir, exist_ok=True)
+        
+        for item in os.listdir(local_path):
+            shutil.copy2(os.path.join(local_path, item), os.path.join(temp_dir, item))
+            
+        try:
+            api.create_branch(
+                repo_id=repo_id,
+                repo_type="model",
+                branch=revision_name,
+                exist_ok=True
+            )
+            print(f"[HF] Created branch/revision '{revision_name}' on repository.")
+        except Exception as branch_err:
+            print(f"[HF] Info: Branch creation failed or exists: {branch_err}")
+
+        print(f"[HF] Uploading staged folder to '{repo_id}' revision '{revision_name}'...")
+        try:
+            api.upload_folder(
+                folder_path=temp_dir,
+                repo_id=repo_id,
+                repo_type="model",
+                revision=revision_name
+            )
+            print(f"[HF] Successfully uploaded revision '{revision_name}' to repository.")
+        except Exception as e:
+            print(f"[HF] Failed to upload revision '{revision_name}': {e}")
+
+    if os.path.exists(temp_dir):
+        shutil.rmtree(temp_dir)
+    print(f"\n[HF] All uploads finished! View your repository at https://huggingface.co/{repo_id}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-name", type=str, default="gpt2_dense_31M")
+    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="active_param_gpt2_31M", help="Hugging Face repository name")
+    parser.add_argument("--upload-token", type=str, default=None, help="Hugging Face API token")
+    args = parser.parse_args()
+
+    if args.upload:
+        upload_pipeline(args.model_name, args.upload_repo, args.upload_token)
+    else:
+        class Tee:
+            def __init__(self, filepath, original_stream):
+                self.file = open(filepath, "a", encoding="utf-8", buffering=1)
+                self.original_stream = original_stream
+
+            def write(self, data):
+                self.original_stream.write(data)
+                self.file.write(data)
+
+            def flush(self):
+                self.original_stream.flush()
+                self.file.flush()
+
+        import sys
+        sys.stdout = Tee("training_terminal.log", sys.stdout)
+        sys.stderr = Tee("training_terminal.log", sys.stderr)
+
+        print("\n=== STARTING NEW TRAINING RUN LOGGING TO training_terminal.log ===")
+        run_pipeline(args.model_name, epochs=args.epochs, skip_eval=args.skip_eval, skip_aoa=args.skip_aoa, skip_glue=args.skip_glue)

upload_all_to_hf.py

@@ -0,0 +1,144 @@
+import os
+import sys
+import shutil
+import re
+import argparse
+from huggingface_hub import HfApi, create_repo
+
+def upload_all(repo_name="active_param_gpt2_31M", model_name="gpt2_dense_fresh"):
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        print("Error: HF_TOKEN environment variable not set!")
+        print("Please set it before running: export HF_TOKEN=your_token")
+        sys.exit(1)
+
+    api = HfApi(token=token)
+    try:
+        user_info = api.whoami()
+        username = user_info["name"]
+        print(f"[HF] Authenticated as: {username}")
+    except Exception as e:
+        print(f"[HF] Authentication failed: {e}")
+        sys.exit(1)
+
+    repo_id = f"{username}/{repo_name}"
+    
+    # Auto-create the repo if it does not 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] Warning during repository creation: {e}")
+
+    local_dir = os.path.dirname(os.path.abspath(__file__))
+
+    # ─────────────────────────────────────────────────────────────
+    # Step 1: Upload scripts, metrics, results, and plots to main branch
+    # ─────────────────────────────────────────────────────────────
+    print(f"\n[HF] Step 1: Uploading codebase, plots, metrics, and results to main branch...")
+    ignore_patterns = [
+        "**/__pycache__/*",
+        "**/*.pyc",
+        "**/hf_cache/*",
+        "**/nltk_data/*",
+        "**/checkpoints/*",
+        "**/babylm_eval_repo/*"
+    ]
+    
+    # Exclude files containing HF Token
+    sensitive_ignores = []
+    token_pattern = re.compile(r"hf_[a-zA-Z0-9]{34}")
+    for root, dirs, files in os.walk(local_dir):
+        if "__pycache__" in root or "checkpoints" in root or "hf_cache" in root or "nltk_data" in root or "babylm_eval_repo" in root:
+            continue
+        for file in files:
+            file_path = os.path.join(root, file)
+            try:
+                with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+                    content = f.read()
+                    if token_pattern.search(content):
+                        rel_path = os.path.relpath(file_path, local_dir)
+                        rel_path_glob = rel_path.replace("\\", "/")
+                        sensitive_ignores.append(rel_path_glob)
+                        print(f"  -> Warning: Sensitive file containing raw token excluded: {rel_path_glob}")
+            except Exception:
+                pass
+
+    all_ignores = ignore_patterns + sensitive_ignores
+
+    try:
+        api.upload_folder(
+            folder_path=local_dir,
+            repo_id=repo_id,
+            repo_type="model",
+            revision="main",
+            ignore_patterns=all_ignores,
+            commit_message="Upload dense baseline codebase, metrics, and plots"
+        )
+        print(f"[HF] Codebase, plots, and results uploaded successfully to main branch.")
+    except Exception as e:
+        print(f"[HF] Failed to upload codebase: {e}")
+
+    # ─────────────────────────────────────────────────────────────
+    # Step 2: Upload checkpoints to their respective revisions/branches
+    # ─────────────────────────────────────────────────────────────
+    checkpoint_dir = os.path.abspath(f"./checkpoints/{model_name}")
+    if not os.path.exists(checkpoint_dir):
+        print(f"[HF] Error: Checkpoint directory not found at: {checkpoint_dir}")
+        return
+
+    revisions = {}
+    if os.path.exists(os.path.join(checkpoint_dir, "main")):
+        revisions["main"] = os.path.join(checkpoint_dir, "main")
+
+    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(checkpoint_dir, ckpt_name)
+        if os.path.exists(ckpt_path):
+            revisions[ckpt_name] = ckpt_path
+
+    print(f"\n[HF] Step 2: Uploading {len(revisions)} checkpoint revisions to their respective branches...")
+    temp_dir = os.path.abspath("./hf_upload_temp")
+    
+    for revision_name, local_path in revisions.items():
+        print(f"\n[HF] Preparing revision '{revision_name}' from folder: {local_path}")
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+        os.makedirs(temp_dir, exist_ok=True)
+        
+        for item in os.listdir(local_path):
+            shutil.copy2(os.path.join(local_path, item), os.path.join(temp_dir, item))
+            
+        try:
+            api.create_branch(
+                repo_id=repo_id,
+                repo_type="model",
+                branch=revision_name,
+                exist_ok=True
+            )
+            print(f"[HF] Created branch/revision '{revision_name}' on repository.")
+        except Exception as branch_err:
+            print(f"[HF] Info: Branch creation failed or exists: {branch_err}")
+
+        print(f"[HF] Uploading staged folder to '{repo_id}' revision '{revision_name}'...")
+        try:
+            api.upload_folder(
+                folder_path=temp_dir,
+                repo_id=repo_id,
+                repo_type="model",
+                revision=revision_name
+            )
+            print(f"[HF] Successfully uploaded revision '{revision_name}' to repository.")
+        except Exception as e:
+            print(f"[HF] Failed to upload revision '{revision_name}': {e}")
+
+    if os.path.exists(temp_dir):
+        shutil.rmtree(temp_dir)
+    print(f"\n[HF] All uploads finished! View your repository at https://huggingface.co/{repo_id}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo-name", type=str, default="active_param_gpt2_31M", help="Hugging Face repository name")
+    parser.add_argument("--model-name", type=str, default="gpt2_dense_fresh", help="Local model checkpoints folder name")
+    args = parser.parse_args()
+    upload_all(args.repo_name, args.model_name)

upload_project_hf.py

@@ -0,0 +1,84 @@
+import os
+import sys
+import re
+import argparse
+from huggingface_hub import HfApi, create_repo
+
+def upload_project(repo_name=None):
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        print("Error: HF_TOKEN environment variable not set!")
+        print("Please set it before running: set HF_TOKEN=your_token")
+        sys.exit(1)
+
+    if not repo_name:
+        repo_name = os.environ.get("HF_REPO_NAME", "active_param_gpt2_31M")
+
+    local_dir = os.path.dirname(os.path.abspath(__file__))
+
+    api = HfApi(token=token)
+    try:
+        user_info = api.whoami()
+        username = user_info["name"]
+        print(f"[HF] Authenticated as: {username}")
+    except Exception as e:
+        print(f"[HF] Authentication failed: {e}")
+        sys.exit(1)
+
+    repo_id = f"{username}/{repo_name}"
+    
+    # Auto-create the repo if it does not 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] Warning during repository creation: {e}")
+
+    print(f"[HF] Uploading GPT-2 46M SwiGLU Dense baseline scripts to '{repo_id}' main branch...")
+    
+    ignore_patterns = [
+        "**/__pycache__/*",
+        "**/*.pyc",
+        "**/hf_cache/*",
+        "**/nltk_data/*",
+        "**/checkpoints/*"
+    ]
+    
+    sensitive_ignores = []
+    token_pattern = re.compile(r"hf_[a-zA-Z0-9]{34}")
+    for root, dirs, files in os.walk(local_dir):
+        if "__pycache__" in root or "checkpoints" in root:
+            continue
+        for file in files:
+            file_path = os.path.join(root, file)
+            try:
+                with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+                    content = f.read()
+                    if token_pattern.search(content):
+                        rel_path = os.path.relpath(file_path, local_dir)
+                        rel_path_glob = rel_path.replace("\\", "/")
+                        sensitive_ignores.append(rel_path_glob)
+                        print(f"  -> Warning: Sensitive file containing raw token excluded: {rel_path_glob}")
+            except Exception:
+                pass
+
+    all_ignores = ignore_patterns + sensitive_ignores
+
+    try:
+        api.upload_folder(
+            folder_path=local_dir,
+            repo_id=repo_id,
+            repo_type="model",
+            revision="main",
+            ignore_patterns=all_ignores,
+            commit_message="Initial commit of GPT-2 46M SwiGLU Dense baseline model architecture and configs"
+        )
+        print(f"\n[HF] Success! Scripts uploaded to: https://huggingface.co/{repo_id}/tree/main")
+    except Exception as e:
+        print(f"[HF] Upload failed: {e}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo-name", type=str, default=None, help="Hugging Face repository name")
+    args = parser.parse_args()
+    upload_project(args.repo_name)