import os import math import time import json import random import inspect import shutil import subprocess import argparse # ───────────────────────────────────────────────────────────── # 1. CORE PIPELINE FUNCTION AND HELPERS # ───────────────────────────────────────────────────────────── def is_zero_shot_cache_valid(cache_dir): if not os.path.exists(cache_dir): return False pred_count = 0 for root, dirs, files in os.walk(cache_dir): if "predictions.json" in files: pred_count += 1 return pred_count >= 3 def is_finetune_cache_valid(cache_dir): if not os.path.exists(cache_dir): return False pred_count = 0 for root, dirs, files in os.walk(cache_dir): if "predictions.json" in files: pred_count += 1 return pred_count >= 3 def run_pipeline(model_name: str, epochs: int = 10, skip_eval: bool = False, skip_aoa: bool = False, skip_glue: bool = False): # Configure persistent cache paths locally to avoid duplicate downloads os.environ["HF_HOME"] = os.path.abspath("./hf_cache") os.environ["NLTK_DATA"] = os.path.abspath("./nltk_data") os.makedirs("./hf_cache", exist_ok=True) os.makedirs("./nltk_data", exist_ok=True) # Programmatic Hugging Face Hub Login if HF_TOKEN is in environment hf_token = os.environ.get("HF_TOKEN") if hf_token: try: from huggingface_hub import login login(token=hf_token) print("[HF] Programmatic login successful using HF_TOKEN.") except Exception as e: print(f"[HF] Warning: Programmatic login failed: {e}") import torch import torch.nn as nn import torch.nn.functional as F from datasets import load_dataset from tokenizers import Tokenizer from transformers import PreTrainedTokenizerFast # Import model architecture from modeling_msit import ( MSITGPTBERTModel, MSITGPTBERTConfig, MSITGPTBERTHFConfig ) # Print GPU details if torch.cuda.is_available(): gpu_name = torch.cuda.get_device_name(0) print(f"\n[GPU] CUDA is available! Using GPU: {gpu_name}\n") else: print("\n[GPU] Warning: CUDA is NOT available! Running on CPU.\n") # ───────────────────────────────────────────────────────────── # GLOBAL HYPERPARAMETERS # ───────────────────────────────────────────────────────────── VOCAB_SIZE = 16384 MASK_TOKEN_ID = 16383 BLOCK_SIZE = 512 BATCH_SIZE = 16 GRAD_ACCUM_STEPS = 1 # grad_acc_step = 1 as requested EPOCHS = epochs LEARNING_RATE = 1e-4 LR_MIN = LEARNING_RATE * 0.05 WARMUP_STEPS = 50 WEIGHT_DECAY = 0.1 GRAD_CLIP = 1.0 NUM_THIN_BLOCKS = 6 EC_CAPACITY_FACTOR = 2.0 CAUSAL_RATIO = 1 / 1 MASK_PROB_START = 0.20 MASK_PROB_END = 0.10 # Output directories locally (persistent in Studio workspace) model_dir = os.path.abspath(f"./checkpoints/{model_name}") os.makedirs(model_dir, exist_ok=True) # ───────────────────────────────────────────────────────────── # Helper: Save Hugging Face Compliant Checkpoint # ───────────────────────────────────────────────────────────── def save_hf_checkpoint(raw_model, checkpoint_dir_name, tokenizer): save_dir = os.path.join(model_dir, checkpoint_dir_name) os.makedirs(save_dir, exist_ok=True) print(f"\n[Checkpoint] Saving Hugging Face format checkpoint to '{save_dir}'...") # A. Convert state dict keys to CausalLM wrapper naming state_dict = raw_model.state_dict() new_state_dict = {} for k, v in state_dict.items(): name = k if name.startswith("_orig_mod."): name = name[10:] if name.startswith("model."): name = name[6:] if name == "lm_head.weight": new_state_dict["lm_head.weight"] = v else: new_state_dict[f"transformer.{name}"] = v torch.save(new_state_dict, os.path.join(save_dir, "pytorch_model.bin")) # B. Copy modeling.py shutil.copy("modeling_msit.py", os.path.join(save_dir, "modeling_msit.py")) # C. Create config.json config_dict = { "auto_map": { "AutoConfig": "modeling_msit.MSITGPTBERTHFConfig", "AutoModel": "modeling_msit.MSITGPTBERTModelWrapper", "AutoModelForCausalLM": "modeling_msit.MSITGPTBERTForCausalLM" }, "vocab_size": VOCAB_SIZE, "block_size": BLOCK_SIZE, "d_model": 384, # 384 d_model as requested "hidden_size": 384, # 384 hidden_size "d_thin": 192, "num_layers": 6, "num_blocks": NUM_THIN_BLOCKS, "capacity_factor": EC_CAPACITY_FACTOR, "dropout": 0.1, "model_type": "msit_gptbert" } with open(os.path.join(save_dir, "config.json"), "w") as f: json.dump(config_dict, f, indent=2) # D. Save tokenizer config files fast_tokenizer = PreTrainedTokenizerFast( tokenizer_object=tokenizer, bos_token="[CLS]", eos_token="[SEP]", unk_token="[UNK]", pad_token="[PAD]", mask_token="[MASK]" ) fast_tokenizer.save_pretrained(save_dir) print(f"[Checkpoint] Checkpoint '{checkpoint_dir_name}' successfully saved.") # ───────────────────────────────────────────────────────────── # Tokenizer Training # ───────────────────────────────────────────────────────────── def build_and_train_tokenizer(texts: list) -> Tokenizer: from tokenizers.models import BPE from tokenizers.trainers import BpeTrainer from tokenizers.pre_tokenizers import Whitespace vocab_path = os.path.join(model_dir, "bpe_vocab_16k.json") if os.path.exists(vocab_path): print(f"[Tokenizer] Loading trained BPE model layout from '{vocab_path}'...") return Tokenizer.from_file(vocab_path) print(f"[Tokenizer] Generating fresh HuggingFace BPE Tokenizer model with {VOCAB_SIZE} slots...") tokenizer = Tokenizer(BPE(unk_token="[UNK]")) tokenizer.pre_tokenizer = Whitespace() trainer = BpeTrainer( vocab_size=VOCAB_SIZE, special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] ) tokenizer.train_from_iterator(texts, trainer) tokenizer.save(vocab_path) print(f"[Tokenizer] Tokenizer training completed and saved to '{vocab_path}'.") return tokenizer # ───────────────────────────────────────────────────────────── # Data Loader Setup # ───────────────────────────────────────────────────────────── class DataLoaderLite: def __init__(self, B: int, T: int, texts: list, tokenizer: Tokenizer, name: str): self.B = B self.T = T print(f"[DataLoader:{name}] Tokenising dataset sequences...") all_ids = [] for t in texts: if t.strip(): encoded = tokenizer.encode(t).ids all_ids.extend(encoded) self.tokens = torch.tensor(all_ids, dtype=torch.long) self.chunk_size = B * T self.n_chunks = (len(self.tokens) - 1) // self.chunk_size self.indices = list(range(self.n_chunks)) self.pos = 0 self._shuffle() print(f"[DataLoader:{name}] Total tokens: {len(self.tokens):,} | Epoch steps: {self.n_chunks:,}") def _shuffle(self): random.shuffle(self.indices) self.pos = 0 def steps_per_epoch(self) -> int: return self.n_chunks def next_batch(self): B, T = self.B, self.T if self.pos >= len(self.indices): self._shuffle() chunk_idx = self.indices[self.pos] self.pos += 1 start_pos = chunk_idx * self.chunk_size temp = self.tokens[start_pos : start_pos + self.chunk_size + 1] x = temp[:-1].view(B, T) y = temp[1:].view(B, T) return x, y # ───────────────────────────────────────────────────────────── # Batch preparation and schedules # ───────────────────────────────────────────────────────────── def get_current_mask_prob(global_step: int, total_steps: int) -> float: ratio = min(1.0, global_step / total_steps) return MASK_PROB_START + ratio * (MASK_PROB_END - MASK_PROB_START) def prepare_causal_batch(x: torch.Tensor, y: torch.Tensor): return x, y, False def prepare_masked_batch(x: torch.Tensor, y: torch.Tensor, mask_prob: float, mask_token_id: int): B, T = x.size() mask = torch.rand(B, T, device=x.device) < mask_prob masked_x = x.clone() masked_x[mask] = mask_token_id targets = torch.full_like(y, -100) targets[mask] = y[mask] return masked_x, targets, True def get_hybrid_batch(train_loader: DataLoaderLite, global_step: int, total_steps: int, device: torch.device): x, y = train_loader.next_batch() x, y = x.to(device), y.to(device) if random.random() < CAUSAL_RATIO: input_ids, targets, bidir = prepare_causal_batch(x, y) else: mask_prob = get_current_mask_prob(global_step, total_steps) input_ids, targets, bidir = prepare_masked_batch(x, y, mask_prob, MASK_TOKEN_ID) return input_ids, targets, bidir def get_lr(it: int, total_steps: int) -> float: if it < WARMUP_STEPS: return LEARNING_RATE * (it + 1) / WARMUP_STEPS if it >= total_steps: return LR_MIN decay_ratio = (it - WARMUP_STEPS) / (total_steps - WARMUP_STEPS) coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) return LR_MIN + coeff * (LEARNING_RATE - LR_MIN) # ───────────────────────────────────────────────────────────── # Dataset Preparation # ───────────────────────────────────────────────────────────── print("\n[Data] Loading BabyLM-2026-Strict-Small ...") ds = load_dataset("BabyLM-community/BabyLM-2026-Strict-Small") all_text = list(ds['train']['text']) tokenizer = build_and_train_tokenizer(all_text) split = int(len(all_text) * 0.95) train_texts = all_text[:split] val_texts = all_text[split:] train_loader = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, train_texts, tokenizer, "train") chunks_per_epoch = train_loader.steps_per_epoch() steps_per_epoch = chunks_per_epoch // GRAD_ACCUM_STEPS total_steps = steps_per_epoch * EPOCHS cfg = MSITGPTBERTConfig() # Will use updated d_model = 384 device = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(42) if torch.cuda.is_available(): torch.cuda.manual_seed(42) random.seed(42) if hasattr(torch, 'set_float32_matmul_precision'): torch.set_float32_matmul_precision('high') model = MSITGPTBERTModel(cfg).to(device) # ───────────────────────────────────────────────────────────── # Training Resume Check # ───────────────────────────────────────────────────────────── words_trained = 0 next_milestone_idx = 0 global_step = 0 milestones = [i * 1_000_000 for i in range(1, 10)] + [i * 10_000_000 for i in range(1, 11)] resume_checkpoint_dir = None for idx in range(len(milestones) - 1, -1, -1): m = milestones[idx] ckpt_name = f"chck_{m // 1_000_000}M" ckpt_path = os.path.join(model_dir, ckpt_name) if os.path.exists(os.path.join(ckpt_path, "pytorch_model.bin")): config_json_path = os.path.join(ckpt_path, "config.json") if os.path.exists(config_json_path): try: with open(config_json_path, "r") as f: saved_config = json.load(f) # Verify d_model matches to avoid weight shape mismatch crashes if saved_config.get("d_model") == 384: resume_checkpoint_dir = ckpt_path next_milestone_idx = idx + 1 words_trained = m global_step = words_trained // (BATCH_SIZE * BLOCK_SIZE) print(f"[Training] Found existing milestone checkpoint '{ckpt_name}'. Resuming from step {global_step:,} ({words_trained:,} tokens trained)...") break else: print(f"[Training] Found checkpoint '{ckpt_name}' but it has mismatch d_model={saved_config.get('d_model')}. Starting fresh.") except Exception as e: pass # Load weights if resuming if resume_checkpoint_dir is not None: print(f"[Model] Loading weights from checkpoint '{resume_checkpoint_dir}'...") state_dict = torch.load(os.path.join(resume_checkpoint_dir, "pytorch_model.bin"), map_location=device) model_state_dict = {} for k, v in state_dict.items(): name = k if name.startswith("transformer."): name = name[12:] model_state_dict[name] = v model.load_state_dict(model_state_dict) # Check if final main model exists main_ckpt_path = os.path.join(model_dir, "main") if os.path.exists(os.path.join(main_ckpt_path, "pytorch_model.bin")): print("\n[Pipeline] Final checkpoint 'main' already exists. Skipping training phase and transitioning directly to evaluations!") else: # torch.compile try: model = torch.compile(model) print("[Model] torch.compile() successfully verified graph optimizations") except Exception as e: print(f"[Model] torch.compile() skipped ({e})") # Optimizer param_dict = {n: p for n, p in model.named_parameters() if p.requires_grad} decay_params = [p for p in param_dict.values() if p.dim() >= 2] nodecay_params = [p for p in param_dict.values() if p.dim() < 2] groups = [ {'params': decay_params, 'weight_decay': WEIGHT_DECAY}, {'params': nodecay_params, 'weight_decay': 0.0}, ] fused_ok = 'fused' in inspect.signature(torch.optim.AdamW).parameters use_fused = fused_ok and ('cuda' in device) optimizer = torch.optim.AdamW(groups, lr=LEARNING_RATE, betas=(0.9, 0.95), eps=1e-8, fused=use_fused) # ───────────────────────────────────────────────────────────── # Training Loop # ───────────────────────────────────────────────────────────── model.train() autocast_ctx = torch.autocast(device_type="cuda" if "cuda" in device else "cpu", dtype=torch.bfloat16, enabled=True) start_epoch = global_step // steps_per_epoch start_chunk = (global_step % steps_per_epoch) * GRAD_ACCUM_STEPS print(f"\n[Training] Starting MSIT-GPTBERT MoEP training for {EPOCHS} epochs...") for epoch in range(start_epoch, EPOCHS): train_loader._shuffle() if epoch == start_epoch and start_chunk > 0: print(f"[Training] Fast-forwarding dataloader to chunk index {start_chunk}...") train_loader.pos = start_chunk optimizer.zero_grad(set_to_none=True) loss_accum = 0.0 start_chunk_idx = start_chunk if epoch == start_epoch else 0 for chunk_step in range(start_chunk_idx, chunks_per_epoch): t0 = time.perf_counter() lr = get_lr(global_step, total_steps) for pg in optimizer.param_groups: pg['lr'] = lr input_ids, targets, bidir = get_hybrid_batch(train_loader, global_step, total_steps, device) words_trained += input_ids.numel() with autocast_ctx: _, loss = model(input_ids, targets, bidirectional=bidir) scaled_loss = loss / GRAD_ACCUM_STEPS loss_accum += scaled_loss.item() scaled_loss.backward() # Optimizer Step if (chunk_step + 1) % GRAD_ACCUM_STEPS == 0: norm = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP) optimizer.step() optimizer.zero_grad(set_to_none=True) if "cuda" in device: torch.cuda.synchronize() dt = (time.perf_counter() - t0) * 1000 mode_tag = "MLM" if bidir else "CLM" mask_p = get_current_mask_prob(global_step, total_steps) current_step = (chunk_step + 1) // GRAD_ACCUM_STEPS print( f"[E{epoch+1:02d} {current_step:>5d}/{steps_per_epoch} G{global_step:>7d}|{mode_tag}] " f"train={loss_accum:.4f} mask={mask_p:.1%} norm={norm:.3f} lr={lr:.2e} dt={dt:6.1f}ms words={words_trained:,}" ) loss_accum = 0.0 global_step += 1 # Check if we passed a milestone for checkpointing if next_milestone_idx < len(milestones) and words_trained >= milestones[next_milestone_idx]: milestone_val = milestones[next_milestone_idx] if milestone_val < 10_000_000: milestone_name = f"chck_{milestone_val // 1_000_000}M" else: milestone_name = f"chck_{(milestone_val // 10_000_000) * 10}M" raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model save_hf_checkpoint(raw_model, milestone_name, tokenizer) next_milestone_idx += 1 # Save final model as 'main' raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model save_hf_checkpoint(raw_model, "main", tokenizer) print("\n[Training] Training phase complete!") if skip_eval: print("[Pipeline] Skipping evaluations phase as requested.") return # ───────────────────────────────────────────────────────────── # 2. RUN EVALUATION PIPELINE # ───────────────────────────────────────────────────────────── local_results_dir = os.path.abspath(f"./results/{model_name}") os.makedirs(local_results_dir, exist_ok=True) local_main_res = os.path.join(local_results_dir, "main") # Delete invalid local caches before we start harvesting/evaluating if not is_zero_shot_cache_valid(os.path.join(local_main_res, "zero_shot")): if os.path.exists(os.path.join(local_main_res, "zero_shot")): print("[Eval] Local zero-shot cache is incomplete or corrupt. Cleaning up...") shutil.rmtree(os.path.join(local_main_res, "zero_shot")) if not is_finetune_cache_valid(os.path.join(local_main_res, "finetune")): if os.path.exists(os.path.join(local_main_res, "finetune")): print("[Eval] Local finetuning cache is incomplete or corrupt. Cleaning up...") shutil.rmtree(os.path.join(local_main_res, "finetune")) # Walk the entire workspace to search for completed evaluations workspace_root = "/teamspace/studios/this_studio" print(f"[Eval] Scanning workspace {workspace_root} to harvest any completed evaluations...") exclude_dirs = {"miniconda3", ".git", ".cache", "hf_cache", "nltk_data", "babylm_eval_repo"} if os.path.exists(workspace_root): for root, dirs, files in os.walk(workspace_root): # Speed up walking by pruning system dirs dirs[:] = [d for d in dirs if d not in exclude_dirs] parts = root.split(os.sep) # Skip checkpoint paths if any(p.startswith("chck_") for p in parts): continue # If we find a directory ending in zero_shot if root.endswith(f"{os.sep}zero_shot"): if is_zero_shot_cache_valid(root): dest = os.path.join(local_main_res, "zero_shot") if not is_zero_shot_cache_valid(dest): if os.path.exists(dest): shutil.rmtree(dest) os.makedirs(os.path.dirname(dest), exist_ok=True) shutil.copytree(root, dest) print(f"[Eval] Successfully harvested zero-shot results from {root}") # If we find a directory ending in finetune if root.endswith(f"{os.sep}finetune"): if is_finetune_cache_valid(root): dest = os.path.join(local_main_res, "finetune") if not is_finetune_cache_valid(dest): if os.path.exists(dest): shutil.rmtree(dest) os.makedirs(os.path.dirname(dest), exist_ok=True) shutil.copytree(root, dest) print(f"[Eval] Successfully harvested GLUE finetuning results from {root}") # Also harvest intermediate checkpoints from the old strict dir if it exists old_strict_dir = os.path.abspath("./babylm_eval_repo/babylm-eval/strict") old_model_results = os.path.join(old_strict_dir, "results", model_name) if os.path.exists(old_model_results): for item in os.listdir(old_model_results): src = os.path.join(old_model_results, item) dst = os.path.join(local_results_dir, item) if os.path.isdir(src) and item != "main" and (not os.path.exists(dst) or not os.listdir(dst)): shutil.copytree(src, dst) print(f"[Eval] Successfully harvested checkpoint '{item}' results.") clone_dir = os.path.abspath("./babylm_eval_repo") if os.path.exists(clone_dir): shutil.rmtree(clone_dir) print("\n[Eval] Cloning evaluation pipeline repository...") subprocess.run(["git", "clone", "https://github.com/atulgithub2/babylm.git", clone_dir], check=True) strict_dir = os.path.join(clone_dir, "babylm-eval", "strict") print("[Eval] Stripping Windows-specific packages from requirements.txt...") req_file_path = os.path.join(strict_dir, "requirements.txt") if os.path.exists(req_file_path): with open(req_file_path, "r") as f: lines = f.readlines() with open(req_file_path, "w") as f: for line in lines: if "pywin" not in line.lower() and "wintypes" not in line.lower(): f.write(line) print("[Eval] Skipping pipeline dependencies installation (as requested)...") print("[Eval] Downloading NLTK tokenizer resources...") import nltk nltk.download('punkt') nltk.download('punkt_tab') print("[Eval] Downloading zero-shot evaluation datasets...") subprocess.run(["python", "-m", "scripts.download_evals"], cwd=strict_dir, check=True) # Unzip EWoK fast ewok_zip = os.path.join(strict_dir, "evaluation_data/fast_eval/ewok_fast.zip") if os.path.exists(ewok_zip): print("[Eval] Unzipping EWoK fast data...") bad_nested_dir = os.path.join(strict_dir, "evaluation_data/fast_eval/evaluation_data") if os.path.exists(bad_nested_dir): shutil.rmtree(bad_nested_dir) subprocess.run(["unzip", "-o", "-P", "BabyLM2025", "evaluation_data/fast_eval/ewok_fast.zip", "-d", "."], cwd=strict_dir, check=True) # Download EWoK full print("[Eval] Downloading and filtering full EWoK dataset...") subprocess.run(["python", "-m", "evaluation_pipeline.ewok.dl_and_filter"], cwd=strict_dir, check=True) # Ensure all scripts are executable print("[Eval] Making evaluation shell scripts executable...") subprocess.run("chmod +x scripts/*.sh", shell=True, cwd=strict_dir, check=True) # ───────────────────────────────────────────────────────────── # A. INTERMEDIATE CHECKPOINTS FAST EVALUATION (First - Normal Flow) # ───────────────────────────────────────────────────────────── print(f"[Eval] Running zero-shot fast evaluations on intermediate checkpoints...") checkpoints = [f"chck_{i}M" for i in range(1, 10)] + [f"chck_{i}M" for i in range(10, 110, 10)] for checkpoint in checkpoints: local_ckpt_res = os.path.join(local_results_dir, checkpoint) # Caching check: if results exist locally, skip evaluations for this checkpoint! if os.path.exists(local_ckpt_res) and os.listdir(local_ckpt_res): print(f"[Eval] Checkpoint '{checkpoint}' already evaluated. Restoring results from cache...") target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint) if os.path.exists(target_results_dir): shutil.rmtree(target_results_dir) os.makedirs(os.path.dirname(target_results_dir), exist_ok=True) shutil.copytree(local_ckpt_res, target_results_dir) else: print(f"[Eval] Evaluating checkpoint '{checkpoint}'...") subprocess.run([ "bash", "scripts/eval_zero_shot_fast.sh", model_dir, checkpoint, "causal", "evaluation_data/fast_eval" ], cwd=strict_dir, check=True) # Immediately cache the result target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint) if os.path.exists(target_results_dir): print(f"[Eval] Caching '{checkpoint}' results...") if os.path.exists(local_ckpt_res): shutil.rmtree(local_ckpt_res) shutil.copytree(target_results_dir, local_ckpt_res) # ───────────────────────────────────────────────────────────── # B. FINAL MODEL 'main' EVALUATION (Second) # ───────────────────────────────────────────────────────────── main_ckpt_path = os.path.join(model_dir, "main") if os.path.exists(main_ckpt_path): # 1. Zero-shot check zero_shot_cache_file = os.path.join(local_main_res, "zero_shot") if is_zero_shot_cache_valid(zero_shot_cache_file): print("[Eval] Full zero-shot evaluation on main already completed and cached. Restoring to all expected targets...") # Restore to all possible directories the pipeline might seek targets = [ os.path.join(strict_dir, "results", "main", "main", "zero_shot"), os.path.join(strict_dir, "results", "main", "zero_shot"), os.path.join(strict_dir, "results", model_name, "main", "zero_shot"), ] for t_zs in targets: if os.path.exists(t_zs): shutil.rmtree(t_zs) os.makedirs(os.path.dirname(t_zs), exist_ok=True) shutil.copytree(zero_shot_cache_file, t_zs) else: print(f"[Eval] Running full zero-shot evaluation on main...") subprocess.run([ "./scripts/eval_zero_shot.sh", main_ckpt_path, "causal", "evaluation_data/full_eval" ], cwd=strict_dir, check=True) # Find generated zero-shot directory and cache it found_zs = None possible_paths = [ os.path.join(strict_dir, "results", "main", "main", "zero_shot"), os.path.join(strict_dir, "results", "main", "zero_shot"), os.path.join(strict_dir, "results", model_name, "main", "zero_shot"), ] for p in possible_paths: if os.path.exists(p) and os.listdir(p): found_zs = p break if found_zs: os.makedirs(local_main_res, exist_ok=True) dest = os.path.join(local_main_res, "zero_shot") if os.path.exists(dest): shutil.rmtree(dest) shutil.copytree(found_zs, dest) print(f"[Eval] Successfully cached zero-shot results from {found_zs}") # 2. GLUE Fine-Tuning check if skip_glue: print("[Eval] Skipping GLUE fine-tuning evaluations as requested by --skip-glue.") else: finetune_cache_file = os.path.join(local_main_res, "finetune") if is_finetune_cache_valid(finetune_cache_file): print("[Eval] GLUE fine-tuning on main already completed and cached. Restoring to all expected targets...") # Restore to all possible directories targets = [ os.path.join(strict_dir, "results", "main", "main", "finetune"), os.path.join(strict_dir, "results", "main", "finetune"), os.path.join(strict_dir, "results", model_name, "main", "finetune"), ] for t_ft in targets: if os.path.exists(t_ft): shutil.rmtree(t_ft) os.makedirs(os.path.dirname(t_ft), exist_ok=True) shutil.copytree(finetune_cache_file, t_ft) else: print(f"[Eval] Running GLUE fine-tuning evaluations on main (enforces batch_size=32 to prevent MoE OOM)...") subprocess.run([ "./scripts/eval_finetuning.sh", "--model_path", main_ckpt_path, "--lr", "3e-5", "--bsz", "32" ], cwd=strict_dir, check=True) # Find generated finetuning directory and cache it found_ft = None possible_paths = [ os.path.join(strict_dir, "results", "main", "main", "finetune"), os.path.join(strict_dir, "results", "main", "finetune"), os.path.join(strict_dir, "results", model_name, "main", "finetune"), ] for p in possible_paths: if os.path.exists(p) and os.listdir(p): found_ft = p break if found_ft: os.makedirs(local_main_res, exist_ok=True) dest = os.path.join(local_main_res, "finetune") if os.path.exists(dest): shutil.rmtree(dest) shutil.copytree(found_ft, dest) print(f"[Eval] Successfully cached finetuning results from {found_ft}") # 3. AoA check if skip_aoa: print("[Eval] Skipping AoA evaluations as requested by --skip-aoa.") else: aoa_cache_file = os.path.join(local_main_res, "aoa") if os.path.exists(aoa_cache_file) and os.listdir(aoa_cache_file): print("[Eval] AoA evaluations already completed and cached. Restoring...") target_main_aoa = os.path.join(strict_dir, "results", model_name, "main", "aoa") if os.path.exists(target_main_aoa): shutil.rmtree(target_main_aoa) os.makedirs(os.path.dirname(target_main_aoa), exist_ok=True) shutil.copytree(aoa_cache_file, target_main_aoa) else: print(f"[Eval] Running AoA metrics evaluations...") subprocess.run([ "./scripts/eval_aoa.sh", model_dir, "causal", "strict-small" ], cwd=strict_dir, check=True) # Cache it target_main_aoa = os.path.join(strict_dir, "results", model_name, "main", "aoa") if os.path.exists(target_main_aoa): os.makedirs(local_main_res, exist_ok=True) dest = os.path.join(local_main_res, "aoa") if os.path.exists(dest): shutil.rmtree(dest) shutil.copytree(target_main_aoa, dest) # ───────────────────────────────────────────────────────────── # C. COLLATE RESULTS # ───────────────────────────────────────────────────────────── print("[Eval] Collating predictions into submission file...") print("[Debug] Listing local cache folder:") if os.path.exists(local_results_dir): for root, dirs, files in os.walk(local_results_dir): rel_path = os.path.relpath(root, local_results_dir) print(f" {rel_path}: {files}") else: print(" Local results dir does not exist!") for item in os.listdir(local_results_dir): src = os.path.join(local_results_dir, item) dst = os.path.join(strict_dir, "results", model_name, item) if os.path.isdir(src): if os.path.exists(dst): shutil.rmtree(dst) os.makedirs(os.path.dirname(dst), exist_ok=True) shutil.copytree(src, dst) print("[Debug] Listing strict results folder before collation:") strict_results_dir = os.path.join(strict_dir, "results") if os.path.exists(strict_results_dir): for root, dirs, files in os.walk(strict_results_dir): rel_path = os.path.relpath(root, strict_results_dir) print(f" {rel_path}: {files}") else: print(" Strict results dir does not exist!") subprocess.run([ "bash", "scripts/collate_preds.sh", model_name, "causal", "strict-small", "--fast" ], cwd=strict_dir, check=True) # Save results to local folder results_src = os.path.join(strict_dir, "results") results_dest = os.path.abspath("./results") if os.path.exists(results_dest): shutil.rmtree(results_dest) shutil.copytree(results_src, results_dest) # Copy final collated json to current folder collated_json = os.path.join(strict_dir, "all_full_preds_and_fast_scores_causal.json") if os.path.exists(collated_json): shutil.copy(collated_json, "./all_full_preds_and_fast_scores_causal.json") print("\n[Eval] Success! Collation completed! Final file is at './all_full_preds_and_fast_scores_causal.json'") print("\n[Eval] Pipeline evaluation run finished.") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model-name", type=str, default="msit_gptbert_model") parser.add_argument("--epochs", type=int, default=10) parser.add_argument("--skip-eval", action="store_true", help="Skip evaluation phase after training") parser.add_argument("--skip-aoa", action="store_true", help="Skip AoA evaluation") parser.add_argument("--skip-glue", action="store_true", help="Skip GLUE fine-tuning") args = parser.parse_args() run_pipeline(args.model_name, epochs=args.epochs, skip_eval=args.skip_eval, skip_aoa=args.skip_aoa, skip_glue=args.skip_glue)