src/lmr/glue_benchmark/glue_benchmark copy 6.py

# lmr/glue_benchmark_grid_full.py
"""
Grid-search GLUE + extra tasks runner (full single-file).

Features:
- LR sweep across bert_lr_candidates
- Random restarts for small/unstable tasks
- Save per-run checkpoints and run_meta.json
- Save all_runs.csv and best_overall per task
- Evaluate best model and compute errorbars for GLUE
- Supports EXTRA_TASKS (boolq/piqa/winogrande + hellaswag/openbookqa/arc) with MC/pair handling
"""

import os
import json
import re
import math
import random
import shutil
import time
from pathlib import Path
from typing import Optional, List, Tuple, Dict, Any

import torch
import numpy as np
import pandas as pd
from datasets import load_dataset
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm

import evaluate

# Project imports: adjust if your package layout differs
try:
    from lmr.checkpointing import Checkpointing
    from lmr.ddp import unwrap_model
except Exception:
    # If these modules are not available, provide lightweight fallbacks to avoid import errors
    Checkpointing = None

    def unwrap_model(m):
        return m

# ---------------------------------------------------------------------
# Tasks config
# ---------------------------------------------------------------------
GLUE_TASKS = {
    "cola": {"type": "classification", "num_labels": 2, "hf_name": "cola"},
    "sst2": {"type": "classification", "num_labels": 2, "hf_name": "sst2"},
    "mrpc": {"type": "classification", "num_labels": 2, "hf_name": "mrpc"},
    "stsb": {"type": "regression",     "num_labels": 1, "hf_name": "stsb"},
    "qqp": {"type": "classification", "num_labels": 2, "hf_name": "qqp"},
    "mnli": {"type": "classification", "num_labels": 3, "hf_name": "mnli"},
    "qnli": {"type": "classification", "num_labels": 2, "hf_name": "qnli"},
    "rte":  {"type": "classification", "num_labels": 2, "hf_name": "rte"},
    "wnli": {"type": "classification", "num_labels": 2, "hf_name": "wnli"},
}

# Extra tasks (BoolQ, PIQA, Winogrande, HellaSwag, OpenBookQA, ARC variants)
EXTRA_TASKS = {
    "boolq": {
        "type": "classification",
        "num_labels": 2,
        "hf_path": "boolq",
        "format": "pair",
    },
    "piqa": {
        "type": "multiple_choice",
        "num_labels": 2,
        "hf_path": "piqa",
        "format": "mc",
    },
    "winogrande": {
        "type": "multiple_choice",
        "num_labels": 2,
        "hf_path": "winogrande",
        "hf_config": "winogrande_xl",
        "format": "mc",
    },
    # Added tasks below
    "hellaswag": {
        "type": "multiple_choice",
        "num_labels": 4,
        "hf_path": "hellaswag",
        "format": "mc",
    },
    "openbookqa": {
        "type": "multiple_choice",
        "num_labels": 4,
        "hf_path": "openbookqa",
        "format": "mc",
    },
    # AI2 ARC splits — align names with common usage
    "arc_easy": {
        "type": "multiple_choice",
        "num_labels": 4,
        "hf_path": "ai2_arc",
        "hf_config": "ARC-Easy",
        "format": "mc",
    },
    "arc_challenge": {
        "type": "multiple_choice",
        "num_labels": 4,
        "hf_path": "ai2_arc",
        "hf_config": "ARC-Challenge",
        "format": "mc",
    },
}

ALL_TASKS = {**GLUE_TASKS, **EXTRA_TASKS}

# Small/unstable tasks for extra random restarts
SMALL_TASKS_RANDOM_RESTARTS = {"cola", "mrpc", "rte", "stsb"}
SMALL_TASKS_RANDOM_RESTARTS_EXTRA = set({"piqa", "boolq", "winogrande", "hellaswag"})  # adjust as desired

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

PREFERRED_METRIC_KEY = {
    "cola": "matthews_correlation",
    "sst2": "accuracy",
    "mrpc": "accuracy",
    "stsb": "pearson",
    "qqp": "accuracy",
    "mnli": "accuracy",
    "qnli": "accuracy",
    "rte": "accuracy",
    "wnli": "accuracy",
    "boolq": "accuracy",
    "piqa": "accuracy",
    "winogrande": "accuracy",
    "hellaswag": "accuracy",
    "openbookqa": "accuracy",
    "arc_easy": "accuracy",
    "arc_challenge": "accuracy",
}

# ---------------------------------------------------------------------
# Repro helpers
# ---------------------------------------------------------------------
def _set_all_seeds(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    try:
        torch.cuda.manual_seed_all(seed)
    except Exception:
        pass
    try:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    except Exception:
        pass

def _json_dump(obj: Any, path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, indent=2, ensure_ascii=False)

def _safe_float(x):
    try:
        if isinstance(x, (np.generic,)):
            return float(x.item())
        return float(x)
    except Exception:
        return None

def _metric_to_scalar(task: str, metric_res: Dict[str, Any], fallback_val_loss: Optional[float] = None) -> float:
    if isinstance(metric_res, dict) and metric_res:
        pref = PREFERRED_METRIC_KEY.get(task)
        if pref is not None and pref in metric_res:
            v = _safe_float(metric_res.get(pref))
            if v is not None and not math.isnan(v):
                return float(v)
        for _, v in metric_res.items():
            fv = _safe_float(v)
            if fv is not None and not math.isnan(fv):
                return float(fv)
    if fallback_val_loss is not None:
        try:
            return -float(fallback_val_loss)
        except Exception:
            pass
    return -1e9

# ---------------------------------------------------------------------
# Task-aware example field extraction (robust)
# ---------------------------------------------------------------------
def _get_text_pair_from_example(task: str, ex: dict):
    """
    Robustly extract (s1, s2) from a HF GLUE/example dict `ex` depending on task.
    Returns (s1:str, s2:Optional[str]) where s2 can be None for single-sentence tasks.
    """
    task_field_map = {
        "cola": ("sentence", None),
        "sst2": ("sentence", None),
        "mrpc": ("sentence1", "sentence2"),
        "stsb": ("sentence1", "sentence2"),
        "qqp": ("question1", "question2"),
        "mnli": ("premise", "hypothesis"),
        "qnli": ("question", "sentence"),
        "rte": ("sentence1", "sentence2"),
        "wnli": ("sentence1", "sentence2"),
    }
    f1, f2 = task_field_map.get(task, (None, None))

    def _try_keys(keys):
        for k in keys:
            if k in ex and ex.get(k) is not None:
                return ex.get(k)
        return None

    s1_candidates = []
    s2_candidates = []

    if f1:
        s1_candidates.append(f1)
    s1_candidates += ["sentence1", "premise", "question", "sentence", "text", "question1"]

    if f2:
        s2_candidates.append(f2)
    s2_candidates += ["sentence2", "hypothesis", "question2", "question1", "text2"]

    s1 = _try_keys(s1_candidates)
    s2 = _try_keys(s2_candidates)

    if s1 is None:
        s1 = ex.get("sentence") or ex.get("premise") or ex.get("question") or ex.get("text")
    if s2 is None:
        s2 = ex.get("sentence2") or ex.get("hypothesis") or ex.get("question2")

    s1 = "" if s1 is None else (s1 if isinstance(s1, str) else str(s1))
    s2 = None if s2 is None else (s2 if isinstance(s2, str) else str(s2))
    return s1, s2

# ---------------------------------------------------------------------
# Tokenization helpers (robust to different tokenizer APIs)
# ---------------------------------------------------------------------
def _pad_and_tensorize(input_ids_list, attention_mask_list, pad_token_id: int):
    max_len = max(len(x) for x in input_ids_list) if input_ids_list else 0
    ids_padded = [x + [pad_token_id] * (max_len - len(x)) for x in input_ids_list]
    mask_padded = [m + [0] * (max_len - len(m)) for m in attention_mask_list]
    input_ids = torch.tensor(ids_padded, dtype=torch.long)
    attention_mask = torch.tensor(mask_padded, dtype=torch.long)
    return input_ids, attention_mask

def _batch_tokenize(tokenizer, texts: List[Tuple[Optional[str], Optional[str]]], max_length: int = 128):
    """
    Robust batch tokenization for a variety of tokenizer APIs.
    - texts: list of (s1, s2) where s2 may be None.
    - Try HF tokenizer(...) first, then various batch methods, then per-example fallback.
    Returns dict with 'input_ids' (list of lists) and 'attention_mask'.
    """
    sanitized = []
    for a, b in texts:
        a_s = "" if a is None else (a if isinstance(a, str) else str(a))
        b_s = None if b is None else (b if isinstance(b, str) else str(b))
        sanitized.append((a_s, b_s))

    # 1) Try HF-like tokenizer(...) first
    try:
        flat = [(a if b is None else (a, b)) for a, b in sanitized]
        enc = tokenizer(flat, truncation=True, padding=False, max_length=max_length)
        if isinstance(enc.get("input_ids", None), torch.Tensor):
            enc["input_ids"] = enc["input_ids"].tolist()
        if isinstance(enc.get("attention_mask", None), torch.Tensor):
            enc["attention_mask"] = enc["attention_mask"].tolist()
        return enc
    except Exception:
        pass

    # 2) Try other batch-like methods
    for method_name in ("batch_encode", "encode_batch", "batch_encode_plus", "encode_batch_pair", "encode_batch_items"):
        fn = getattr(tokenizer, method_name, None)
        if fn is None:
            continue
        try:
            try:
                enc = fn(sanitized, max_length=max_length, truncation=True, padding=False)
            except TypeError:
                enc = fn(sanitized)
            if isinstance(enc.get("input_ids", None), torch.Tensor):
                enc["input_ids"] = enc["input_ids"].tolist()
            if isinstance(enc.get("attention_mask", None), torch.Tensor):
                enc["attention_mask"] = enc["attention_mask"].tolist()
            return enc
        except Exception:
            continue

    # 3) Fallback per-example
    input_ids_list = []
    attention_mask_list = []
    for a, b in sanitized:
        try:
            if b is None:
                try:
                    single = tokenizer.encode(a)
                except TypeError:
                    single = tokenizer.encode([a])
            else:
                single = None
                try:
                    single = tokenizer.encode((a, b))
                except Exception:
                    try:
                        single = tokenizer.encode(a, b)
                    except Exception:
                        single = tokenizer(a if b is None else (a, b))

            if isinstance(single, dict):
                ids = single.get("input_ids") or single.get("ids") or []
                mask = single.get("attention_mask") or single.get("mask") or [1] * len(ids)
            elif isinstance(single, torch.Tensor):
                ids = single.tolist()
                mask = [1] * len(ids)
            elif isinstance(single, list):
                ids = single
                mask = [1] * len(ids)
            else:
                tmp = tokenizer(a if b is None else (a, b))
                if isinstance(tmp, dict):
                    ids = tmp.get("input_ids") or tmp.get("ids") or []
                    mask = tmp.get("attention_mask") or tmp.get("mask") or [1] * len(ids)
                elif torch.is_tensor(tmp):
                    ids = tmp.tolist()
                    mask = [1] * len(ids)
                else:
                    ids = list(tmp)
                    mask = [1] * len(ids)

            if len(ids) > max_length:
                ids = ids[:max_length]
                mask = mask[:max_length]

            input_ids_list.append(ids)
            attention_mask_list.append(mask)
        except Exception as e:
            snippet = (a[:80] + "...") if a else "<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 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

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
    import pdb
    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


# ---------------------------
# 新增：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
    # import pdb
    # pdb.set_trace()
    for ex in raw_split:
        total += 1
        ctx, opts, lab = _extract_mc_example(task, ex)
        if not opts:
            skipped_no_opts += 1
            continue
        if expect_num_choices is not None and len(opts) != expect_num_choices:
            skipped_inconsistent += 1
            continue
        if lab is None:
            skipped_bad_label += 1
            continue
        try:
            lab_i = int(lab)
        except Exception:
            skipped_bad_label += 1
            continue
        if lab_i < 0 or lab_i >= len(opts):
            skipped_bad_label += 1
            continue

        contexts.append(ctx if ctx is not None else "")
        options.append(opts)
        labels.append(lab_i)
        kept += 1

    print(f"[tokenize_generic_mc][{task}] total={total} kept={kept} skipped_no_opts={skipped_no_opts} skipped_inconsistent={skipped_inconsistent} skipped_bad_label={skipped_bad_label}")

    if len(contexts) == 0:
        return torch.zeros((0, 1, 1), dtype=torch.long), torch.zeros((0, 1, 1), dtype=torch.long), torch.tensor([], dtype=torch.long)

    pad_id = getattr(tokenizer, "pad_token_id", None)
    if pad_id is None:
        try:
            pad_id = tokenizer.token_to_id("[PAD]")
        except Exception:
            pad_id = 0

    num_choices = len(options[0])

    # Build flat pairs for batch tokenization and assert tokenizer returns expected length
    input_ids_rows = []
    attention_rows = []

    for i in range(0, len(contexts), batch_tokenize_size):
        chunk_ctx = contexts[i:i+batch_tokenize_size]
        chunk_opts = options[i:i+batch_tokenize_size]
        flat_pairs = []
        for c, opts in zip(chunk_ctx, chunk_opts):
            for o in opts:
                flat_pairs.append((c, o))

        enc = tokenizer(flat_pairs, truncation=True, padding=False, max_length=max_length)
        ids_flat = enc.get("input_ids")
        masks_flat = enc.get("attention_mask") or enc.get("mask") or enc.get("masks")

        if isinstance(ids_flat, torch.Tensor):
            ids_flat = ids_flat.tolist()
        if isinstance(masks_flat, torch.Tensor):
            masks_flat = masks_flat.tolist()

        expected = len(chunk_ctx) * num_choices
        if len(ids_flat) != expected:
            # raise explicit error so user sees where tokenizer mismatched
            raise RuntimeError(f"[tokenize_generic_mc][{task}] tokenizer returned {len(ids_flat)} items but expected {expected} (i={i}). "
                               f"First flat_pairs sample: {flat_pairs[0] if len(flat_pairs)>0 else None}")

        per_example = []
        per_mask_example = []
        idx = 0
        for _ in range(len(chunk_ctx)):
            row = []
            row_mask = []
            for _ in range(num_choices):
                row.append(ids_flat[idx])
                row_mask.append(masks_flat[idx])
                idx += 1
            per_example.append(row)
            per_mask_example.append(row_mask)

        input_ids_rows.extend(per_example)
        attention_rows.extend(per_mask_example)

    max_len = max(len(seq) for row in input_ids_rows for seq in row) if input_ids_rows else 1
    input_ids_padded = [
        [ seq + [pad_id] * (max_len - len(seq)) for seq in row ]
        for row in input_ids_rows
    ]
    attention_padded = [
        [ mask + [0] * (max_len - len(mask)) for mask in row ]
        for row in attention_rows
    ]

    input_ids_t = torch.tensor(input_ids_padded, dtype=torch.long)  # (N, C, L)
    attention_t = torch.tensor(attention_padded, dtype=torch.long)
    labels_t = torch.tensor(labels, dtype=torch.long)
    return input_ids_t, attention_t, labels_t

# ---------------------------
# train_full_finetune_extra（去除对关键异常的吞掉，遇到异常直接抛出）
# ---------------------------
def train_full_finetune_extra(task: str, tokenizer, model, raw_train, raw_val,
                              device: str = "cuda", epochs: int = 3, batch_size: int = 16,
                              lr: float = 2e-5, weight_decay: float = 0.01, warmup_steps: int = 100,
                              max_length: int = 128, grad_accum_steps: int = 1, out_checkpoint_dir: Optional[str] = None,
                              expect_num_choices: Optional[int] = None):
    """
    Fine-tune for EXTRA_TASKS with a safe projection fallback.
    This variant does not silence errors: tokenizer/model forward issues will raise and surface to the caller.

    For multiple-choice tasks we flatten (B, C, L) -> (B*C, L) before feeding into the transformer,
    then reshape logits back to (B, C) and compute CrossEntropyLoss against labels (B,).
    """
    import torch.nn as nn
    cfg = EXTRA_TASKS[task]
    device = torch.device(device if torch.cuda.is_available() else "cpu")
    model.to(device)

    is_mc = cfg["format"] == "mc"
    if is_mc:
        # Tokenize to (N, C, L)
        train_ids, train_mask, train_labels = _tokenize_generic_mc_split_to_tensors(task, tokenizer, raw_train, max_length=max_length, batch_tokenize_size=256, expect_num_choices=expect_num_choices)
        val_ids, val_mask, val_labels = _tokenize_generic_mc_split_to_tensors(task, tokenizer, raw_val, max_length=max_length, batch_tokenize_size=256, expect_num_choices=expect_num_choices)

        # Basic checks
        if train_ids.ndim != 3:
            raise RuntimeError(f"[ExtraTrain][{task}] train_ids expected 3 dims (N,C,L) got {train_ids.ndim}")
        if val_ids.ndim != 3 and len(val_ids) > 0:
            raise RuntimeError(f"[ExtraTrain][{task}] val_ids expected 3 dims (N,C,L) got {val_ids.ndim}")

        train_ds = TensorDataset(train_ids, train_mask, train_labels)
        val_ds = TensorDataset(val_ids, val_mask, val_labels)
    else:
        # pair tasks: reuse existing pair tokenizer (expect to be defined)
        train_ids, train_mask, train_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_train, cfg, max_length=max_length)
        val_ids, val_mask, val_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_val, cfg, max_length=max_length)
        train_ds = TensorDataset(train_ids, train_mask, train_labels)
        val_ds = TensorDataset(val_ids, val_mask, val_labels)

    # If train_ds is empty, bail early with explicit error
    if len(train_ds) == 0:
        raise RuntimeError(f"No training samples after tokenization for task={task}; aborting.")

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, pin_memory=True)
    val_loader = DataLoader(val_ds, batch_size=max(64, batch_size), shuffle=False, pin_memory=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    total_steps = max(1, (len(train_loader) // max(1, grad_accum_steps)) * epochs)
    from transformers import get_cosine_schedule_with_warmup
    scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)

    # For MC tasks we will use CE on (B,C) vs labels (B,)
    ce_loss_fn = torch.nn.CrossEntropyLoss()
    mse_loss_fn = torch.nn.MSELoss()

    model.train()
    global_step = 0
    final_metric_res = {}

    expected_num_labels = cfg["num_labels"]

    for epoch in range(epochs):
        running_loss = 0.0
        for step, batch in enumerate(tqdm(train_loader, desc=f"[ExtraTrain] {task} epoch {epoch+1}")):
            if is_mc:
                ids_b, mask_b, labs_b = batch  # ids_b: (B, C, L)
                ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device)

                # Ensure labels shape matches batch
                if labs_b.dim() > 1:
                    labs_b = labs_b.view(-1)
                B = ids_b.size(0)
                if labs_b.numel() != B:
                    raise RuntimeError(f"[ExtraTrain][FATAL] label count ({labs_b.numel()}) != batch size ({B}) at step {step}; aborting.")

                # Flatten choices to (B*C, L)
                Bf, Cf, Lf = ids_b.shape
                flat_ids = ids_b.view(Bf * Cf, Lf).to(device)
                flat_mask = mask_b.view(Bf * Cf, Lf).to(device) if mask_b is not None else None

                # Forward pass (robust)
                out_flat = None
                try:
                    if flat_mask is not None:
                        out_flat = model(input_ids=flat_ids, attention_mask=flat_mask, labels=None)
                    else:
                        out_flat = model(input_ids=flat_ids, labels=None)
                except TypeError:
                    out_flat = model(flat_ids)

                flat_logits = getattr(out_flat, "logits", None)
                if flat_logits is None and isinstance(out_flat, (tuple, list)):
                    flat_logits = out_flat[0]
                if flat_logits is None:
                    raise RuntimeError("Model did not return logits for MC forward")

                # Normalize flat_logits -> choice_scores (B, C)
                if flat_logits.ndim == 1:
                    choice_scores = flat_logits.view(Bf, Cf)
                elif flat_logits.ndim == 2:
                    D = flat_logits.shape[1]
                    if D == 1:
                        choice_scores = flat_logits.view(Bf, Cf).squeeze(-1)
                    else:
                        # learned projector from D -> 1, reuse if present
                        if not hasattr(model, "_mc_projector") or model._mc_projector.weight.shape[1] != D:
                            model._mc_projector = nn.Linear(D, 1).to(device)
                        proj = model._mc_projector(flat_logits)  # (B*C, 1)
                        choice_scores = proj.view(Bf, Cf).squeeze(-1)
                elif flat_logits.ndim == 3:
                    # (B*C, S, D), reduce to scalar per choice
                    choice_scores = flat_logits.mean(dim=tuple(range(1, flat_logits.ndim))).view(Bf, Cf)
                else:
                    # fallback: flatten last dims then mean
                    choice_scores = flat_logits.reshape(flat_logits.size(0), -1).mean(dim=1).view(Bf, Cf)

                if choice_scores.ndim == 1:
                    choice_scores = choice_scores.unsqueeze(0)

                if choice_scores.shape != (Bf, Cf):
                    raise RuntimeError(f"Unexpected choice_scores shape {choice_scores.shape} at step {step}")

                # Project if number of classes doesn't match expected labels
                if choice_scores.ndim != 2:
                    raise RuntimeError(f"Unexpected logits ndim {choice_scores.ndim} from model at train step {step}")

                # ensure final logits align (B, num_labels) and labels (B,)
                if Cf != expected_num_labels:
                    # if dataset choices count differs, project to expected_num_labels
                    exist_dim = choice_scores.shape[1]
                    model._mc_choice_to_label = nn.Linear(exist_dim, expected_num_labels).to(device)
                    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)
                    used_logits = model._mc_projector(logits)
                else:
                    used_logits = logits
                if cfg["type"] == "classification":
                    loss = ce_loss_fn(used_logits, labs_b.long())
                else:
                    # regression
                    if used_logits.ndim == 2 and used_logits.shape[1] == 1:
                        preds = used_logits.squeeze(1)
                    elif used_logits.ndim == 2:
                        preds = used_logits.mean(dim=1)
                    else:
                        preds = used_logits
                    loss = mse_loss_fn(preds, labs_b.float())

            loss = loss / max(1, grad_accum_steps)
            loss.backward()

            if (step + 1) % max(1, grad_accum_steps) == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
                global_step += 1
            running_loss += float(loss.item()) * (ids_b.size(0) if is_mc else ids_b.size(0))

        # validation (no exception swallowing)
        model.eval()
        all_logits = []
        all_labels = []
        with torch.no_grad():
            for batch in tqdm(val_loader, desc=f"[ExtraVal] {task} epoch {epoch+1}", leave=False):
                if is_mc:
                    ids_b, mask_b, labs_b = batch
                    if labs_b.dim() > 1:
                        labs_b = labs_b.view(-1)
                    B = ids_b.size(0)
                    if labs_b.numel() != B:
                        raise RuntimeError(f"[ExtraVal][FATAL] val labels count ({labs_b.numel()}) != batch size ({B})")
                    ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device)

                    Bf, Cf, Lf = ids_b.shape
                    flat_ids = ids_b.view(Bf*Cf, Lf).to(device)
                    flat_mask = mask_b.view(Bf*Cf, Lf).to(device) if mask_b is not None else None
                    out_flat = None
                    try:
                        if flat_mask is not None:
                            out_flat = model(input_ids=flat_ids, attention_mask=flat_mask, labels=None)
                        else:
                            out_flat = model(input_ids=flat_ids, labels=None)
                    except TypeError:
                        out_flat = model(flat_ids)

                    flat_logits = getattr(out_flat, "logits", None)
                    if flat_logits is None and isinstance(out_flat, (tuple, list)):
                        flat_logits = out_flat[0]
                    if flat_logits is None:
                        raise RuntimeError("Model did not return logits during MC validation fallback")
                    if flat_logits.ndim == 2 and flat_logits.shape[1] == 1:
                        logits = flat_logits.view(Bf, Cf)
                    elif flat_logits.ndim == 2:
                        D = flat_logits.shape[1]
                        if D == 1:
                            logits = flat_logits.view(Bf, Cf).squeeze(-1)
                        else:
                            if not hasattr(model, "_mc_projector") or model._mc_projector.weight.shape[1] != D:
                                model._mc_projector = nn.Linear(D, 1).to(device)
                            proj = model._mc_projector(flat_logits)
                            logits = proj.view(Bf, Cf).squeeze(-1)
                    elif flat_logits.ndim == 1:
                        logits = flat_logits.view(Bf, Cf)
                    elif flat_logits.ndim == 3:
                        logits = flat_logits.mean(dim=tuple(range(1, flat_logits.ndim))).view(Bf, Cf)
                    else:
                        logits = flat_logits.reshape(flat_logits.size(0), -1).mean(dim=1).view(Bf, Cf)

                    if logits.shape[1] != expected_num_labels:
                        exist_dim = logits.shape[1]
                        model._mc_choice_to_label = nn.Linear(exist_dim, expected_num_labels).to(device)
                        used_logits = model._mc_choice_to_label(logits)
                    else:
                        used_logits = logits

                    all_logits.append(used_logits.detach().cpu().numpy())
                    all_labels.append(labs_b.detach().cpu().numpy())
                else:
                    ids_b, mask_b, labs_b = batch
                    ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device)
                    out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                    logits = getattr(out, "logits", None)
                    if logits is None and isinstance(out, (tuple, list)):
                        logits = out[0]
                    if logits.ndim == 2 and logits.shape[1] != expected_num_labels:
                        exist_dim = logits.shape[1]
                        model._mc_projector = nn.Linear(exist_dim, expected_num_labels).to(device)
                        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


# ---------------------------------------------------------------------
# Grid-runner: LR sweep + restarts + checkpointing + eval
# ---------------------------------------------------------------------
def run_glue_benchmark(
    config,
    tokenizer,
    model,
    checkpointing: Optional[Checkpointing] = None,
    out_dir: str = "glue_outputs_grid",
):
    tasks = getattr(config, "glue_tasks", None)
    if tasks is None:
        tasks = list(ALL_TASKS.keys())
    tasks = ['openbookqa','openbookqa','arc_easy','arc_challenge']  
    # tasks = ["piqa", "boolq", "winogrande"]
    #     "hellaswag": "accuracy",
    # "openbookqa": "accuracy",
    # "arc_easy": "accuracy",
    # "arc_challenge": "accuracy",
    batch_size = getattr(config, "batch_size", 64)
    max_length = getattr(config, "max_length", 128)
    device = getattr(config, "device", "cuda")
    auto_train = getattr(config, "auto_train", True)

    train_epochs = getattr(config, "train_epochs", 3)
    train_epochs_per_task = getattr(config, "train_epochs_per_task", {}) or {
        t: train_epochs for t in tasks
    }

    train_batch_size = getattr(config, "train_batch_size", 32)
    train_warmup_steps = getattr(config, "train_warmup_steps", 100)
    train_weight_decay = getattr(config, "train_weight_decay", 0.01)
    train_grad_accum_steps = getattr(config, "train_grad_accum_steps", 1)

    lr_candidates = getattr(config, "bert_lr_candidates", BERT_LR_CANDIDATES)
    random_restarts_small = int(getattr(config, "random_restarts_small", 1))
    base_seed = int(getattr(config, "base_seed", 543211))

    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    checkpoint_out_root = out_dir / "checkpoints"
    checkpoint_out_root.mkdir(parents=True, exist_ok=True)

    rows = []

    # save original model state
    try:
        original_state = unwrap_model(model).state_dict()
    except Exception:
        original_state = model.state_dict()

    def _reset_model_to_original():
        try:
            unwrap_model(model).load_state_dict(original_state, strict=False)
        except Exception:
            model.load_state_dict(original_state, strict=False)

    for task in tasks:
        print(f"\n==== Grid-running task: {task} ====")
        epochs_this_task = int(train_epochs_per_task.get(task, train_epochs))

        # load dataset
        if task in EXTRA_TASKS:
            cfg = EXTRA_TASKS[task]
            if "hf_config" in cfg:
                ds = load_dataset(cfg["hf_path"], cfg["hf_config"])
            else:
                ds = load_dataset(cfg["hf_path"])
            train_raw = ds.get("train")
            val_raw = ds.get("validation") or ds.get("test")
        else:
            hf = load_dataset("glue", GLUE_TASKS[task]["hf_name"])
            train_raw = hf["train"]
            val_raw = hf["validation_matched"] if task == "mnli" else hf["validation"]

        task_ckpt_root = checkpoint_out_root / task
        task_ckpt_root.mkdir(parents=True, exist_ok=True)

        all_run_records = []
        best_run = {"score": None, "lr": None, "restart": None, "seed": None, "best_ckpt_path": None}

        small_flag = (task in SMALL_TASKS_RANDOM_RESTARTS) or (task in SMALL_TASKS_RANDOM_RESTARTS_EXTRA)
        restarts_per_lr = random_restarts_small if small_flag else 1

        if auto_train:
            for lr in lr_candidates:
                for restart_idx in range(restarts_per_lr):
                    seed = (
                        base_seed
                        + (abs(hash(task)) % 10000) * 1000
                        + restart_idx * 10
                        + (int(round(lr * 1e7)) % 1000)
                    )
                    print(f"[SWEEP] task={task} lr={lr:g} restart={restart_idx} seed={seed}")

                    _reset_model_to_original()

                    run_dir = task_ckpt_root / f"lr_{lr:g}" / f"restart_{restart_idx}"
                    run_dir.mkdir(parents=True, exist_ok=True)

                    if task in EXTRA_TASKS:
                        det = detect_num_choices(val_raw, task)
                        expect_num = det.get("detected_num_choices")
                        model, metric_res = train_full_finetune_extra(
                            task=task,
                            tokenizer=tokenizer,
                            model=model,
                            raw_train=train_raw,
                            raw_val=val_raw,
                            device=device,
                            epochs=epochs_this_task,
                            batch_size=train_batch_size,
                            lr=float(lr),
                            weight_decay=train_weight_decay,
                            warmup_steps=train_warmup_steps,
                            max_length=max_length,
                            grad_accum_steps=train_grad_accum_steps,
                            out_checkpoint_dir=str(run_dir),
                            expect_num_choices=expect_num,
                        )
                        sc = _metric_to_scalar(task, metric_res)
                        best_epoch = None
                    else:
                        model, metric_res, sc, best_epoch = train_full_finetune(
                            task=task,
                            tokenizer=tokenizer,
                            model=model,
                            raw_train=train_raw,
                            raw_val=val_raw,
                            device=device,
                            epochs=epochs_this_task,
                            batch_size=train_batch_size,
                            lr=float(lr),
                            weight_decay=train_weight_decay,
                            warmup_steps=train_warmup_steps,
                            max_length=max_length,
                            grad_accum_steps=train_grad_accum_steps,
                            out_checkpoint_dir=str(run_dir),
                            seed=seed,
                        )

                    rec = {
                        "task": task,
                        "lr": float(lr),
                        "restart": int(restart_idx),
                        "seed": int(seed),
                        "dev_score": float(sc) if sc is not None else None,
                        "run_dir": str(run_dir),
                        "best_ckpt_path": str(run_dir / "best_finetuned.pt"),
                    }
                    all_run_records.append(rec)

                    if best_run["score"] is None or (sc is not None and sc > best_run["score"]):
                        best_run.update(
                            {
                                "score": float(sc),
                                "lr": float(lr),
                                "restart": int(restart_idx),
                                "seed": int(seed),
                                "best_ckpt_path": rec["best_ckpt_path"],
                            }
                        )

            pd.DataFrame(all_run_records).to_csv(task_ckpt_root / "all_runs.csv", index=False)

            best_overall_dir = task_ckpt_root / "best_overall"
            best_overall_dir.mkdir(parents=True, exist_ok=True)
            if best_run["best_ckpt_path"] and os.path.exists(best_run["best_ckpt_path"]):
                shutil.copy2(best_run["best_ckpt_path"], best_overall_dir / "best_finetuned.pt")
            _json_dump(best_run, best_overall_dir / "best_meta.json")

            _reset_model_to_original()
            model_for_eval = _load_finetuned_checkpoint_for_task(
                task, model, str(best_overall_dir / "best_finetuned.pt")
            )
        else:
            model_for_eval = model

        task_out_dir = out_dir / task
        task_out_dir.mkdir(parents=True, exist_ok=True)

        if task in EXTRA_TASKS:
            det = detect_num_choices(val_raw, task)
            expect_num = det.get("detected_num_choices")
            metric_res = run_extra_task(
                task=task,
                tokenizer=tokenizer,
                model=model_for_eval,
                device=device,
                batch_size=batch_size,
                max_length=max_length,
                output_dir=str(task_out_dir),
                expect_num_choices=expect_num,
            )
            rows.append({"task": task, "split": "validation", "metrics": json.dumps(metric_res)})
        else:
            res = run_glue_task(
                task=task,
                tokenizer=tokenizer,
                model=model_for_eval,
                device=device,
                batch_size=batch_size,
                max_length=max_length,
                output_dir=str(task_out_dir),
                compute_errorbar=True,
            )
            for split, pack in res.items():
                rows.append(
                    {
                        "task": task,
                        "split": split,
                        "metrics": json.dumps(pack["metrics"]),
                        "errorbar": json.dumps(pack["errorbar"]),
                    }
                )

    summary_csv = out_dir / "glue_summary.csv"
    pd.DataFrame(rows).to_csv(summary_csv, index=False)
    print(f"[Grid] Summary saved to: {summary_csv}")
    return pd.DataFrame(rows)


# ---------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------
if __name__ == "__main__":
    print(
        "This module is intended to be imported and used by providing tokenizer/model objects.\n"
        "Example:\n"
        "  tok = AutoTokenizer.from_pretrained('bert-base-uncased')\n"
        "  model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')\n"
        "  run_glue_benchmark(cfg, tok, model)\n"
    )