almanach/benchmark-in-a-haystack

1from abc import ABC, abstractmethod2import random3from datasets import load_dataset4 5class Benchmark(ABC):6    @abstractmethod7    def load_samples(self, count=5, subjects=None):8        pass9 10    @abstractmethod11    def format_sample(self, sample, subject=None):12        pass13 14class MMLUBenchmark(Benchmark):15    dataset = "cais/mmlu"16    split = "test"17    format_template = "Subject: {subject}\nQuestion: {question}\n{choices}\nAnswer: {answer}"18 19    def load_samples(self, count=5, subjects=None):20        samples = []21        if not subjects:22            raise ValueError("MMLU requires subjects")23        for subject in subjects:24            dataset = load_dataset(self.dataset, subject, split=self.split)25            for idx in range(count):26                samples.append({27                    "subject": subject,28                    "data": dataset[idx],29                    "benchmark_type": "mmlu"30                })31        return samples32 33    def format_sample(self, sample, subject=None):34        data = sample["data"]35        question = data["question"]36        answer = chr(65 + data["answer"])37        choices = "\n".join([f"{chr(65+j)}. {choice}" for j, choice in enumerate(data["choices"])])38        subject = subject or sample.get("subject")39        return self.format_template.format(subject=subject, question=question, choices=choices, answer=answer)40 41class GSM8KBenchmark(Benchmark):42    dataset = "openai/gsm8k"43    name = "main"44    split = "test"45    format_template = "Math Problem: {question}\n\nSolution: {answer}"46 47    def load_samples(self, count=5, subjects=None):48        dataset = load_dataset(self.dataset, name=self.name, split=self.split)49        indices = random.sample(range(len(dataset)), count)50        return [{"data": dataset[i], "benchmark_type": "gsm8k"} for i in indices]51 52    def format_sample(self, sample, subject=None):53        data = sample["data"]54        return self.format_template.format(question=data["question"], answer=data["answer"])55 56class GPQABenchmark(Benchmark):57    dataset = "hendrydong/gpqa_diamond"58    split = "test"59    format_template = "Problem:\n{problem}\n\nSolution:\n{solution}"60 61    def load_samples(self, count=5, subjects=None):62        dataset = load_dataset(self.dataset, split=self.split)63        indices = random.sample(range(len(dataset)), count)64        return [{"data": dataset[i], "benchmark_type": "gpqa"} for i in indices]65 66    def format_sample(self, sample, subject=None):67        data = sample["data"]68        return self.format_template.format(problem=data["problem"], solution=data["solution"])69 70class ARCChallengeBenchmark(Benchmark):71    dataset = "allenai/ai2_arc"72    config = "ARC-Challenge"73    split = "test"74    format_template = "Question: {question}\n{choices}\nAnswer: {answer}"75 76    def load_samples(self, count=5, subjects=None):77        dataset = load_dataset(self.dataset, self.config, split=self.split)78        indices = random.sample(range(len(dataset)), min(count, len(dataset)))79        return [{"data": dataset[i], "benchmark_type": "arc_challenge"} for i in indices]80 81    def format_sample(self, sample, subject=None):82        data = sample["data"]83        choices = "\n".join([f"{label}. {text}" for label, text in zip(data['choices']['label'], data['choices']['text'])])84        return self.format_template.format(question=data["question"], choices=choices, answer=data["answerKey"])85 86class ARCEasyBenchmark(Benchmark):87    dataset = "allenai/ai2_arc"88    config = "ARC-Easy"89    split = "test"90    format_template = "Question: {question}\n{choices}\nAnswer: {answer}"91 92    def load_samples(self, count=5, subjects=None):93        dataset = load_dataset(self.dataset, self.config, split=self.split)94        indices = random.sample(range(len(dataset)), min(count, len(dataset)))95        return [{"data": dataset[i], "benchmark_type": "arc_easy"} for i in indices]96 97    def format_sample(self, sample, subject=None):98        data = sample["data"]99        choices = "\n".join([f"{label}. {text}" for label, text in zip(data['choices']['label'], data['choices']['text'])])100        return self.format_template.format(question=data["question"], choices=choices, answer=data["answerKey"])101 102class HellaSwagBenchmark(Benchmark):103    dataset = "Rowan/hellaswag"104    split = "validation"105    format_template = "Context: {context}\n\nChoose the most plausible continuation:\n{endings}\nAnswer: {answer}"106 107    def load_samples(self, count=5, subjects=None):108        dataset = load_dataset(self.dataset, split=self.split)109        indices = random.sample(range(len(dataset)), min(count, len(dataset)))110        return [{"data": dataset[i], "benchmark_type": "hellaswag"} for i in indices]111 112    def format_sample(self, sample, subject=None):113        data = sample["data"]114        endings = "\n".join([f"{chr(65+i)}. {ending}" for i, ending in enumerate(data['endings'])])115        answer = chr(65 + int(data['label']))116        return self.format_template.format(context=data["ctx"], endings=endings, answer=answer)117 118class PIQABenchmark(Benchmark):119    dataset = "gimmaru/piqa"120    split = "validation"121    format_template = "Goal: {goal}\n\nWhich solution is better?\nA. {sol1}\nB. {sol2}\nAnswer: {answer}"122 123    def load_samples(self, count=5, subjects=None):124        dataset = load_dataset(self.dataset, split=self.split)125        indices = random.sample(range(len(dataset)), min(count, len(dataset)))126        return [{"data": dataset[i], "benchmark_type": "piqa"} for i in indices]127 128    def format_sample(self, sample, subject=None):129        data = sample["data"]130        answer = chr(65 + data['label'])131        return self.format_template.format(goal=data["goal"], sol1=data["sol1"], sol2=data["sol2"], answer=answer)132 133class TruthfulQABenchmark(Benchmark):134    dataset = "truthful_qa"135    config = "generation"136    split = "validation"137    format_template = "Question: {question}\n\nBest Answer: {best_answer}\n\nCorrect Answers:\n{correct_answers}"138 139    def load_samples(self, count=5, subjects=None):140        dataset = load_dataset(self.dataset, self.config, split=self.split)141        indices = random.sample(range(len(dataset)), min(count, len(dataset)))142        return [{"data": dataset[i], "benchmark_type": "truthfulqa"} for i in indices]143 144    def format_sample(self, sample, subject=None):145        data = sample["data"]146        correct_answers = "\n".join([f"- {ans}" for ans in data['correct_answers']])147        return self.format_template.format(148            question=data["question"], 149            best_answer=data["best_answer"],150            correct_answers=correct_answers151        )152 153# Registry for easy extensibility154BENCHMARKS = {155    "mmlu": MMLUBenchmark(),156    "gsm8k": GSM8KBenchmark(),157    "gpqa": GPQABenchmark(),158    "arc_challenge": ARCChallengeBenchmark(),159    "arc_easy": ARCEasyBenchmark(),160    "hellaswag": HellaSwagBenchmark(),161    "piqa": PIQABenchmark(),162    "truthfulqa": TruthfulQABenchmark(),163}164