# # 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({"openbookqa", "arc_easy", "arc_challenge","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（替换版：用替换 classifier 输出 1 的方法） # # --------------------------- # 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 replacing the model's classification head to output 1 scalar: # - Replace model.classifier / model.lm_head / model.score / similar with Linear(hidden, 1) # - For MC tasks: flatten (B, C, L) -> (B*C, L), forward, obtain (B*C,1) -> reshape (B, C) -> CE loss # - Ensures new head parameters are created before optimizer so they are optimized. # This variant surfaces errors (no silent swallowing). # """ # 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" # # --------------------------- # # Tokenize splits # # --------------------------- # if is_mc: # 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 # ) # 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: # 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 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) # # --------------------------- # # Replace model head to output scalar per example (for MC we'll apply it to each choice) # # --------------------------- # def _infer_hidden_size_from_model(m): # # try config.hidden_size first # hidden = None # try: # if hasattr(m, "config") and hasattr(m.config, "hidden_size"): # hidden = int(m.config.hidden_size) # except Exception: # hidden = None # if hidden is not None: # return hidden # # try to inspect state_dict looking for embedding/proj dimensions # try: # un = unwrap_model(m) # 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: # return int(v.shape[1]) # 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: # return int(cand) # except Exception: # pass # return None # # Only replace head for MC and for pair tasks we keep existing pipeline (but we still may want to ensure classifier shape) # if is_mc: # # detect suitable hidden size # hidden_size = _infer_hidden_size_from_model(model) # if hidden_size is None: # raise RuntimeError("Cannot infer hidden size from model; please set model.config.hidden_size or pass hidden_size explicitly.") # # Candidate head attribute names to replace # head_attrs = ["classifier", "lm_head", "score", "classifier_head", "head"] # replaced = False # for attr in head_attrs: # if hasattr(model, attr): # try: # old = getattr(model, attr) # # create new linear head that maps hidden -> 1 # new_head = nn.Linear(hidden_size, 1).to(device) # # init reasonably # try: # nn.init.xavier_uniform_(new_head.weight) # if new_head.bias is not None: # nn.init.zeros_(new_head.bias) # except Exception: # pass # setattr(model, attr, new_head) # print(f"[ExtraTrain][{task}] Replaced model.{attr} with Linear({hidden_size},1)") # replaced = True # break # except Exception: # # ignore and try next # pass # if not replaced: # # fallback: attach as attribute _mc_choice_head # model._mc_choice_head = nn.Linear(hidden_size, 1).to(device) # try: # nn.init.xavier_uniform_(model._mc_choice_head.weight) # if model._mc_choice_head.bias is not None: # nn.init.zeros_(model._mc_choice_head.bias) # except Exception: # pass # print(f"[ExtraTrain][{task}] Attached model._mc_choice_head = Linear({hidden_size},1)") # # --------------------------- # # Create optimizer AFTER head replacement so head params 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) # ce_loss_fn = torch.nn.CrossEntropyLoss() # mse_loss_fn = torch.nn.MSELoss() # model.train() # global_step = 0 # final_metric_res = {} # expected_num_labels = cfg["num_labels"] # # training loop # 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) # 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 (B,C,L) -> (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 # 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: # # Some models accept single positional tensor # out_flat = model(flat_ids) # # get a pooled representation or logits from out_flat # flat_logits = getattr(out_flat, "logits", None) # pooled = None # if flat_logits is None: # # try pooler_output or last_hidden_state # pooled = getattr(out_flat, "pooler_output", None) # if pooled is None: # lh = getattr(out_flat, "last_hidden_state", None) # if lh is not None: # pooled = lh[:, 0, :] # else: # # If model already returns logits of shape (B*C, D) or (B*C,1), we can still feed them to head or reduce # pass # # compute scalar per flat example using the replaced head # # Case 1: model already returned logits as scalar (B*C,1) or (B*C,) # scalar_flat = None # if flat_logits is not None: # # flat_logits could be (B*C,) or (B*C,1) or (B*C,D) # if flat_logits.ndim == 1: # scalar_flat = flat_logits.view(-1, 1) # elif flat_logits.ndim == 2 and flat_logits.shape[1] == 1: # scalar_flat = flat_logits.view(-1, 1) # elif flat_logits.ndim == 2: # # project via the replaced head (we expect a head that maps hidden->1, so treat flat_logits as features) # # If we replaced model.classifier, it's likely expecting hidden vectors, but we have flat_logits features: # # attempt to use model._mc_choice_head if present; else take mean over dims then linear # feat = flat_logits # if hasattr(model, "_mc_choice_head"): # # ensure dimensions align: if feat dim matches head.in_features, use it # try: # in_f = model._mc_choice_head.in_features # if feat.shape[1] == in_f: # scalar_flat = model._mc_choice_head(feat) # else: # # reduce feat to hidden via mean # reduced = feat.mean(dim=1) # scalar_flat = model._mc_choice_head(reduced.unsqueeze(1) if len(reduced.shape)==1 else reduced) # except Exception: # reduced = feat.mean(dim=1) # scalar_flat = model._mc_choice_head(reduced.unsqueeze(1) if len(reduced.shape)==1 else reduced) # else: # # fallback: mean over feature dims to scalar # scalar_flat = feat.mean(dim=1, keepdim=True) # elif flat_logits.ndim == 3: # # reduce seq dim and feature dim to scalar # scalar_flat = flat_logits.mean(dim=tuple(range(1, flat_logits.ndim))).view(-1, 1) # else: # scalar_flat = flat_logits.reshape(flat_logits.size(0), -1).mean(dim=1, keepdim=True) # elif pooled is not None: # # pooled is (B*C, hidden), apply the replaced head # # find which head attribute we replaced # head_found = None # for attr in ("classifier", "lm_head", "score", "classifier_head", "head", "_mc_choice_head"): # if hasattr(model, attr): # head_found = getattr(model, attr) # break # if head_found is None: # raise RuntimeError("[ExtraTrain] no replacement head found on model to project pooled -> scalar") # # head_found may expect input dim hidden_size # try: # scalar_flat = head_found(pooled) # except Exception as e: # # if mismatch, try a linear mapping by reducing pooled # if pooled.ndim == 2: # # If head expects 1-d input, ensure shape matches # try: # scalar_flat = head_found(pooled) # except Exception: # scalar_flat = pooled.mean(dim=1, keepdim=True) # else: # scalar_flat = pooled.mean(dim=1, keepdim=True) # else: # raise RuntimeError("[ExtraTrain] cannot obtain features or logits from model forward to compute scalar per choice") # # Now scalar_flat should be (B*C, 1) -> reshape to (B, C) # if scalar_flat is None: # raise RuntimeError("[ExtraTrain] scalar_flat is None after processing model output") # if scalar_flat.ndim == 1: # scalar_flat = scalar_flat.view(-1, 1) # choice_scores = scalar_flat.view(Bf, Cf) # # If dataset's number of choices (Cf) differs from expected_num_labels, project to expected # if Cf != expected_num_labels: # # create a mapping layer if not exist (choice_count -> expected_num_labels) # if not hasattr(model, "_mc_choice_to_label"): # model._mc_choice_to_label = nn.Linear(Cf, expected_num_labels).to(device) # # register params in optimizer # try: # optimizer.add_param_group({'params': model._mc_choice_to_label.parameters()}) # except Exception: # pass # print(f"[ExtraTrain][{task}] created _mc_choice_to_label: {Cf}->{expected_num_labels}") # 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] # model._mc_projector = nn.Linear(exist_dim, expected_num_labels).to(device) # try: # optimizer.add_param_group({'params': model._mc_projector.parameters()}) # except Exception: # pass # used_logits = model._mc_projector(logits) # else: # used_logits = logits # if cfg["type"] == "classification": # loss = ce_loss_fn(used_logits, labs_b.long()) # else: # 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 # 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 # 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 # 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) # pooled = None # if flat_logits is None: # pooled = getattr(out_flat, "pooler_output", None) # if pooled is None: # lh = getattr(out_flat, "last_hidden_state", None) # if lh is not None: # pooled = lh[:, 0, :] # scalar_flat = None # if flat_logits is not None: # if flat_logits.ndim == 1: # scalar_flat = flat_logits.view(-1, 1) # elif flat_logits.ndim == 2 and flat_logits.shape[1] == 1: # scalar_flat = flat_logits.view(-1, 1) # elif flat_logits.ndim == 2: # if hasattr(model, "_mc_choice_head"): # try: # in_f = model._mc_choice_head.in_features # if flat_logits.shape[1] == in_f: # scalar_flat = model._mc_choice_head(flat_logits) # else: # scalar_flat = flat_logits.mean(dim=1, keepdim=True) # except Exception: # scalar_flat = flat_logits.mean(dim=1, keepdim=True) # else: # scalar_flat = flat_logits.mean(dim=1, keepdim=True) # elif flat_logits.ndim == 3: # scalar_flat = flat_logits.mean(dim=tuple(range(1, flat_logits.ndim))).view(-1, 1) # else: # scalar_flat = flat_logits.reshape(flat_logits.size(0), -1).mean(dim=1, keepdim=True) # elif pooled is not None: # head_found = None # for attr in ("classifier", "lm_head", "score", "classifier_head", "head", "_mc_choice_head"): # if hasattr(model, attr): # head_found = getattr(model, attr) # break # if head_found is None: # raise RuntimeError("[ExtraVal] no replacement head found on model to project pooled -> scalar") # try: # scalar_flat = head_found(pooled) # except Exception: # scalar_flat = pooled.mean(dim=1, keepdim=True) # else: # raise RuntimeError("Model did not return logits/poolable outputs during MC validation.") # choice_scores = scalar_flat.view(Bf, Cf) # if choice_scores.shape[1] != expected_num_labels: # exist_dim = choice_scores.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()}) # except Exception: # pass # used_logits = model._mc_choice_to_label(choice_scores) # else: # used_logits = choice_scores # 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()}) # except Exception: # pass # 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 = ['hellaswag','openbookqa','arc_easy','arc_challenge'] # # tasks = tasks[::-1] # # ['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(): # pass # # 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") # checkpointing.load_model_states("recent") # 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" # )