Add benchmark harness: tasks.py - Standard NLP task definitions

benchmark/tasks.py

@@ -0,0 +1,336 @@
+"""
+Standard NLP benchmark task definitions.
+
+Each task loads from HuggingFace datasets and formats examples for
+log-likelihood scoring.
+
+Supported tasks:
+  - HellaSwag (commonsense NLI, 4-choice)
+  - ARC-Easy / ARC-Challenge (science QA, 3-5 choices)
+  - PIQA (physical intuition, 2-choice)
+  - WinoGrande (coreference, 2-choice)
+  - MMLU (multi-domain knowledge, 4-choice)
+  - HaluEval-QA (hallucination detection, 2-choice)
+"""
+
+import random
+from abc import ABC, abstractmethod
+from typing import List, Dict, Tuple, Optional, Any
+from datasets import load_dataset
+
+
+class BenchmarkTask(ABC):
+    """Base class for benchmark tasks."""
+    
+    name: str = "base"
+    num_few_shot: int = 0
+    
+    @abstractmethod
+    def load_examples(self, n: Optional[int] = None, seed: int = 42) -> List[Dict]:
+        """
+        Load and format examples.
+        
+        Returns list of dicts, each with:
+          - "context": str — the prompt/context
+          - "continuations": List[str] — possible completions
+          - "gold_idx": int — index of the correct continuation
+        """
+        ...
+    
+    def format_few_shot(self, examples: List[Dict], train_examples: List[Dict]) -> List[Dict]:
+        """Prepend few-shot examples to each test example's context."""
+        if not train_examples or self.num_few_shot == 0:
+            return examples
+        
+        # Build few-shot prefix
+        shots = train_examples[:self.num_few_shot]
+        prefix = ""
+        for shot in shots:
+            prefix += shot["context"] + shot["continuations"][shot["gold_idx"]] + "\n\n"
+        
+        for ex in examples:
+            ex["context"] = prefix + ex["context"]
+        
+        return examples
+
+
+class HellaSwag(BenchmarkTask):
+    """
+    HellaSwag: Can a Machine Really Finish Your Sentence?
+    4-choice commonsense NLI. Dataset: Rowan/hellaswag
+    """
+    name = "hellaswag"
+    num_few_shot = 5
+    
+    def load_examples(self, n: Optional[int] = None, seed: int = 42) -> List[Dict]:
+        ds = load_dataset("Rowan/hellaswag", split="validation")
+        
+        if n is not None:
+            ds = ds.shuffle(seed=seed).select(range(min(n, len(ds))))
+        
+        examples = []
+        few_shot_ds = load_dataset("Rowan/hellaswag", split="train")
+        few_shot_ds = few_shot_ds.shuffle(seed=seed).select(range(self.num_few_shot))
+        
+        train_examples = []
+        for row in few_shot_ds:
+            ctx = row["ctx"]
+            endings = row["endings"]
+            gold = int(row["label"])
+            train_examples.append({
+                "context": ctx,
+                "continuations": endings,
+                "gold_idx": gold,
+            })
+        
+        for row in ds:
+            ctx = row["ctx"]
+            endings = row["endings"]
+            gold = int(row["label"])
+            examples.append({
+                "context": ctx,
+                "continuations": endings,
+                "gold_idx": gold,
+            })
+        
+        return self.format_few_shot(examples, train_examples)
+
+
+class ARC(BenchmarkTask):
+    """
+    AI2 Reasoning Challenge. Dataset: allenai/ai2_arc
+    Subsets: ARC-Easy, ARC-Challenge
+    """
+    name = "arc"
+    num_few_shot = 5
+    
+    def __init__(self, subset: str = "ARC-Easy"):
+        self.subset = subset
+        self.name = f"arc-{'easy' if 'Easy' in subset else 'challenge'}"
+    
+    def load_examples(self, n: Optional[int] = None, seed: int = 42) -> List[Dict]:
+        ds = load_dataset("allenai/ai2_arc", self.subset, split="test")
+        
+        if n is not None:
+            ds = ds.shuffle(seed=seed).select(range(min(n, len(ds))))
+        
+        # Few-shot from train
+        train_ds = load_dataset("allenai/ai2_arc", self.subset, split="train")
+        train_ds = train_ds.shuffle(seed=seed).select(range(self.num_few_shot))
+        
+        def format_row(row):
+            question = row["question"]
+            choices = row["choices"]
+            labels = choices["label"]
+            texts = choices["text"]
+            answer_key = row["answerKey"]
+            
+            # Map answer key to index
+            gold_idx = labels.index(answer_key) if answer_key in labels else 0
+            
+            # Format as "Question: ...\nA) ... B) ...\nAnswer:"
+            choice_str = " ".join(f"{l}) {t}" for l, t in zip(labels, texts))
+            context = f"Question: {question}\n{choice_str}\nAnswer:"
+            
+            continuations = [f" {t}" for t in texts]
+            
+            return {
+                "context": context,
+                "continuations": continuations,
+                "gold_idx": gold_idx,
+            }
+        
+        train_examples = [format_row(row) for row in train_ds]
+        examples = [format_row(row) for row in ds]
+        
+        return self.format_few_shot(examples, train_examples)
+
+
+class PIQA(BenchmarkTask):
+    """
+    Physical Intuition QA. 2-choice.
+    Dataset: gimmaru/piqa (parquet mirror — ybisk/piqa loading script no longer supported)
+    """
+    name = "piqa"
+    num_few_shot = 0  # No train split in mirror; use 0-shot
+    
+    def load_examples(self, n: Optional[int] = None, seed: int = 42) -> List[Dict]:
+        ds = load_dataset("gimmaru/piqa", split="validation")
+        
+        if n is not None:
+            ds = ds.shuffle(seed=seed).select(range(min(n, len(ds))))
+        
+        def format_row(row):
+            goal = row["goal"]
+            sol1 = row["sol1"]
+            sol2 = row["sol2"]
+            gold = row["label"]  # 0 or 1
+            
+            context = f"Goal: {goal}\nSolution 1: {sol1}\nSolution 2: {sol2}\nThe better solution is Solution"
+            continuations = [" 1", " 2"]
+            
+            return {
+                "context": context,
+                "continuations": continuations,
+                "gold_idx": gold,
+            }
+        
+        examples = [format_row(row) for row in ds]
+        return examples
+
+
+class WinoGrande(BenchmarkTask):
+    """
+    WinoGrande: Winograd-style coreference. 2-choice.
+    Dataset: allenai/winogrande (winogrande_xl)
+    """
+    name = "winogrande"
+    num_few_shot = 5
+    
+    def load_examples(self, n: Optional[int] = None, seed: int = 42) -> List[Dict]:
+        ds = load_dataset("allenai/winogrande", "winogrande_xl", split="validation")
+        
+        if n is not None:
+            ds = ds.shuffle(seed=seed).select(range(min(n, len(ds))))
+        
+        train_ds = load_dataset("allenai/winogrande", "winogrande_xl", split="train")
+        train_ds = train_ds.shuffle(seed=seed).select(range(self.num_few_shot))
+        
+        def format_row(row):
+            sentence = row["sentence"]
+            option1 = row["option1"]
+            option2 = row["option2"]
+            answer = int(row["answer"]) - 1  # 1-indexed -> 0-indexed
+            
+            # Replace _ with each option
+            sent1 = sentence.replace("_", option1)
+            sent2 = sentence.replace("_", option2)
+            
+            context = f"Which makes more sense?\nA) {sent1}\nB) {sent2}\nAnswer:"
+            continuations = [" A", " B"]
+            
+            return {
+                "context": context,
+                "continuations": continuations,
+                "gold_idx": answer,
+            }
+        
+        train_examples = [format_row(row) for row in train_ds]
+        examples = [format_row(row) for row in ds]
+        
+        return self.format_few_shot(examples, train_examples)
+
+
+class MMLU(BenchmarkTask):
+    """
+    Massive Multitask Language Understanding. 4-choice.
+    Dataset: cais/mmlu (all subjects)
+    """
+    name = "mmlu"
+    num_few_shot = 5
+    
+    def __init__(self, subject: Optional[str] = None):
+        """If subject is None, sample across all subjects."""
+        self.subject = subject
+        if subject:
+            self.name = f"mmlu-{subject}"
+    
+    def load_examples(self, n: Optional[int] = None, seed: int = 42) -> List[Dict]:
+        if self.subject:
+            ds = load_dataset("cais/mmlu", self.subject, split="test")
+            train_ds = load_dataset("cais/mmlu", self.subject, split="validation")
+        else:
+            ds = load_dataset("cais/mmlu", "all", split="test")
+            train_ds = load_dataset("cais/mmlu", "all", split="validation")
+        
+        if n is not None:
+            ds = ds.shuffle(seed=seed).select(range(min(n, len(ds))))
+        
+        train_ds = train_ds.shuffle(seed=seed).select(range(min(self.num_few_shot, len(train_ds))))
+        
+        def format_row(row):
+            question = row["question"]
+            choices = row["choices"]
+            answer = row["answer"]  # 0-3
+            
+            labels = ["A", "B", "C", "D"]
+            choice_str = "\n".join(f"{l}) {c}" for l, c in zip(labels, choices))
+            context = f"Question: {question}\n{choice_str}\nAnswer:"
+            continuations = [f" {l}" for l in labels]
+            
+            return {
+                "context": context,
+                "continuations": continuations,
+                "gold_idx": answer,
+            }
+        
+        train_examples = [format_row(row) for row in train_ds]
+        examples = [format_row(row) for row in ds]
+        
+        return self.format_few_shot(examples, train_examples)
+
+
+class HaluEval(BenchmarkTask):
+    """
+    HaluEval: Hallucination Evaluation.
+    Dataset: pminervini/HaluEval (qa_samples)
+    
+    Tests whether the model can identify hallucinated answers.
+    """
+    name = "halueval"
+    num_few_shot = 2
+    
+    def load_examples(self, n: Optional[int] = None, seed: int = 42) -> List[Dict]:
+        ds = load_dataset("pminervini/HaluEval", "qa_samples", split="data")
+        
+        if n is not None:
+            ds = ds.shuffle(seed=seed).select(range(min(n, len(ds))))
+        
+        examples = []
+        for row in ds:
+            question = row["question"]
+            knowledge = row.get("knowledge", "")
+            right_answer = row.get("right_answer", "")
+            hallucinated_answer = row.get("hallucinated_answer", "")
+            
+            if not right_answer or not hallucinated_answer:
+                continue
+            
+            # Randomly order the options
+            rng = random.Random(seed + len(examples))
+            options = [(right_answer, 0), (hallucinated_answer, 1)]
+            if rng.random() > 0.5:
+                options = options[::-1]
+            
+            # Gold is the correct (non-hallucinated) answer
+            gold_idx = 0 if options[0][1] == 0 else 1
+            
+            context_parts = [f"Question: {question}"]
+            if knowledge:
+                context_parts.insert(0, f"Knowledge: {knowledge[:300]}")
+            context_parts.append(f"Answer A: {options[0][0][:200]}")
+            context_parts.append(f"Answer B: {options[1][0][:200]}")
+            context_parts.append("Which answer is correct? Answer:")
+            
+            context = "\n".join(context_parts)
+            continuations = [" A", " B"]
+            
+            examples.append({
+                "context": context,
+                "continuations": continuations,
+                "gold_idx": gold_idx,
+            })
+        
+        return examples
+
+
+# Task registry for easy lookup
+TASK_REGISTRY = {
+    "hellaswag": HellaSwag,
+    "arc-easy": lambda: ARC("ARC-Easy"),
+    "arc-challenge": lambda: ARC("ARC-Challenge"),
+    "piqa": PIQA,
+    "winogrande": WinoGrande,
+    "mmlu": MMLU,
+    "halueval": HaluEval,
+}