# lmr/glue_benchmark_grid_full.py """ Grid-search GLUE + extra tasks runner (full single-file). Features: - LR sweep across bert_lr_candidates - Random restarts for small/unstable tasks - Save per-run checkpoints and run_meta.json - Save all_runs.csv and best_overall per task - Evaluate best model and compute errorbars for GLUE - Supports EXTRA_TASKS (boolq/piqa/winogrande + hellaswag/openbookqa/arc) with MC/pair handling """ import os import json import re import math import random import shutil import time from pathlib import Path from typing import Optional, List, Tuple, Dict, Any import torch import numpy as np import pandas as pd from datasets import load_dataset from torch.utils.data import DataLoader, TensorDataset from tqdm import tqdm import evaluate # Project imports: adjust if your package layout differs try: from lmr.checkpointing import Checkpointing from lmr.ddp import unwrap_model except Exception: # If these modules are not available, provide lightweight fallbacks to avoid import errors Checkpointing = None def unwrap_model(m): return m # --------------------------------------------------------------------- # Tasks config # --------------------------------------------------------------------- GLUE_TASKS = { "cola": {"type": "classification", "num_labels": 2, "hf_name": "cola"}, "sst2": {"type": "classification", "num_labels": 2, "hf_name": "sst2"}, "mrpc": {"type": "classification", "num_labels": 2, "hf_name": "mrpc"}, "stsb": {"type": "regression", "num_labels": 1, "hf_name": "stsb"}, "qqp": {"type": "classification", "num_labels": 2, "hf_name": "qqp"}, "mnli": {"type": "classification", "num_labels": 3, "hf_name": "mnli"}, "qnli": {"type": "classification", "num_labels": 2, "hf_name": "qnli"}, "rte": {"type": "classification", "num_labels": 2, "hf_name": "rte"}, "wnli": {"type": "classification", "num_labels": 2, "hf_name": "wnli"}, } # Extra tasks (BoolQ, PIQA, Winogrande, HellaSwag, OpenBookQA, ARC variants) EXTRA_TASKS = { "boolq": { "type": "classification", "num_labels": 2, "hf_path": "boolq", "format": "pair", }, "piqa": { "type": "multiple_choice", "num_labels": 2, "hf_path": "piqa", "format": "mc", }, "winogrande": { "type": "multiple_choice", "num_labels": 2, "hf_path": "winogrande", "hf_config": "winogrande_xl", "format": "mc", }, # Added tasks below "hellaswag": { "type": "multiple_choice", "num_labels": 4, "hf_path": "hellaswag", "format": "mc", }, "openbookqa": { "type": "multiple_choice", "num_labels": 4, "hf_path": "openbookqa", "format": "mc", }, # AI2 ARC splits — align names with common usage "arc_easy": { "type": "multiple_choice", "num_labels": 4, "hf_path": "ai2_arc", "hf_config": "ARC-Easy", "format": "mc", }, "arc_challenge": { "type": "multiple_choice", "num_labels": 4, "hf_path": "ai2_arc", "hf_config": "ARC-Challenge", "format": "mc", }, } ALL_TASKS = {**GLUE_TASKS, **EXTRA_TASKS} # Small/unstable tasks for extra random restarts SMALL_TASKS_RANDOM_RESTARTS = {"cola", "mrpc", "rte", "stsb"} SMALL_TASKS_RANDOM_RESTARTS_EXTRA = set({"piqa", "boolq", "winogrande", "hellaswag"}) # adjust as desired BERT_LR_CANDIDATES = [2e-5, 3e-5, 4e-5, 5e-5] PREFERRED_METRIC_KEY = { "cola": "matthews_correlation", "sst2": "accuracy", "mrpc": "accuracy", "stsb": "pearson", "qqp": "accuracy", "mnli": "accuracy", "qnli": "accuracy", "rte": "accuracy", "wnli": "accuracy", "boolq": "accuracy", "piqa": "accuracy", "winogrande": "accuracy", "hellaswag": "accuracy", "openbookqa": "accuracy", "arc_easy": "accuracy", "arc_challenge": "accuracy", } # --------------------------------------------------------------------- # Repro helpers # --------------------------------------------------------------------- def _set_all_seeds(seed: int): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) try: torch.cuda.manual_seed_all(seed) except Exception: pass try: torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False except Exception: pass def _json_dump(obj: Any, path: Path): path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", encoding="utf-8") as f: json.dump(obj, f, indent=2, ensure_ascii=False) def _safe_float(x): try: if isinstance(x, (np.generic,)): return float(x.item()) return float(x) except Exception: return None def _metric_to_scalar(task: str, metric_res: Dict[str, Any], fallback_val_loss: Optional[float] = None) -> float: if isinstance(metric_res, dict) and metric_res: pref = PREFERRED_METRIC_KEY.get(task) if pref is not None and pref in metric_res: v = _safe_float(metric_res.get(pref)) if v is not None and not math.isnan(v): return float(v) for _, v in metric_res.items(): fv = _safe_float(v) if fv is not None and not math.isnan(fv): return float(fv) if fallback_val_loss is not None: try: return -float(fallback_val_loss) except Exception: pass return -1e9 # --------------------------------------------------------------------- # Task-aware example field extraction (robust) # --------------------------------------------------------------------- def _get_text_pair_from_example(task: str, ex: dict): """ Robustly extract (s1, s2) from a HF GLUE/example dict `ex` depending on task. Returns (s1:str, s2:Optional[str]) where s2 can be None for single-sentence tasks. """ task_field_map = { "cola": ("sentence", None), "sst2": ("sentence", None), "mrpc": ("sentence1", "sentence2"), "stsb": ("sentence1", "sentence2"), "qqp": ("question1", "question2"), "mnli": ("premise", "hypothesis"), "qnli": ("question", "sentence"), "rte": ("sentence1", "sentence2"), "wnli": ("sentence1", "sentence2"), } f1, f2 = task_field_map.get(task, (None, None)) def _try_keys(keys): for k in keys: if k in ex and ex.get(k) is not None: return ex.get(k) return None s1_candidates = [] s2_candidates = [] if f1: s1_candidates.append(f1) s1_candidates += ["sentence1", "premise", "question", "sentence", "text", "question1"] if f2: s2_candidates.append(f2) s2_candidates += ["sentence2", "hypothesis", "question2", "question1", "text2"] s1 = _try_keys(s1_candidates) s2 = _try_keys(s2_candidates) if s1 is None: s1 = ex.get("sentence") or ex.get("premise") or ex.get("question") or ex.get("text") if s2 is None: s2 = ex.get("sentence2") or ex.get("hypothesis") or ex.get("question2") s1 = "" if s1 is None else (s1 if isinstance(s1, str) else str(s1)) s2 = None if s2 is None else (s2 if isinstance(s2, str) else str(s2)) return s1, s2 # --------------------------------------------------------------------- # Tokenization helpers (robust to different tokenizer APIs) # --------------------------------------------------------------------- def _pad_and_tensorize(input_ids_list, attention_mask_list, pad_token_id: int): max_len = max(len(x) for x in input_ids_list) if input_ids_list else 0 ids_padded = [x + [pad_token_id] * (max_len - len(x)) for x in input_ids_list] mask_padded = [m + [0] * (max_len - len(m)) for m in attention_mask_list] input_ids = torch.tensor(ids_padded, dtype=torch.long) attention_mask = torch.tensor(mask_padded, dtype=torch.long) return input_ids, attention_mask def _batch_tokenize(tokenizer, texts: List[Tuple[Optional[str], Optional[str]]], max_length: int = 128): """ Robust batch tokenization for a variety of tokenizer APIs. - texts: list of (s1, s2) where s2 may be None. - Try HF tokenizer(...) first, then various batch methods, then per-example fallback. Returns dict with 'input_ids' (list of lists) and 'attention_mask'. """ sanitized = [] for a, b in texts: a_s = "" if a is None else (a if isinstance(a, str) else str(a)) b_s = None if b is None else (b if isinstance(b, str) else str(b)) sanitized.append((a_s, b_s)) # 1) Try HF-like tokenizer(...) first try: flat = [(a if b is None else (a, b)) for a, b in sanitized] enc = tokenizer(flat, truncation=True, padding=False, max_length=max_length) if isinstance(enc.get("input_ids", None), torch.Tensor): enc["input_ids"] = enc["input_ids"].tolist() if isinstance(enc.get("attention_mask", None), torch.Tensor): enc["attention_mask"] = enc["attention_mask"].tolist() return enc except Exception: pass # 2) Try other batch-like methods for method_name in ("batch_encode", "encode_batch", "batch_encode_plus", "encode_batch_pair", "encode_batch_items"): fn = getattr(tokenizer, method_name, None) if fn is None: continue try: try: enc = fn(sanitized, max_length=max_length, truncation=True, padding=False) except TypeError: enc = fn(sanitized) if isinstance(enc.get("input_ids", None), torch.Tensor): enc["input_ids"] = enc["input_ids"].tolist() if isinstance(enc.get("attention_mask", None), torch.Tensor): enc["attention_mask"] = enc["attention_mask"].tolist() return enc except Exception: continue # 3) Fallback per-example input_ids_list = [] attention_mask_list = [] for a, b in sanitized: try: if b is None: try: single = tokenizer.encode(a) except TypeError: single = tokenizer.encode([a]) else: single = None try: single = tokenizer.encode((a, b)) except Exception: try: single = tokenizer.encode(a, b) except Exception: single = tokenizer(a if b is None else (a, b)) if isinstance(single, dict): ids = single.get("input_ids") or single.get("ids") or [] mask = single.get("attention_mask") or single.get("mask") or [1] * len(ids) elif isinstance(single, torch.Tensor): ids = single.tolist() mask = [1] * len(ids) elif isinstance(single, list): ids = single mask = [1] * len(ids) else: tmp = tokenizer(a if b is None else (a, b)) if isinstance(tmp, dict): ids = tmp.get("input_ids") or tmp.get("ids") or [] mask = tmp.get("attention_mask") or tmp.get("mask") or [1] * len(ids) elif torch.is_tensor(tmp): ids = tmp.tolist() mask = [1] * len(ids) else: ids = list(tmp) mask = [1] * len(ids) if len(ids) > max_length: ids = ids[:max_length] mask = mask[:max_length] input_ids_list.append(ids) attention_mask_list.append(mask) except Exception as e: snippet = (a[:80] + "...") if a else "" raise RuntimeError(f"Tokenizer fallback encode failed for example '{snippet}': {e}") return {"input_ids": input_ids_list, "attention_mask": attention_mask_list} # --------------------------------------------------------------------- # Postprocess preds to the right shapes/types (fixes metric mismatches) # --------------------------------------------------------------------- def _postprocess_predictions(task: str, logits_np: np.ndarray, cfg_task: dict): """ Take logits (N, C) or (N,) or (N,1) and produce preds array ready for evaluate.compute: - classification -> 1D ints (class indices or binary 0/1) - regression -> 1D floats (for stsb typically 0..5) """ ttype = cfg_task["type"] num_labels = cfg_task["num_labels"] if logits_np is None or logits_np.size == 0: return np.array([]) # If logits are shape (N, ) -> treat as single score per example (binary/regression) if logits_np.ndim == 1: if ttype == "classification": preds = (logits_np > 0.5).astype(int) else: preds = logits_np.astype(float) return preds # If logits shape (N, 1) if logits_np.ndim == 2 and logits_np.shape[1] == 1: col = logits_np[:, 0] if ttype == "classification": preds = (col > 0.5).astype(int) else: preds = col.astype(float) return preds # If logits shape (N, C) if logits_np.ndim == 2 and logits_np.shape[1] >= 1: if ttype == "classification": preds = np.argmax(logits_np, axis=-1).astype(int) return preds else: if logits_np.shape[1] == 1: preds = logits_np[:, 0].astype(float) else: preds = logits_np.mean(axis=1).astype(float) if task == "stsb": preds = np.clip(preds, 0.0, 5.0) return preds return logits_np.ravel() # --------------------------------------------------------------------- # Model wrapping helper (robust) # --------------------------------------------------------------------- def make_wrapped_model_if_needed(model, hidden_size: Optional[int], num_labels: int, force_num_labels: Optional[int] = None): """ Robust wrapper factory with resilient hidden_size inference. Returns (model_or_wrapper, wrapped_flag) """ import torch.nn as nn base_model = model def _detect_head_dim(m): try: if hasattr(m, "classifier") and isinstance(getattr(m, "classifier"), nn.Linear): return getattr(m, "classifier").out_features if hasattr(m, "lm_head") and isinstance(getattr(m, "lm_head"), nn.Linear): return getattr(m, "lm_head").out_features if hasattr(m, "get_output_embeddings"): out_emb = m.get_output_embeddings() if out_emb is not None: if isinstance(out_emb, nn.Embedding): return out_emb.embedding_dim if hasattr(out_emb, "embedding_dim") else out_emb.num_embeddings if isinstance(out_emb, nn.Linear): return out_emb.out_features except Exception: pass return None if force_num_labels is None: head_dim = _detect_head_dim(base_model) if head_dim is not None and head_dim == num_labels: return base_model, False inferred_hidden = hidden_size if inferred_hidden is None: try: cand = getattr(base_model, "config", None) if cand is not None and hasattr(cand, "hidden_size"): inferred_hidden = int(cand.hidden_size) except Exception: inferred_hidden = None if inferred_hidden is None: try: un = unwrap_model(base_model) sd = un.state_dict() for k, v in sd.items(): if re.search(r"embed|embedding|word_embeddings|token_embedding|embed_tokens", k, re.I): if hasattr(v, "shape") and len(v.shape) == 2: inferred_hidden = int(v.shape[1]) break if re.search(r"q_proj|k_proj|v_proj|o_proj|dense|fc|linear|proj", k, re.I): if hasattr(v, "shape") and len(v.shape) == 2: cand = max(v.shape) if 1 < cand < 1_000_000: inferred_hidden = int(cand) break except Exception: inferred_hidden = None if inferred_hidden is None: raise RuntimeError( "Cannot infer hidden_size for wrapped classifier head. " "Please set `model.config.hidden_size` or pass `hidden_size` explicitly." ) class _WrappedModel(nn.Module): def __init__(self, base, hidden_size, num_labels): super().__init__() self.base = base self.classifier = nn.Linear(hidden_size, num_labels) self.logits_projector = None def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs): try: out = self.base(input_ids=input_ids, attention_mask=attention_mask, **kwargs) except TypeError: out = self.base(input_ids) last_hidden = getattr(out, "last_hidden_state", None) if last_hidden is not None: pooled = last_hidden[:, 0, :] logits = self.classifier(pooled) return type("Out", (), {"logits": logits, "loss": None}) if isinstance(out, (tuple, list)) and len(out) > 0: cand = out[0] if torch.is_tensor(cand): if cand.ndim == 3: pooled = cand[:, 0, :] logits = self.classifier(pooled) return type("Out", (), {"logits": logits, "loss": None}) if cand.ndim == 2 and cand.shape[1] == num_labels: return type("Out", (), {"logits": cand, "loss": None}) logits = getattr(out, "logits", None) if logits is not None: if logits.ndim == 2 and logits.shape[1] == num_labels: return type("Out", (), {"logits": logits, "loss": getattr(out, "loss", None)}) exist_dim = logits.shape[1] if self.logits_projector is None or self.logits_projector.weight.shape[1] != exist_dim: self.logits_projector = nn.Linear(exist_dim, num_labels).to(logits.device) projected = self.logits_projector(logits) return type("Out", (), {"logits": projected, "loss": getattr(out, "loss", None)}) hidden_states = getattr(out, "hidden_states", None) if hidden_states is not None: last_hidden = hidden_states[-1] if isinstance(hidden_states, (list, tuple)) else hidden_states if torch.is_tensor(last_hidden) and last_hidden.ndim == 3: pooled = last_hidden[:, 0, :] logits = self.classifier(pooled) return type("Out", (), {"logits": logits, "loss": None}) raise RuntimeError("Wrapped base model did not return recognizable hidden states or logits") return _WrappedModel(base_model, inferred_hidden, num_labels), True # --------------------------------------------------------------------- # Tokenize HF split to tensors (for finetune) # --------------------------------------------------------------------- def _tokenize_hf_split_to_tensors(task: str, tokenizer, raw_split, cfg_task, max_length=128, batch_tokenize_size=512): texts = [] labels = [] empty_s1 = 0 empty_s2 = 0 for ex in raw_split: s1, s2 = _get_text_pair_from_example(task, ex) texts.append((s1, s2)) labels.append(ex.get("label") if "label" in ex else -100) if not s1 or (isinstance(s1, str) and s1.strip() == ""): empty_s1 += 1 if s2 is not None and (not s2 or (isinstance(s2, str) and s2.strip() == "")): empty_s2 += 1 total = len(texts) print( f"[tokenize] task={task} samples={total} empty_s1={empty_s1} empty_s2={empty_s2} " f"({(empty_s1/total if total>0 else 0):.2%}, {(empty_s2/total if total>0 else 0):.2%})" ) input_ids_all = [] attention_all = [] for i in range(0, len(texts), batch_tokenize_size): enc = _batch_tokenize(tokenizer, texts[i:i+batch_tokenize_size], max_length=max_length) ids = enc.get("input_ids") masks = enc.get("attention_mask") or enc.get("mask") or enc.get("masks") if isinstance(ids, torch.Tensor): ids = ids.tolist() if isinstance(masks, torch.Tensor): masks = masks.tolist() input_ids_all.extend(ids) attention_all.extend(masks) pad_id = getattr(tokenizer, "pad_token_id", None) if pad_id is None: try: pad_id = tokenizer.token_to_id("[PAD]") except Exception: pad_id = 0 input_ids_t, attention_mask_t = _pad_and_tensorize(input_ids_all, attention_all, pad_id) labels_t = torch.tensor(labels, dtype=torch.long if cfg_task["type"] == "classification" else torch.float) return input_ids_t, attention_mask_t, labels_t def _normalize_label(label): """Robustly normalize various label encodings to an int index or None. Handles: int, float-like strings, single-letter answers ('A','b'), empty strings, None. """ if label is None: return None # If it's already int-like if isinstance(label, (int, np.integer)): return int(label) # string handling if isinstance(label, str): s = label.strip() if s == "": return None # single-letter like 'A'/'b' if len(s) == 1 and s.isalpha(): return ord(s.upper()) - ord("A") # Try int try: return int(s) except Exception: pass # Try float then cast to int if reasonable (e.g. '1.0') try: f = float(s) # only accept if it's integer-valued (e.g. 1.0 -> 1) if abs(f - round(f)) < 1e-6: return int(round(f)) # otherwise treat as None (can't map to choice index) return None except Exception: return None # other numeric-like (np types) try: return int(label) except Exception: return None def _extract_mc_example(task: str, ex: dict): """ Robust extractor for multiple-choice examples across a range of HF dataset schemas. Returns (context, options_list, label_index_or_None). """ # detect label raw value first (don't int() it yet) label_raw = None if "label" in ex: label_raw = ex.get("label") if label_raw is None: label_raw = ex.get("answerKey") or ex.get("answer") or ex.get("correct") or ex.get("gold") # normalize into int index or None label = _normalize_label(label_raw) # 1) 'choices' list (strings or dicts) if "choices" in ex and ex["choices"] is not None: ch = ex["choices"] if isinstance(ch, list) and len(ch) > 0: opts = [] for c in ch: if isinstance(c, dict): opts.append(c.get("text") or c.get("label") or c.get("choice") or str(c)) else: opts.append(str(c)) ctx = ex.get("context") or ex.get("question") or ex.get("story") or ex.get("sentence") or ex.get("passage") return ctx, opts, label # 2) 'endings' pattern (hellaswag) if "endings" in ex and isinstance(ex["endings"], list) and len(ex["endings"]) > 0: ctx = ex.get("context") or ex.get("article") or ex.get("sentence") or ex.get("story") or ex.get("paragraph") opts = [str(x) for x in ex["endings"]] return ctx, opts, label # 3) explicit option fields like 'choice1','choice2' or 'option1'.. opts = [] for prefix in ("choice", "option", "ending", "answer"): i = 1 found = False while True: key = f"{prefix}{i}" if key in ex: opts.append(str(ex[key])) found = True i += 1 else: break if found: ctx = ex.get("question") or ex.get("context") or ex.get("passage") or ex.get("sentence") return ctx, opts, label # 4) common QA fields: 'question' + 'choices' (where choices might be list of dicts) if "question" in ex: ctx = ex["question"] if "choices" in ex: ch = ex["choices"] if isinstance(ch, list) and len(ch) > 0: opts = [] for c in ch: if isinstance(c, dict): opts.append(c.get("text") or c.get("choice") or str(c)) else: opts.append(str(c)) return ctx, opts, label # 5) ai2_arc / openbookqa style: search for list-like values for k, v in ex.items(): if isinstance(v, list) and 2 <= len(v) <= 10 and all(isinstance(x, (str, dict)) for x in v): opts = [x.get("text") if isinstance(x, dict) and x.get("text") else str(x) for x in v] ctx = ex.get("goal") or ex.get("question") or ex.get("context") or ex.get("passage") or "" return ctx, opts, label # 6) fallback: collect fields that look like options candidate_opts = [] for k in sorted(ex.keys()): if any(tok in k.lower() for tok in ("option", "choice", "ending", "answer", "alt", "sol")): candidate_opts.append(str(ex[k])) if candidate_opts: ctx = ex.get("question") or ex.get("context") or "" return ctx, candidate_opts, label # last resort ctx = ex.get("question") or ex.get("context") or ex.get("passage") or "" return ctx, [], label def _tokenize_generic_mc_split_to_tensors(task: str, tokenizer, raw_split, max_length=128, batch_tokenize_size=256): """ Generic MC tokenizer that uses _extract_mc_example to normalize different HF schemas. Returns (input_ids_t (N,C,L), attention_t (N,C,L), labels_t (N,)) """ contexts = [] options = [] labels = [] num_choices = None for ex in raw_split: ctx, opts, lab = _extract_mc_example(task, ex) if not opts: # skip if no options recognized continue if num_choices is None: num_choices = len(opts) if len(opts) != num_choices: # inconsistent number of choices; skip example continue contexts.append(ctx if ctx is not None else "") options.append(opts) labels.append(-1 if lab is None else int(lab)) if len(contexts) == 0: return torch.zeros((0, 1, 1), dtype=torch.long), torch.zeros((0, 1, 1), dtype=torch.long), torch.tensor([], dtype=torch.long) input_ids_rows = [] attention_rows = [] pad_id = getattr(tokenizer, "pad_token_id", None) if pad_id is None: try: pad_id = tokenizer.token_to_id("[PAD]") except Exception: pad_id = 0 for i in range(0, len(contexts), batch_tokenize_size): chunk_ctx = contexts[i:i+batch_tokenize_size] chunk_opts = options[i:i+batch_tokenize_size] flat_pairs = [] for c, opts in zip(chunk_ctx, chunk_opts): for o in opts: flat_pairs.append((c, o)) enc = _batch_tokenize(tokenizer, flat_pairs, max_length=max_length) ids_flat = enc.get("input_ids") masks_flat = enc.get("attention_mask") or enc.get("mask") or enc.get("masks") if isinstance(ids_flat, torch.Tensor): ids_flat = ids_flat.tolist() if isinstance(masks_flat, torch.Tensor): masks_flat = masks_flat.tolist() per_example = [] per_mask_example = [] idx = 0 for _ in chunk_ctx: row = [] row_mask = [] for _ in range(num_choices): row.append(ids_flat[idx]) row_mask.append(masks_flat[idx]) idx += 1 per_example.append(row) per_mask_example.append(row_mask) input_ids_rows.extend(per_example) attention_rows.extend(per_mask_example) max_len = max(len(seq) for row in input_ids_rows for seq in row) if input_ids_rows else 1 input_ids_padded = [ [ seq + [pad_id] * (max_len - len(seq)) for seq in row ] for row in input_ids_rows ] attention_padded = [ [ mask + [0] * (max_len - len(mask)) for mask in row ] for row in attention_rows ] input_ids_t = torch.tensor(input_ids_padded, dtype=torch.long) # (N, C, L) attention_t = torch.tensor(attention_padded, dtype=torch.long) labels_t = torch.tensor(labels, dtype=torch.long) return input_ids_t, attention_t, labels_t # --------------------------------------------------------------------- # Multiple-choice tokenizer entry (keeps fast paths for known tasks, else generic) # --------------------------------------------------------------------- def _tokenize_mc_split_to_tensors(task: str, tokenizer, raw_split, max_length=128, batch_tokenize_size=256): """ Build tensors for multiple choice tasks: returns input_ids tensor shape (N, num_choices, L), attention_mask tensor same, labels tensor (N,) Uses fast paths for known tasks (piqa, winogrande), otherwise uses generic parser. """ # fast paths if task == "piqa": contexts = [] options = [] labels = [] for ex in raw_split: try: ctx = ex.get("goal") or ex.get("question") or ex.get("context") or "" opts = [ex["sol1"], ex["sol2"]] lab = int(ex["label"]) except Exception: # fallback to generic return _tokenize_generic_mc_split_to_tensors(task, tokenizer, raw_split, max_length=max_length, batch_tokenize_size=batch_tokenize_size) contexts.append(ctx) options.append(opts) labels.append(lab) # then pack like generic elif task == "winogrande": contexts = [] options = [] labels = [] for ex in raw_split: try: ctx = ex.get("sentence") or ex.get("context") or ex.get("question") or "" opts = [ex["option1"], ex["option2"]] lab = int(ex.get("answer", 1)) - 1 except Exception: return _tokenize_generic_mc_split_to_tensors(task, tokenizer, raw_split, max_length=max_length, batch_tokenize_size=batch_tokenize_size) contexts.append(ctx) options.append(opts) labels.append(lab) else: # fallback to generic parser which supports hellaswag, openbookqa, arc, etc. return _tokenize_generic_mc_split_to_tensors(task, tokenizer, raw_split, max_length=max_length, batch_tokenize_size=batch_tokenize_size) # From here pack contexts/options/labels into tensors (same logic as generic) if len(contexts) == 0: return torch.zeros((0, 1, 1), dtype=torch.long), torch.zeros((0, 1, 1), dtype=torch.long), torch.tensor([], dtype=torch.long) input_ids_rows = [] attention_rows = [] pad_id = getattr(tokenizer, "pad_token_id", None) if pad_id is None: try: pad_id = tokenizer.token_to_id("[PAD]") except Exception: pad_id = 0 num_choices = len(options[0]) for i in range(0, len(contexts), batch_tokenize_size): chunk_ctx = contexts[i:i+batch_tokenize_size] chunk_opts = options[i:i+batch_tokenize_size] flat_pairs = [] for c, opts in zip(chunk_ctx, chunk_opts): for o in opts: flat_pairs.append((c, o)) enc = _batch_tokenize(tokenizer, flat_pairs, max_length=max_length) ids_flat = enc.get("input_ids") masks_flat = enc.get("attention_mask") or enc.get("mask") or enc.get("masks") if isinstance(ids_flat, torch.Tensor): ids_flat = ids_flat.tolist() if isinstance(masks_flat, torch.Tensor): masks_flat = masks_flat.tolist() per_example = [] per_mask_example = [] idx = 0 for _ in chunk_ctx: row = [] row_mask = [] for _ in range(num_choices): row.append(ids_flat[idx]) row_mask.append(masks_flat[idx]) idx += 1 per_example.append(row) per_mask_example.append(row_mask) input_ids_rows.extend(per_example) attention_rows.extend(per_mask_example) max_len = max(len(seq) for row in input_ids_rows for seq in row) if input_ids_rows else 1 input_ids_padded = [ [ seq + [pad_id] * (max_len - len(seq)) for seq in row ] for row in input_ids_rows ] attention_padded = [ [ mask + [0] * (max_len - len(mask)) for mask in row ] for row in attention_rows ] input_ids_t = torch.tensor(input_ids_padded, dtype=torch.long) # (N, C, L) attention_t = torch.tensor(attention_padded, dtype=torch.long) labels_t = torch.tensor(labels, dtype=torch.long) return input_ids_t, attention_t, labels_t # --------------------------------------------------------------------- # Pairwise tokenizer for BoolQ (if not already present) # --------------------------------------------------------------------- def _tokenize_pair_split_to_tensors(task: str, tokenizer, raw_split, max_length=128, batch_tokenize_size=512): texts = [] labels = [] for ex in raw_split: if task == "boolq": a = ex.get("passage") or ex.get("context") or ex.get("article") or "" b = ex.get("question") or ex.get("query") or "" lab = int(ex.get("answer") or ex.get("label") or 0) else: # If unknown, try generic pair fields a = ex.get("passage") or ex.get("context") or ex.get("article") or "" b = ex.get("question") or ex.get("query") or "" lab = int(ex.get("answer") or ex.get("label") or 0) texts.append((a, b)) labels.append(lab) input_ids_all = [] attention_all = [] for i in range(0, len(texts), batch_tokenize_size): chunk = texts[i:i+batch_tokenize_size] enc = _batch_tokenize(tokenizer, chunk, max_length=max_length) ids = enc.get("input_ids") masks = enc.get("attention_mask") or enc.get("mask") or enc.get("masks") if isinstance(ids, torch.Tensor): ids = ids.tolist() if isinstance(masks, torch.Tensor): masks = masks.tolist() input_ids_all.extend(ids) attention_all.extend(masks) pad_id = getattr(tokenizer, "pad_token_id", None) if pad_id is None: try: pad_id = tokenizer.token_to_id("[PAD]") except Exception: pad_id = 0 input_ids_t, attention_mask_t = _pad_and_tensorize(input_ids_all, attention_all, pad_id) labels_t = torch.tensor(labels, dtype=torch.long) return input_ids_t, attention_mask_t, labels_t # --------------------------------------------------------------------- # Postprocess preds helper already defined above (_postprocess_predictions) # --------------------------------------------------------------------- # --------------------------------------------------------------------- # Training for GLUE tasks (full fine-tune) # --------------------------------------------------------------------- def train_full_finetune( task: str, tokenizer, model, raw_train, raw_val, device: str = "cuda", epochs: int = 3, batch_size: int = 32, lr: float = 2e-5, weight_decay: float = 0.01, warmup_steps: int = 100, max_length: int = 128, grad_accum_steps: int = 1, out_checkpoint_dir: Optional[str] = None, seed: Optional[int] = None, ): cfg_task = GLUE_TASKS[task] device_t = torch.device(device if torch.cuda.is_available() else "cpu") if seed is not None: _set_all_seeds(int(seed)) hidden_size = None if hasattr(model, "config") and hasattr(model.config, "hidden_size"): try: hidden_size = int(model.config.hidden_size) except Exception: hidden_size = None model, wrapped_flag = make_wrapped_model_if_needed( model, hidden_size, cfg_task["num_labels"], force_num_labels=cfg_task["num_labels"] ) model.to(device_t) train_ids, train_mask, train_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_train, cfg_task, max_length=max_length) val_ids, val_mask, val_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_val, cfg_task, max_length=max_length) train_ds = TensorDataset(train_ids, train_mask, train_labels) val_ds = TensorDataset(val_ids, val_mask, val_labels) g = torch.Generator() if seed is not None: g.manual_seed(int(seed)) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, pin_memory=True, generator=g) val_loader = DataLoader(val_ds, batch_size=max(64, batch_size), shuffle=False, pin_memory=True) optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay) total_steps = max(1, (len(train_loader) // max(1, grad_accum_steps)) * epochs) try: from transformers import get_cosine_schedule_with_warmup scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps) except Exception: scheduler = None loss_fn = torch.nn.CrossEntropyLoss() if cfg_task["type"] == "classification" else torch.nn.MSELoss() best_metric_res: Dict[str, Any] = {} best_score: Optional[float] = None best_epoch = -1 model.train() for epoch in range(epochs): for step, batch in enumerate(tqdm(train_loader, desc=f"Train {task} epoch {epoch+1} (lr={lr:g})")): ids_b, mask_b, labs_b = batch ids_b = ids_b.to(device_t) mask_b = mask_b.to(device_t) labs_b = labs_b.to(device_t) out = model(input_ids=ids_b, attention_mask=mask_b, labels=None) logits = getattr(out, "logits", None) if logits is None: if isinstance(out, (tuple, list)): logits = out[0] else: raise RuntimeError("Model did not return logits during finetune") if cfg_task["type"] == "classification": loss = loss_fn(logits, labs_b.long()) else: if logits.ndim == 2 and logits.shape[1] == 1: preds = logits.squeeze(1) elif logits.ndim == 2: preds = logits.mean(dim=1) else: preds = logits loss = loss_fn(preds, labs_b.float()) loss = loss / max(1, grad_accum_steps) loss.backward() if (step + 1) % max(1, grad_accum_steps) == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() if scheduler is not None: scheduler.step() optimizer.zero_grad() # validation model.eval() tot_val_loss = 0.0 all_logits = [] all_labels = [] with torch.no_grad(): for ids_b, mask_b, labs_b in tqdm(val_loader, desc=f"Validate {task} epoch {epoch+1}", leave=False): ids_b = ids_b.to(device_t) mask_b = mask_b.to(device_t) labs_b = labs_b.to(device_t) out = model(input_ids=ids_b, attention_mask=mask_b, labels=None) logits = getattr(out, "logits", None) if logits is None: if isinstance(out, (tuple, list)): logits = out[0] else: raise RuntimeError("Model did not return logits during validation") if cfg_task["type"] == "classification": l = loss_fn(logits, labs_b.long()) else: if logits.ndim == 2 and logits.shape[1] == 1: preds = logits.squeeze(1) elif logits.ndim == 2: preds = logits.mean(dim=1) else: preds = logits l = loss_fn(preds, labs_b.float()) tot_val_loss += l.item() * ids_b.size(0) all_logits.append(logits.detach().cpu().numpy()) all_labels.append(labs_b.detach().cpu().numpy()) model.train() all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg_task["num_labels"])) all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,)) preds = _postprocess_predictions(task, all_logits, cfg_task) metric = evaluate.load("glue", cfg_task["hf_name"]) try: metric_res = metric.compute(predictions=preds.tolist(), references=all_labels.tolist()) except Exception: try: metric_res = metric.compute(predictions=preds, references=all_labels) except Exception as e: metric_res = {"error": str(e)} avg_val_loss = tot_val_loss / len(val_ds) if len(val_ds) > 0 else float("nan") score = _metric_to_scalar(task, metric_res, fallback_val_loss=avg_val_loss) print(f"[FT] {task} epoch {epoch+1} lr={lr:g} val_loss={avg_val_loss:.6f} metric={metric_res} score={score:.6f}") if best_score is None or float(score) > float(best_score): best_score = float(score) best_metric_res = metric_res best_epoch = epoch + 1 if out_checkpoint_dir: outp = Path(out_checkpoint_dir) outp.mkdir(parents=True, exist_ok=True) best_fname = outp / "best_finetuned.pt" try: sd = unwrap_model(model).state_dict() except Exception: sd = model.state_dict() torch.save(sd, str(best_fname)) print(f"[FT] Saved best checkpoint (epoch {best_epoch}) to: {best_fname}") if out_checkpoint_dir: outp = Path(out_checkpoint_dir) outp.mkdir(parents=True, exist_ok=True) fname = outp / "finetuned.pt" try: sd = unwrap_model(model).state_dict() except Exception: sd = model.state_dict() torch.save(sd, str(fname)) print(f"[FT] Saved finetuned model to: {fname}") meta = { "task": task, "lr": lr, "seed": seed, "epochs": epochs, "batch_size": batch_size, "grad_accum_steps": grad_accum_steps, "warmup_steps": warmup_steps, "weight_decay": weight_decay, "max_length": max_length, "wrapped_flag": bool(wrapped_flag), "best_epoch": int(best_epoch), "best_score": float(best_score) if best_score is not None else None, "best_metrics": best_metric_res, } _json_dump(meta, outp / "run_meta.json") print(f"[FT] Best validation for task '{task}' (lr={lr:g}, seed={seed}): epoch={best_epoch}, score={best_score}, metrics={best_metric_res}") return model, best_metric_res, float(best_score) if best_score is not None else -1e9, best_epoch # --------------------------------------------------------------------- # Load finetuned checkpoint for eval (wrapped/unwrapped) # --------------------------------------------------------------------- def _load_finetuned_checkpoint_for_task(task: str, base_model, checkpoint_path: str): cfg = GLUE_TASKS.get(task) or EXTRA_TASKS.get(task) sd = torch.load(checkpoint_path, map_location="cpu") keys = list(sd.keys()) if isinstance(sd, dict) else [] looks_wrapped = any(k.startswith("base.") for k in keys) or any(k.startswith("classifier.") for k in keys) hidden_size = None if hasattr(base_model, "config") and hasattr(base_model.config, "hidden_size"): try: hidden_size = int(base_model.config.hidden_size) except Exception: hidden_size = None if looks_wrapped and cfg is not None: wrapped_model, _ = make_wrapped_model_if_needed( base_model, hidden_size, cfg["num_labels"], force_num_labels=cfg["num_labels"] ) try: unwrap_model(wrapped_model).load_state_dict(sd, strict=False) except Exception: try: wrapped_model.load_state_dict(sd, strict=False) except Exception: pass return wrapped_model try: unwrap_model(base_model).load_state_dict(sd, strict=False) except Exception: try: base_model.load_state_dict(sd, strict=False) except Exception: pass return base_model # --------------------------------------------------------------------- # Error bar evaluation (5 folds -> 5 leave-one-fold-out subsets) # --------------------------------------------------------------------- def _fivefold_indices(n: int): idx = np.arange(n) folds = np.array_split(idx, 5) return [f.tolist() for f in folds] def _errorbar_subsets_from_folds(folds: List[List[int]]): assert len(folds) == 5 combos = [ ("0123", [0,1,2,3]), ("1234", [1,2,3,4]), ("0124", [0,1,2,4]), ("0234", [0,2,3,4]), ("0134", [0,1,3,4]), ] subsets = [] for name, keep in combos: inds = [] for k in keep: inds.extend(folds[k]) subsets.append({"name": name, "indices": inds}) return subsets def _compute_metric_for_indices(task: str, cfg: dict, metric_obj, preds_all: np.ndarray, labels_all: np.ndarray, indices: List[int]): if len(indices) == 0: return {"error": "empty_indices"} p = preds_all[indices] y = labels_all[indices] if cfg["type"] == "classification" or cfg.get("type") == "multiple_choice": preds_out = p.astype(int).tolist() refs_out = y.astype(int).tolist() else: preds_out = p.astype(float).tolist() refs_out = y.astype(float).tolist() try: return metric_obj.compute(predictions=preds_out, references=refs_out) except Exception: try: return metric_obj.compute(predictions=np.array(preds_out), references=np.array(refs_out)) except Exception as e: return {"error": str(e)} def _compute_errorbar(task: str, cfg: dict, metric_obj, preds_all: np.ndarray, labels_all: np.ndarray): n = int(len(labels_all)) folds = _fivefold_indices(n) subsets = _errorbar_subsets_from_folds(folds) pref = PREFERRED_METRIC_KEY.get(task) subset_scores = [] scores = [] for s in subsets: m = _compute_metric_for_indices(task, cfg, metric_obj, preds_all, labels_all, s["indices"]) sc = _metric_to_scalar(task, m, fallback_val_loss=None) subset_scores.append({"subset": s["name"], "score": float(sc), "metrics": m}) scores.append(float(sc)) arr = np.array(scores, dtype=float) mean = float(np.mean(arr)) if len(arr) else float("nan") std = float(np.std(arr, ddof=1)) if len(arr) > 1 else 0.0 stderr = float(std / math.sqrt(len(arr))) if len(arr) > 0 else float("nan") return { "preferred_key": pref, "subset_scores": subset_scores, "mean": mean, "std": std, "stderr": stderr, } # --------------------------------------------------------------------- # Evaluation (returns metrics + also writes preds/results) # --------------------------------------------------------------------- def run_glue_task( task: str, tokenizer, model, checkpointing: Optional[Checkpointing] = None, device: str = "cuda", batch_size: int = 64, max_length: int = 128, output_dir: str = "glue_output", compute_errorbar: bool = False, ): assert task in GLUE_TASKS, f"Unknown GLUE task: {task}" cfg = GLUE_TASKS[task] hf = load_dataset("glue", cfg["hf_name"]) if task == "mnli": val_splits = ["validation_matched", "validation_mismatched"] else: val_splits = ["validation"] results_by_split = {} for split in val_splits: raw = hf[split] print(f"[GLUE] Task={task} split={split} samples={len(raw)}") texts = [] labels = [] for ex in raw: s1, s2 = _get_text_pair_from_example(task, ex) texts.append((s1, s2)) labels.append(ex.get("label") if "label" in ex else -100) BATCH = 512 input_ids_all = [] attention_all = [] for i in range(0, len(texts), BATCH): enc = _batch_tokenize(tokenizer, texts[i:i+BATCH], max_length=max_length) ids = enc.get("input_ids") masks = enc.get("attention_mask") or enc.get("mask") or enc.get("masks") if isinstance(ids, torch.Tensor): ids = ids.tolist() if isinstance(masks, torch.Tensor): masks = masks.tolist() input_ids_all.extend(ids) attention_all.extend(masks) pad_id = getattr(tokenizer, "pad_token_id", None) if pad_id is None: try: pad_id = tokenizer.token_to_id("[PAD]") except Exception: pad_id = 0 input_ids, attention_mask = _pad_and_tensorize(input_ids_all, attention_all, pad_id) labels_t = torch.tensor(labels, dtype=torch.long if cfg["type"] == "classification" else torch.float) ds = TensorDataset(input_ids, attention_mask, labels_t) loader = DataLoader(ds, batch_size=batch_size, shuffle=False, pin_memory=True) if checkpointing is not None: try: checkpointing.load_model_states("recent") except Exception: pass device_t = torch.device(device if torch.cuda.is_available() else "cpu") model.to(device_t) model.eval() hidden_size = None if hasattr(model, "config") and hasattr(model.config, "hidden_size"): try: hidden_size = int(model.config.hidden_size) except Exception: hidden_size = None force = 1 if cfg["type"] == "regression" else cfg["num_labels"] wrapped_model, _ = make_wrapped_model_if_needed(model, hidden_size, cfg["num_labels"], force_num_labels=force) wrapped_model.to(device_t) wrapped_model.eval() all_logits = [] all_labels = [] with torch.no_grad(): for batch in tqdm(loader, desc=f"Eval {task}:{split}"): ids_b, mask_b, labels_b = batch ids_b = ids_b.to(device_t) mask_b = mask_b.to(device_t) out = wrapped_model(input_ids=ids_b, attention_mask=mask_b, labels=None) logits = getattr(out, "logits", None) if logits is None: if isinstance(out, (tuple, list)): logits = out[0] else: raise RuntimeError("Model forward did not return logits") all_logits.append(logits.detach().cpu().numpy()) all_labels.append(labels_b.detach().cpu().numpy()) all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg["num_labels"])) all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,)) preds = _postprocess_predictions(task, all_logits, cfg) metric = evaluate.load("glue", cfg["hf_name"]) metric_res = _compute_metric_for_indices(task, cfg, metric, preds, all_labels, list(range(len(all_labels)))) errorbar_res = None if compute_errorbar: errorbar_res = _compute_errorbar(task, cfg, metric, preds, all_labels) os.makedirs(output_dir, exist_ok=True) out_json = Path(output_dir) / f"{task}_{split}_results.json" with open(out_json, "w", encoding="utf-8") as f: json.dump({"task": task, "split": split, "metrics": metric_res}, f, indent=2) if errorbar_res is not None: out_eb = Path(output_dir) / f"{task}_{split}_errorbar.json" with open(out_eb, "w", encoding="utf-8") as f: json.dump({"task": task, "split": split, "errorbar": errorbar_res}, f, indent=2) csv_p = Path(output_dir) / f"{task}_{split}_preds.csv" pd.DataFrame({"pred": preds.tolist(), "label": all_labels.tolist()}).to_csv(csv_p, index=False) results_by_split[split] = {"metrics": metric_res, "errorbar": errorbar_res} return results_by_split # --------------------------------------------------------------------- # Train full fine-tune for EXTRA tasks (pair + mc) # --------------------------------------------------------------------- def train_full_finetune_extra(task: str, tokenizer, model, raw_train, raw_val, device: str = "cuda", epochs: int = 3, batch_size: int = 16, lr: float = 2e-5, weight_decay: float = 0.01, warmup_steps: int = 100, max_length: int = 128, grad_accum_steps: int = 1, out_checkpoint_dir: Optional[str] = None): """ Fine-tune for EXTRA_TASKS (boolq/piqa/winogrande/hellaswag/openbookqa/arc) Expects the provided `model` to be compatible with HF's MultipleChoice or SequenceClassification APIs. """ cfg = EXTRA_TASKS[task] device = torch.device(device if torch.cuda.is_available() else "cpu") model.to(device) is_mc = cfg["format"] == "mc" if is_mc: train_ids, train_mask, train_labels = _tokenize_mc_split_to_tensors(task, tokenizer, raw_train, max_length=max_length) val_ids, val_mask, val_labels = _tokenize_mc_split_to_tensors(task, tokenizer, raw_val, max_length=max_length) train_ds = TensorDataset(train_ids, train_mask, train_labels) val_ds = TensorDataset(val_ids, val_mask, val_labels) else: train_ids, train_mask, train_labels = _tokenize_pair_split_to_tensors(task, tokenizer, raw_train, max_length=max_length) val_ids, val_mask, val_labels = _tokenize_pair_split_to_tensors(task, tokenizer, raw_val, max_length=max_length) train_ds = TensorDataset(train_ids, train_mask, train_labels) val_ds = TensorDataset(val_ids, val_mask, val_labels) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, pin_memory=True) val_loader = DataLoader(val_ds, batch_size=max(64, batch_size), shuffle=False, pin_memory=True) optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay) total_steps = max(1, (len(train_loader) // max(1, grad_accum_steps)) * epochs) try: from transformers import get_cosine_schedule_with_warmup scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps) except Exception: scheduler = None loss_fn = torch.nn.CrossEntropyLoss() model.train() global_step = 0 final_metric_res = {} for epoch in range(epochs): running_loss = 0.0 for step, batch in enumerate(tqdm(train_loader, desc=f"[ExtraTrain] {task} epoch {epoch+1}")): if is_mc: ids_b, mask_b, labs_b = batch # ids_b: (B, C, L) ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device) out = None logits = None try: out = model(input_ids=ids_b, attention_mask=mask_b, labels=None) logits = getattr(out, "logits", None) if logits is None and isinstance(out, (list, tuple)): logits = out[0] except Exception: B, C, L = ids_b.shape flat_ids = ids_b.view(B*C, L).to(device) flat_mask = mask_b.view(B*C, L).to(device) out_flat = model(input_ids=flat_ids, attention_mask=flat_mask) flat_logits = getattr(out_flat, "logits", None) if flat_logits is None and isinstance(out_flat, (tuple, list)): flat_logits = out_flat[0] if flat_logits is None: raise RuntimeError("Model did not return logits for MC fallback") if flat_logits.ndim == 2 and flat_logits.shape[1] == 1: logits = flat_logits.view(B, C) else: logits = flat_logits.view(B, C, -1).mean(dim=-1) loss = loss_fn(logits, labs_b.long()) else: ids_b, mask_b, labs_b = batch ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device) out = model(input_ids=ids_b, attention_mask=mask_b, labels=None) logits = getattr(out, "logits", None) if logits is None: if isinstance(out, (tuple, list)): logits = out[0] else: raise RuntimeError("Model did not return logits for pair task") loss = loss_fn(logits, labs_b.long()) loss = loss / max(1, grad_accum_steps) loss.backward() if (step + 1) % max(1, grad_accum_steps) == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() if scheduler is not None: scheduler.step() optimizer.zero_grad() global_step += 1 running_loss += loss.item() * (ids_b.size(0) if not is_mc else ids_b.size(0)) # validation model.eval() all_logits = [] all_labels = [] with torch.no_grad(): for batch in tqdm(val_loader, desc=f"[ExtraVal] {task} epoch {epoch+1}", leave=False): if is_mc: ids_b, mask_b, labs_b = batch ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device) out = None try: out = model(input_ids=ids_b, attention_mask=mask_b, labels=None) logits = getattr(out, "logits", None) if logits is None and isinstance(out, (tuple, list)): logits = out[0] except Exception: B, C, L = ids_b.shape flat_ids = ids_b.view(B*C, L).to(device) flat_mask = mask_b.view(B*C, L).to(device) out_flat = model(input_ids=flat_ids, attention_mask=flat_mask) flat_logits = getattr(out_flat, "logits", None) if flat_logits is None and isinstance(out_flat, (tuple, list)): flat_logits = out_flat[0] if flat_logits is None: raise RuntimeError("Model did not return logits during MC validation fallback") if flat_logits.ndim == 2 and flat_logits.shape[1] == 1: logits = flat_logits.view(B, C) else: logits = flat_logits.view(B, C, -1).mean(dim=-1) all_logits.append(logits.detach().cpu().numpy()) all_labels.append(labs_b.detach().cpu().numpy()) else: ids_b, mask_b, labs_b = batch ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device) out = model(input_ids=ids_b, attention_mask=mask_b, labels=None) logits = getattr(out, "logits", None) if logits is None: if isinstance(out, (tuple, list)): logits = out[0] else: raise RuntimeError("Model did not return logits during pair validation") all_logits.append(logits.detach().cpu().numpy()) all_labels.append(labs_b.detach().cpu().numpy()) model.train() all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg["num_labels"])) all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,)) if cfg["type"] == "classification" or cfg["type"] == "multiple_choice": preds = np.argmax(all_logits, axis=-1) if all_logits.size else np.array([]) else: preds = _postprocess_predictions(task, all_logits, {"type": cfg["type"], "num_labels": cfg["num_labels"]}) try: metric = evaluate.load("accuracy") metric_res = metric.compute(predictions=preds.tolist(), references=all_labels.tolist()) except Exception as e: metric_res = {"error": str(e)} print(f"[Extra FT] {task} epoch {epoch+1} metric={metric_res}") final_metric_res = metric_res if out_checkpoint_dir: outp = Path(out_checkpoint_dir) outp.mkdir(parents=True, exist_ok=True) fname = outp / "finetuned_extra.pt" try: sd = unwrap_model(model).state_dict() except Exception: sd = model.state_dict() torch.save(sd, str(fname)) print(f"[Extra FT] Saved finetuned model to: {fname}") return model, final_metric_res # --------------------------------------------------------------------- # Evaluation-only runner for EXTRA tasks # --------------------------------------------------------------------- # Evaluation-only runner for EXTRA tasks # --------------------------------------------------------------------- def run_extra_task(task: str, tokenizer, model, checkpointing: Optional[Checkpointing] = None, device: str = "cuda", batch_size: int = 64, max_length: int = 128, output_dir: str = "extra_output", prefer_test_if_available: bool = True): """ Evaluate an EXTRA_TASK on HF dataset. Behavior: - If the dataset provides a 'test' split we will prefer it (unless it is unlabeled). - Otherwise use 'validation' or other labeled splits. - Returns metric dict (usually accuracy) and writes preds + metrics to output_dir. """ assert task in EXTRA_TASKS, f"Unknown extra task: {task}" cfg = EXTRA_TASKS[task] # load hf dataset robustly (handle hf_config if present) try: if "hf_config" in cfg: ds = load_dataset(cfg["hf_path"], cfg["hf_config"]) else: ds = load_dataset(cfg["hf_path"]) except Exception as e: raise RuntimeError(f"Failed to load HF dataset for task={task}: {e}") # prefer test split if available and (prefer_test_if_available True). # But if test is unlabeled (no label fields), we may fall back to validation. chosen_split = None candidate_order = [] # explicit preference order: test, validation, validation_matched, validation_unlabeled, train candidate_order = ["test", "validation", "validation_matched", "validation_unlabeled", "train"] available_splits = list(ds.keys()) if hasattr(ds, "keys") else [] # If prefer_test_if_available try to pick test first for cand in candidate_order: if cand in ds: # check if split has at least one example and at least one of common label keys split_ds = ds[cand] try: first = next(iter(split_ds), None) except Exception: first = None has_label = False if first is not None: if any(k in first for k in ("label", "answer", "answerKey", "correct", "gold")): has_label = True # choose test even if has_label False (many datasets publish test with labels in HF) if cand == "test" and cand in available_splits: chosen_split = "test" break if cand == "validation" and has_label: chosen_split = "validation" break if chosen_split is None and cand in available_splits: chosen_split = cand if chosen_split is None: # fallback to first available chosen_split = available_splits[0] val = ds.get(chosen_split) print(f"[Extra Eval] Task={task} using split='{chosen_split}' samples={len(val)}") device_t = torch.device(device if torch.cuda.is_available() else "cpu") model.to(device_t) model.eval() # Build tensors depending on format if cfg["format"] == "mc": ids_t, mask_t, labels_t = _tokenize_mc_split_to_tensors(task, tokenizer, val, max_length=max_length) # ids_t shape: (N, C, L) or (0,...) ds_t = TensorDataset(ids_t, mask_t, labels_t) loader = DataLoader(ds_t, batch_size=batch_size, shuffle=False, pin_memory=True) else: ids_t, mask_t, labels_t = _tokenize_pair_split_to_tensors(task, tokenizer, val, max_length=max_length) ds_t = TensorDataset(ids_t, mask_t, labels_t) loader = DataLoader(ds_t, batch_size=batch_size, shuffle=False, pin_memory=True) if checkpointing is not None: try: checkpointing.load_model_states("recent") except Exception: pass all_logits = [] all_labels = [] with torch.no_grad(): for batch in tqdm(loader, desc=f"Eval {task}"): if cfg["format"] == "mc": ids_b, mask_b, labs_b = batch ids_b = ids_b.to(device_t); mask_b = mask_b.to(device_t) # Try direct MC forward, else flatten fallback try: out = model(input_ids=ids_b, attention_mask=mask_b, labels=None) logits = getattr(out, "logits", None) if logits is None and isinstance(out, (list, tuple)): logits = out[0] except Exception: B, C, L = ids_b.shape flat_ids = ids_b.view(B*C, L).to(device_t) flat_mask = mask_b.view(B*C, L).to(device_t) out_flat = model(input_ids=flat_ids, attention_mask=flat_mask) flat_logits = getattr(out_flat, "logits", None) if flat_logits is None and isinstance(out_flat, (list, tuple)): flat_logits = out_flat[0] if flat_logits is None: raise RuntimeError("Model did not return logits for MC fallback") if flat_logits.ndim == 2 and flat_logits.shape[1] == 1: logits = flat_logits.view(B, C) else: logits = flat_logits.view(B, C, -1).mean(dim=-1) all_logits.append(logits.detach().cpu().numpy()) all_labels.append(labs_b.detach().cpu().numpy()) else: ids_b, mask_b, labs_b = batch ids_b = ids_b.to(device_t); mask_b = mask_b.to(device_t) out = model(input_ids=ids_b, attention_mask=mask_b, labels=None) logits = getattr(out, "logits", None) if logits is None and isinstance(out, (list, tuple)): logits = out[0] all_logits.append(logits.detach().cpu().numpy()) all_labels.append(labs_b.detach().cpu().numpy()) all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg["num_labels"])) all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,)) if cfg["format"] == "mc" or cfg["type"] == "multiple_choice" or cfg["type"] == "classification": preds = np.argmax(all_logits, axis=-1).astype(int) if all_logits.size else np.array([]) else: preds = _postprocess_predictions(task, all_logits, {"type": cfg["type"], "num_labels": cfg["num_labels"]}) # compute metric (accuracy for most MC tasks) try: metric = evaluate.load("accuracy") metric_res = metric.compute(predictions=preds.tolist(), references=all_labels.tolist()) except Exception as e: metric_res = {"error": str(e)} # write outputs os.makedirs(output_dir, exist_ok=True) out_json = Path(output_dir) / f"{task}_{chosen_split}_results.json" with open(out_json, "w", encoding="utf-8") as f: json.dump({"task": task, "split": chosen_split, "metrics": metric_res}, f, indent=2) csv_p = Path(output_dir) / f"{task}_{chosen_split}_preds.csv" pd.DataFrame({"pred": preds.tolist(), "label": all_labels.tolist()}).to_csv(csv_p, index=False) print(f"[Extra Eval] {task} split={chosen_split} metric={metric_res}") return metric_res # --------------------------------------------------------------------- # Grid-runner: LR sweep + restarts + checkpointing + eval # --------------------------------------------------------------------- def run_glue_benchmark(config, tokenizer, model, checkpointing: Optional[Checkpointing] = None, out_dir: str = "glue_outputs_grid"): """ Grid-search runner for GLUE + EXTRA tasks. Config attributes supported (defaults will be used if missing): - glue_tasks: list of tasks (GLUE or EXTRA) - batch_size, max_length, device - auto_train (bool) - train_epochs_per_task (dict) - bert_lr_candidates (list) - random_restarts_small (int) - base_seed (int) - train_batch_size, train_warmup_steps, train_weight_decay, train_grad_accum_steps """ tasks = getattr(config, "glue_tasks", None) if tasks is None: # default to a sensible subset; you can pass config with glue_tasks list tasks = ["rte", "cola", "mnli"] # allow passing a single string if isinstance(tasks, str): tasks = [t.strip() for t in tasks.split(",") if t.strip()] tasks = [ "arc_easy", "arc_challenge","hellaswag","openbookqa"] # "piqa": 3, "winogrande": 3, "boolq": 3, "hellaswag": 3, "openbookqa": 3, "arc_easy": 3, "arc_challenge": 3 # sanity-check tasks are in ALL_TASKS tasks = [t for t in tasks if t in ALL_TASKS] if not tasks: raise RuntimeError("No valid tasks found in config.glue_tasks (must be in GLUE_TASKS or EXTRA_TASKS).") batch_size = int(getattr(config, "batch_size", 64)) max_length = int(getattr(config, "max_length", 128)) device = getattr(config, "device", "cuda") auto_train = bool(getattr(config, "auto_train", True)) train_epochs = int(getattr(config, "train_epochs", 3)) train_epochs_per_task = getattr(config, "train_epochs_per_task", {}) if not train_epochs_per_task: train_epochs_per_task = { "cola": 5, "mrpc": 3, "rte": 5, "stsb": 3, "sst2": 3, "qqp": 3, "qnli": 3, "mnli": 3, "wnli": 5, "piqa": 3, "winogrande": 3, "boolq": 3, "hellaswag": 3, "openbookqa": 3, "arc_easy": 3, "arc_challenge": 3 } train_batch_size = int(getattr(config, "train_batch_size", 32)) train_warmup_steps = int(getattr(config, "train_warmup_steps", 100)) train_weight_decay = float(getattr(config, "train_weight_decay", 0.01)) train_grad_accum_steps = int(getattr(config, "train_grad_accum_steps", 1)) lr_candidates = getattr(config, "bert_lr_candidates", BERT_LR_CANDIDATES) random_restarts_small = int(getattr(config, "random_restarts_small", 1)) base_seed = int(getattr(config, "base_seed", 543211)) out_dir = Path(out_dir); out_dir.mkdir(parents=True, exist_ok=True) checkpoint_out_root = out_dir / "checkpoints"; checkpoint_out_root.mkdir(parents=True, exist_ok=True) # try to save original model state so we can reset between runs original_state = None try: original_state = unwrap_model(model).state_dict() except Exception: try: original_state = model.state_dict() except Exception: original_state = None def _reset_model_to_original(): if original_state is None: return try: unwrap_model(model).load_state_dict(original_state, strict=False) except Exception: try: model.load_state_dict(original_state, strict=False) except Exception: pass summary_rows = [] for task in tasks: assert (task in GLUE_TASKS) or (task in EXTRA_TASKS), f"Unknown task '{task}'" print(f"\n==== Grid-running task: {task} ====") epochs_this_task = int(train_epochs_per_task.get(task, train_epochs)) print(f"[Grid] Epochs for task '{task}': {epochs_this_task}") # load data splits if task in EXTRA_TASKS: cfg = EXTRA_TASKS[task] try: if "hf_config" in cfg: ds = load_dataset(cfg["hf_path"], cfg["hf_config"]) else: ds = load_dataset(cfg["hf_path"]) except Exception as e: raise RuntimeError(f"Failed to load dataset for extra task {task}: {e}") # pick train/val splits train_raw = ds.get("train") # prefer validation if available, else test val_raw = ds.get("validation") or ds.get("test") or ds.get("validation_matched") if val_raw is None: # pick first available split as val val_raw = next(iter(ds.values())) else: hf = load_dataset("glue", GLUE_TASKS[task]["hf_name"]) train_raw = hf["train"] val_raw = hf["validation_matched"] if task == "mnli" else hf["validation"] task_ckpt_root = checkpoint_out_root / task task_ckpt_root.mkdir(parents=True, exist_ok=True) all_run_records = [] best_run = { "score": None, "metrics": None, "lr": None, "restart": None, "seed": None, "best_epoch": None, "run_dir": None, "best_ckpt_path": None } small_flag = (task in SMALL_TASKS_RANDOM_RESTARTS) or (task in SMALL_TASKS_RANDOM_RESTARTS_EXTRA) restarts_per_lr = int(random_restarts_small) if small_flag and auto_train else 1 if auto_train: print(f"[Grid] Auto-training. LRs={lr_candidates}. Restarts/LR={restarts_per_lr} (small={small_flag}).") for lr in lr_candidates: for restart_idx in range(restarts_per_lr): seed = base_seed + (abs(hash(task)) % 10000) * 1000 + int(restart_idx) * 10 + (int(round(lr * 1e7)) % 1000) print(f"\n[SWEEP] task={task} lr={lr:g} restart={restart_idx}/{restarts_per_lr-1} seed={seed}") _reset_model_to_original() if checkpointing is not None: try: checkpointing.load_model_states("recent") except Exception: pass run_dir = task_ckpt_root / f"lr_{lr:g}" / f"restart_{restart_idx}" run_dir.mkdir(parents=True, exist_ok=True) try: if task in EXTRA_TASKS: # train extra model, metric_res = train_full_finetune_extra( task=task, tokenizer=tokenizer, model=model, raw_train=train_raw, raw_val=val_raw, device=device, epochs=epochs_this_task, batch_size=train_batch_size, lr=float(lr), weight_decay=train_weight_decay, warmup_steps=train_warmup_steps, max_length=max_length, grad_accum_steps=train_grad_accum_steps, out_checkpoint_dir=str(run_dir) ) sc = _metric_to_scalar(task, metric_res, fallback_val_loss=None) best_epoch = None else: model, metric_res, score, best_epoch = train_full_finetune( task=task, tokenizer=tokenizer, model=model, raw_train=train_raw, raw_val=val_raw, device=device, epochs=epochs_this_task, batch_size=train_batch_size, lr=float(lr), weight_decay=train_weight_decay, warmup_steps=train_warmup_steps, max_length=max_length, grad_accum_steps=train_grad_accum_steps, out_checkpoint_dir=str(run_dir), seed=int(seed) ) sc = float(score) # update run_meta.json with lr/restart/seed meta_path = Path(run_dir) / "run_meta.json" meta = {} if meta_path.exists(): try: meta = json.loads(meta_path.read_text(encoding="utf-8")) except Exception: meta = {} meta.update({"restart": int(restart_idx), "seed": int(seed), "lr": float(lr)}) _json_dump(meta, meta_path) rec = { "task": task, "lr": float(lr), "restart": int(restart_idx), "seed": int(seed), "epochs": int(epochs_this_task), "dev_score": float(sc) if sc is not None else None, "dev_metrics": json.dumps(metric_res), "run_dir": str(run_dir), "best_ckpt_path": str(Path(run_dir) / "best_finetuned.pt") if (Path(run_dir) / "best_finetuned.pt").exists() else None, "final_ckpt_path": str(Path(run_dir) / "finetuned.pt") if (Path(run_dir) / "finetuned.pt").exists() else None, } all_run_records.append(rec) if best_run["score"] is None or (sc is not None and float(sc) > float(best_run["score"])): best_run.update({ "score": float(sc) if sc is not None else None, "metrics": metric_res, "lr": float(lr), "restart": int(restart_idx), "seed": int(seed), "best_epoch": int(best_epoch) if best_epoch is not None else None, "run_dir": str(run_dir), "best_ckpt_path": rec["best_ckpt_path"], }) except Exception as e: print(f"[WARN] Training run failed for {task} lr={lr} restart={restart_idx}: {e}") all_run_records.append({ "task": task, "lr": float(lr), "restart": int(restart_idx), "seed": int(seed), "epochs": int(epochs_this_task), "dev_score": None, "dev_metrics": json.dumps({"error": str(e)}), "run_dir": str(run_dir), "best_ckpt_path": None, "final_ckpt_path": None, }) # Save all_runs.csv for this task all_runs_csv = task_ckpt_root / "all_runs.csv" pd.DataFrame(all_run_records).to_csv(all_runs_csv, index=False) print(f"[Grid] Saved all runs summary to: {all_runs_csv}") # Save best_overall best_overall_dir = task_ckpt_root / "best_overall" best_overall_dir.mkdir(parents=True, exist_ok=True) if best_run.get("best_ckpt_path") and best_run["best_ckpt_path"] and os.path.exists(best_run["best_ckpt_path"]): try: shutil.copy2(best_run["best_ckpt_path"], best_overall_dir / "best_finetuned.pt") except Exception: pass _json_dump(best_run, best_overall_dir / "best_meta.json") print(f"[Grid] Best run for task='{task}': lr={best_run.get('lr')}, restart={best_run.get('restart')}, seed={best_run.get('seed')}, dev_score={best_run.get('score')}") # Reset model and load best for evaluation _reset_model_to_original() model_for_eval = model try: best_ckpt = best_overall_dir / "best_finetuned.pt" if best_ckpt.exists(): model_for_eval = _load_finetuned_checkpoint_for_task(task, model, str(best_ckpt)) except Exception as e: print(f"[WARN] Failed to load best_overall checkpoint for eval; using current model. err={e}") model_for_eval = model # Evaluate and write summary rows task_out_dir = out_dir / task; task_out_dir.mkdir(parents=True, exist_ok=True) if task in EXTRA_TASKS: metric_res = run_extra_task(task=task, tokenizer=tokenizer, model=model_for_eval, checkpointing=None, device=device, batch_size=batch_size, max_length=max_length, output_dir=str(task_out_dir)) summary_rows.append({"task": task, "split": "selected", "epochs": epochs_this_task, "metrics": json.dumps(metric_res), "selected_lr": best_run.get("lr")}) else: res = run_glue_task(task=task, tokenizer=tokenizer, model=model_for_eval, checkpointing=None, device=device, batch_size=batch_size, max_length=max_length, output_dir=str(task_out_dir), compute_errorbar=True) for split, pack in res.items(): metrics = pack["metrics"] eb = pack["errorbar"] summary_rows.append({ "task": task, "split": split, "epochs": epochs_this_task, "selected_lr": best_run.get("lr"), "selected_restart": best_run.get("restart"), "selected_seed": best_run.get("seed"), "selected_dev_score": best_run.get("score"), "eval_metrics": json.dumps(metrics), "errorbar_mean": (eb["mean"] if eb else None), "errorbar_std": (eb["std"] if eb else None), "errorbar_stderr": (eb["stderr"] if eb else None), "errorbar_detail": json.dumps(eb) if eb else None, }) summary_csv = out_dir / "glue_summary.csv" pd.DataFrame(summary_rows).to_csv(summary_csv, index=False) print(f"\n[Grid] Summary saved to: {summary_csv}") return pd.DataFrame(summary_rows) # --------------------------------------------------------------------- # CLI shim (optional) # --------------------------------------------------------------------- if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument("--tasks", type=str, default="sst2", help="comma separated tasks (supports glue and extra tasks: boolq,piqa,winogrande,hellaswag,openbookqa,arc_easy,arc_challenge)") parser.add_argument("--batch_size", type=int, default=64) parser.add_argument("--max_length", type=int, default=128) parser.add_argument("--device", type=str, default="cuda") parser.add_argument("--out_dir", type=str, default="glue_outputs_grid") args = parser.parse_args() print("This module is intended to be invoked from your project's main which provides tokenizer/model/checkpointing.") print(f"LI args tasks={args.tasks} batch_size={args.batch_size} max_length={args.max_length} device={args.device} out_dir={args.out_dir}") # Example usage (pseudo): # from transformers import AutoTokenizer, AutoModelForSequenceClassification # tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") # model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2) # cfg = type("C", (), {"glue_tasks": args.tasks.split(","), "batch_size": args.batch_size, "max_length": args.max_length, "device": args.device}) # run_glue_benchmark(cfg, tokenizer, model, checkpointing=None, out_dir=args.out_dir)