src/lmr/glue_benchmark/glue_benchmark copy 3.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):
        try:
            if hasattr(m, "classifier") and isinstance(getattr(m, "classifier"), nn.Linear):
                return getattr(m, "classifier").out_features
            if hasattr(m, "lm_head") and isinstance(getattr(m, "lm_head"), nn.Linear):
                return getattr(m, "lm_head").out_features
            if hasattr(m, "get_output_embeddings"):
                out_emb = m.get_output_embeddings()
                if out_emb is not None:
                    if isinstance(out_emb, nn.Embedding):
                        return out_emb.embedding_dim if hasattr(out_emb, "embedding_dim") else out_emb.num_embeddings
                    if isinstance(out_emb, nn.Linear):
                        return out_emb.out_features
        except Exception:
            pass
        return None

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

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

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

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

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

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

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

            if isinstance(out, (tuple, list)) and len(out) > 0:
                cand = out[0]
                if torch.is_tensor(cand):
                    if cand.ndim == 3:
                        pooled = cand[:, 0, :]
                        logits = self.classifier(pooled)
                        return type("Out", (), {"logits": logits, "loss": None})
                    if cand.ndim == 2 and cand.shape[1] == num_labels:
                        return type("Out", (), {"logits": cand, "loss": None})

            logits = getattr(out, "logits", None)
            if logits is not None:
                if logits.ndim == 2 and logits.shape[1] == num_labels:
                    return type("Out", (), {"logits": logits, "loss": getattr(out, "loss", None)})
                exist_dim = logits.shape[1]
                if self.logits_projector is None or self.logits_projector.weight.shape[1] != exist_dim:
                    self.logits_projector = nn.Linear(exist_dim, num_labels).to(logits.device)
                projected = self.logits_projector(logits)
                return type("Out", (), {"logits": projected, "loss": getattr(out, "loss", None)})

            hidden_states = getattr(out, "hidden_states", None)
            if hidden_states is not None:
                last_hidden = hidden_states[-1] if isinstance(hidden_states, (list, tuple)) else hidden_states
                if torch.is_tensor(last_hidden) and last_hidden.ndim == 3:
                    pooled = last_hidden[:, 0, :]
                    logits = self.classifier(pooled)
                    return type("Out", (), {"logits": logits, "loss": None})

            raise RuntimeError("Wrapped base model did not return recognizable hidden states or logits")

    return _WrappedModel(base_model, inferred_hidden, num_labels), True

# ---------------------------------------------------------------------
# Tokenize HF split to tensors (for finetune)
# ---------------------------------------------------------------------
def _tokenize_hf_split_to_tensors(task: str, tokenizer, raw_split, cfg_task, max_length=128, batch_tokenize_size=512):
    texts = []
    labels = []
    empty_s1 = 0
    empty_s2 = 0

    for ex in raw_split:
        s1, s2 = _get_text_pair_from_example(task, ex)
        texts.append((s1, s2))
        labels.append(ex.get("label") if "label" in ex else -100)
        if not s1 or (isinstance(s1, str) and s1.strip() == ""):
            empty_s1 += 1
        if s2 is not None and (not s2 or (isinstance(s2, str) and s2.strip() == "")):
            empty_s2 += 1

    total = len(texts)
    print(
        f"[tokenize] task={task} samples={total} empty_s1={empty_s1} empty_s2={empty_s2} "
        f"({(empty_s1/total if total>0 else 0):.2%}, {(empty_s2/total if total>0 else 0):.2%})"
    )

    input_ids_all = []
    attention_all = []
    for i in range(0, len(texts), batch_tokenize_size):
        enc = _batch_tokenize(tokenizer, texts[i:i+batch_tokenize_size], max_length=max_length)
        ids = enc.get("input_ids")
        masks = enc.get("attention_mask") or enc.get("mask") or enc.get("masks")
        if isinstance(ids, torch.Tensor):
            ids = ids.tolist()
        if isinstance(masks, torch.Tensor):
            masks = masks.tolist()
        input_ids_all.extend(ids)
        attention_all.extend(masks)

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

    input_ids_t, attention_mask_t = _pad_and_tensorize(input_ids_all, attention_all, pad_id)
    labels_t = torch.tensor(labels, dtype=torch.long if cfg_task["type"] == "classification" else torch.float)
    return input_ids_t, attention_mask_t, labels_t
def _normalize_label(label):
    """Robustly normalize various label encodings to an int index or None.
    Handles: int, float-like strings, single-letter answers ('A','b'), empty strings, None.
    """
    if label is None:
        return None
    # If it's already int-like
    if isinstance(label, (int, np.integer)):
        return int(label)
    # string handling
    if isinstance(label, str):
        s = label.strip()
        if s == "":
            return None
        # single-letter like 'A'/'b'
        if len(s) == 1 and s.isalpha():
            return ord(s.upper()) - ord("A")
        # Try int
        try:
            return int(s)
        except Exception:
            pass
        # Try float then cast to int if reasonable (e.g. '1.0')
        try:
            f = float(s)
            # only accept if it's integer-valued (e.g. 1.0 -> 1)
            if abs(f - round(f)) < 1e-6:
                return int(round(f))
            # otherwise treat as None (can't map to choice index)
            return None
        except Exception:
            return None
    # other numeric-like (np types)
    try:
        return int(label)
    except Exception:
        return None


def _extract_mc_example(task: str, ex: dict):
    """
    Robust extractor for multiple-choice examples across a range of HF dataset schemas.
    Returns (context, options_list, label_index_or_None).
    """
    # detect label raw value first (don't int() it yet)
    label_raw = None
    if "label" in ex:
        label_raw = ex.get("label")
    if label_raw is None:
        label_raw = ex.get("answerKey") or ex.get("answer") or ex.get("correct") or ex.get("gold")

    # normalize into int index or None
    label = _normalize_label(label_raw)

    # 1) 'choices' list (strings or dicts)
    if "choices" in ex and ex["choices"] is not None:
        ch = ex["choices"]
        if isinstance(ch, list) and len(ch) > 0:
            opts = []
            for c in ch:
                if isinstance(c, dict):
                    opts.append(c.get("text") or c.get("label") or c.get("choice") or str(c))
                else:
                    opts.append(str(c))
            ctx = ex.get("context") or ex.get("question") or ex.get("story") or ex.get("sentence") or ex.get("passage")
            return ctx, opts, label

    # 2) 'endings' pattern (hellaswag)
    if "endings" in ex and isinstance(ex["endings"], list) and len(ex["endings"]) > 0:
        ctx = ex.get("context") or ex.get("article") or ex.get("sentence") or ex.get("story") or ex.get("paragraph")
        opts = [str(x) for x in ex["endings"]]
        return ctx, opts, label

    # 3) explicit option fields like 'choice1','choice2' or 'option1'..
    opts = []
    for prefix in ("choice", "option", "ending", "answer"):
        i = 1
        found = False
        while True:
            key = f"{prefix}{i}"
            if key in ex:
                opts.append(str(ex[key]))
                found = True
                i += 1
            else:
                break
        if found:
            ctx = ex.get("question") or ex.get("context") or ex.get("passage") or ex.get("sentence")
            return ctx, opts, label

    # 4) common QA fields: 'question' + 'choices' (where choices might be list of dicts)
    if "question" in ex:
        ctx = ex["question"]
        if "choices" in ex:
            ch = ex["choices"]
            if isinstance(ch, list) and len(ch) > 0:
                opts = []
                for c in ch:
                    if isinstance(c, dict):
                        opts.append(c.get("text") or c.get("choice") or str(c))
                    else:
                        opts.append(str(c))
                return ctx, opts, label

    # 5) ai2_arc / openbookqa style: search for list-like values
    for k, v in ex.items():
        if isinstance(v, list) and 2 <= len(v) <= 10 and all(isinstance(x, (str, dict)) for x in v):
            opts = [x.get("text") if isinstance(x, dict) and x.get("text") else str(x) for x in v]
            ctx = ex.get("goal") or ex.get("question") or ex.get("context") or ex.get("passage") or ""
            return ctx, opts, label

    # 6) fallback: collect fields that look like options
    candidate_opts = []
    for k in sorted(ex.keys()):
        if any(tok in k.lower() for tok in ("option", "choice", "ending", "answer", "alt", "sol")):
            candidate_opts.append(str(ex[k]))
    if candidate_opts:
        ctx = ex.get("question") or ex.get("context") or ""
        return ctx, candidate_opts, label

    # last resort
    ctx = ex.get("question") or ex.get("context") or ex.get("passage") or ""
    return ctx, [], label

def _tokenize_generic_mc_split_to_tensors(task: str, tokenizer, raw_split, max_length=128, batch_tokenize_size=256):
    """
    Generic MC tokenizer that uses _extract_mc_example to normalize different HF schemas.
    Returns (input_ids_t (N,C,L), attention_t (N,C,L), labels_t (N,))
    """
    contexts = []
    options = []
    labels = []
    num_choices = None

    for ex in raw_split:
        ctx, opts, lab = _extract_mc_example(task, ex)
        if not opts:
            # skip if no options recognized
            continue
        if num_choices is None:
            num_choices = len(opts)
        if len(opts) != num_choices:
            # inconsistent number of choices; skip example
            continue
        contexts.append(ctx if ctx is not None else "")
        options.append(opts)
        labels.append(-1 if lab is None else int(lab))

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

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

    for i in range(0, len(contexts), batch_tokenize_size):
        chunk_ctx = contexts[i:i+batch_tokenize_size]
        chunk_opts = options[i:i+batch_tokenize_size]
        flat_pairs = []
        for c, opts in zip(chunk_ctx, chunk_opts):
            for o in opts:
                flat_pairs.append((c, o))
        enc = _batch_tokenize(tokenizer, flat_pairs, max_length=max_length)
        ids_flat = enc.get("input_ids")
        masks_flat = enc.get("attention_mask") or enc.get("mask") or enc.get("masks")
        if isinstance(ids_flat, torch.Tensor):
            ids_flat = ids_flat.tolist()
        if isinstance(masks_flat, torch.Tensor):
            masks_flat = masks_flat.tolist()

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

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

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

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

# ---------------------------------------------------------------------
# Multiple-choice tokenizer entry (keeps fast paths for known tasks, else generic)
# ---------------------------------------------------------------------
def _tokenize_mc_split_to_tensors(task: str, tokenizer, raw_split, max_length=128, batch_tokenize_size=256):
    """
    Build tensors for multiple choice tasks:
      returns input_ids tensor shape (N, num_choices, L), attention_mask tensor same, labels tensor (N,)
    Uses fast paths for known tasks (piqa, winogrande), otherwise uses generic parser.
    """
    # fast paths
    if task == "piqa":
        contexts = []
        options = []
        labels = []
        for ex in raw_split:
            try:
                ctx = ex.get("goal") or ex.get("question") or ex.get("context") or ""
                opts = [ex["sol1"], ex["sol2"]]
                lab = int(ex["label"])
            except Exception:
                # fallback to generic
                return _tokenize_generic_mc_split_to_tensors(task, tokenizer, raw_split, max_length=max_length, batch_tokenize_size=batch_tokenize_size)
            contexts.append(ctx)
            options.append(opts)
            labels.append(lab)
        # then pack like generic
    elif task == "winogrande":
        contexts = []
        options = []
        labels = []
        for ex in raw_split:
            try:
                ctx = ex.get("sentence") or ex.get("context") or ex.get("question") or ""
                opts = [ex["option1"], ex["option2"]]
                lab = int(ex.get("answer", 1)) - 1
            except Exception:
                return _tokenize_generic_mc_split_to_tensors(task, tokenizer, raw_split, max_length=max_length, batch_tokenize_size=batch_tokenize_size)
            contexts.append(ctx)
            options.append(opts)
            labels.append(lab)
    else:
        # fallback to generic parser which supports hellaswag, openbookqa, arc, etc.
        return _tokenize_generic_mc_split_to_tensors(task, tokenizer, raw_split, max_length=max_length, batch_tokenize_size=batch_tokenize_size)

    # From here pack contexts/options/labels into tensors (same logic as generic)
    if len(contexts) == 0:
        return torch.zeros((0, 1, 1), dtype=torch.long), torch.zeros((0, 1, 1), dtype=torch.long), torch.tensor([], dtype=torch.long)

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

    num_choices = len(options[0])
    for i in range(0, len(contexts), batch_tokenize_size):
        chunk_ctx = contexts[i:i+batch_tokenize_size]
        chunk_opts = options[i:i+batch_tokenize_size]
        flat_pairs = []
        for c, opts in zip(chunk_ctx, chunk_opts):
            for o in opts:
                flat_pairs.append((c, o))
        enc = _batch_tokenize(tokenizer, flat_pairs, max_length=max_length)
        ids_flat = enc.get("input_ids")
        masks_flat = enc.get("attention_mask") or enc.get("mask") or enc.get("masks")
        if isinstance(ids_flat, torch.Tensor):
            ids_flat = ids_flat.tolist()
        if isinstance(masks_flat, torch.Tensor):
            masks_flat = masks_flat.tolist()

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

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

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

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

# ---------------------------------------------------------------------
# Pairwise tokenizer for BoolQ (if not already present)
# ---------------------------------------------------------------------
def _tokenize_pair_split_to_tensors(task: str, tokenizer, raw_split, max_length=128, batch_tokenize_size=512):
    texts = []
    labels = []
    for ex in raw_split:
        if task == "boolq":
            a = ex.get("passage") or ex.get("context") or ex.get("article") or ""
            b = ex.get("question") or ex.get("query") or ""
            lab = int(ex.get("answer") or ex.get("label") or 0)
        else:
            # If unknown, try generic pair fields
            a = ex.get("passage") or ex.get("context") or ex.get("article") or ""
            b = ex.get("question") or ex.get("query") or ""
            lab = int(ex.get("answer") or ex.get("label") or 0)
        texts.append((a, b))
        labels.append(lab)

    input_ids_all = []
    attention_all = []
    for i in range(0, len(texts), batch_tokenize_size):
        chunk = texts[i:i+batch_tokenize_size]
        enc = _batch_tokenize(tokenizer, chunk, max_length=max_length)
        ids = enc.get("input_ids")
        masks = enc.get("attention_mask") or enc.get("mask") or enc.get("masks")
        if isinstance(ids, torch.Tensor):
            ids = ids.tolist()
        if isinstance(masks, torch.Tensor):
            masks = masks.tolist()
        input_ids_all.extend(ids)
        attention_all.extend(masks)

    pad_id = getattr(tokenizer, "pad_token_id", None)
    if pad_id is None:
        try:
            pad_id = tokenizer.token_to_id("[PAD]")
        except Exception:
            pad_id = 0
    input_ids_t, attention_mask_t = _pad_and_tensorize(input_ids_all, attention_all, pad_id)
    labels_t = torch.tensor(labels, dtype=torch.long)
    return input_ids_t, attention_mask_t, labels_t

# ---------------------------------------------------------------------
# Postprocess preds helper already defined above (_postprocess_predictions)
# ---------------------------------------------------------------------

# ---------------------------------------------------------------------
# Training for GLUE tasks (full fine-tune)
# ---------------------------------------------------------------------
def train_full_finetune(
    task: str,
    tokenizer,
    model,
    raw_train,
    raw_val,
    device: str = "cuda",
    epochs: int = 3,
    batch_size: int = 32,
    lr: float = 2e-5,
    weight_decay: float = 0.01,
    warmup_steps: int = 100,
    max_length: int = 128,
    grad_accum_steps: int = 1,
    out_checkpoint_dir: Optional[str] = None,
    seed: Optional[int] = None,
):
    cfg_task = GLUE_TASKS[task]
    device_t = torch.device(device if torch.cuda.is_available() else "cpu")

    if seed is not None:
        _set_all_seeds(int(seed))

    hidden_size = None
    if hasattr(model, "config") and hasattr(model.config, "hidden_size"):
        try:
            hidden_size = int(model.config.hidden_size)
        except Exception:
            hidden_size = None

    model, wrapped_flag = make_wrapped_model_if_needed(
        model, hidden_size, cfg_task["num_labels"], force_num_labels=cfg_task["num_labels"]
    )
    model.to(device_t)

    train_ids, train_mask, train_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_train, cfg_task, max_length=max_length)
    val_ids, val_mask, val_labels = _tokenize_hf_split_to_tensors(task, tokenizer, raw_val, cfg_task, max_length=max_length)

    train_ds = TensorDataset(train_ids, train_mask, train_labels)
    val_ds = TensorDataset(val_ids, val_mask, val_labels)

    g = torch.Generator()
    if seed is not None:
        g.manual_seed(int(seed))

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

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

    loss_fn = torch.nn.CrossEntropyLoss() if cfg_task["type"] == "classification" else torch.nn.MSELoss()

    best_metric_res: Dict[str, Any] = {}
    best_score: Optional[float] = None
    best_epoch = -1

    model.train()

    for epoch in range(epochs):
        for step, batch in enumerate(tqdm(train_loader, desc=f"Train {task} epoch {epoch+1} (lr={lr:g})")):
            ids_b, mask_b, labs_b = batch
            ids_b = ids_b.to(device_t)
            mask_b = mask_b.to(device_t)
            labs_b = labs_b.to(device_t)

            out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
            logits = getattr(out, "logits", None)
            if logits is None:
                if isinstance(out, (tuple, list)):
                    logits = out[0]
                else:
                    raise RuntimeError("Model did not return logits during finetune")

            if cfg_task["type"] == "classification":
                loss = loss_fn(logits, labs_b.long())
            else:
                if logits.ndim == 2 and logits.shape[1] == 1:
                    preds = logits.squeeze(1)
                elif logits.ndim == 2:
                    preds = logits.mean(dim=1)
                else:
                    preds = logits
                loss = loss_fn(preds, labs_b.float())

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

            if (step + 1) % max(1, grad_accum_steps) == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                if scheduler is not None:
                    scheduler.step()
                optimizer.zero_grad()

        # validation
        model.eval()
        tot_val_loss = 0.0
        all_logits = []
        all_labels = []
        with torch.no_grad():
            for ids_b, mask_b, labs_b in tqdm(val_loader, desc=f"Validate {task} epoch {epoch+1}", leave=False):
                ids_b = ids_b.to(device_t)
                mask_b = mask_b.to(device_t)
                labs_b = labs_b.to(device_t)

                out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                logits = getattr(out, "logits", None)
                if logits is None:
                    if isinstance(out, (tuple, list)):
                        logits = out[0]
                    else:
                        raise RuntimeError("Model did not return logits during validation")

                if cfg_task["type"] == "classification":
                    l = loss_fn(logits, labs_b.long())
                else:
                    if logits.ndim == 2 and logits.shape[1] == 1:
                        preds = logits.squeeze(1)
                    elif logits.ndim == 2:
                        preds = logits.mean(dim=1)
                    else:
                        preds = logits
                    l = loss_fn(preds, labs_b.float())

                tot_val_loss += l.item() * ids_b.size(0)
                all_logits.append(logits.detach().cpu().numpy())
                all_labels.append(labs_b.detach().cpu().numpy())

        model.train()

        all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg_task["num_labels"]))
        all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,))
        preds = _postprocess_predictions(task, all_logits, cfg_task)

        metric = evaluate.load("glue", cfg_task["hf_name"])
        try:
            metric_res = metric.compute(predictions=preds.tolist(), references=all_labels.tolist())
        except Exception:
            try:
                metric_res = metric.compute(predictions=preds, references=all_labels)
            except Exception as e:
                metric_res = {"error": str(e)}

        avg_val_loss = tot_val_loss / len(val_ds) if len(val_ds) > 0 else float("nan")
        score = _metric_to_scalar(task, metric_res, fallback_val_loss=avg_val_loss)

        print(f"[FT] {task} epoch {epoch+1} lr={lr:g} val_loss={avg_val_loss:.6f} metric={metric_res} score={score:.6f}")

        if best_score is None or float(score) > float(best_score):
            best_score = float(score)
            best_metric_res = metric_res
            best_epoch = epoch + 1

            if out_checkpoint_dir:
                outp = Path(out_checkpoint_dir)
                outp.mkdir(parents=True, exist_ok=True)
                best_fname = outp / "best_finetuned.pt"
                try:
                    sd = unwrap_model(model).state_dict()
                except Exception:
                    sd = model.state_dict()
                torch.save(sd, str(best_fname))
                print(f"[FT] Saved best checkpoint (epoch {best_epoch}) to: {best_fname}")

    if out_checkpoint_dir:
        outp = Path(out_checkpoint_dir)
        outp.mkdir(parents=True, exist_ok=True)
        fname = outp / "finetuned.pt"
        try:
            sd = unwrap_model(model).state_dict()
        except Exception:
            sd = model.state_dict()
        torch.save(sd, str(fname))
        print(f"[FT] Saved finetuned model to: {fname}")

        meta = {
            "task": task,
            "lr": lr,
            "seed": seed,
            "epochs": epochs,
            "batch_size": batch_size,
            "grad_accum_steps": grad_accum_steps,
            "warmup_steps": warmup_steps,
            "weight_decay": weight_decay,
            "max_length": max_length,
            "wrapped_flag": bool(wrapped_flag),
            "best_epoch": int(best_epoch),
            "best_score": float(best_score) if best_score is not None else None,
            "best_metrics": best_metric_res,
        }
        _json_dump(meta, outp / "run_meta.json")

    print(f"[FT] Best validation for task '{task}' (lr={lr:g}, seed={seed}): epoch={best_epoch}, score={best_score}, metrics={best_metric_res}")
    return model, best_metric_res, float(best_score) if best_score is not None else -1e9, best_epoch

# ---------------------------------------------------------------------
# Load finetuned checkpoint for eval (wrapped/unwrapped)
# ---------------------------------------------------------------------
def _load_finetuned_checkpoint_for_task(task: str, base_model, checkpoint_path: str):
    cfg = GLUE_TASKS.get(task) or EXTRA_TASKS.get(task)
    sd = torch.load(checkpoint_path, map_location="cpu")
    keys = list(sd.keys()) if isinstance(sd, dict) else []
    looks_wrapped = any(k.startswith("base.") for k in keys) or any(k.startswith("classifier.") for k in keys)

    hidden_size = None
    if hasattr(base_model, "config") and hasattr(base_model.config, "hidden_size"):
        try:
            hidden_size = int(base_model.config.hidden_size)
        except Exception:
            hidden_size = None

    if looks_wrapped and cfg is not None:
        wrapped_model, _ = make_wrapped_model_if_needed(
            base_model, hidden_size, cfg["num_labels"], force_num_labels=cfg["num_labels"]
        )
        try:
            unwrap_model(wrapped_model).load_state_dict(sd, strict=False)
        except Exception:
            try:
                wrapped_model.load_state_dict(sd, strict=False)
            except Exception:
                pass
        return wrapped_model

    try:
        unwrap_model(base_model).load_state_dict(sd, strict=False)
    except Exception:
        try:
            base_model.load_state_dict(sd, strict=False)
        except Exception:
            pass
    return base_model

# ---------------------------------------------------------------------
# Error bar evaluation (5 folds -> 5 leave-one-fold-out subsets)
# ---------------------------------------------------------------------
def _fivefold_indices(n: int):
    idx = np.arange(n)
    folds = np.array_split(idx, 5)
    return [f.tolist() for f in folds]

def _errorbar_subsets_from_folds(folds: List[List[int]]):
    assert len(folds) == 5
    combos = [
        ("0123", [0,1,2,3]),
        ("1234", [1,2,3,4]),
        ("0124", [0,1,2,4]),
        ("0234", [0,2,3,4]),
        ("0134", [0,1,3,4]),
    ]
    subsets = []
    for name, keep in combos:
        inds = []
        for k in keep:
            inds.extend(folds[k])
        subsets.append({"name": name, "indices": inds})
    return subsets

def _compute_metric_for_indices(task: str, cfg: dict, metric_obj, preds_all: np.ndarray, labels_all: np.ndarray, indices: List[int]):
    if len(indices) == 0:
        return {"error": "empty_indices"}
    p = preds_all[indices]
    y = labels_all[indices]
    if cfg["type"] == "classification" or cfg.get("type") == "multiple_choice":
        preds_out = p.astype(int).tolist()
        refs_out = y.astype(int).tolist()
    else:
        preds_out = p.astype(float).tolist()
        refs_out = y.astype(float).tolist()
    try:
        return metric_obj.compute(predictions=preds_out, references=refs_out)
    except Exception:
        try:
            return metric_obj.compute(predictions=np.array(preds_out), references=np.array(refs_out))
        except Exception as e:
            return {"error": str(e)}

def _compute_errorbar(task: str, cfg: dict, metric_obj, preds_all: np.ndarray, labels_all: np.ndarray):
    n = int(len(labels_all))
    folds = _fivefold_indices(n)
    subsets = _errorbar_subsets_from_folds(folds)

    pref = PREFERRED_METRIC_KEY.get(task)
    subset_scores = []
    scores = []

    for s in subsets:
        m = _compute_metric_for_indices(task, cfg, metric_obj, preds_all, labels_all, s["indices"])
        sc = _metric_to_scalar(task, m, fallback_val_loss=None)
        subset_scores.append({"subset": s["name"], "score": float(sc), "metrics": m})
        scores.append(float(sc))

    arr = np.array(scores, dtype=float)
    mean = float(np.mean(arr)) if len(arr) else float("nan")
    std = float(np.std(arr, ddof=1)) if len(arr) > 1 else 0.0
    stderr = float(std / math.sqrt(len(arr))) if len(arr) > 0 else float("nan")

    return {
        "preferred_key": pref,
        "subset_scores": subset_scores,
        "mean": mean,
        "std": std,
        "stderr": stderr,
    }

# ---------------------------------------------------------------------
# Evaluation (returns metrics + also writes preds/results)
# ---------------------------------------------------------------------
def run_glue_task(
    task: str,
    tokenizer,
    model,
    checkpointing: Optional[Checkpointing] = None,
    device: str = "cuda",
    batch_size: int = 64,
    max_length: int = 128,
    output_dir: str = "glue_output",
    compute_errorbar: bool = False,
):
    assert task in GLUE_TASKS, f"Unknown GLUE task: {task}"
    cfg = GLUE_TASKS[task]

    hf = load_dataset("glue", cfg["hf_name"])
    if task == "mnli":
        val_splits = ["validation_matched", "validation_mismatched"]
    else:
        val_splits = ["validation"]

    results_by_split = {}
    for split in val_splits:
        raw = hf[split]
        print(f"[GLUE] Task={task} split={split} samples={len(raw)}")

        texts = []
        labels = []
        for ex in raw:
            s1, s2 = _get_text_pair_from_example(task, ex)
            texts.append((s1, s2))
            labels.append(ex.get("label") if "label" in ex else -100)

        BATCH = 512
        input_ids_all = []
        attention_all = []
        for i in range(0, len(texts), BATCH):
            enc = _batch_tokenize(tokenizer, texts[i:i+BATCH], max_length=max_length)
            ids = enc.get("input_ids")
            masks = enc.get("attention_mask") or enc.get("mask") or enc.get("masks")
            if isinstance(ids, torch.Tensor):
                ids = ids.tolist()
            if isinstance(masks, torch.Tensor):
                masks = masks.tolist()
            input_ids_all.extend(ids)
            attention_all.extend(masks)

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

        input_ids, attention_mask = _pad_and_tensorize(input_ids_all, attention_all, pad_id)
        labels_t = torch.tensor(labels, dtype=torch.long if cfg["type"] == "classification" else torch.float)

        ds = TensorDataset(input_ids, attention_mask, labels_t)
        loader = DataLoader(ds, batch_size=batch_size, shuffle=False, pin_memory=True)

        if checkpointing is not None:
            try:
                checkpointing.load_model_states("recent")
            except Exception:
                pass

        device_t = torch.device(device if torch.cuda.is_available() else "cpu")
        model.to(device_t)
        model.eval()

        hidden_size = None
        if hasattr(model, "config") and hasattr(model.config, "hidden_size"):
            try:
                hidden_size = int(model.config.hidden_size)
            except Exception:
                hidden_size = None

        force = 1 if cfg["type"] == "regression" else cfg["num_labels"]
        wrapped_model, _ = make_wrapped_model_if_needed(model, hidden_size, cfg["num_labels"], force_num_labels=force)
        wrapped_model.to(device_t)
        wrapped_model.eval()

        all_logits = []
        all_labels = []
        with torch.no_grad():
            for batch in tqdm(loader, desc=f"Eval {task}:{split}"):
                ids_b, mask_b, labels_b = batch
                ids_b = ids_b.to(device_t)
                mask_b = mask_b.to(device_t)
                out = wrapped_model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                logits = getattr(out, "logits", None)
                if logits is None:
                    if isinstance(out, (tuple, list)):
                        logits = out[0]
                    else:
                        raise RuntimeError("Model forward did not return logits")
                all_logits.append(logits.detach().cpu().numpy())
                all_labels.append(labels_b.detach().cpu().numpy())

        all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg["num_labels"]))
        all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,))

        preds = _postprocess_predictions(task, all_logits, cfg)

        metric = evaluate.load("glue", cfg["hf_name"])
        metric_res = _compute_metric_for_indices(task, cfg, metric, preds, all_labels, list(range(len(all_labels))))

        errorbar_res = None
        if compute_errorbar:
            errorbar_res = _compute_errorbar(task, cfg, metric, preds, all_labels)

        os.makedirs(output_dir, exist_ok=True)
        out_json = Path(output_dir) / f"{task}_{split}_results.json"
        with open(out_json, "w", encoding="utf-8") as f:
            json.dump({"task": task, "split": split, "metrics": metric_res}, f, indent=2)

        if errorbar_res is not None:
            out_eb = Path(output_dir) / f"{task}_{split}_errorbar.json"
            with open(out_eb, "w", encoding="utf-8") as f:
                json.dump({"task": task, "split": split, "errorbar": errorbar_res}, f, indent=2)

        csv_p = Path(output_dir) / f"{task}_{split}_preds.csv"
        pd.DataFrame({"pred": preds.tolist(), "label": all_labels.tolist()}).to_csv(csv_p, index=False)

        results_by_split[split] = {"metrics": metric_res, "errorbar": errorbar_res}

    return results_by_split

# ---------------------------------------------------------------------
# Train full fine-tune for EXTRA tasks (pair + mc)
# ---------------------------------------------------------------------
def train_full_finetune_extra(task: str, tokenizer, model, raw_train, raw_val,
                              device: str = "cuda", epochs: int = 3, batch_size: int = 16,
                              lr: float = 2e-5, weight_decay: float = 0.01, warmup_steps: int = 100,
                              max_length: int = 128, grad_accum_steps: int = 1, out_checkpoint_dir: Optional[str] = None):
    """
    Fine-tune for EXTRA_TASKS (boolq/piqa/winogrande/hellaswag/openbookqa/arc)
    Expects the provided `model` to be compatible with HF's MultipleChoice or SequenceClassification APIs.
    """
    cfg = EXTRA_TASKS[task]
    device = torch.device(device if torch.cuda.is_available() else "cpu")

    model.to(device)

    is_mc = cfg["format"] == "mc"
    if is_mc:
        train_ids, train_mask, train_labels = _tokenize_mc_split_to_tensors(task, tokenizer, raw_train, max_length=max_length)
        val_ids, val_mask, val_labels = _tokenize_mc_split_to_tensors(task, tokenizer, raw_val, max_length=max_length)

        train_ds = TensorDataset(train_ids, train_mask, train_labels)
        val_ds = TensorDataset(val_ids, val_mask, val_labels)
    else:
        train_ids, train_mask, train_labels = _tokenize_pair_split_to_tensors(task, tokenizer, raw_train, max_length=max_length)
        val_ids, val_mask, val_labels = _tokenize_pair_split_to_tensors(task, tokenizer, raw_val, max_length=max_length)
        train_ds = TensorDataset(train_ids, train_mask, train_labels)
        val_ds = TensorDataset(val_ids, val_mask, val_labels)

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

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

    loss_fn = torch.nn.CrossEntropyLoss()

    model.train()
    global_step = 0
    final_metric_res = {}
    for epoch in range(epochs):
        running_loss = 0.0
        for step, batch in enumerate(tqdm(train_loader, desc=f"[ExtraTrain] {task} epoch {epoch+1}")):
            if is_mc:
                ids_b, mask_b, labs_b = batch  # ids_b: (B, C, L)
                ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device)
                out = None
                logits = None
                try:
                    out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                    logits = getattr(out, "logits", None)
                    if logits is None and isinstance(out, (list, tuple)):
                        logits = out[0]
                except Exception:
                    B, C, L = ids_b.shape
                    flat_ids = ids_b.view(B*C, L).to(device)
                    flat_mask = mask_b.view(B*C, L).to(device)
                    out_flat = model(input_ids=flat_ids, attention_mask=flat_mask)
                    flat_logits = getattr(out_flat, "logits", None)
                    if flat_logits is None and isinstance(out_flat, (tuple, list)):
                        flat_logits = out_flat[0]
                    if flat_logits is None:
                        raise RuntimeError("Model did not return logits for MC fallback")
                    if flat_logits.ndim == 2 and flat_logits.shape[1] == 1:
                        logits = flat_logits.view(B, C)
                    else:
                        logits = flat_logits.view(B, C, -1).mean(dim=-1)

                loss = loss_fn(logits, labs_b.long())
            else:
                ids_b, mask_b, labs_b = batch
                ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device)
                out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                logits = getattr(out, "logits", None)
                if logits is None:
                    if isinstance(out, (tuple, list)):
                        logits = out[0]
                    else:
                        raise RuntimeError("Model did not return logits for pair task")
                loss = loss_fn(logits, labs_b.long())

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

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

        # validation
        model.eval()
        all_logits = []
        all_labels = []
        with torch.no_grad():
            for batch in tqdm(val_loader, desc=f"[ExtraVal] {task} epoch {epoch+1}", leave=False):
                if is_mc:
                    ids_b, mask_b, labs_b = batch
                    ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device)
                    out = None
                    try:
                        out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                        logits = getattr(out, "logits", None)
                        if logits is None and isinstance(out, (tuple, list)):
                            logits = out[0]
                    except Exception:
                        B, C, L = ids_b.shape
                        flat_ids = ids_b.view(B*C, L).to(device)
                        flat_mask = mask_b.view(B*C, L).to(device)
                        out_flat = model(input_ids=flat_ids, attention_mask=flat_mask)
                        flat_logits = getattr(out_flat, "logits", None)
                        if flat_logits is None and isinstance(out_flat, (tuple, list)):
                            flat_logits = out_flat[0]
                        if flat_logits is None:
                            raise RuntimeError("Model did not return logits during MC validation fallback")
                        if flat_logits.ndim == 2 and flat_logits.shape[1] == 1:
                            logits = flat_logits.view(B, C)
                        else:
                            logits = flat_logits.view(B, C, -1).mean(dim=-1)
                    all_logits.append(logits.detach().cpu().numpy())
                    all_labels.append(labs_b.detach().cpu().numpy())
                else:
                    ids_b, mask_b, labs_b = batch
                    ids_b = ids_b.to(device); mask_b = mask_b.to(device); labs_b = labs_b.to(device)
                    out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                    logits = getattr(out, "logits", None)
                    if logits is None:
                        if isinstance(out, (tuple, list)):
                            logits = out[0]
                        else:
                            raise RuntimeError("Model did not return logits during pair validation")
                    all_logits.append(logits.detach().cpu().numpy())
                    all_labels.append(labs_b.detach().cpu().numpy())
        model.train()

        all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg["num_labels"]))
        all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,))

        if cfg["type"] == "classification" or cfg["type"] == "multiple_choice":
            preds = np.argmax(all_logits, axis=-1) if all_logits.size else np.array([])
        else:
            preds = _postprocess_predictions(task, all_logits, {"type": cfg["type"], "num_labels": cfg["num_labels"]})

        try:
            metric = evaluate.load("accuracy")
            metric_res = metric.compute(predictions=preds.tolist(), references=all_labels.tolist())
        except Exception as e:
            metric_res = {"error": str(e)}

        print(f"[Extra FT] {task} epoch {epoch+1} metric={metric_res}")
        final_metric_res = metric_res

    if out_checkpoint_dir:
        outp = Path(out_checkpoint_dir)
        outp.mkdir(parents=True, exist_ok=True)
        fname = outp / "finetuned_extra.pt"
        try:
            sd = unwrap_model(model).state_dict()
        except Exception:
            sd = model.state_dict()
        torch.save(sd, str(fname))
        print(f"[Extra FT] Saved finetuned model to: {fname}")

    return model, final_metric_res

# ---------------------------------------------------------------------
# Evaluation-only runner for EXTRA tasks

# ---------------------------------------------------------------------
# Evaluation-only runner for EXTRA tasks
# ---------------------------------------------------------------------
def run_extra_task(task: str,
                   tokenizer,
                   model,
                   checkpointing: Optional[Checkpointing] = None,
                   device: str = "cuda",
                   batch_size: int = 64,
                   max_length: int = 128,
                   output_dir: str = "extra_output",
                   prefer_test_if_available: bool = True):
    """
    Evaluate an EXTRA_TASK on HF dataset. Behavior:
      - If the dataset provides a 'test' split we will prefer it (unless it is unlabeled).
      - Otherwise use 'validation' or other labeled splits.
      - Returns metric dict (usually accuracy) and writes preds + metrics to output_dir.
    """
    assert task in EXTRA_TASKS, f"Unknown extra task: {task}"
    cfg = EXTRA_TASKS[task]

    # load hf dataset robustly (handle hf_config if present)
    try:
        if "hf_config" in cfg:
            ds = load_dataset(cfg["hf_path"], cfg["hf_config"])
        else:
            ds = load_dataset(cfg["hf_path"])
    except Exception as e:
        raise RuntimeError(f"Failed to load HF dataset for task={task}: {e}")

    # prefer test split if available and (prefer_test_if_available True).
    # But if test is unlabeled (no label fields), we may fall back to validation.
    chosen_split = None
    candidate_order = []
    # explicit preference order: test, validation, validation_matched, validation_unlabeled, train
    candidate_order = ["test", "validation", "validation_matched", "validation_unlabeled", "train"]
    available_splits = list(ds.keys()) if hasattr(ds, "keys") else []
    # If prefer_test_if_available try to pick test first
    for cand in candidate_order:
        if cand in ds:
            # check if split has at least one example and at least one of common label keys
            split_ds = ds[cand]
            try:
                first = next(iter(split_ds), None)
            except Exception:
                first = None
            has_label = False
            if first is not None:
                if any(k in first for k in ("label", "answer", "answerKey", "correct", "gold")):
                    has_label = True
            # choose test even if has_label False (many datasets publish test with labels in HF)
            if cand == "test" and cand in available_splits:
                chosen_split = "test"
                break
            if cand == "validation" and has_label:
                chosen_split = "validation"
                break
            if chosen_split is None and cand in available_splits:
                chosen_split = cand

    if chosen_split is None:
        # fallback to first available
        chosen_split = available_splits[0]

    val = ds.get(chosen_split)
    print(f"[Extra Eval] Task={task} using split='{chosen_split}' samples={len(val)}")

    device_t = torch.device(device if torch.cuda.is_available() else "cpu")
    model.to(device_t)
    model.eval()

    # Build tensors depending on format
    if cfg["format"] == "mc":
        ids_t, mask_t, labels_t = _tokenize_mc_split_to_tensors(task, tokenizer, val, max_length=max_length)
        # ids_t shape: (N, C, L) or (0,...)
        ds_t = TensorDataset(ids_t, mask_t, labels_t)
        loader = DataLoader(ds_t, batch_size=batch_size, shuffle=False, pin_memory=True)
    else:
        ids_t, mask_t, labels_t = _tokenize_pair_split_to_tensors(task, tokenizer, val, max_length=max_length)
        ds_t = TensorDataset(ids_t, mask_t, labels_t)
        loader = DataLoader(ds_t, batch_size=batch_size, shuffle=False, pin_memory=True)

    if checkpointing is not None:
        try:
            checkpointing.load_model_states("recent")
        except Exception:
            pass

    all_logits = []
    all_labels = []
    with torch.no_grad():
        for batch in tqdm(loader, desc=f"Eval {task}"):
            if cfg["format"] == "mc":
                ids_b, mask_b, labs_b = batch
                ids_b = ids_b.to(device_t); mask_b = mask_b.to(device_t)
                # Try direct MC forward, else flatten fallback
                try:
                    out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                    logits = getattr(out, "logits", None)
                    if logits is None and isinstance(out, (list, tuple)):
                        logits = out[0]
                except Exception:
                    B, C, L = ids_b.shape
                    flat_ids = ids_b.view(B*C, L).to(device_t)
                    flat_mask = mask_b.view(B*C, L).to(device_t)
                    out_flat = model(input_ids=flat_ids, attention_mask=flat_mask)
                    flat_logits = getattr(out_flat, "logits", None)
                    if flat_logits is None and isinstance(out_flat, (list, tuple)):
                        flat_logits = out_flat[0]
                    if flat_logits is None:
                        raise RuntimeError("Model did not return logits for MC fallback")
                    if flat_logits.ndim == 2 and flat_logits.shape[1] == 1:
                        logits = flat_logits.view(B, C)
                    else:
                        logits = flat_logits.view(B, C, -1).mean(dim=-1)
                all_logits.append(logits.detach().cpu().numpy())
                all_labels.append(labs_b.detach().cpu().numpy())
            else:
                ids_b, mask_b, labs_b = batch
                ids_b = ids_b.to(device_t); mask_b = mask_b.to(device_t)
                out = model(input_ids=ids_b, attention_mask=mask_b, labels=None)
                logits = getattr(out, "logits", None)
                if logits is None and isinstance(out, (list, tuple)):
                    logits = out[0]
                all_logits.append(logits.detach().cpu().numpy())
                all_labels.append(labs_b.detach().cpu().numpy())

    all_logits = np.concatenate(all_logits, axis=0) if all_logits else np.zeros((0, cfg["num_labels"]))
    all_labels = np.concatenate(all_labels, axis=0) if all_labels else np.zeros((0,))

    if cfg["format"] == "mc" or cfg["type"] == "multiple_choice" or cfg["type"] == "classification":
        preds = np.argmax(all_logits, axis=-1).astype(int) if all_logits.size else np.array([])
    else:
        preds = _postprocess_predictions(task, all_logits, {"type": cfg["type"], "num_labels": cfg["num_labels"]})

    # compute metric (accuracy for most MC tasks)
    try:
        metric = evaluate.load("accuracy")
        metric_res = metric.compute(predictions=preds.tolist(), references=all_labels.tolist())
    except Exception as e:
        metric_res = {"error": str(e)}

    # write outputs
    os.makedirs(output_dir, exist_ok=True)
    out_json = Path(output_dir) / f"{task}_{chosen_split}_results.json"
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump({"task": task, "split": chosen_split, "metrics": metric_res}, f, indent=2)

    csv_p = Path(output_dir) / f"{task}_{chosen_split}_preds.csv"
    pd.DataFrame({"pred": preds.tolist(), "label": all_labels.tolist()}).to_csv(csv_p, index=False)

    print(f"[Extra Eval] {task} split={chosen_split} metric={metric_res}")
    return metric_res

# ---------------------------------------------------------------------
# Grid-runner: LR sweep + restarts + checkpointing + eval
# ---------------------------------------------------------------------
def run_glue_benchmark(config, tokenizer, model, checkpointing: Optional[Checkpointing] = None, out_dir: str = "glue_outputs_grid"):
    """
    Grid-search runner for GLUE + EXTRA tasks.

    Config attributes supported (defaults will be used if missing):
      - glue_tasks: list of tasks (GLUE or EXTRA)
      - batch_size, max_length, device
      - auto_train (bool)
      - train_epochs_per_task (dict)
      - bert_lr_candidates (list)
      - random_restarts_small (int)
      - base_seed (int)
      - train_batch_size, train_warmup_steps, train_weight_decay, train_grad_accum_steps
    """
    tasks = getattr(config, "glue_tasks", None)
    if tasks is None:
        # default to a sensible subset; you can pass config with glue_tasks list
        tasks = ["rte", "cola", "mnli"]
    # allow passing a single string
    if isinstance(tasks, str):
        tasks = [t.strip() for t in tasks.split(",") if t.strip()]
    tasks =  [ "arc_easy", "arc_challenge","hellaswag","openbookqa"]
#  "piqa": 3, "winogrande": 3, "boolq": 3, "hellaswag": 3, "openbookqa": 3, "arc_easy": 3, "arc_challenge": 3

    # sanity-check tasks are in ALL_TASKS
    tasks = [t for t in tasks if t in ALL_TASKS]
    if not tasks:
        raise RuntimeError("No valid tasks found in config.glue_tasks (must be in GLUE_TASKS or EXTRA_TASKS).")

    batch_size = int(getattr(config, "batch_size", 64))
    max_length = int(getattr(config, "max_length", 128))
    device = getattr(config, "device", "cuda")
    auto_train = bool(getattr(config, "auto_train", True))
    train_epochs = int(getattr(config, "train_epochs", 3))
    train_epochs_per_task = getattr(config, "train_epochs_per_task", {})
    if not train_epochs_per_task:
        train_epochs_per_task = {
            "cola": 5, "mrpc": 3, "rte": 5, "stsb": 3, "sst2": 3, "qqp": 3, "qnli": 3, "mnli": 3, "wnli": 5,
            "piqa": 3, "winogrande": 3, "boolq": 3, "hellaswag": 3, "openbookqa": 3, "arc_easy": 3, "arc_challenge": 3
        }

    train_batch_size = int(getattr(config, "train_batch_size", 32))
    train_warmup_steps = int(getattr(config, "train_warmup_steps", 100))
    train_weight_decay = float(getattr(config, "train_weight_decay", 0.01))
    train_grad_accum_steps = int(getattr(config, "train_grad_accum_steps", 1))
    lr_candidates = getattr(config, "bert_lr_candidates", BERT_LR_CANDIDATES)
    random_restarts_small = int(getattr(config, "random_restarts_small", 1))
    base_seed = int(getattr(config, "base_seed", 543211))

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

    # try to save original model state so we can reset between runs
    original_state = None
    try:
        original_state = unwrap_model(model).state_dict()
    except Exception:
        try:
            original_state = model.state_dict()
        except Exception:
            original_state = None

    def _reset_model_to_original():
        if original_state is None:
            return
        try:
            unwrap_model(model).load_state_dict(original_state, strict=False)
        except Exception:
            try:
                model.load_state_dict(original_state, strict=False)
            except Exception:
                pass

    summary_rows = []

    for task in tasks:
        assert (task in GLUE_TASKS) or (task in EXTRA_TASKS), f"Unknown task '{task}'"
        print(f"\n==== Grid-running task: {task} ====")
        epochs_this_task = int(train_epochs_per_task.get(task, train_epochs))
        print(f"[Grid] Epochs for task '{task}': {epochs_this_task}")

        # load data splits
        if task in EXTRA_TASKS:
            cfg = EXTRA_TASKS[task]
            try:
                if "hf_config" in cfg:
                    ds = load_dataset(cfg["hf_path"], cfg["hf_config"])
                else:
                    ds = load_dataset(cfg["hf_path"])
            except Exception as e:
                raise RuntimeError(f"Failed to load dataset for extra task {task}: {e}")
            # pick train/val splits
            train_raw = ds.get("train")
            # prefer validation if available, else test
            val_raw = ds.get("validation") or ds.get("test") or ds.get("validation_matched")
            if val_raw is None:
                # pick first available split as val
                val_raw = next(iter(ds.values()))
        else:
            hf = load_dataset("glue", GLUE_TASKS[task]["hf_name"])
            train_raw = hf["train"]
            val_raw = hf["validation_matched"] if task == "mnli" else hf["validation"]

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

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

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

        if auto_train:
            print(f"[Grid] Auto-training. LRs={lr_candidates}. Restarts/LR={restarts_per_lr} (small={small_flag}).")
            for lr in lr_candidates:
                for restart_idx in range(restarts_per_lr):
                    seed = base_seed + (abs(hash(task)) % 10000) * 1000 + int(restart_idx) * 10 + (int(round(lr * 1e7)) % 1000)
                    print(f"\n[SWEEP] task={task} lr={lr:g} restart={restart_idx}/{restarts_per_lr-1} seed={seed}")
                    _reset_model_to_original()
                    if checkpointing is not None:
                        try:
                            checkpointing.load_model_states("recent")
                        except Exception:
                            pass

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

                    try:
                        if task in EXTRA_TASKS:
                            # train extra
                            model, metric_res = train_full_finetune_extra(
                                task=task, tokenizer=tokenizer, model=model,
                                raw_train=train_raw, raw_val=val_raw,
                                device=device, epochs=epochs_this_task,
                                batch_size=train_batch_size, lr=float(lr),
                                weight_decay=train_weight_decay, warmup_steps=train_warmup_steps,
                                max_length=max_length, grad_accum_steps=train_grad_accum_steps,
                                out_checkpoint_dir=str(run_dir)
                            )
                            sc = _metric_to_scalar(task, metric_res, fallback_val_loss=None)
                            best_epoch = None
                        else:
                            model, metric_res, score, best_epoch = train_full_finetune(
                                task=task, tokenizer=tokenizer, model=model,
                                raw_train=train_raw, raw_val=val_raw,
                                device=device, epochs=epochs_this_task,
                                batch_size=train_batch_size, lr=float(lr),
                                weight_decay=train_weight_decay, warmup_steps=train_warmup_steps,
                                max_length=max_length, grad_accum_steps=train_grad_accum_steps,
                                out_checkpoint_dir=str(run_dir), seed=int(seed)
                            )
                            sc = float(score)
                        # update run_meta.json with lr/restart/seed
                        meta_path = Path(run_dir) / "run_meta.json"
                        meta = {}
                        if meta_path.exists():
                            try:
                                meta = json.loads(meta_path.read_text(encoding="utf-8"))
                            except Exception:
                                meta = {}
                        meta.update({"restart": int(restart_idx), "seed": int(seed), "lr": float(lr)})
                        _json_dump(meta, meta_path)

                        rec = {
                            "task": task,
                            "lr": float(lr),
                            "restart": int(restart_idx),
                            "seed": int(seed),
                            "epochs": int(epochs_this_task),
                            "dev_score": float(sc) if sc is not None else None,
                            "dev_metrics": json.dumps(metric_res),
                            "run_dir": str(run_dir),
                            "best_ckpt_path": str(Path(run_dir) / "best_finetuned.pt") if (Path(run_dir) / "best_finetuned.pt").exists() else None,
                            "final_ckpt_path": str(Path(run_dir) / "finetuned.pt") if (Path(run_dir) / "finetuned.pt").exists() else None,
                        }
                        all_run_records.append(rec)

                        if best_run["score"] is None or (sc is not None and float(sc) > float(best_run["score"])):
                            best_run.update({
                                "score": float(sc) if sc is not None else None,
                                "metrics": metric_res,
                                "lr": float(lr),
                                "restart": int(restart_idx),
                                "seed": int(seed),
                                "best_epoch": int(best_epoch) if best_epoch is not None else None,
                                "run_dir": str(run_dir),
                                "best_ckpt_path": rec["best_ckpt_path"],
                            })
                    except Exception as e:
                        print(f"[WARN] Training run failed for {task} lr={lr} restart={restart_idx}: {e}")
                        all_run_records.append({
                            "task": task, "lr": float(lr), "restart": int(restart_idx), "seed": int(seed),
                            "epochs": int(epochs_this_task), "dev_score": None,
                            "dev_metrics": json.dumps({"error": str(e)}), "run_dir": str(run_dir),
                            "best_ckpt_path": None, "final_ckpt_path": None,
                        })

        # Save all_runs.csv for this task
        all_runs_csv = task_ckpt_root / "all_runs.csv"
        pd.DataFrame(all_run_records).to_csv(all_runs_csv, index=False)
        print(f"[Grid] Saved all runs summary to: {all_runs_csv}")

        # Save best_overall
        best_overall_dir = task_ckpt_root / "best_overall"
        best_overall_dir.mkdir(parents=True, exist_ok=True)
        if best_run.get("best_ckpt_path") and best_run["best_ckpt_path"] and os.path.exists(best_run["best_ckpt_path"]):
            try:
                shutil.copy2(best_run["best_ckpt_path"], best_overall_dir / "best_finetuned.pt")
            except Exception:
                pass
        _json_dump(best_run, best_overall_dir / "best_meta.json")
        print(f"[Grid] Best run for task='{task}': lr={best_run.get('lr')}, restart={best_run.get('restart')}, seed={best_run.get('seed')}, dev_score={best_run.get('score')}")

        # Reset model and load best for evaluation
        _reset_model_to_original()
        model_for_eval = model
        try:
            best_ckpt = best_overall_dir / "best_finetuned.pt"
            if best_ckpt.exists():
                model_for_eval = _load_finetuned_checkpoint_for_task(task, model, str(best_ckpt))
        except Exception as e:
            print(f"[WARN] Failed to load best_overall checkpoint for eval; using current model. err={e}")
            model_for_eval = model

        # Evaluate and write summary rows
        task_out_dir = out_dir / task; task_out_dir.mkdir(parents=True, exist_ok=True)
        if task in EXTRA_TASKS:
            metric_res = run_extra_task(task=task, tokenizer=tokenizer, model=model_for_eval,
                                        checkpointing=None, device=device, batch_size=batch_size,
                                        max_length=max_length, output_dir=str(task_out_dir))
            summary_rows.append({"task": task, "split": "selected", "epochs": epochs_this_task, "metrics": json.dumps(metric_res), "selected_lr": best_run.get("lr")})
        else:
            res = run_glue_task(task=task, tokenizer=tokenizer, model=model_for_eval,
                                checkpointing=None, device=device, batch_size=batch_size,
                                max_length=max_length, output_dir=str(task_out_dir), compute_errorbar=True)
            for split, pack in res.items():
                metrics = pack["metrics"]
                eb = pack["errorbar"]
                summary_rows.append({
                    "task": task,
                    "split": split,
                    "epochs": epochs_this_task,
                    "selected_lr": best_run.get("lr"),
                    "selected_restart": best_run.get("restart"),
                    "selected_seed": best_run.get("seed"),
                    "selected_dev_score": best_run.get("score"),
                    "eval_metrics": json.dumps(metrics),
                    "errorbar_mean": (eb["mean"] if eb else None),
                    "errorbar_std": (eb["std"] if eb else None),
                    "errorbar_stderr": (eb["stderr"] if eb else None),
                    "errorbar_detail": json.dumps(eb) if eb else None,
                })

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

# ---------------------------------------------------------------------
# CLI shim (optional)
# ---------------------------------------------------------------------
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--tasks", type=str, default="sst2", help="comma separated tasks (supports glue and extra tasks: boolq,piqa,winogrande,hellaswag,openbookqa,arc_easy,arc_challenge)")
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--max_length", type=int, default=128)
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--out_dir", type=str, default="glue_outputs_grid")
    args = parser.parse_args()

    print("This module is intended to be invoked from your project's main which provides tokenizer/model/checkpointing.")
    print(f"LI args tasks={args.tasks} batch_size={args.batch_size} max_length={args.max_length} device={args.device} out_dir={args.out_dir}")
    # Example usage (pseudo):
    # from transformers import AutoTokenizer, AutoModelForSequenceClassification
    # tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    # model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
    # cfg = type("C", (), {"glue_tasks": args.tasks.split(","), "batch_size": args.batch_size, "max_length": args.max_length, "device": args.device})
    # run_glue_benchmark(cfg, tokenizer, model, checkpointing=None, out_dir=args.out_dir)