src/lmr/glue_benchmark/benchmark_with_mc.py

# """
# Extended GLUE + extra MC/QA/RC benchmark runner with grid-search (LR sweep + restarts).

# Features:

# BERT-style LR sweep for GLUE tasks (and extended tasks)

# Small-task multiple random restarts (seeded) for robustness

# Full fine-tune training & evaluation

# 5-fold leave-one-fold-out error bars for GLUE eval

# Support for EXTRA_TASKS: boolq, piqa, winogrande, openbookqa, hellaswag,
# arc_challenge (ARC-Challenge), arc_easy (ARC-Easy), race_middle, race_high

# Robust tokenization helpers that support many tokenizer APIs

# Multiple-choice support for variable numbers of choices (2, 4, ...)

# Usage:
# Integrate with your main script which provides tokenizer/model/checkpointing.
# Example:
# run_glue_benchmark(config.benchmark, tokenizer, model, checkpointing, out_dir="glue_outputs")

# Notes:

# This file aims to be comprehensive. It's long by design.

# It expects datasets, transformers, evaluate, torch, numpy, pandas, tqdm to be installed.
# """

# 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)

# from lmr.checkpointing import Checkpointing
# from lmr.ddp import unwrap_model

# ---------------------------------------------------------------------
# GLUE config + EXTRA tasks
# ---------------------------------------------------------------------

# 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"},
# }

# Preferred metric key per task (for selecting best run / best epoch)

# PREFERRED_METRIC_KEY = {
# "cola": "matthews_correlation",
# "sst2": "accuracy",
# "mrpc": "accuracy",
# "stsb": "pearson",
# "qqp": "accuracy",
# "mnli": "accuracy",
# "qnli": "accuracy",
# "rte": "accuracy",
# "wnli": "accuracy",
# }

# SMALL_TASKS_RANDOM_RESTARTS = {"cola", "mrpc", "rte", "stsb"}

# BERT_LR_CANDIDATES = [2e-5, 3e-5, 4e-5, 5e-5]

# Extra tasks (multiple-choice and pairwise)

# 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 QA / MC / RC datasets
# "openbookqa": {
# "type": "multiple_choice",
# "num_labels": 4,
# "hf_path": "allenai/openbookqa",
# "format": "mc",
# },
# "hellaswag": {
# "type": "multiple_choice",
# "num_labels": 4,
# "hf_path": "hellaswag",
# "format": "mc",
# },
# # ARC: ai2_arc has configs "ARC-Challenge" and "ARC-Easy"
# "arc_challenge": {
# "type": "multiple_choice",
# "num_labels": 4,
# "hf_path": "allenai/ai2_arc",
# "hf_config": "ARC-Challenge",
# "format": "mc",
# },
# "arc_easy": {
# "type": "multiple_choice",
# "num_labels": 4,
# "hf_path": "allenai/ai2_arc",
# "hf_config": "ARC-Easy",
# "format": "mc",
# },
# # RACE reading comprehension
# "race_middle": {
# "type": "multiple_choice",
# "num_labels": 4,
# "hf_path": "race",
# "hf_config": "middle",
# "format": "mc",
# },
# "race_high": {
# "type": "multiple_choice",
# "num_labels": 4,
# "hf_path": "race",
# "hf_config": "high",
# "format": "mc",
# },
# }

# ALL_TASKS = {**GLUE_TASKS, **EXTRA_TASKS}

# ---------------------------------------------------------------------
# Repro helpers
# ---------------------------------------------------------------------

# def _set_all_seeds(seed: int):
# random.seed(seed)
# np.random.seed(seed)
# torch.manual_seed(seed)
# torch.cuda.manual_seed_all(seed)
# torch.backends.cudnn.deterministic = True
# torch.backends.cudnn.benchmark = False

# 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)
# # fallback: first numeric entry
# 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 example extraction
# ---------------------------------------------------------------------

# 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) where s2 can be None for single-sentence tasks.
# """
# # Known per-task fields
# 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)
# ---------------------------------------------------------------------

# 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):
# 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))

# 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

# 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

# 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 "<empty>"
#         raise RuntimeError(f"Tokenizer fallback encode failed for example '{snippet}': {e}")

# return {"input_ids": input_ids_list, "attention_mask": attention_mask_list}

# ---------------------------------------------------------------------
# Postprocess preds
# ---------------------------------------------------------------------

# def _postprocess_predictions(task: str, logits_np: np.ndarray, cfg_task: dict):
# ttype = cfg_task["type"]

# if logits_np is None or logits_np.size == 0:
#     return np.array([])

# if logits_np.ndim == 1:
#     if ttype == "classification":
#         return (logits_np > 0.5).astype(int)
#     return logits_np.astype(float)

# if logits_np.ndim == 2 and logits_np.shape[1] == 1:
#     col = logits_np[:, 0]
#     if ttype == "classification":
#         return (col > 0.5).astype(int)
#     return col.astype(float)

# if logits_np.ndim == 2:
#     if ttype == "classification":
#         return np.argmax(logits_np, axis=-1).astype(int)
#     preds = logits_np[:, 0].astype(float) if logits_np.shape[1] == 1 else 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):
# import torch.nn as nn

# base_model = model

# def _detect_head_dim(m):
#     try:
#         if hasattr(m, "classifier") and isinstance(getattr(m, "classifier"), nn.Linear):
#             return getattr(m, "classifier").out_features
#         if hasattr(m, "lm_head") and isinstance(getattr(m, "lm_head"), nn.Linear):
#             return getattr(m, "lm_head").out_features
#         if hasattr(m, "get_output_embeddings"):
#             out_emb = m.get_output_embeddings()
#             if out_emb is not None:
#                 if isinstance(out_emb, nn.Embedding):
#                     return out_emb.embedding_dim if hasattr(out_emb, "embedding_dim") else out_emb.num_embeddings
#                 if isinstance(out_emb, nn.Linear):
#                     return out_emb.out_features
#     except Exception:
#         pass
#     return None

# if force_num_labels is None:
#     head_dim = _detect_head_dim(base_model)
#     if head_dim is not None and head_dim == num_labels:
#         return base_model, False

# inferred_hidden = hidden_size
# if inferred_hidden is None:
#     try:
#         cand = getattr(base_model, "config", None)
#         if cand is not None and hasattr(cand, "hidden_size"):
#             inferred_hidden = int(cand.hidden_size)
#     except Exception:
#         inferred_hidden = None

# if inferred_hidden is None:
#     try:
#         un = unwrap_model(base_model)
#         sd = un.state_dict()
#         for k, v in sd.items():
#             if re.search(r"embed|embedding|word_embeddings|token_embedding|embed_tokens", k, re.I):
#                 if hasattr(v, "shape") and len(v.shape) == 2:
#                     inferred_hidden = int(v.shape[1])
#                     break
#             if re.search(r"q_proj|k_proj|v_proj|o_proj|dense|fc|linear|proj", k, re.I):
#                 if hasattr(v, "shape") and len(v.shape) == 2:
#                     cand = max(v.shape)
#                     if 1 < cand < 1_000_000:
#                         inferred_hidden = int(cand)
#                         break
#     except Exception:
#         inferred_hidden = None

# if inferred_hidden is None:
#     raise RuntimeError(
#         "Cannot infer hidden_size for wrapped classifier head. "
#         "Please set `model.config.hidden_size` or pass `hidden_size` explicitly."
#     )

# class _WrappedModel(nn.Module):
#     def __init__(self, base, hidden_size, num_labels):
#         super().__init__()
#         self.base = base
#         self.classifier = nn.Linear(hidden_size, num_labels)
#         self.logits_projector = None

#     def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
#         try:
#             out = self.base(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
#         except TypeError:
#             out = self.base(input_ids)

#         last_hidden = getattr(out, "last_hidden_state", None)
#         if last_hidden is not None:
#             pooled = last_hidden[:, 0, :]
#             logits = self.classifier(pooled)
#             return type("Out", (), {"logits": logits, "loss": None})

#         if isinstance(out, (tuple, list)) and len(out) > 0:


# ChatGPT can make mistakes. OpenAI doesn't use Duke University workspace data to train its models.