# lmr/glue_benchmark.py """ Standalone GLUE benchmark runner for your project. Updated fixes: - task-aware example extraction (fixes MNLI/QQP/other empty-text problems) - robust tokenizer handling + debug empty-sample logging - robust wrapper and postprocessing for predictions (fixes STS-B / QQP metric mismatches) - auto-train per-task (full fine-tune) saving under out_dir/checkpoints/ """ import os import json import re import time from pathlib import Path from typing import Optional, List, Tuple import torch import numpy as np import pandas as pd from datasets import load_dataset from torch.utils.data import DataLoader, TensorDataset from tqdm import tqdm import evaluate # Project imports (adjust if your package layout differs) from lmr.checkpointing import Checkpointing from lmr.ddp import unwrap_model # --------------------------------------------------------------------- # GLUE config # --------------------------------------------------------------------- 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"}, } # --------------------------------------------------------------------- # Task-aware example field extraction (fixes empty-text issues) # --------------------------------------------------------------------- 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. """ # Known per-task fields task_field_map = { "cola": ("sentence", None), "sst2": ("sentence", None), "mrpc": ("sentence1", "sentence2"), "stsb": ("sentence1", "sentence2"), "qqp": ("question1", "question2"), "mnli": ("premise", "hypothesis"), "qnli": ("question", "sentence"), "rte": ("sentence1", "sentence2"), "wnli": ("sentence1", "sentence2"), } f1, f2 = task_field_map.get(task, (None, None)) def _try_keys(keys): for k in keys: if k in ex and ex.get(k) is not None: return ex.get(k) return None # Compose candidate key lists for robustness 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) # Final fallback: try fields common in some GLUE subsets if s1 is None: s1 = ex.get("sentence") or ex.get("premise") or ex.get("question") or ex.get("text") if s2 is None: # allow None for single-sentence tasks s2 = ex.get("sentence2") or ex.get("hypothesis") or ex.get("question2") # Normalize types 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) # normalize to python lists 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() # debug tag # print("[tokenizer] used hf-style batch tokenizer") 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: # try pair 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)) # interpret return 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: # try tokenizer(...) convenience 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) # truncation if len(ids) > max_length: ids = ids[:max_length] mask = mask[:max_length] input_ids_list.append(ids) attention_mask_list.append(mask) except Exception as e: snippet = (a[:80] + "...") if a else "" raise RuntimeError(f"Tokenizer fallback encode failed for example '{snippet}': {e}") return {"input_ids": input_ids_list, "attention_mask": attention_mask_list} # --------------------------------------------------------------------- # Postprocess preds to the right shapes/types (fixes metric mismatches) # --------------------------------------------------------------------- def _postprocess_predictions(task: str, logits_np: np.ndarray, cfg_task: dict): """ Take logits (N, C) or (N,) or (N,1) and produce preds array ready for evaluate.compute: - classification -> 1D ints (class indices or binary 0/1) - regression -> 1D floats (for stsb typically 0..5) """ ttype = cfg_task["type"] num_labels = cfg_task["num_labels"] if logits_np is None or logits_np.size == 0: return np.array([]) # If logits are shape (N, ) -> treat as single score per example (binary/regression) if logits_np.ndim == 1: if ttype == "classification": # binary: threshold at 0.5 for scores in [0,1] or sign threshold 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": # argmax -> class index preds = np.argmax(logits_np, axis=-1).astype(int) return preds else: # regression: if multiple dims, average or take first if logits_np.shape[1] == 1: preds = logits_np[:, 0].astype(float) else: preds = logits_np.mean(axis=1).astype(float) # clamp STS-B to 0..5 if it's that task (defensive) if task == "stsb": preds = np.clip(preds, 0.0, 5.0) return preds # fallback 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 import re 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 existing head already matches desired num_labels -> reuse 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 # infer hidden_size if not provided 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 # try unwrap and inspect params/state_dict shapes if inferred_hidden is None: try: un = unwrap_model(base_model) sd = un.state_dict() # search for embedding weight shapes 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 cand > 1 and cand < 1000000: 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` (e.g. model.config.hidden_size = 1024) " "or pass `hidden_size` explicitly when calling make_wrapped_model_if_needed." ) 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 several base calling conventions try: out = self.base(input_ids=input_ids, attention_mask=attention_mask, **kwargs) except TypeError: out = self.base(input_ids) # last_hidden_state 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}) # tuple/list case 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}) # out.logits present 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)}) # project if dims mismatch 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 attribute hidden_states = getattr(out, "hidden_states", None) if hidden_states is not None: if isinstance(hidden_states, (list, tuple)): last_hidden = hidden_states[-1] else: last_hidden = 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 dataset split into tensors (used for finetune) # --------------------------------------------------------------------- def _tokenize_hf_split_to_tensors(task: str, tokenizer, raw_split, cfg_task, max_length=128, batch_tokenize_size=512): """ Convert HF dataset split to tensors (input_ids tensor, attention_mask tensor, labels tensor) """ 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): batch_texts = texts[i:i+batch_tokenize_size] enc = _batch_tokenize(tokenizer, batch_texts, 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 # --------------------------------------------------------------------- # Train full fine-tune (entire model) for a GLUE task # --------------------------------------------------------------------- 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): """ Fine-tune the full model on the task train set, validate on raw_val. Saves final state_dict to out_checkpoint_dir/finetuned.pt if provided. Returns the fine-tuned model (in-place) and a dict with final validation metrics. """ cfg_task = GLUE_TASKS[task] device = torch.device(device if torch.cuda.is_available() else "cpu") # Wrap/create classifier if necessary (force correct output dim) 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 wrapped_model, wrapped_flag = make_wrapped_model_if_needed(model, hidden_size, cfg_task["num_labels"], force_num_labels=cfg_task["num_labels"]) model = wrapped_model model.to(device) # Prepare tensors (task-aware) 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) 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 & scheduler 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() 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"Train {task} epoch {epoch+1}")): 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 finetune") # compute loss if cfg_task["type"] == "classification": # logits shape (B, C) loss = loss_fn(logits, labs_b.long()) else: # regression 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() global_step += 1 running_loss += loss.item() * ids_b.size(0) # validation at epoch end 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); 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 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) # compute metric with resilient handling 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)} print(f"[FT] {task} epoch {epoch+1} val_loss={tot_val_loss/len(val_ds):.6f} metric={metric_res}") final_metric_res = metric_res # Save final full model state_dict if desired 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}") return model, final_metric_res # --------------------------------------------------------------------- # Main per-task runner (evaluation only) # --------------------------------------------------------------------- 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"): """ Run a single GLUE task evaluation. Returns metric dict per split. """ assert task in GLUE_TASKS, f"Unknown GLUE task: {task}" cfg = GLUE_TASKS[task] # load HF dataset hf = load_dataset("glue", cfg["hf_name"]) # choose split(s): mnli has two val splits 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)}") # prepare texts and labels (task-aware) 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 tokenize robustly BATCH = 512 input_ids_all = [] attention_all = [] for i in range(0, len(texts), BATCH): batch_texts = texts[i:i+BATCH] enc = _batch_tokenize(tokenizer, batch_texts, 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 and tensorize 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) # optionally load checkpoint (recent pretrained) if provided if checkpointing is not None: try: checkpointing.load_model_states("recent") except Exception: pass device = torch.device(device if torch.cuda.is_available() else "cpu") model.to(device) 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 # decide wrapper force logic force = None if cfg["type"] == "regression": force = 1 else: try: import torch.nn as nn existing_dim = None if hasattr(model, "classifier") and isinstance(getattr(model, "classifier"), nn.Linear): existing_dim = getattr(model, "classifier").out_features elif hasattr(model, "lm_head") and isinstance(getattr(model, "lm_head"), nn.Linear): existing_dim = getattr(model, "lm_head").out_features if existing_dim is not None and existing_dim != cfg["num_labels"]: force = cfg["num_labels"] except Exception: force = cfg["num_labels"] wrapped_model, wrapped = make_wrapped_model_if_needed(model, hidden_size, cfg["num_labels"], force_num_labels=force) wrapped_model.to(device) 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) mask_b = mask_b.to(device) 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") logits_np = logits.detach().cpu().numpy() all_logits.append(logits_np) 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) # ensure preds shape compatible with evaluate if cfg["type"] == "classification": preds_out = preds.astype(int).tolist() refs_out = all_labels.astype(int).tolist() else: preds_out = preds.astype(float).tolist() refs_out = all_labels.astype(float).tolist() metric = evaluate.load("glue", cfg["hf_name"]) try: metric_res = metric.compute(predictions=preds_out, references=refs_out) except Exception: # second attempt with numpy arrays (some metrics accept np) try: metric_res = metric.compute(predictions=np.array(preds_out), references=np.array(refs_out)) except Exception as e: metric_res = {"error": str(e)} 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) csv_p = Path(output_dir) / f"{task}_{split}_preds.csv" pd.DataFrame({"pred": preds_out, "label": refs_out}).to_csv(csv_p, index=False) results_by_split[split] = metric_res return results_by_split # --------------------------------------------------------------------- # Top-level runner: multiple tasks + (optional) auto-train # --------------------------------------------------------------------- def run_glue_benchmark(config, tokenizer, model, checkpointing: Optional[Checkpointing] = None, out_dir: str = "glue_outputs"): """ config: object/dict with fields: - glue_tasks: list of task names (e.g. ["sst2","mnli"]) - batch_size: int - max_length: int - device: "cuda" or "cpu" (optional) - auto_train: bool (if True, train per-task before eval) - train_epochs, train_batch_size, train_lr, train_warmup_steps, train_weight_decay, train_grad_accum_steps """ tasks = getattr( config, "glue_tasks", ["sst2", "cola", "mrpc", "stsb", "qqp", "mnli", "qnli", "rte", "wnli"]) # "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"}, # ["wnli", "rte", "stsb", "mrpc", "cola", "sst2", "qnli", "qqp", "mnli"] # tasks = [tasks[2]] # tasks = getattr(config, "glue_tasks", [ "qqp", "mnli", "qnli"]) batch_size = getattr(config, "batch_size", 64) max_length = getattr(config, "max_length", 128) device = getattr(config, "device", "cuda") # training hyperparams for auto_train auto_train = getattr(config, "auto_train", True) train_epochs = getattr(config, "train_epochs", 3) train_batch_size = getattr(config, "train_batch_size", 32) train_lr = getattr(config, "train_lr", 2e-5) 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) 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 = [] # retain original pretrained model state so each task starts from same point original_state = None try: original_state = unwrap_model(model).state_dict() except Exception: try: original_state = model.state_dict() except Exception: original_state = None for task in tasks: print(f"\n==== Running GLUE task: {task} ====") # load hf dataset splits hf = load_dataset("glue", GLUE_TASKS[task]["hf_name"]) train_raw = hf["train"] if task == "mnli": val_raw = hf["validation_matched"] else: val_raw = hf["validation"] # If auto_train: reload pretrained checkpoint and finetune whole model if auto_train: print(f"[GLUE] Auto-training enabled. Loading pretrained checkpoint (if any) then fine-tuning for task '{task}'") # reset model params to original pretrained snapshot (so each task starts same) if original_state is not None: try: m_unwrap = unwrap_model(model) m_unwrap.load_state_dict(original_state, strict=False) except Exception: try: model.load_state_dict(original_state, strict=False) except Exception: pass # also try checkpointing loader if checkpointing is not None: try: ckdir = checkpointing.load_model_states("recent") # If checkpointing returned path, try to load weights robustly (if trainer helper available) # Trainer's load_only_model_weights may exist elsewhere; we avoid circular import here. except Exception: pass task_ckpt_dir = checkpoint_out_root / task task_ckpt_dir.mkdir(parents=True, exist_ok=True) model, metric_res = train_full_finetune( task, tokenizer, model, train_raw, val_raw, device=device, epochs=train_epochs, batch_size=train_batch_size, lr=train_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(task_ckpt_dir) ) print(f"[GLUE] Finished fine-tuning for task {task}. Val metric: {metric_res}") else: if checkpointing is not None: try: checkpointing.load_model_states("recent") except Exception: pass # Evaluate (uses the current model in memory) task_out_dir = out_dir / task task_out_dir.mkdir(parents=True, exist_ok=True) res = run_glue_task(task, tokenizer, model, checkpointing=None, device=device, batch_size=batch_size, max_length=max_length, output_dir=str(task_out_dir)) for split, metrics in res.items(): metric_str = json.dumps(metrics) rows.append({"task": task, "split": split, "metrics": metric_str}) summary_csv = out_dir / "glue_summary.csv" pd.DataFrame(rows).to_csv(summary_csv, index=False) print(f"\nGLUE summary saved to: {summary_csv}") return pd.DataFrame(rows) # --------------------------------------------------------------------- # CLI for quick testing (optional) # --------------------------------------------------------------------- if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument("--tasks", type=str, default="sst2", help="comma separated glue tasks") 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") args = parser.parse_args() print("This module is intended to be invoked from your project's main, which provides tokenizer/model/checkpointing.") print(f"Example usage in your main: run_glue_benchmark(config.benchmark, tokenizer, model, checkpointing, out_dir={args.out_dir})")