# 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): 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 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: cand = getattr(base_model, "config", None) if cand is not None and hasattr(cand, "hidden_size"): try: inferred_hidden = int(cand.hidden_size) except Exception: inferred_hidden = None if inferred_hidden is None: 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 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): out = None 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 # --------------------------------------------------------------------- # Generic multiple-choice extractor & tokenizer # --------------------------------------------------------------------- # Improved unwrapping + robust context assembly and MC example extractor def _maybe_unwrap(v): """ Improved unwrap: - numpy / torch scalars -> Python scalars - single-element lists/tuples -> unwrap - dicts are left as-is (they may represent structured choices) - numeric strings -> int when appropriate (helps labels like '3') - lists of strings left as-is (caller may join) """ try: import numpy as _np import torch as _torch except Exception: _np = None _torch = None if v is None: return None # torch tensor scalar or 0-d array if _torch is not None and isinstance(v, _torch.Tensor): if v.ndim == 0: return v.item() if v.numel() == 1: return v.view(-1).tolist()[0] return v # numpy scalar / array if _np is not None and isinstance(v, _np.ndarray): if v.shape == () or v.size == 1: return v.flatten().tolist()[0] return v.tolist() # dict likely meaningful (choices dict) -> keep as-is if isinstance(v, dict): return v # list/tuple with single element -> unwrap to element if isinstance(v, (list, tuple)) and len(v) == 1: return _maybe_unwrap(v[0]) # string that represents an integer -> convert to int (helps labels like '3') if isinstance(v, str): s = v.strip() if s.isdigit(): try: return int(s) except Exception: pass return v return v def _assemble_context_from_example(ex: dict): """ Build a single string context from many possible fields. Priority / heuristics: 1. If 'ctx' or 'context' present and non-empty -> use it. 2. If 'ctx_a' and 'ctx_b' (or 'ctxA'/'ctxB') present -> join them. 3. If 'premise' and 'hypothesis' present -> join them (suitable for MNLI-like). 4. Use question-oriented fields if that is the best we can do: question_stem, question, stem. 5. Fallback: join a selection of textual fields in a sensible order. Converts lists/tuples of strings to joined text. """ def _to_text(x): if x is None: return "" if isinstance(x, (list, tuple)): # attempt to join list of strings or dicts parts = [] for it in x: if isinstance(it, dict): txt = it.get("text") or it.get("choice") or it.get("label") or str(it) parts.append(str(txt)) else: parts.append(str(it)) return " ".join([p.strip() for p in parts if p is not None and str(p).strip() != ""]) if isinstance(x, dict): # pick likely textual fields for k in ("text", "content", "question", "sentence", "passage", "context"): if k in x and x[k]: return _to_text(x[k]) return str(x) return str(x) # common direct fields keys = {k.lower(): v for k, v in ex.items()} # 1) explicit ctx / context for k in ("ctx", "context"): if k in keys and keys[k]: return _to_text(keys[k]).strip() # 2) ctx_a + ctx_b variants a_keys = ("ctx_a", "ctxA", "context_a", "contextA", "ctxa") b_keys = ("ctx_b", "ctxB", "context_b", "contextB", "ctxb") a_val = None b_val = None for ka in a_keys: if ka in keys and keys[ka]: a_val = keys[ka] break for kb in b_keys: if kb in keys and keys[kb]: b_val = keys[kb] break if a_val is not None or b_val is not None: parts = [] if a_val is not None: parts.append(_to_text(a_val)) if b_val is not None: parts.append(_to_text(b_val)) return " ".join([p.strip() for p in parts if p and p.strip() != ""]).strip() # 3) premise + hypothesis if "premise" in keys or "hypothesis" in keys: p = keys.get("premise") h = keys.get("hypothesis") parts = [] if p: parts.append(_to_text(p)) if h: parts.append(_to_text(h)) return " ".join([p.strip() for p in parts if p and p.strip() != ""]).strip() # 4) question stem / question / stem for k in ("question_stem", "questionStem", "question", "stem", "prompt", "goal"): if k.lower() in keys and keys[k.lower()]: return _to_text(keys[k.lower()]).strip() # 5) passage/article/story/sentence for k in ("passage", "article", "story", "sentence", "paragraph"): if k in keys and keys[k]: return _to_text(keys[k]).strip() # 6) try to build a concat from any textual fields in a sensible order fallback_fields = [ "context", "question", "passage", "article", "story", "sentence", "prompt", "goal", "stem", "query", "description", "narration" ] parts = [] for f in fallback_fields: if f in keys and keys[f]: parts.append(_to_text(keys[f])) combined = " ".join([p.strip() for p in parts if p and p.strip() != ""]).strip() if combined: return combined # final fallback: try first non-empty string-ish field for k, v in ex.items(): if isinstance(v, str) and v.strip(): return v.strip() if isinstance(v, (list, tuple)) and len(v) > 0 and all(isinstance(x, str) for x in v): return " ".join(v).strip() return "" def _extract_mc_example(task: str, ex: dict): """ Robust extractor for multiple-choice examples across HF dataset schemas. Returns (context:str, options_list:List[str], label_index_or_None) Improvements over earlier versions: - Uses _maybe_unwrap on fields - Assembles context robustly via _assemble_context_from_example (handles ctx_a+ctx_b, premise+hypothesis, etc.) - Accepts choices as: * dict {'text': [...], 'label': [...]} (OpenBookQA / ARC) * list of dicts or strings * 'endings' list (HellaSwag) * explicit fields choice1/option1 etc. - Normalizes label forms: 'A'..'D' -> 0..3, '3' -> int(3), torch/numpy scalars handled """ def _map_label_to_int(label, opts): if label is None: return None # unwrap containers if isinstance(label, (list, tuple)) and len(label) > 0: label = label[0] try: import numpy as _np import torch as _torch except Exception: _np = None _torch = None if _torch is not None and isinstance(label, _torch.Tensor): try: return int(label.item()) except Exception: pass if _np is not None and isinstance(label, _np.ndarray): if label.size == 1: return int(label.flatten().tolist()[0]) if isinstance(label, str): s = label.strip() if len(s) == 1 and s.isalpha(): return ord(s.upper()) - ord("A") if s.isdigit(): try: return int(s) except Exception: pass # if label equals one of option texts, return that index if opts: for i, o in enumerate(opts): if isinstance(o, str) and s == o: return i return None if isinstance(label, (int, np.integer)): return int(label) return None # shallow normalized copy ex_norm = {} for k, v in ex.items(): try: ex_norm[k] = _maybe_unwrap(v) except Exception: ex_norm[k] = v ex = ex_norm # assemble context early ctx = _assemble_context_from_example(ex) # 1) OpenBookQA / ARC-style: 'choices' is dict with 'text' list if "choices" in ex and isinstance(ex["choices"], dict): chd = ex["choices"] texts = chd.get("text") or chd.get("texts") or chd.get("choice") or None if isinstance(texts, (list, tuple)) and len(texts) > 0: opts = [str(x) for x in texts] lab = ex.get("answerKey") or ex.get("answer") or ex.get("correctAnswer") or ex.get("label") lab_idx = _map_label_to_int(lab, opts) return ctx, opts, lab_idx # 2) HellaSwag style 'endings' if "endings" in ex and isinstance(ex["endings"], list) and len(ex["endings"]) > 0: opts = [str(x) for x in ex["endings"]] lab = ex.get("label") or ex.get("answerKey") or ex.get("answer") lab_idx = _map_label_to_int(lab, opts) return ctx, opts, lab_idx # 3) 'choices' as 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("choice") or c.get("label") or str(c)) else: opts.append(str(c)) lab = ex.get("answerKey") or ex.get("answer") or ex.get("label") or ex.get("correct") lab_idx = _map_label_to_int(lab, opts) return ctx, opts, lab_idx # 4) 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: val = _maybe_unwrap(ex.get(key)) opts.append(str(val)) found = True i += 1 else: break if found: lab = ex.get("label") or ex.get("answerKey") or ex.get("answer") lab_idx = _map_label_to_int(lab, opts) return ctx, opts, lab_idx # 5) fallback: look for list-valued fields that look like options for k, v in ex.items(): if isinstance(v, (list, tuple)) 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] lab = ex.get("answer") or ex.get("label") or ex.get("answerKey") lab_idx = _map_label_to_int(lab, opts) return ctx, opts, lab_idx # 6) collect candidate option-fields heuristically candidate_opts = [] for k in sorted(ex.keys()): if any(tok in k.lower() for tok in ("option", "choice", "ending", "answer", "alt", "sol", "choices")): val = _maybe_unwrap(ex.get(k)) candidate_opts.append(str(val)) if candidate_opts: lab = ex.get("label") or ex.get("answer") or ex.get("answerKey") lab_idx = _map_label_to_int(lab, candidate_opts) return ctx, candidate_opts, lab_idx # last resort: no options discovered lab = ex.get("label") or ex.get("answerKey") or ex.get("answer") lab_idx = _map_label_to_int(lab, []) return ctx, [], lab_idx # --------------------------- # 新增：detect_num_choices # --------------------------- def detect_num_choices(raw_split, task: str, max_samples: int = 200): max_s = min(len(raw_split), max_samples) # 强制拿出前 max_s 个示例为「标量形式」：select 确保每个示例是 dict of scalars 而不是子 Dataset subset = raw_split.select(range(max_s)) if hasattr(raw_split, "select") else raw_split[:max_s] counts = {} total = 0 kept = 0 bad_label = 0 for ex in subset: total += 1 try: ctx, opts, lab = _extract_mc_example(task, ex) except Exception: counts.setdefault("no_opts", 0) counts["no_opts"] += 1 continue if not opts: counts.setdefault("no_opts", 0) counts["no_opts"] += 1 else: kept += 1 counts.setdefault(len(opts), 0) counts[len(opts)] += 1 if lab is None or (isinstance(lab, int) and lab < 0): bad_label += 1 numeric = {k: v for k, v in counts.items() if isinstance(k, int)} detected = None if numeric: detected = max(numeric.items(), key=lambda x: x[1])[0] summary = { "task": task, "total_scanned": total, "kept": kept, "counts": counts, "detected_num_choices": detected, "skipped_bad_label": bad_label, } print(f"[detect_num_choices][{task}] scanned={total} kept={kept} detected={detected} counts={counts} bad_label={bad_label}") return summary # --------------------------- # 更新：_tokenize_generic_mc_split_to_tensors # --------------------------- def _tokenize_generic_mc_split_to_tensors(task: str, tokenizer, raw_split, max_length=128, batch_tokenize_size=256, expect_num_choices: Optional[int] = None): """ Generic MC tokenizer with optional expect_num_choices enforcement. Returns (input_ids_t (N,C,L), attention_t (N,C,L), labels_t (N,)), and prints debug counts. If expect_num_choices is provided, examples whose number of options != expect_num_choices are considered 'inconsistent' and skipped. This version does not swallow exceptions: tokenizer errors and unexpected shapes will raise. """ contexts = [] options = [] labels = [] total = 0 skipped_no_opts = 0 skipped_inconsistent = 0 skipped_bad_label = 0 kept = 0 for ex in raw_split: ctx, opts, lab = _extract_mc_example(task, ex) total += 1 # print(task) # print(ex) # print(ctx) # print(opts) # print(lab) # import pdb # pdb.set_trace() # print(ex) # print(ctx) # print(opts) # print(lab) if not opts: skipped_no_opts += 1 continue if expect_num_choices is not None and len(opts) != expect_num_choices: skipped_inconsistent += 1 continue if lab is None: skipped_bad_label += 1 continue try: lab_i = int(lab) except Exception: skipped_bad_label += 1 continue if lab_i < 0 or lab_i >= len(opts): skipped_bad_label += 1 continue contexts.append(ctx if ctx is not None else "") options.append(opts) labels.append(lab_i) kept += 1 print(f"[tokenize_generic_mc][{task}] total={total} kept={kept} skipped_no_opts={skipped_no_opts} skipped_inconsistent={skipped_inconsistent} skipped_bad_label={skipped_bad_label}") 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) 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]) # Build flat pairs for batch tokenization and assert tokenizer returns expected length input_ids_rows = [] attention_rows = [] 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 = tokenizer(flat_pairs, truncation=True, padding=False, 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() expected = len(chunk_ctx) * num_choices if len(ids_flat) != expected: # raise explicit error so user sees where tokenizer mismatched raise RuntimeError(f"[tokenize_generic_mc][{task}] tokenizer returned {len(ids_flat)} items but expected {expected} (i={i}). " f"First flat_pairs sample: {flat_pairs[0] if len(flat_pairs)>0 else None}") per_example = [] per_mask_example = [] idx = 0 for _ in range(len(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 # --------------------------- # train_full_finetune_extra（修补版） # --------------------------- 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, expect_num_choices: Optional[int] = None): """ Fine-tune for EXTRA_TASKS with a safe projection fallback. This variant does not silence errors: tokenizer/model forward issues will raise and surface to the caller. For multiple-choice tasks we flatten (B, C, L) -> (B*C, L) before feeding into the transformer, then reshape logits back to (B, C) and compute CrossEntropyLoss against labels (B,). This patched version ensures any linear/projector layers created for MC handling are included in the optimizer (either by pre-creating before optimizer creation, or by adding param groups when created at runtime). """ import torch.nn as nn 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: # Tokenize to (N, C, L) train_ids, train_mask, train_labels = _tokenize_generic_mc_split_to_tensors( task, tokenizer, raw_train, max_length=max_length, batch_tokenize_size=256, expect_num_choices=expect_num_choices ) val_ids, val_mask, val_labels = _tokenize_generic_mc_split_to_tensors( task, tokenizer, raw_val, max_length=max_length, batch_tokenize_size=256, expect_num_choices=expect_num_choices ) # Basic checks if train_ids.ndim != 3: raise RuntimeError(f"[ExtraTrain][{task}] train_ids expected 3 dims (N,C,L) got {train_ids.ndim}") if val_ids.ndim != 3 and len(val_ids) > 0: raise RuntimeError(f"[ExtraTrain][{task}] val_ids expected 3 dims (N,C,L) got {val_ids.ndim}") train_ds = TensorDataset(train_ids, train_mask, train_labels) val_ds = TensorDataset(val_ids, val_mask, val_labels) else: # pair tasks: reuse existing pair tokenizer (expect to be defined) train_ids, train_mask, train_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_train, cfg, max_length=max_length) val_ids, val_mask, val_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_val, cfg, max_length=max_length) train_ds = TensorDataset(train_ids, train_mask, train_labels) val_ds = TensorDataset(val_ids, val_mask, val_labels) # If train_ds is empty, bail early with explicit error if len(train_ds) == 0: raise RuntimeError(f"No training samples after tokenization for task={task}; aborting.") 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) # ------------------------------------------------------- # PRE-CREATE any stable projectors BEFORE optimizer creation # This ensures their parameters are part of model.parameters() and thus # included in the optimizer when it's built. # ------------------------------------------------------- expected_num_labels = cfg["num_labels"] if is_mc: # Try to infer num_choices from training tensor if possible try: num_choices_train = int(train_ids.shape[1]) except Exception: num_choices_train = None if num_choices_train is not None and num_choices_train != expected_num_labels: # Create a stable choice->label mapping so optimizer will include it if not hasattr(model, "_mc_choice_to_label"): model._mc_choice_to_label = nn.Linear(num_choices_train, expected_num_labels).to(device) # note: optimizer not created yet, so this will be captured by model.parameters() print(f"[ExtraTrain][{task}] pre-created _mc_choice_to_label: {num_choices_train} -> {expected_num_labels}") # Create optimizer AFTER any pre-created layers so they are included 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) 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) # For MC tasks we will use CE on (B,C) vs labels (B,) ce_loss_fn = torch.nn.CrossEntropyLoss() mse_loss_fn = torch.nn.MSELoss() # Debug: report optimizer param count try: tot_p = sum(p.numel() for g in optimizer.param_groups for p in g['params']) print(f"[ExtraTrain][{task}] optimizer param count: {tot_p}") except Exception: pass model.train() global_step = 0 final_metric_res = {} # runtime _mc_projector (maps feature-dim D -> 1) may be unknown until first forward. # If created later, we must add its params to optimizer via add_param_group. 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) # Ensure labels shape matches batch if labs_b.dim() > 1: labs_b = labs_b.view(-1) B = ids_b.size(0) if labs_b.numel() != B: raise RuntimeError(f"[ExtraTrain][FATAL] label count ({labs_b.numel()}) != batch size ({B}) at step {step}; aborting.") # Flatten choices to (B*C, L) Bf, Cf, Lf = ids_b.shape flat_ids = ids_b.view(Bf * Cf, Lf).to(device) flat_mask = mask_b.view(Bf * Cf, Lf).to(device) if mask_b is not None else None # Forward pass (robust) out_flat = None try: if flat_mask is not None: out_flat = model(input_ids=flat_ids, attention_mask=flat_mask, labels=None) else: out_flat = model(input_ids=flat_ids, labels=None) except TypeError: out_flat = model(flat_ids) 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 forward") # Normalize flat_logits -> choice_scores (B, C) if flat_logits.ndim == 1: choice_scores = flat_logits.view(Bf, Cf) elif flat_logits.ndim == 2: D = flat_logits.shape[1] if D == 1: choice_scores = flat_logits.view(Bf, Cf).squeeze(-1) else: # learned projector from D -> 1, reuse if present; if not, create once and register with optimizer need_new_proj = False if not hasattr(model, "_mc_projector"): need_new_proj = True else: # check compatible input dim try: cur_in = getattr(model._mc_projector, "in_features", None) if cur_in is None: cur_in = model._mc_projector.weight.shape[1] if cur_in != D: need_new_proj = True except Exception: need_new_proj = True if need_new_proj: # create projector and add to optimizer param group immediately model._mc_projector = nn.Linear(D, 1).to(device) try: optimizer.add_param_group({'params': model._mc_projector.parameters()}) print(f"[ExtraTrain][{task}] created and registered _mc_projector with D={D}") except Exception as e: print(f"[WARN][{task}] failed to add_param_group for _mc_projector: {e}") proj = model._mc_projector(flat_logits) # (B*C, 1) choice_scores = proj.view(Bf, Cf).squeeze(-1) elif flat_logits.ndim == 3: # (B*C, S, D), reduce to scalar per choice choice_scores = flat_logits.mean(dim=tuple(range(1, flat_logits.ndim))).view(Bf, Cf) else: # fallback: flatten last dims then mean choice_scores = flat_logits.reshape(flat_logits.size(0), -1).mean(dim=1).view(Bf, Cf) if choice_scores.ndim == 1: choice_scores = choice_scores.unsqueeze(0) if choice_scores.shape != (Bf, Cf): raise RuntimeError(f"Unexpected choice_scores shape {choice_scores.shape} at step {step}") # ensure final logits align (B, num_labels) and labels (B,) if Cf != expected_num_labels: # if dataset choices count differs, project to expected_num_labels exist_dim = choice_scores.shape[1] # Prefer pre-created stable _mc_choice_to_label if available if not hasattr(model, "_mc_choice_to_label"): # create and register to optimizer model._mc_choice_to_label = nn.Linear(exist_dim, expected_num_labels).to(device) try: optimizer.add_param_group({'params': model._mc_choice_to_label.parameters()}) print(f"[ExtraTrain][{task}] created and registered _mc_choice_to_label: {exist_dim} -> {expected_num_labels}") except Exception as e: print(f"[WARN][{task}] failed to add_param_group for _mc_choice_to_label: {e}") used_logits = model._mc_choice_to_label(choice_scores) else: used_logits = choice_scores if used_logits.ndim != 2 or used_logits.size(0) != labs_b.size(0): raise RuntimeError(f"Logits/labels shape mismatch used_logits={tuple(used_logits.shape)} labels={tuple(labs_b.shape)} at step {step}") loss = ce_loss_fn(used_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 and isinstance(out, (tuple, list)): logits = out[0] if logits is None: raise RuntimeError("Model did not return logits for pair task") if logits.ndim == 2 and logits.shape[1] != expected_num_labels: exist_dim = logits.shape[1] if not hasattr(model, "_mc_projector"): model._mc_projector = nn.Linear(exist_dim, expected_num_labels).to(device) try: optimizer.add_param_group({'params': model._mc_projector.parameters()}) print(f"[ExtraTrain][{task}] created and registered _mc_projector for pair-task: {exist_dim} -> {expected_num_labels}") except Exception as e: print(f"[WARN][{task}] failed to add_param_group for pair _mc_projector: {e}") used_logits = model._mc_projector(logits) else: used_logits = logits if cfg["type"] == "classification": loss = ce_loss_fn(used_logits, labs_b.long()) else: # regression if used_logits.ndim == 2 and used_logits.shape[1] == 1: preds = used_logits.squeeze(1) elif used_logits.ndim == 2: preds = used_logits.mean(dim=1) else: preds = used_logits loss = mse_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() scheduler.step() optimizer.zero_grad() global_step += 1 running_loss += float(loss.item()) * (ids_b.size(0) if is_mc else ids_b.size(0)) # validation (no exception swallowing) model.eval() all_logits = [] all_labels = [] print(running_loss) 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 if labs_b.dim() > 1: labs_b = labs_b.view(-1) B = ids_b.size(0) if labs_b.numel() != B: raise RuntimeError(f"[ExtraVal][FATAL] val labels count ({labs_b.numel()}) != batch size ({B})") ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device) Bf, Cf, Lf = ids_b.shape flat_ids = ids_b.view(Bf*Cf, Lf).to(device) flat_mask = mask_b.view(Bf*Cf, Lf).to(device) if mask_b is not None else None out_flat = None try: if flat_mask is not None: out_flat = model(input_ids=flat_ids, attention_mask=flat_mask, labels=None) else: out_flat = model(input_ids=flat_ids, labels=None) except TypeError: out_flat = model(flat_ids) 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(Bf, Cf) elif flat_logits.ndim == 2: D = flat_logits.shape[1] if D == 1: logits = flat_logits.view(Bf, Cf).squeeze(-1) else: # ensure projector exists (create+register if not) need_new_proj = False if not hasattr(model, "_mc_projector"): need_new_proj = True else: try: cur_in = getattr(model._mc_projector, "in_features", None) if cur_in is None: cur_in = model._mc_projector.weight.shape[1] if cur_in != D: need_new_proj = True except Exception: need_new_proj = True if need_new_proj: model._mc_projector = nn.Linear(D, 1).to(device) try: optimizer.add_param_group({'params': model._mc_projector.parameters()}) print(f"[ExtraVal][{task}] created and registered _mc_projector with D={D}") except Exception as e: print(f"[WARN][{task}] failed to add_param_group for _mc_projector during val: {e}") proj = model._mc_projector(flat_logits) logits = proj.view(Bf, Cf).squeeze(-1) elif flat_logits.ndim == 1: logits = flat_logits.view(Bf, Cf) elif flat_logits.ndim == 3: logits = flat_logits.mean(dim=tuple(range(1, flat_logits.ndim))).view(Bf, Cf) else: logits = flat_logits.reshape(flat_logits.size(0), -1).mean(dim=1).view(Bf, Cf) if logits.shape[1] != expected_num_labels: exist_dim = logits.shape[1] if not hasattr(model, "_mc_choice_to_label"): model._mc_choice_to_label = nn.Linear(exist_dim, expected_num_labels).to(device) try: optimizer.add_param_group({'params': model._mc_choice_to_label.parameters()}) print(f"[ExtraVal][{task}] created and registered _mc_choice_to_label: {exist_dim} -> {expected_num_labels}") except Exception as e: print(f"[WARN][{task}] failed to add_param_group for _mc_choice_to_label during val: {e}") used_logits = model._mc_choice_to_label(logits) else: used_logits = logits all_logits.append(used_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 and isinstance(out, (tuple, list)): logits = out[0] if logits.ndim == 2 and logits.shape[1] != expected_num_labels: exist_dim = logits.shape[1] if not hasattr(model, "_mc_projector"): model._mc_projector = nn.Linear(exist_dim, expected_num_labels).to(device) try: optimizer.add_param_group({'params': model._mc_projector.parameters()}) print(f"[ExtraVal][{task}] created and registered _mc_projector for pair-task val: {exist_dim} -> {expected_num_labels}") except Exception as e: print(f"[WARN][{task}] failed to add_param_group for pair _mc_projector during val: {e}") used_logits = model._mc_projector(logits) else: used_logits = logits all_logits.append(used_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"]}) metric = evaluate.load("accuracy") metric_res = metric.compute(predictions=preds.tolist(), references=all_labels.tolist()) 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" sd = unwrap_model(model).state_dict() torch.save(sd, str(fname)) print(f"[Extra FT] Saved finetuned model to: {fname}") return model, final_metric_res # --------------------------- # 辅助：打印 HellaSwag 样例（用来“给我看一下”） # --------------------------- def print_hellaswag_samples(n: int = 5): """ 加载 hellaswag 数据集并打印前 n 个样例经 _extract_mc_example 解析后的 context/options/label。 运行时会在 stdout 中输出，方便你检查数据格式。 """ ds = load_dataset("hellaswag") split = ds.get("validation") or ds.get("train") or ds.get("test") if split is None: print("[print_hellaswag_samples] no split") return for i, ex in enumerate(split[:n]): ctx, opts, lab = _extract_mc_example("hellaswag", ex) print(f"=== sample {i} ===") print("context:", (ctx[:400] + "...") if ctx is not None and len(str(ctx))>400 else ctx) print("num_options:", len(opts)) for j, o in enumerate(opts): print(f" [{j}] {o[:200]}{'...' if len(str(o))>200 else ''}") print("label:", lab) print() # --------------------------- # _tokenize_mc_split_to_tensors (general MC tokenizer) # --------------------------- def _tokenize_mc_split_to_tensors(task: str, tokenizer, raw_split, max_length=128, batch_tokenize_size=256): """ Generalized multiple-choice tokenizer that supports: - PIQA (goal/sol1/sol2) - Winogrande (sentence/option1/option2) - ARC-like (question + choices, answerKey) - RACE/other formats with 'options'/'choices' (list of strings or list of dicts with 'text') Returns: input_ids: (N, C, L) attention_mask: (N, C, L) labels: (N,) with -1 for unknown/no-label examples This function performs strict checks and will raise on inconsistencies (no silent catches). """ contexts = [] options = [] labels = [] def _extract_choices_text(choices_field): out = [] for c in choices_field: if isinstance(c, str): out.append(c) elif isinstance(c, dict): txt = c.get("text") or c.get("content") or c.get("choice") or c.get("label") out.append(txt if txt is not None else str(c)) else: out.append(str(c)) return out def _map_answer_to_index(ans, choice_texts): if ans is None: return -1 if isinstance(ans, (list, tuple)): if len(ans) == 0: return -1 ans = ans[0] if isinstance(ans, str): s = ans.strip() if len(s) == 1 and s.isalpha(): idx = ord(s.upper()) - ord("A") if 0 <= idx < len(choice_texts): return idx try: ni = int(s) if 0 <= ni < len(choice_texts): return ni if 1 <= ni <= len(choice_texts): return ni - 1 except Exception: pass for i, ct in enumerate(choice_texts): if isinstance(ct, str) and s == ct: return i for i, ct in enumerate(choice_texts): if isinstance(ct, str) and s in ct: return i low = s.lower() if low in ("true", "t", "yes", "y"): return 1 if low in ("false", "f", "no", "n"): return 0 return -1 if isinstance(ans, (int, np.integer)): ai = int(ans) if 0 <= ai < len(choice_texts): return ai if 1 <= ai <= len(choice_texts): return ai - 1 return -1 return -1 num_choices = None for ex in raw_split: ctx = None opts = None lab = None # PIQA if "goal" in ex and ("sol1" in ex or "sol2" in ex): ctx = ex.get("goal") opts = [ex.get("sol1"), ex.get("sol2")] lab = ex.get("label") or ex.get("answer") or ex.get("answerKey") if isinstance(lab, (str, np.str_)) and str(lab).isdigit(): lab = int(lab) # Winogrande elif "sentence" in ex and ("option1" in ex or "option2" in ex): ctx = ex.get("sentence") opts = [ex.get("option1"), ex.get("option2")] lab = ex.get("answer") or ex.get("label") or ex.get("answerKey") if isinstance(lab, str) and lab.isdigit(): lab = int(lab) - 1 # convert 1-based -> 0-based # ARC / RACE-style: question + choices (choices may be list of strings or list of dicts) elif "question" in ex and ("choices" in ex or "options" in ex): ctx = ex.get("question") choices_field = ex.get("choices") or ex.get("options") opts = _extract_choices_text(choices_field) lab = ex.get("answerKey") or ex.get("answer") or ex.get("correct_answer") or ex.get("label") # Generic fallback: try known pairs elif "context" in ex and "options" in ex: ctx = ex.get("context") opts = _extract_choices_text(ex.get("options")) lab = ex.get("label") or ex.get("answer") else: if "question" in ex: ctx = ex.get("question") elif "query" in ex: ctx = ex.get("query") elif "prompt" in ex: ctx = ex.get("prompt") else: for k, v in ex.items(): if isinstance(v, str) and len(v) > 0: ctx = v break for k, v in ex.items(): if isinstance(v, (list, tuple)) and len(v) > 1: if all(isinstance(x, (str, dict)) for x in v): opts = _extract_choices_text(v) break lab = ex.get("answer") or ex.get("label") or ex.get("answerKey") or ex.get("correct_answer") if opts is None: opts = [""] if num_choices is None: num_choices = len(opts) else: if len(opts) != num_choices: if len(opts) < num_choices: opts = opts + [""] * (num_choices - len(opts)) else: opts = opts[:num_choices] lab_idx = _map_answer_to_index(lab, opts) contexts.append(ctx) options.append(opts) labels.append(int(lab_idx) if lab_idx is not None else -1) if num_choices is None: 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 if c is not None else "", o if o is not None else "")) enc = tokenizer(flat_pairs, truncation=True, padding=False, 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 range(len(chunk_ctx)): row = [] row_mask = [] for _ in range(num_choices): if idx >= len(ids_flat): # this will raise later if shapes inconsistent; better to raise now raise RuntimeError(f"[tokenize_mc] tokenizer returned too few items at batch starting index {i} (got {len(ids_flat)}, needed at least {len(chunk_ctx)*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 # --------------------------------------------------------------------- # 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}") # save final checkpoint 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": float(lr), "seed": int(seed) if seed is not None else None, "epochs": int(epochs), "batch_size": int(batch_size), "grad_accum_steps": int(grad_accum_steps), "warmup_steps": int(warmup_steps), "weight_decay": float(weight_decay), "max_length": int(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}' " f"(lr={lr:g}, seed={seed}): epoch={best_epoch}, " f"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): if not checkpoint_path or not os.path.exists(checkpoint_path): return base_model 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: wrapped_model.load_state_dict(sd, strict=False) return wrapped_model try: unwrap_model(base_model).load_state_dict(sd, strict=False) except Exception: base_model.load_state_dict(sd, strict=False) return base_model # --------------------------------------------------------------------- # Evaluation only for EXTRA_TASKS (no training) # --------------------------------------------------------------------- def run_extra_task( task: str, tokenizer, model, device: str = "cuda", batch_size: int = 32, max_length: int = 128, output_dir: Optional[str] = None, expect_num_choices: Optional[int] = None, ): """ Evaluation-only runner for EXTRA_TASKS. Uses the SAME MC flatten->reshape logic as train_full_finetune_extra. Reports accuracy on validation split (lm-eval style). """ cfg = EXTRA_TASKS[task] device_t = torch.device(device if torch.cuda.is_available() else "cpu") model.to(device_t) model.eval() # load dataset if "hf_config" in cfg: ds = load_dataset(cfg["hf_path"], cfg["hf_config"]) else: ds = load_dataset(cfg["hf_path"]) raw_eval = ds.get("validation") or ds.get("test") if raw_eval is None: raise RuntimeError(f"[run_extra_task] No validation/test split for task={task}") is_mc = cfg["format"] == "mc" if is_mc: ids, masks, labels = _tokenize_generic_mc_split_to_tensors( task, tokenizer, raw_eval, max_length=max_length, expect_num_choices=expect_num_choices, ) if ids.ndim != 3: raise RuntimeError(f"[run_extra_task][{task}] ids must be (N,C,L)") ds_eval = TensorDataset(ids, masks, labels) else: ids, masks, labels = _tokenize_hf_split_to_tensors( task, tokenizer, raw_eval, cfg, max_length=max_length, ) ds_eval = TensorDataset(ids, masks, labels) loader = DataLoader(ds_eval, batch_size=batch_size, shuffle=False) all_preds = [] all_refs = [] with torch.no_grad(): for batch in tqdm(loader, desc=f"[ExtraEval] {task}"): if is_mc: 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) B, C, L = ids_b.shape flat_ids = ids_b.view(B * C, L) flat_mask = mask_b.view(B * C, L) try: out = model(input_ids=flat_ids, attention_mask=flat_mask, labels=None) except TypeError: out = model(flat_ids) flat_logits = getattr(out, "logits", None) if flat_logits is None and isinstance(out, (tuple, list)): flat_logits = out[0] if flat_logits is None: raise RuntimeError("Model did not return logits during MC eval") # reduce to (B, C) if flat_logits.ndim == 1: choice_scores = flat_logits.view(B, C) elif flat_logits.ndim == 2: D = flat_logits.shape[1] if D == 1: choice_scores = flat_logits.view(B, C).squeeze(-1) else: if not hasattr(model, "_mc_projector") or model._mc_projector.weight.shape[1] != D: model._mc_projector = torch.nn.Linear(D, 1).to(device_t) proj = model._mc_projector(flat_logits) choice_scores = proj.view(B, C).squeeze(-1) else: choice_scores = flat_logits.mean(dim=tuple(range(1, flat_logits.ndim))).view(B, C) preds = torch.argmax(choice_scores, dim=1) 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, (tuple, list)): logits = out[0] preds = torch.argmax(logits, dim=-1) all_preds.append(preds.cpu().numpy()) all_refs.append(labs_b.cpu().numpy()) all_preds = np.concatenate(all_preds, axis=0) all_refs = np.concatenate(all_refs, axis=0) acc = float((all_preds == all_refs).mean()) metric_res = {"accuracy": acc} print(f"[ExtraEval] {task} accuracy={acc:.4f}") if output_dir: Path(output_dir).mkdir(parents=True, exist_ok=True) with open(Path(output_dir) / "eval_metrics.json", "w") as f: json.dump(metric_res, f, indent=2) 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", ): tasks = getattr(config, "glue_tasks", None) if tasks is None: tasks = list(ALL_TASKS.keys()) tasks = ['openbookqa','openbookqa','arc_easy','arc_challenge'] # ['hellaswag','openbookqa','arc_easy','arc_challenge'] # tasks = tasks[::-1] # tasks = ["piqa", "boolq", "winogrande"] # "hellaswag": "accuracy", # "openbookqa": "accuracy", # "arc_easy": "accuracy", # "arc_challenge": "accuracy", batch_size = getattr(config, "batch_size", 64) max_length = getattr(config, "max_length", 128) device = getattr(config, "device", "cuda") auto_train = getattr(config, "auto_train", True) train_epochs = getattr(config, "train_epochs", 3) train_epochs_per_task = getattr(config, "train_epochs_per_task", {}) or { t: train_epochs for t in tasks } train_batch_size = getattr(config, "train_batch_size", 32) train_warmup_steps = getattr(config, "train_warmup_steps", 100) train_weight_decay = getattr(config, "train_weight_decay", 0.01) train_grad_accum_steps = 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) rows = [] # save original model state try: original_state = unwrap_model(model).state_dict() except Exception: original_state = model.state_dict() def _reset_model_to_original(): try: unwrap_model(model).load_state_dict(original_state, strict=False) except Exception: model.load_state_dict(original_state, strict=False) for task in tasks: print(f"\n==== Grid-running task: {task} ====") epochs_this_task = int(train_epochs_per_task.get(task, train_epochs)) # load dataset if task in EXTRA_TASKS: cfg = EXTRA_TASKS[task] if "hf_config" in cfg: ds = load_dataset(cfg["hf_path"], cfg["hf_config"]) else: ds = load_dataset(cfg["hf_path"]) train_raw = ds.get("train") val_raw = ds.get("validation") or ds.get("test") 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, "lr": None, "restart": None, "seed": None, "best_ckpt_path": None} small_flag = (task in SMALL_TASKS_RANDOM_RESTARTS) or (task in SMALL_TASKS_RANDOM_RESTARTS_EXTRA) restarts_per_lr = random_restarts_small if small_flag else 1 if auto_train: for lr in lr_candidates: for restart_idx in range(restarts_per_lr): seed = ( base_seed + (abs(hash(task)) % 10000) * 1000 + restart_idx * 10 + (int(round(lr * 1e7)) % 1000) ) print(f"[SWEEP] task={task} lr={lr:g} restart={restart_idx} seed={seed}") _reset_model_to_original() run_dir = task_ckpt_root / f"lr_{lr:g}" / f"restart_{restart_idx}" run_dir.mkdir(parents=True, exist_ok=True) if checkpointing is not None: try: checkpointing.load_model_states("recent") except Exception: pass if task in EXTRA_TASKS: det = detect_num_choices(val_raw, task) expect_num = det.get("detected_num_choices") 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), expect_num_choices=expect_num, ) sc = _metric_to_scalar(task, metric_res) best_epoch = None else: model, metric_res, sc, 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=seed, ) rec = { "task": task, "lr": float(lr), "restart": int(restart_idx), "seed": int(seed), "dev_score": float(sc) if sc is not None else None, "run_dir": str(run_dir), "best_ckpt_path": str(run_dir / "best_finetuned.pt"), } all_run_records.append(rec) if best_run["score"] is None or (sc is not None and sc > best_run["score"]): best_run.update( { "score": float(sc), "lr": float(lr), "restart": int(restart_idx), "seed": int(seed), "best_ckpt_path": rec["best_ckpt_path"], } ) pd.DataFrame(all_run_records).to_csv(task_ckpt_root / "all_runs.csv", index=False) best_overall_dir = task_ckpt_root / "best_overall" best_overall_dir.mkdir(parents=True, exist_ok=True) if best_run["best_ckpt_path"] and os.path.exists(best_run["best_ckpt_path"]): shutil.copy2(best_run["best_ckpt_path"], best_overall_dir / "best_finetuned.pt") _json_dump(best_run, best_overall_dir / "best_meta.json") _reset_model_to_original() model_for_eval = _load_finetuned_checkpoint_for_task( task, model, str(best_overall_dir / "best_finetuned.pt") ) else: model_for_eval = model task_out_dir = out_dir / task task_out_dir.mkdir(parents=True, exist_ok=True) if task in EXTRA_TASKS: det = detect_num_choices(val_raw, task) expect_num = det.get("detected_num_choices") metric_res = run_extra_task( task=task, tokenizer=tokenizer, model=model_for_eval, device=device, batch_size=batch_size, max_length=max_length, output_dir=str(task_out_dir), expect_num_choices=expect_num, ) rows.append({"task": task, "split": "validation", "metrics": json.dumps(metric_res)}) else: res = run_glue_task( task=task, tokenizer=tokenizer, model=model_for_eval, 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(): rows.append( { "task": task, "split": split, "metrics": json.dumps(pack["metrics"]), "errorbar": json.dumps(pack["errorbar"]), } ) summary_csv = out_dir / "glue_summary.csv" pd.DataFrame(rows).to_csv(summary_csv, index=False) print(f"[Grid] Summary saved to: {summary_csv}") return pd.DataFrame(rows) # --------------------------------------------------------------------- # CLI # --------------------------------------------------------------------- if __name__ == "__main__": print( "This module is intended to be imported and used by providing tokenizer/model objects.\n" "Example:\n" " tok = AutoTokenizer.from_pretrained('bert-base-uncased')\n" " model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')\n" " run_glue_benchmark(cfg, tok, model)\n" )