Theskywalker07/full_model_1
119
1import os
2import math
3import time
4import json
5import random
6import inspect
7import shutil
8import subprocess
9import argparse
10
11# ─────────────────────────────────────────────────────────────
12# 1. CORE PIPELINE FUNCTION AND HELPERS
13# ─────────────────────────────────────────────────────────────
14
15def is_zero_shot_cache_valid(cache_dir):
16 if not os.path.exists(cache_dir):
17 return False
18 pred_count = 0
19 for root, dirs, files in os.walk(cache_dir):
20 if "predictions.json" in files:
21 pred_count += 1
22 return pred_count >= 3
23
24def is_finetune_cache_valid(cache_dir):
25 if not os.path.exists(cache_dir):
26 return False
27 pred_count = 0
28 for root, dirs, files in os.walk(cache_dir):
29 if "predictions.json" in files:
30 pred_count += 1
31 return pred_count >= 3
32
33def verify_eval_run(path_to_check, description):
34 print(f"[Eval] Verifying {description} path: {path_to_check}...")
35 if not os.path.exists(path_to_check):
36 raise RuntimeError(f"CRITICAL ERROR: {description} directory was NOT created at: {path_to_check}")
37
38 # Check if predictions.json or results.txt or surprisal.json exists and is non-empty
39 found_valid = False
40 for root, dirs, files in os.walk(path_to_check):
41 for f in files:
42 if f in ["predictions.json", "results.txt", "surprisal.json"]:
43 file_path = os.path.join(root, f)
44 if os.path.getsize(file_path) > 0:
45 found_valid = True
46 break
47 if found_valid:
48 break
49
50 if not found_valid:
51 raise RuntimeError(f"CRITICAL ERROR: {description} completed but no valid prediction/result files were written in: {path_to_check}")
52 print(f"[Eval] Success! Verified {description} results are stored correctly.")
53
54def run_pipeline(model_name: str, epochs: int = 10, skip_eval: bool = False, skip_aoa: bool = True, skip_glue: bool = False):
55 # Configure persistent cache paths locally to avoid duplicate downloads
56 os.environ["HF_HOME"] = os.path.abspath("./hf_cache")
57 os.environ["NLTK_DATA"] = os.path.abspath("./nltk_data")
58 os.makedirs("./hf_cache", exist_ok=True)
59 os.makedirs("./nltk_data", exist_ok=True)
60
61 # Programmatic Hugging Face Hub Login if HF_TOKEN is in environment
62 hf_token = os.environ.get("HF_TOKEN")
63
64 if hf_token:
65 try:
66 from huggingface_hub import login
67 login(token=hf_token)
68 print("[HF] Programmatic login successful using HF_TOKEN.")
69 except Exception as e:
70 print(f"[HF] Warning: Programmatic login failed: {e}")
71
72 import torch
73 import torch.nn as nn
74 import torch.nn.functional as F
75 from datasets import load_dataset
76 from tokenizers import Tokenizer
77 from transformers import PreTrainedTokenizerFast
78
79 # Import model architecture
80 from modeling_xpertgpt import (
81 XpertGPTModel,
82 XpertGPTModelConfig,
83 XpertGPTConfig
84 )
85
86 # Print GPU details
87 if torch.cuda.is_available():
88 gpu_name = torch.cuda.get_device_name(0)
89 print(f"\n[GPU] CUDA is available! Using GPU: {gpu_name}\n")
90 else:
91 print("\n[GPU] Warning: CUDA is NOT available! Running on CPU.\n")
92
93 # ─────────────────────────────────────────────────────────────
94 # GLOBAL HYPERPARAMETERS
95 # ─────────────────────────────────────────────────────────────
96 VOCAB_SIZE = 16384
97 MASK_TOKEN_ID = 16383
98 BLOCK_SIZE = 512
99 BATCH_SIZE = 16
100 GRAD_ACCUM_STEPS = 1 # grad_acc_step = 1
101 EPOCHS = epochs
102 LEARNING_RATE = 3e-4 # lr = 3e-4
103 LR_MIN = LEARNING_RATE * 0.05
104 WARMUP_STEPS = 800 # warmup_steps = 800
105 WEIGHT_DECAY = 0.1
106 GRAD_CLIP = 1.0
107
108 NUM_THIN_BLOCKS = 4
109 EC_CAPACITY_FACTOR = 2.0
110 CAUSAL_RATIO = 1 / 1
111
112 MASK_PROB_START = 0.20
113 MASK_PROB_END = 0.10
114
115 # Output directories locally
116 model_dir = os.path.abspath(f"./checkpoints/{model_name}")
117 os.makedirs(model_dir, exist_ok=True)
118 local_results_dir = os.path.abspath(f"./results/{model_name}")
119 os.makedirs(local_results_dir, exist_ok=True)
120
121
122
123 # ─────────────────────────────────────────────────────────────
124 # Helper: Save Hugging Face Compliant Checkpoint
125 # ─────────────────────────────────────────────────────────────
126 def save_hf_checkpoint(raw_model, checkpoint_dir_name, tokenizer):
127 save_dir = os.path.join(model_dir, checkpoint_dir_name)
128 os.makedirs(save_dir, exist_ok=True)
129 print(f"\n[Checkpoint] Saving Hugging Face format checkpoint to '{save_dir}'...")
130
131 # A. Convert state dict keys to CausalLM wrapper naming
132 state_dict = raw_model.state_dict()
133 new_state_dict = {}
134 for k, v in state_dict.items():
135 name = k
136 if name.startswith("_orig_mod."):
137 name = name[10:]
138 if name.startswith("model."):
139 name = name[6:]
140
141 if name == "lm_head.weight":
142 new_state_dict["lm_head.weight"] = v
143 else:
144 new_state_dict[f"transformer.{name}"] = v
145
146 torch.save(new_state_dict, os.path.join(save_dir, "pytorch_model.bin"))
147
148 # B. Copy modeling.py and configuration.py
149 shutil.copy("modeling_xpertgpt.py", os.path.join(save_dir, "modeling_xpertgpt.py"))
150 shutil.copy("configuration_xpertgpt.py", os.path.join(save_dir, "configuration_xpertgpt.py"))
151
152 # C. Create config.json
153 config_dict = {
154 "auto_map": {
155 "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
156 "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
157 "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
158 },
159 "vocab_size": VOCAB_SIZE,
160 "block_size": BLOCK_SIZE,
161 "d_model": 256, # d_model = 256
162 "hidden_size": 256, # hidden_size = 256
163 "d_thin": 384,
164 "num_layers": 6,
165 "num_blocks": NUM_THIN_BLOCKS,
166 "capacity_factor": EC_CAPACITY_FACTOR,
167 "dropout": 0.1,
168 "model_type": "xpertgpt"
169 }
170 with open(os.path.join(save_dir, "config.json"), "w") as f:
171 json.dump(config_dict, f, indent=2)
172
173 # D. Save tokenizer config files
174 fast_tokenizer = PreTrainedTokenizerFast(
175 tokenizer_object=tokenizer,
176 bos_token="[CLS]",
177 eos_token="[SEP]",
178 unk_token="[UNK]",
179 pad_token="[PAD]",
180 mask_token="[MASK]"
181 )
182 # E. Copy loss metrics
183 main_loss_csv = os.path.join(model_dir, "loss_history.csv")
184 main_loss_json = os.path.join(model_dir, "loss_metrics.json")
185 if os.path.exists(main_loss_csv):
186 shutil.copy2(main_loss_csv, os.path.join(save_dir, "loss_history.csv"))
187 if os.path.exists(main_loss_json):
188 shutil.copy2(main_loss_json, os.path.join(save_dir, "loss_metrics.json"))
189
190 print(f"[Checkpoint] Checkpoint '{checkpoint_dir_name}' successfully saved.")
191
192 # ─────────────────────────────────────────────────────────────
193 # Tokenizer Training
194 # ─────────────────────────────────────────────────────────────
195 def build_and_train_tokenizer(texts: list) -> Tokenizer:
196 from tokenizers.models import BPE
197 from tokenizers.trainers import BpeTrainer
198 from tokenizers.pre_tokenizers import Whitespace
199
200 vocab_path = os.path.join(model_dir, "bpe_vocab_16k.json")
201 if os.path.exists(vocab_path):
202 print(f"[Tokenizer] Loading trained BPE model layout from '{vocab_path}'...")
203 return Tokenizer.from_file(vocab_path)
204
205 print(f"[Tokenizer] Generating fresh HuggingFace BPE Tokenizer model with {VOCAB_SIZE} slots...")
206 tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
207 tokenizer.pre_tokenizer = Whitespace()
208
209 trainer = BpeTrainer(
210 vocab_size=VOCAB_SIZE,
211 special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
212 )
213 tokenizer.train_from_iterator(texts, trainer)
214 tokenizer.save(vocab_path)
215 print(f"[Tokenizer] Tokenizer training completed and saved to '{vocab_path}'.")
216 return tokenizer
217
218 # ─────────────────────────────────────────────────────────────
219 # Data Loader Setup
220 # ─────────────────────────────────────────────────────────────
221 class DataLoaderLite:
222 def __init__(self, B: int, T: int, texts: list, tokenizer: Tokenizer, name: str):
223 self.B = B
224 self.T = T
225
226 print(f"[DataLoader:{name}] Tokenising dataset sequences...")
227 all_ids = []
228 for t in texts:
229 if t.strip():
230 encoded = tokenizer.encode(t).ids
231 all_ids.extend(encoded)
232
233 self.tokens = torch.tensor(all_ids, dtype=torch.long)
234 self.chunk_size = B * T
235 self.n_chunks = (len(self.tokens) - 1) // self.chunk_size
236 self.indices = list(range(self.n_chunks))
237 self.pos = 0
238 self._shuffle()
239
240 print(f"[DataLoader:{name}] Total tokens: {len(self.tokens):,} | Epoch steps: {self.n_chunks:,}")
241
242 def _shuffle(self):
243 random.shuffle(self.indices)
244 self.pos = 0
245
246 def steps_per_epoch(self) -> int:
247 return self.n_chunks
248
249 def next_batch(self):
250 B, T = self.B, self.T
251 if self.pos >= len(self.indices):
252 self._shuffle()
253
254 chunk_idx = self.indices[self.pos]
255 self.pos += 1
256
257 start_pos = chunk_idx * self.chunk_size
258 temp = self.tokens[start_pos : start_pos + self.chunk_size + 1]
259
260 x = temp[:-1].view(B, T)
261 y = temp[1:].view(B, T)
262 return x, y
263
264 # ─────────────────────────────────────────────────────────────
265 # Batch preparation and schedules
266 # ─────────────────────────────────────────────────────────────
267 def get_current_mask_prob(global_step: int, total_steps: int) -> float:
268 ratio = min(1.0, global_step / total_steps)
269 return MASK_PROB_START + ratio * (MASK_PROB_END - MASK_PROB_START)
270
271 def prepare_causal_batch(x: torch.Tensor, y: torch.Tensor):
272 return x, y, False
273
274 def prepare_masked_batch(x: torch.Tensor, y: torch.Tensor, mask_prob: float, mask_token_id: int):
275 B, T = x.size()
276 mask = torch.rand(B, T, device=x.device) < mask_prob
277 masked_x = x.clone()
278 masked_x[mask] = mask_token_id
279
280 targets = torch.full_like(y, -100)
281 targets[mask] = y[mask]
282
283 return masked_x, targets, True
284
285 def get_hybrid_batch(train_loader: DataLoaderLite, global_step: int, total_steps: int, device: torch.device):
286 x, y = train_loader.next_batch()
287 x, y = x.to(device), y.to(device)
288
289 if random.random() < CAUSAL_RATIO:
290 input_ids, targets, bidir = prepare_causal_batch(x, y)
291 else:
292 mask_prob = get_current_mask_prob(global_step, total_steps)
293 input_ids, targets, bidir = prepare_masked_batch(x, y, mask_prob, MASK_TOKEN_ID)
294
295 return input_ids, targets, bidir
296
297 def get_lr(it: int, total_steps: int) -> float:
298 if it < WARMUP_STEPS:
299 return LEARNING_RATE * (it + 1) / WARMUP_STEPS
300 if it >= total_steps:
301 return LR_MIN
302 decay_ratio = (it - WARMUP_STEPS) / (total_steps - WARMUP_STEPS)
303 coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
304 return LR_MIN + coeff * (LEARNING_RATE - LR_MIN)
305
306 # ─────────────────────────────────────────────────────────────
307 # Dataset Preparation
308 # ─────────────────────────────────────────────────────────────
309 print("\n[Data] Loading BabyLM-2026-Strict-Small ...")
310 ds = load_dataset("BabyLM-community/BabyLM-2026-Strict-Small")
311 all_text = list(ds['train']['text'])
312
313 tokenizer = build_and_train_tokenizer(all_text)
314
315 split = int(len(all_text) * 0.95)
316 train_texts = all_text[:split]
317 val_texts = all_text[split:]
318
319 train_loader = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, train_texts, tokenizer, "train")
320 val_loader = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, val_texts, tokenizer, "val")
321
322 def evaluate_val_loss(model_to_eval, v_loader, dev, autocast_context, num_batches=30):
323 raw_m = model_to_eval._orig_mod if hasattr(model_to_eval, '_orig_mod') else model_to_eval
324 raw_m.eval()
325 total_v_loss = 0.0
326 eval_steps = min(num_batches, v_loader.steps_per_epoch())
327 with torch.no_grad():
328 for _ in range(eval_steps):
329 x_v, y_v = v_loader.next_batch()
330 x_v, y_v = x_v.to(dev), y_v.to(dev)
331 with autocast_context:
332 _, loss_v = raw_m(x_v, y_v, bidirectional=False)
333 total_v_loss += loss_v.item()
334 raw_m.train()
335 return total_v_loss / max(1, eval_steps)
336
337 chunks_per_epoch = train_loader.steps_per_epoch()
338 steps_per_epoch = chunks_per_epoch // GRAD_ACCUM_STEPS
339 total_steps = steps_per_epoch * EPOCHS
340
341 cfg = XpertGPTModelConfig()
342 device = "cuda" if torch.cuda.is_available() else "cpu"
343 torch.manual_seed(42)
344 if torch.cuda.is_available():
345 torch.cuda.manual_seed(42)
346 random.seed(42)
347 if hasattr(torch, 'set_float32_matmul_precision'):
348 torch.set_float32_matmul_precision('high')
349
350 model = XpertGPTModel(cfg).to(device)
351
352 # ─────────────────────────────────────────────────────────────
353 # Training Resume Check
354 # ─────────────────────────────────────────────────────────────
355 words_trained = 0
356 next_milestone_idx = 0
357 global_step = 0
358
359 milestones = sorted(list(set([i * 1_000_000 for i in range(1, 11)] + [i * 10_000_000 for i in range(1, 11)])))
360
361 resume_checkpoint_dir = None
362 for idx in range(len(milestones) - 1, -1, -1):
363 m = milestones[idx]
364 ckpt_name = f"chck_{m // 1_000_000}M"
365 ckpt_path = os.path.join(model_dir, ckpt_name)
366 if os.path.exists(os.path.join(ckpt_path, "pytorch_model.bin")):
367 config_json_path = os.path.join(ckpt_path, "config.json")
368 if os.path.exists(config_json_path):
369 try:
370 with open(config_json_path, "r") as f:
371 saved_config = json.load(f)
372 if saved_config.get("d_model") == 256:
373 resume_checkpoint_dir = ckpt_path
374 next_milestone_idx = idx + 1
375 words_trained = m
376 global_step = words_trained // (BATCH_SIZE * BLOCK_SIZE)
377 print(f"[Training] Found existing milestone checkpoint '{ckpt_name}'. Resuming from step {global_step:,} ({words_trained:,} tokens trained)...")
378 break
379 else:
380 print(f"[Training] Found checkpoint '{ckpt_name}' but it has mismatch d_model={saved_config.get('d_model')}. Starting fresh.")
381 except Exception as e:
382 pass
383
384 # Load weights if resuming
385 if resume_checkpoint_dir is not None:
386 print(f"[Model] Loading weights from checkpoint '{resume_checkpoint_dir}'...")
387 state_dict = torch.load(os.path.join(resume_checkpoint_dir, "pytorch_model.bin"), map_location=device)
388 model_state_dict = {}
389 for k, v in state_dict.items():
390 name = k
391 if name.startswith("transformer."):
392 name = name[12:]
393 model_state_dict[name] = v
394 model.load_state_dict(model_state_dict)
395
396 # Check if final main model exists
397 main_ckpt_path = os.path.join(model_dir, "main")
398 if os.path.exists(os.path.join(main_ckpt_path, "pytorch_model.bin")):
399 print("\n[Pipeline] Final checkpoint 'main' already exists. Skipping training phase and transitioning directly to evaluations!")
400 else:
401 # torch.compile
402 try:
403 model = torch.compile(model)
404 print("[Model] torch.compile() successfully verified graph optimizations")
405 except Exception as e:
406 print(f"[Model] torch.compile() skipped ({e})")
407
408 # Optimizer
409 param_dict = {n: p for n, p in model.named_parameters() if p.requires_grad}
410 decay_params = [p for p in param_dict.values() if p.dim() >= 2]
411 nodecay_params = [p for p in param_dict.values() if p.dim() < 2]
412 groups = [
413 {'params': decay_params, 'weight_decay': WEIGHT_DECAY},
414 {'params': nodecay_params, 'weight_decay': 0.0},
415 ]
416 fused_ok = 'fused' in inspect.signature(torch.optim.AdamW).parameters
417 use_fused = fused_ok and ('cuda' in device)
418 optimizer = torch.optim.AdamW(groups, lr=LEARNING_RATE, betas=(0.9, 0.95), eps=1e-8, fused=use_fused)
419
420 # ─────────────────────────────────────────────────────────────
421 # Training Loop & Loss Metrics Logging
422 # ─────────────────────────────────────────────────────────────
423 model.train()
424 autocast_ctx = torch.autocast(device_type="cuda" if "cuda" in device else "cpu", dtype=torch.bfloat16, enabled=True)
425
426 loss_history_csv = os.path.join(model_dir, "loss_history.csv")
427 loss_history_json = os.path.join(model_dir, "loss_metrics.json")
428 os.makedirs(model_dir, exist_ok=True)
429 if not os.path.exists(loss_history_csv):
430 with open(loss_history_csv, "w", encoding="utf-8") as f_csv:
431 f_csv.write("global_step,epoch,words_trained,train_loss,val_loss,lr\n")
432
433 loss_records = []
434 if os.path.exists(loss_history_json):
435 try:
436 with open(loss_history_json, "r", encoding="utf-8") as f_json:
437 loss_records = json.load(f_json)
438 except Exception:
439 loss_records = []
440
441 start_epoch = global_step // steps_per_epoch
442 start_chunk = (global_step % steps_per_epoch) * GRAD_ACCUM_STEPS
443
444 print(f"\n[Training] Starting XpertGPT MoEP training for {EPOCHS} epochs...")
445 for epoch in range(start_epoch, EPOCHS):
446 train_loader._shuffle()
447 if epoch == start_epoch and start_chunk > 0:
448 print(f"[Training] Fast-forwarding dataloader to chunk index {start_chunk}...")
449 train_loader.pos = start_chunk
450
451 optimizer.zero_grad(set_to_none=True)
452 loss_accum = 0.0
453
454 start_chunk_idx = start_chunk if epoch == start_epoch else 0
455 for chunk_step in range(start_chunk_idx, chunks_per_epoch):
456 t0 = time.perf_counter()
457
458 lr = get_lr(global_step, total_steps)
459 for pg in optimizer.param_groups:
460 pg['lr'] = lr
461
462 input_ids, targets, bidir = get_hybrid_batch(train_loader, global_step, total_steps, device)
463 words_trained += input_ids.numel()
464
465 with autocast_ctx:
466 _, loss = model(input_ids, targets, bidirectional=bidir)
467 scaled_loss = loss / GRAD_ACCUM_STEPS
468 loss_accum += scaled_loss.item()
469 scaled_loss.backward()
470
471 # Optimizer Step
472 if (chunk_step + 1) % GRAD_ACCUM_STEPS == 0:
473 norm = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
474 optimizer.step()
475 optimizer.zero_grad(set_to_none=True)
476 if "cuda" in device:
477 torch.cuda.synchronize()
478
479 dt = (time.perf_counter() - t0) * 1000
480 mode_tag = "MLM" if bidir else "CLM"
481 mask_p = get_current_mask_prob(global_step, total_steps)
482 current_step = (chunk_step + 1) // GRAD_ACCUM_STEPS
483
484 # Evaluate validation loss every 50 steps (and at step 1)
485 val_loss = None
486 if (global_step + 1) == 1 or (global_step + 1) % 50 == 0:
487 val_loss = evaluate_val_loss(model, val_loader, device, autocast_ctx)
488
489 val_str = f" val={val_loss:.4f}" if val_loss is not None else ""
490 print(
491 f"[E{epoch+1:02d} {current_step:>5d}/{steps_per_epoch} G{global_step+1:>7d}|{mode_tag}] "
492 f"train={loss_accum:.4f}{val_str} mask={mask_p:.1%} norm={norm:.3f} lr={lr:.2e} dt={dt:6.1f}ms words={words_trained:,}"
493 )
494
495 # Log to CSV and JSON
496 rec = {
497 "global_step": global_step + 1,
498 "epoch": epoch + 1,
499 "words_trained": words_trained,
500 "train_loss": round(loss_accum, 5),
501 "val_loss": round(val_loss, 5) if val_loss is not None else None,
502 "lr": lr
503 }
504 loss_records.append(rec)
505
506 val_csv_str = f"{val_loss:.5f}" if val_loss is not None else ""
507 with open(loss_history_csv, "a", encoding="utf-8") as f_csv:
508 f_csv.write(f"{global_step+1},{epoch+1},{words_trained},{loss_accum:.5f},{val_csv_str},{lr}\n")
509
510 with open(loss_history_json, "w", encoding="utf-8") as f_json:
511 json.dump(loss_records, f_json, indent=2)
512
513 loss_accum = 0.0
514 global_step += 1
515
516 # Check if we passed a milestone for checkpointing
517 if next_milestone_idx < len(milestones) and words_trained >= milestones[next_milestone_idx]:
518 milestone_val = milestones[next_milestone_idx]
519 if milestone_val < 10_000_000:
520 milestone_name = f"chck_{milestone_val // 1_000_000}M"
521 else:
522 milestone_name = f"chck_{(milestone_val // 10_000_000) * 10}M"
523
524 raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
525 save_hf_checkpoint(raw_model, milestone_name, tokenizer)
526 next_milestone_idx += 1
527
528 # Save final model as 'main'
529 raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
530 save_hf_checkpoint(raw_model, "main", tokenizer)
531 print("\n[Training] Training phase complete!")
532
533 if skip_eval:
534 print("[Pipeline] Skipping evaluations phase as requested.")
535 return
536
537 # ─────────────────────────────────────────────────────────────
538 # 2. RUN EVALUATION PIPELINE
539 # ─────────────────────────────────────────────────────────────
540 local_results_dir = os.path.abspath(f"./results/{model_name}")
541 os.makedirs(local_results_dir, exist_ok=True)
542 local_main_res = os.path.join(local_results_dir, "main")
543
544 # Ensure clone_dir exists and has the global_piqa files
545 clone_parent = os.path.abspath("./babylm_eval_repo")
546
547 # Self-healing check for global_piqa presence
548 has_global_piqa = False
549 for potential_strict in [os.path.join(clone_parent, "babylm-eval", "strict"), os.path.join(clone_parent, "strict")]:
550 if os.path.exists(os.path.join(potential_strict, "evaluation_pipeline", "global_piqa")):
551 has_global_piqa = True
552 break
553
554 if not has_global_piqa:
555 print("[Eval] Cloned repository does not contain global_piqa tasks.")
556 print("[Eval] Deleting and cloning official main branch...")
557 if os.path.exists(clone_parent):
558 shutil.rmtree(clone_parent)
559 subprocess.run([
560 "git", "clone", "-b", "main",
561 "https://github.com/babylm-org/babylm-eval.git",
562 clone_parent
563 ], check=True)
564
565 # Determine strict_dir path dynamically
566 if os.path.exists(os.path.join(clone_parent, "strict")):
567 strict_dir = os.path.join(clone_parent, "strict")
568 else:
569 strict_dir = os.path.join(clone_parent, "babylm-eval", "strict")
570
571 print(f"[Eval] Using strict directory: {strict_dir}")
572 os.environ["PYTHONPATH"] = strict_dir
573
574 def patch_evaluation_run_script(strict_dir):
575 import pathlib
576 run_file = os.path.join(strict_dir, "evaluation_pipeline", "sentence_zero_shot", "run.py")
577 if not os.path.exists(run_file):
578 print(f"[GlobalPIQA] Warning: {run_file} not found. Cannot patch.")
579 return
580
581 print(f"[GlobalPIQA] Patching local checkpoint loader in {run_file}...")
582 with open(run_file, "r") as f:
583 content = f.read()
584
585 # Check if already patched
586 if "Local checkpoint directory patch" in content:
587 print("[GlobalPIQA] Script already patched.")
588 return
589
590 target_str = """def main():
591 args = _parse_arguments()
592 if args.images_path is not None:
593 assert args.batch_size == 1, "Multimodal only works in batch size 1!"
594 dataset = args.data_path.stem
595 args.model_name = pathlib.Path(args.model_path_or_name).stem
596 if args.revision_name is None:
597 revision_name = "main"
598 else:
599 revision_name = args.revision_name"""
600
601 patch_str = """def main():
602 args = _parse_arguments()
603 if args.images_path is not None:
604 assert args.batch_size == 1, "Multimodal only works in batch size 1!"
605 dataset = args.data_path.stem
606
607 # Local checkpoint directory patch
608 import os
609 model_path = args.model_path_or_name
610 args.model_name = pathlib.Path(model_path).stem
611 revision_name = args.revision_name if args.revision_name else "main"
612
613 if os.path.isdir(model_path):
614 target_revision = args.revision_name if args.revision_name else "main"
615 if os.path.exists(os.path.join(model_path, target_revision)):
616 args.model_path_or_name = os.path.join(model_path, target_revision)
617 args.revision_name = None"""
618
619 if target_str in content:
620 new_content = content.replace(target_str, patch_str)
621 with open(run_file, "w") as f:
622 f.write(new_content)
623 print("[GlobalPIQA] Successfully patched run.py")
624 else:
625 print("[GlobalPIQA] Warning: Could not find target pattern in run.py. Manual patch may be needed.")
626
627 # Patch sentence zero shot loader inside cloned repo
628 patch_evaluation_run_script(strict_dir)
629
630 print("[Eval] Stripping Windows-specific packages from requirements.txt...")
631 req_file_path = os.path.join(strict_dir, "requirements.txt")
632 if os.path.exists(req_file_path):
633 with open(req_file_path, "r") as f:
634 lines = f.readlines()
635 with open(req_file_path, "w") as f:
636 for line in lines:
637 if "pywin" not in line.lower() and "wintypes" not in line.lower():
638 f.write(line)
639
640 print("[Eval] Verifying and installing evaluation dependencies programmatically...")
641 required_packages = {
642 "nltk": "nltk",
643 "pandas": "pandas",
644 "statsmodels": "statsmodels",
645 "sklearn": "scikit-learn",
646 "scipy": "scipy"
647 }
648 for pkg_import, pkg_install in required_packages.items():
649 try:
650 __import__(pkg_import)
651 except ImportError:
652 print(f"[Eval] Package '{pkg_install}' not found. Installing it programmatically...")
653 import sys
654 subprocess.run([sys.executable, "-m", "pip", "install", pkg_install], check=True)
655
656 print("[Eval] Downloading NLTK tokenizer resources...")
657 import nltk
658 nltk.download('punkt', download_dir=os.environ["NLTK_DATA"])
659 nltk.download('punkt_tab', download_dir=os.environ["NLTK_DATA"])
660
661 # Ensure standard zero-shot datasets are downloaded
662 blimp_fast_dir = os.path.join(strict_dir, "evaluation_data", "fast_eval", "blimp_fast")
663 if not os.path.exists(blimp_fast_dir) or not os.listdir(blimp_fast_dir):
664 print("[Eval] Standard zero-shot datasets not found. Downloading...")
665 subprocess.run(["python", "-m", "scripts.download_evals"], cwd=strict_dir, check=True)
666
667 # Unzip EWoK fast
668 ewok_zip = os.path.join(strict_dir, "evaluation_data/fast_eval/ewok_fast.zip")
669 if os.path.exists(ewok_zip):
670 print("[Eval] Unzipping EWoK fast data...")
671 bad_nested_dir = os.path.join(strict_dir, "evaluation_data/fast_eval/evaluation_data")
672 if os.path.exists(bad_nested_dir):
673 shutil.rmtree(bad_nested_dir)
674 subprocess.run(["unzip", "-o", "-P", "BabyLM2025", "evaluation_data/fast_eval/ewok_fast.zip", "-d", "."], cwd=strict_dir, check=True)
675
676 # Download EWoK full
677 print("[Eval] Downloading and filtering full EWoK dataset...")
678 subprocess.run(["python", "-m", "evaluation_pipeline.ewok.dl_and_filter"], cwd=strict_dir, check=True)
679
680 # Download GlobalPIQA dataset
681 global_piqa_parallel_dir = os.path.join(strict_dir, "evaluation_data", "fast_eval", "global_piqa_parallel")
682 if not os.path.exists(global_piqa_parallel_dir) or not os.listdir(global_piqa_parallel_dir):
683 print("[Eval] GlobalPIQA dataset not found. Downloading...")
684 subprocess.run(["python", "evaluation_pipeline/global_piqa/dl.py"], cwd=strict_dir, check=True)
685
686 # Ensure all scripts are executable
687 print("[Eval] Making evaluation shell scripts executable...")
688 subprocess.run("chmod +x scripts/*.sh", shell=True, cwd=strict_dir, check=True)
689
690 def run_task_with_cache(checkpoint, task, output_subpath, cmd):
691 # Determine paths
692 local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", task, output_subpath)
693 if task == "reading":
694 local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", "reading")
695 elif task == "comps":
696 local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", "comps", "comps")
697
698 target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", task, output_subpath)
699 if task == "reading":
700 target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", "reading")
701 elif task == "comps":
702 target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", "comps", "comps")
703
704 # If cached, copy it over
705 cache_file = os.path.join(local_cache_path, "predictions.json")
706
707 # Self-healing: invalidate old unfiltered entity_tracking caches
708 if task == "entity_tracking" and os.path.exists(cache_file):
709 try:
710 import json
711 with open(cache_file, "r") as f:
712 preds = json.load(f)
713 is_valid_cache = True
714 for k, v in preds.items():
715 if len(v.get("predictions", [])) in [605, 606, 607, 615, 529, 156, 187, 159]:
716 is_valid_cache = False
717 break
718 if not is_valid_cache:
719 print(f"[Eval] Cached entity_tracking for '{checkpoint}' has incorrect old sizes. Invalidate and re-run fresh...")
720 shutil.rmtree(local_cache_path, ignore_errors=True)
721 except Exception:
722 pass
723
724 if os.path.exists(cache_file):
725 print(f"[Eval] Task '{task}' ({output_subpath}) for checkpoint '{checkpoint}' is cached. Restoring...")
726 if os.path.exists(target_results_dir):
727 shutil.rmtree(target_results_dir)
728 os.makedirs(target_results_dir, exist_ok=True)
729 for item in os.listdir(local_cache_path):
730 s = os.path.join(local_cache_path, item)
731 d = os.path.join(target_results_dir, item)
732 if os.path.isdir(s):
733 shutil.copytree(s, d)
734 else:
735 shutil.copy2(s, d)
736 return
737
738 print(f"[Eval] Running task '{task}' ({output_subpath}) for checkpoint '{checkpoint}'...")
739 subprocess.run(cmd, cwd=strict_dir, check=True)
740
741 # Relocate from actual_model_basename, main, or checkpoint subfolder if needed
742 actual_model_basename = os.path.basename(model_dir.rstrip("/"))
743 possible_actual_path = os.path.join(strict_dir, "results", actual_model_basename, checkpoint, "zero_shot", "causal", task, output_subpath)
744 if task == "reading":
745 possible_actual_path = os.path.join(strict_dir, "results", actual_model_basename, checkpoint, "zero_shot", "causal", "reading")
746 elif task == "comps":
747 possible_actual_path = os.path.join(strict_dir, "results", actual_model_basename, checkpoint, "zero_shot", "causal", "comps", "comps")
748
749 possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", task, output_subpath)
750 if task == "reading":
751 possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", "reading")
752 elif task == "comps":
753 possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", "comps", "comps")
754
755 possible_local_reading_path = os.path.join(strict_dir, "results", checkpoint, "main", "zero_shot", "causal", "reading")
756
757 for p_path in [possible_actual_path, possible_main_path, possible_local_reading_path]:
758 if os.path.exists(p_path) and p_path != target_results_dir:
759 print(f"[Eval] Relocating results from {p_path} to {target_results_dir}...")
760 if os.path.exists(target_results_dir):
761 shutil.rmtree(target_results_dir)
762 os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
763 shutil.move(p_path, target_results_dir)
764 break
765
766 # Verify
767 verify_eval_run(target_results_dir, f"{checkpoint} {task} ({output_subpath})")
768
769 # Save to local cache
770 if os.path.exists(local_cache_path):
771 shutil.rmtree(local_cache_path)
772 os.makedirs(local_cache_path, exist_ok=True)
773 for item in os.listdir(target_results_dir):
774 s = os.path.join(target_results_dir, item)
775 d = os.path.join(local_cache_path, item)
776 if os.path.isdir(s):
777 shutil.copytree(s, d)
778 else:
779 shutil.copy2(s, d)
780
781 def run_finetune_task_with_cache(task, cmd):
782 local_cache_path = os.path.join("./results", model_name, "main", "finetune", task)
783 target_results_dir = os.path.join(strict_dir, "results", model_name, "main", "finetune", task)
784
785 if os.path.exists(os.path.join(local_cache_path, "predictions.json")):
786 print(f"[Eval] GLUE task '{task}' is cached. Restoring...")
787 if os.path.exists(target_results_dir):
788 shutil.rmtree(target_results_dir)
789 os.makedirs(target_results_dir, exist_ok=True)
790 for item in os.listdir(local_cache_path):
791 s = os.path.join(local_cache_path, item)
792 d = os.path.join(target_results_dir, item)
793 if os.path.isdir(s):
794 shutil.copytree(s, d)
795 else:
796 shutil.copy2(s, d)
797 return
798
799 print(f"[Eval] Running GLUE task '{task}'...")
800 subprocess.run(cmd, cwd=strict_dir, check=True)
801
802 # Relocate from results/main/main if needed
803 possible_main_path = os.path.join(strict_dir, "results", "main", "main", "finetune", task)
804 if os.path.exists(possible_main_path) and possible_main_path != target_results_dir:
805 print(f"[Eval] Relocating results from {possible_main_path} to {target_results_dir}...")
806 if os.path.exists(target_results_dir):
807 shutil.rmtree(target_results_dir)
808 os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
809 shutil.move(possible_main_path, target_results_dir)
810
811 # Verify
812 verify_eval_run(target_results_dir, f"GLUE task {task}")
813
814 # Cache locally
815 if os.path.exists(local_cache_path):
816 shutil.rmtree(local_cache_path)
817 os.makedirs(local_cache_path, exist_ok=True)
818 for item in os.listdir(target_results_dir):
819 s = os.path.join(target_results_dir, item)
820 d = os.path.join(local_cache_path, item)
821 if os.path.isdir(s):
822 shutil.copytree(s, d)
823 else:
824 shutil.copy2(s, d)
825
826 def run_aoa_with_cache(cmd):
827 local_cache_path = os.path.join("./results", model_name, "main", "aoa")
828 target_results_dir = os.path.join(strict_dir, "results", model_name, "main", "aoa")
829
830 if os.path.exists(os.path.join(local_cache_path, "aoa_score.json")) or os.path.exists(os.path.join(local_cache_path, "surprisal.json")):
831 print(f"[Eval] AoA task is cached. Restoring...")
832 if os.path.exists(target_results_dir):
833 shutil.rmtree(target_results_dir)
834 os.makedirs(target_results_dir, exist_ok=True)
835 for item in os.listdir(local_cache_path):
836 s = os.path.join(local_cache_path, item)
837 d = os.path.join(target_results_dir, item)
838 if os.path.isdir(s):
839 shutil.copytree(s, d)
840 else:
841 shutil.copy2(s, d)
842 return
843
844 print("[Eval] Running AoA task...")
845 subprocess.run(cmd, cwd=strict_dir, check=True)
846
847 # Relocate from results/main/main if needed
848 possible_main_path = os.path.join(strict_dir, "results", "main", "main", "aoa")
849 if os.path.exists(possible_main_path) and possible_main_path != target_results_dir:
850 print(f"[Eval] Relocating results from {possible_main_path} to {target_results_dir}...")
851 if os.path.exists(target_results_dir):
852 shutil.rmtree(target_results_dir)
853 os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
854 shutil.move(possible_main_path, target_results_dir)
855
856 # Verify
857 verify_eval_run(target_results_dir, "AoA task")
858
859 # Cache locally
860 if os.path.exists(local_cache_path):
861 shutil.rmtree(local_cache_path)
862 os.makedirs(local_cache_path, exist_ok=True)
863 for item in os.listdir(target_results_dir):
864 s = os.path.join(target_results_dir, item)
865 d = os.path.join(local_cache_path, item)
866 if os.path.isdir(s):
867 shutil.copytree(s, d)
868 else:
869 shutil.copy2(s, d)
870
871 # ─────────────────────────────────────────────────────────────
872 # B. FINAL MODEL 'main' FULL ZERO-SHOT EVALUATION
873 # ─────────────────────────────────────────────────────────────
874 main_ckpt_path = os.path.join(model_dir, "main")
875 if os.path.exists(main_ckpt_path):
876 print(f"[Eval] Running full zero-shot evaluation on main...")
877 # blimp filtered
878 run_task_with_cache(
879 "main", "blimp", "blimp_filtered",
880 ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/blimp_filtered", "--save_predictions"]
881 )
882 # supplement filtered
883 run_task_with_cache(
884 "main", "blimp", "supplement_filtered",
885 ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/supplement_filtered", "--save_predictions"]
886 )
887 # ewok filtered
888 run_task_with_cache(
889 "main", "ewok", "ewok_filtered",
890 ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "ewok", "--data_path", "evaluation_data/full_eval/ewok_filtered", "--save_predictions"]
891 )
892 # entity tracking
893 run_task_with_cache(
894 "main", "entity_tracking", "entity_tracking",
895 ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "entity_tracking", "--data_path", "evaluation_data/full_eval/entity_tracking", "--save_predictions"]
896 )
897 # comps
898 run_task_with_cache(
899 "main", "comps", "comps",
900 ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "comps", "--data_path", "evaluation_data/full_eval/comps", "--save_predictions"]
901 )
902 # reading
903 run_task_with_cache(
904 "main", "reading", "reading",
905 ["python", "-m", "evaluation_pipeline.reading.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--data_path", "evaluation_data/full_eval/reading/reading_data.csv"]
906 )
907 # global piqa parallel
908 run_task_with_cache(
909 "main", "global_piqa_parallel", "global_piqa_parallel",
910 ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "global_piqa_parallel", "--data_path", "evaluation_data/full_eval/global_piqa_parallel", "--save_predictions"]
911 )
912 # global piqa nonparallel
913 run_task_with_cache(
914 "main", "global_piqa_nonparallel", "global_piqa_nonparallel",
915 ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "global_piqa_nonparallel", "--data_path", "evaluation_data/full_eval/global_piqa_nonparallel", "--save_predictions"]
916 )
917
918 # ─────────────────────────────────────────────────────────────
919 # C. FINAL MODEL GLUE AND AOA EVALUATION
920 # ─────────────────────────────────────────────────────────────
921 if os.path.exists(main_ckpt_path):
922 # 1. GLUE fine-tuning
923 if skip_glue:
924 print("[Eval] Skipping GLUE fine-tuning evaluations as requested.")
925 else:
926 print("[Eval] Running GLUE fine-tuning evaluations on main task-by-task...")
927 glue_tasks = {
928 "boolq": ["boolq", "16", "10"],
929 "multirc": ["multirc", "16", "10"],
930 "rte": ["rte", "32", "10"],
931 "wsc": ["wsc", "32", "30"],
932 "mrpc": ["mrpc", "32", "10"],
933 "qqp": ["qqp", "32", "10"],
934 "mnli": ["mnli", "32", "10"]
935 }
936 for task_name, (task, bsz, max_epochs) in glue_tasks.items():
937 num_labels = "3" if task == "mnli" else "2"
938 metric_for_valid = "accuracy"
939 if task in ["mrpc", "qqp"]:
940 metric_for_valid = "f1"
941 metrics = ["accuracy"]
942 if task != "mnli":
943 metrics = ["accuracy", "f1", "mcc"]
944
945 cmd = [
946 "python", "-m", "evaluation_pipeline.finetune.run",
947 "--model_name_or_path", main_ckpt_path,
948 "--train_data", f"evaluation_data/full_eval/glue_filtered/{task}.train.jsonl",
949 "--valid_data", f"evaluation_data/full_eval/glue_filtered/{task}.valid.jsonl",
950 "--predict_data", f"evaluation_data/full_eval/glue_filtered/{task}.valid.jsonl",
951 "--task", task,
952 "--num_labels", num_labels,
953 "--batch_size", bsz,
954 "--learning_rate", "3e-5",
955 "--num_epochs", max_epochs,
956 "--sequence_length", "512",
957 "--results_dir", "results",
958 "--save",
959 "--save_dir", "models",
960 "--metric_for_valid", metric_for_valid,
961 "--seed", "42",
962 "--verbose",
963 "--padding_side", "left",
964 "--take_final"
965 ]
966 cmd.append("--metrics")
967 cmd.extend(metrics)
968
969 run_finetune_task_with_cache(task_name, cmd)
970
971 # 2. AoA
972 if skip_aoa:
973 print("[Eval] Skipping AoA evaluations as requested.")
974 else:
975 run_aoa_with_cache([
976 "python", "-m", "evaluation_pipeline.AoA_word.run",
977 "--model_name", model_dir,
978 "--backend", "causal",
979 "--track_name", "strict-small",
980 "--word_path", "evaluation_data/full_eval/aoa/cdi_childes.json",
981 "--output_dir", "results"
982 ])
983
984 # ─────────────────────────────────────────────────────────────
985 # A. INTERMEDIATE CHECKPOINTS FAST EVALUATION
986 # ─────────────────────────────────────────────────────────────
987 print(f"[Eval] Running zero-shot fast evaluations on intermediate checkpoints...")
988 checkpoints = [f"chck_{i}M" for i in range(1, 10)] + [f"chck_{i}M" for i in range(10, 110, 10)]
989
990 eval_model_path = model_dir
991
992 for checkpoint in checkpoints:
993 ckpt_full_path = os.path.join(model_dir, checkpoint)
994 if not os.path.exists(ckpt_full_path):
995 continue
996
997 # blimp fast
998 run_task_with_cache(
999 checkpoint, "blimp", "blimp_fast",
1000 ["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]
1001 )
1002 # supplement fast
1003 run_task_with_cache(
1004 checkpoint, "blimp", "supplement_fast",
1005 ["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]
1006 )
1007 # ewok fast
1008 run_task_with_cache(
1009 checkpoint, "ewok", "ewok_fast",
1010 ["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]
1011 )
1012 # entity tracking fast
1013 run_task_with_cache(
1014 checkpoint, "entity_tracking", "entity_tracking_fast",
1015 ["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]
1016 )
1017 # reading fast
1018 run_task_with_cache(
1019 checkpoint, "reading", "reading",
1020 ["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"]
1021 )
1022 # global piqa parallel fast
1023 run_task_with_cache(
1024 checkpoint, "global_piqa_parallel", "global_piqa_parallel",
1025 ["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]
1026 )
1027 # global piqa nonparallel fast
1028 run_task_with_cache(
1029 checkpoint, "global_piqa_nonparallel", "global_piqa_nonparallel",
1030 ["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]
1031 )
1032
1033 # ─────────────────────────────────────────────────────────────
1034 # D. COLLATE RESULTS AND CLEANUP
1035 # ─────────────────────────────────────────────────────────────
1036 print("[Eval] Collating predictions into submission file...")
1037 # Clean collation destination in evaluation repo
1038 collate_results_dir = os.path.join(strict_dir, "results", model_name)
1039 if os.path.exists(collate_results_dir):
1040 shutil.rmtree(collate_results_dir)
1041 os.makedirs(os.path.dirname(collate_results_dir), exist_ok=True)
1042
1043 # Copy from local cache results to strict results for collation
1044 shutil.copytree(local_results_dir, collate_results_dir)
1045
1046 # Run collation
1047 subprocess.run([
1048 "python", "-m", "evaluation_pipeline.collate_preds",
1049 "--model_path_or_name", model_name,
1050 "--backend", "causal",
1051 "--track", "strict-small",
1052 "--fast"
1053 ], cwd=strict_dir, check=True)
1054
1055 # Save results to local folder
1056 results_src = os.path.join(strict_dir, "results")
1057 results_dest = os.path.abspath("./results")
1058 if os.path.exists(results_dest):
1059 shutil.rmtree(results_dest)
1060 shutil.copytree(results_src, results_dest)
1061
1062 # Copy final collated json to current folder
1063 collated_json = os.path.join(strict_dir, "all_full_preds_and_fast_scores_causal.json")
1064 if os.path.exists(collated_json):
1065 shutil.copy(collated_json, "./all_full_preds_and_fast_scores_causal.json")
1066 print("\n[Eval] Success! Collation completed! Final file is at './all_full_preds_and_fast_scores_causal.json'")
1067
1068 print("\n[Eval] Pipeline evaluation run finished.")
1069
1070def upload_pipeline(model_name, repo_name, token=None):
1071 import os
1072 import shutil
1073 import json
1074 import hashlib
1075 from huggingface_hub import HfApi, create_repo
1076
1077 if not token:
1078 token = os.environ.get("HF_TOKEN")
1079
1080 api = HfApi(token=token)
1081 try:
1082 user_info = api.whoami()
1083 username = user_info["name"]
1084 print(f"[HF] Authenticated successfully as user: {username}")
1085 except Exception as e:
1086 print(f"[HF] Authentication failed. Error: {e}")
1087 return
1088
1089 repo_id = f"{username}/{repo_name}"
1090 print(f"[HF] Target Repository ID: {repo_id}")
1091
1092 # Create the repository if it doesn't exist
1093 try:
1094 create_repo(repo_id=repo_id, repo_type="model", token=token, exist_ok=True)
1095 print(f"[HF] Repository '{repo_id}' is ready.")
1096 except Exception as e:
1097 print(f"[HF] Failed to verify or create repository. Error: {e}")
1098 return
1099
1100 checkpoint_dir = os.path.abspath(f"./checkpoints/{model_name}")
1101
1102 # Resolve the main checkpoint directory using self-healing rules
1103 revisions = {}
1104 if os.path.exists("pytorch_model.bin") and os.path.exists("config.json"):
1105 print("[HF] Detected weight and config files in the current working directory. Using current folder as 'main' checkpoint.")
1106 revisions = {"main": os.getcwd()}
1107 elif os.path.exists(os.path.join(checkpoint_dir, "main")):
1108 revisions = {"main": os.path.join(checkpoint_dir, "main")}
1109 elif os.path.exists(os.path.abspath("./checkpoints/msit_gptbert_fresh/main")):
1110 print("[HF] Using fallback checkpoint folder './checkpoints/msit_gptbert_fresh/main'...")
1111 revisions = {"main": os.path.abspath("./checkpoints/msit_gptbert_fresh/main")}
1112 elif os.path.exists(os.path.abspath("./checkpoints/main")):
1113 revisions = {"main": os.path.abspath("./checkpoints/main")}
1114 else:
1115 print(f"[HF] Error: Could not locate the 'main' checkpoint weights. Checked: {checkpoint_dir}/main, current directory, and fallbacks.")
1116 return
1117
1118 # Check for intermediate checkpoints relative to the main checkpoint's parent folder
1119 main_dir = revisions["main"]
1120 parent_dir = os.path.dirname(main_dir)
1121
1122 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]:
1123 ckpt_name = f"chck_{m}M"
1124 ckpt_path = os.path.join(parent_dir, ckpt_name)
1125 if os.path.exists(ckpt_path):
1126 revisions[ckpt_name] = ckpt_path
1127 else:
1128 # Check if there is a .pt file in the parent folder
1129 pt_path = os.path.join(parent_dir, f"{ckpt_name}.pt")
1130 if os.path.exists(pt_path):
1131 print(f"[HF] Found legacy checkpoint file '{ckpt_name}.pt'. Converting to HF format for upload...")
1132 import torch
1133 from transformers import AutoTokenizer
1134 from tokenizers import Tokenizer
1135 try:
1136 from modeling_xpertgpt import XpertGPTForCausalLM, XpertGPTConfig
1137 cfg = XpertGPTConfig(d_model=256, d_thin=384, num_layers=6, num_blocks=4)
1138 model_to_save = XpertGPTForCausalLM(cfg)
1139 sd = torch.load(pt_path, map_location="cpu")
1140 clean_sd = {}
1141 for k, v in sd.items():
1142 new_k = k.replace("module.", "")
1143 clean_sd[new_k] = v
1144 model_to_save.load_state_dict(clean_sd)
1145
1146 vocab_path = os.path.join(parent_dir, "bpe_vocab_16k.json")
1147 if not os.path.exists(vocab_path):
1148 vocab_path = os.path.join(main_dir, "bpe_vocab_16k.json")
1149 if os.path.exists(vocab_path):
1150 tok = Tokenizer.from_file(vocab_path)
1151 else:
1152 tok = None
1153
1154 os.makedirs(ckpt_path, exist_ok=True)
1155 state_dict = model_to_save.state_dict()
1156 new_state_dict = {}
1157 for k, v in state_dict.items():
1158 name = k
1159 if name.startswith("_orig_mod."):
1160 name = name[10:]
1161 if name.startswith("model."):
1162 name = name[6:]
1163 if name == "lm_head.weight":
1164 new_state_dict["lm_head.weight"] = v
1165 else:
1166 new_state_dict[f"transformer.{name}"] = v
1167 torch.save(new_state_dict, os.path.join(ckpt_path, "pytorch_model.bin"))
1168
1169 config_dict = {
1170 "auto_map": {
1171 "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
1172 "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
1173 "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
1174 },
1175 "vocab_size": 16384,
1176 "block_size": 512,
1177 "d_model": 256,
1178 "hidden_size": 256,
1179 "d_thin": 384,
1180 "num_layers": 6,
1181 "num_blocks": 4,
1182 "capacity_factor": 2.0,
1183 "dropout": 0.1,
1184 "model_type": "xpertgpt",
1185 "num_hidden_layers": 6
1186 }
1187 with open(os.path.join(ckpt_path, "config.json"), "w") as f:
1188 json.dump(config_dict, f, indent=2)
1189
1190 if tok:
1191 from transformers import PreTrainedTokenizerFast
1192 fast_tokenizer = PreTrainedTokenizerFast(
1193 tokenizer_object=tok,
1194 bos_token="[CLS]",
1195 eos_token="[SEP]",
1196 unk_token="[UNK]",
1197 pad_token="[PAD]",
1198 mask_token="[MASK]"
1199 )
1200 fast_tokenizer.save_pretrained(ckpt_path)