Upload GPT-2 Strict-Small 16k BPE model, checkpoints, and logs

run_lightning_eval.py

@@ -1,9 +1,10 @@
 """
 Lightning AI Evaluation Runner for BabyLM 2026 Strict-Small Track
-Follows exact order requested:
+Fixes PYTHONPATH and argument formatting for babylm-eval.
+Order:
 1. Full Zero-Shot Evaluation
-2. Fine-Tuning Evaluation
-3. Fast Evaluation & Score Collating / Zipping
+2. Fine-Tuning (GLUE Sequence Classification) Evaluation
+3. Fast Evaluation & Prediction Collating
 """
 
 import os
@@ -11,9 +12,9 @@ import sys
 import shutil
 import subprocess
 
-def run_cmd(cmd, cwd=None):
+def run_cmd(cmd, cwd=None, env=None):
     print(f"\n[RUNNING] {cmd}")
-    res = subprocess.run(cmd, shell=True, cwd=cwd)
+    res = subprocess.run(cmd, shell=True, cwd=cwd, env=env)
     print(f"[STATUS] Exit Code: {res.returncode}")
     return res.returncode
 
@@ -43,23 +44,21 @@ def main():
     if not os.path.exists(eval_repo_dir):
         run_cmd("git clone https://github.com/babylm-org/babylm-eval.git", cwd=base_dir)
 
+    # Prepare environment with PYTHONPATH pointing to babylm-eval root and subfolders
+    custom_env = os.environ.copy()
+    pythonpath = eval_repo_dir
+    for root, dirs, _ in os.walk(eval_repo_dir):
+        if not root.startswith(os.path.join(eval_repo_dir, ".git")):
+            pythonpath += ":" + root
+    custom_env["PYTHONPATH"] = pythonpath + ":" + custom_env.get("PYTHONPATH", "")
+
     # 2. Install dependencies via requirements.txt
     print("\n[Setup] Installing dependencies...")
     req_file = find_script(eval_repo_dir, ["requirements.txt"])
     if req_file:
-        run_cmd(f"pip install -r {req_file}", cwd=os.path.dirname(req_file))
+        run_cmd(f"pip install -r {req_file}", cwd=os.path.dirname(req_file), env=custom_env)
     else:
-        run_cmd("pip install evaluate scikit-learn lm-eval datasets transformers", cwd=eval_repo_dir)
-
-    # List structure of babylm-eval for debugging transparency
-    print("\n[Inspection] Discovered files inside babylm-eval:")
-    found_files = []
-    for root, _, files in os.walk(eval_repo_dir):
-        for f in files:
-            rel = os.path.relpath(os.path.join(root, f), eval_repo_dir)
-            if not rel.startswith(".git"):
-                found_files.append(rel)
-                print(f"  - {rel}")
+        run_cmd("pip install evaluate scikit-learn lm-eval datasets transformers", cwd=eval_repo_dir, env=custom_env)
 
     # ------------------------------------------------------------
     # STEP 1: FULL ZERO-SHOT EVALUATION FIRST
@@ -68,39 +67,32 @@ def main():
     print(" STEP 1: FULL ZERO-SHOT EVALUATION")
     print("="*65)
 
-    zeroshot_py = find_script(eval_repo_dir, ["eval_pipeline.py", "eval_zero_shot.py", "zeroshot.py", "eval_model.py"])
-    zeroshot_sh = find_script(eval_repo_dir, ["zeroshot_model.sh", "run_zeroshot.sh"])
-
-    if zeroshot_py:
-        script_dir = os.path.dirname(zeroshot_py)
-        script_name = os.path.basename(zeroshot_py)
-        run_cmd(f"python {script_name} --model_path {checkpoint_dir} --task causal --type zeroshot", cwd=script_dir)
-    elif zeroshot_sh:
-        script_dir = os.path.dirname(zeroshot_sh)
-        script_name = os.path.basename(zeroshot_sh)
-        run_cmd(f"bash {script_name} {checkpoint_dir} causal strict_small", cwd=script_dir)
+    eval_model_py = find_script(eval_repo_dir, ["eval_model.py", "eval_pipeline.py"])
+    if eval_model_py:
+        script_dir = os.path.dirname(eval_model_py)
+        script_name = os.path.basename(eval_model_py)
+        run_cmd(f"python {script_name} --model_path {checkpoint_dir} --task causal", cwd=script_dir, env=custom_env)
     else:
         print("[Notice] Running fallback zero-shot lm-eval on checkpoints...")
-        run_cmd(f"lm_eval --model hf --model_args pretrained={checkpoint_dir}/main --tasks blimp --batch_size 16", cwd=eval_repo_dir)
+        run_cmd(f"lm_eval --model hf --model_args pretrained={checkpoint_dir}/main --tasks blimp --batch_size 16", cwd=eval_repo_dir, env=custom_env)
 
     # ------------------------------------------------------------
-    # STEP 2: FINE-TUNING (GLUE / CLASSIFICATION) EVALUATION SECOND
+    # STEP 2: FINE-TUNING (GLUE CLASSIFICATION) EVALUATION SECOND
     # ------------------------------------------------------------
     print("\n" + "="*65)
     print(" STEP 2: FINE-TUNING (GLUE CLASSIFICATION) EVALUATION")
     print("="*65)
 
-    finetune_py = find_script(eval_repo_dir, ["finetune_all.py", "eval_glue.py", "finetune.py", "eval_pipeline.py"])
     finetune_sh = find_script(eval_repo_dir, ["finetune_model.sh", "run_finetune.sh"])
-
-    if finetune_py:
-        script_dir = os.path.dirname(finetune_py)
-        script_name = os.path.basename(finetune_py)
-        run_cmd(f"python {script_name} --model_path {checkpoint_dir} --task causal --type finetune", cwd=script_dir)
+    if eval_model_py:
+        script_dir = os.path.dirname(eval_model_py)
+        script_name = os.path.basename(eval_model_py)
+        # Run eval_model.py for finetune tasks
+        run_cmd(f"python {script_name} --model_path {checkpoint_dir} --task causal --finetune", cwd=script_dir, env=custom_env)
     elif finetune_sh:
         script_dir = os.path.dirname(finetune_sh)
         script_name = os.path.basename(finetune_sh)
-        run_cmd(f"bash {script_name} {checkpoint_dir} causal strict_small", cwd=script_dir)
+        run_cmd(f"bash {script_name} --model_path {checkpoint_dir} --task causal", cwd=script_dir, env=custom_env)
 
     # ------------------------------------------------------------
     # STEP 3: FAST EVALUATION & SCORE COLLATING THIRD
@@ -110,16 +102,16 @@ def main():
     print("="*65)
 
     collate_sh = find_script(eval_repo_dir, ["collate_preds.sh", "collate_results.sh"])
-    collate_py = find_script(eval_repo_dir, ["collate_results.py", "collate.py", "eval_pipeline.py"])
-
     if collate_sh:
         script_dir = os.path.dirname(collate_sh)
         script_name = os.path.basename(collate_sh)
-        run_cmd(f"bash {script_name} gpt2_strict_small causal strict_small", cwd=script_dir)
-    elif collate_py:
-        script_dir = os.path.dirname(collate_py)
-        script_name = os.path.basename(collate_py)
-        run_cmd(f"python {script_name} --model_path {checkpoint_dir} --task causal --fast", cwd=script_dir)
+        # Run collate script with PYTHONPATH properly exported
+        run_cmd(f"bash {script_name} gpt2_strict_small causal strict_small", cwd=script_dir, env=custom_env)
+    
+    if eval_model_py:
+        script_dir = os.path.dirname(eval_model_py)
+        script_name = os.path.basename(eval_model_py)
+        run_cmd(f"python {script_name} --model_path {checkpoint_dir} --task causal --fast", cwd=script_dir, env=custom_env)
 
     # ------------------------------------------------------------
     # PACKAGING & ZIPPING RESULTS FOR LEADERBOARD SUBMISSION
@@ -130,6 +122,12 @@ def main():
             results_dir = os.path.join(root, "results")
             break
 
+    if not results_dir or not os.path.exists(results_dir):
+        for root, dirs, _ in os.walk(base_dir):
+            if "results" in dirs and "babylm" in root:
+                results_dir = os.path.join(root, "results")
+                break
+
     if results_dir and os.path.exists(results_dir):
         zip_path = os.path.join(base_dir, "results.zip")
         print(f"\n[Packaging] Zipping results directory '{results_dir}' to {zip_path}...")