#!/usr/bin/env python3

import os

import json

import random

from typing import List, Dict, Any

import datasets

from datasets import Dataset, DatasetDict

from tqdm import tqdm

import numpy as np

from huggingface_hub import HfApi

# Configure random seed for reproducibility

random.seed(42)

np.random.seed(42)

# Configuration

DATASET_NAME = "codelion/optillmbench"

NUM_SAMPLES = 500 # Total samples in the benchmark

SPLIT_RATIO = {"train": 0.8, "test": 0.2} # 80-20 split

SOURCES = {

"gsm8k": {

"name": "gsm8k",

"subset": "main",

"samples": 125,

"field_map": {

"question": "question",

"answer": "answer"

}

},

"boolq": {

"name": "boolq",

"subset": None,

"samples": 125,

"field_map": {

"question": "question",

"passage": "passage",

"answer": "answer"

}

},

"mmlu_math": {

"name": "cais/mmlu",

"subset": "all",

"samples": 125,

"field_map": {

"question": "question",

"choices": "choices",

"answer": "answer"

}

},

"aqua_rat": {

"name": "aqua_rat",

"subset": None,

"samples": 125,

"field_map": {

"question": "question",

"answer": "correct"

}

}

}

def select_challenging_examples(

dataset: datasets.Dataset,

category: str,

num_samples: int,

field_map: Dict[str, str]

) -> List[Dict[str, Any]]:

"""Select challenging examples from the dataset"""

examples = []

# Get all available examples

all_examples = dataset["train"] if "train" in dataset else dataset["validation"]

# Shuffle to randomize selection

shuffled_indices = list(range(len(all_examples)))

random.shuffle(shuffled_indices)

# Select examples

for idx in shuffled_indices:

example = all_examples[idx]

try:

if category == "gsm8k":

question = str(example[field_map["question"]])

answer = str(example[field_map["answer"]])

# Select only multi-step problems

if answer.count("=") < 3:

continue

elif category == "boolq":

passage = str(example[field_map["passage"]])

q = str(example[field_map["question"]])

question = f"Context: {passage}\nQuestion: {q}"

answer = "Yes" if example[field_map["answer"]] else "No"

elif category == "mmlu_math":

question = str(example[field_map["question"]])

choices = example[field_map["choices"]]

answer_index = int(example[field_map["answer"]]) # Convert answer to integer

# Ensure answer index is within bounds

if 0 <= answer_index < len(choices):

answer = choices[answer_index]

else:

print(f"Warning: Answer index '{answer_index}' is out of range for choices: {choices}")

continue # Skip this example if answer index is invalid

# Format choices

choices_text = "\n".join([f"{i}. {choice}" for i, choice in enumerate(choices)])

question = f"{question}\nChoices:\n{choices_text}"

elif category == "aqua_rat":

question = str(example[field_map["question"]])

answer = str(example[field_map["answer"]])

# Ensure non-trivial multiple-choice math problems

if len(question.split()) < 12:

continue

# General filtering

if len(question.split()) < 10: # Ensure substantial questions

continue

examples.append(format_question(category, question, answer))

if len(examples) >= num_samples:

break

except Exception as e:

print(f"Error processing example from {category}: {str(e)}")

continue

return examples

def clean_text(text: str) -> str:

"""Clean text by removing extra whitespace and normalizing newlines"""

return " ".join(text.replace("\r", "\n").split())

def format_question(category: str, question: str, answer: str) -> Dict[str, Any]:

"""Format a question for the benchmark dataset"""

# Basic sanity checks

if not question or not answer:

raise ValueError(f"Empty question or answer in {category}")

return {

"id": f"{category}_{random.getrandbits(32):08x}",

"category": category,

"question": clean_text(question),

"answer": clean_text(answer),

"metadata": {

"source": SOURCES[category]["name"],

"type": category,

"difficulty": "challenging" # All examples are chosen to be challenging

}

}

def load_source_dataset(config: Dict[str, Any]) -> datasets.Dataset:

"""Load a source dataset with error handling"""

try:

dataset = datasets.load_dataset(

config["name"],

config.get("subset")

)

return dataset

except Exception as e:

print(f"Error loading dataset {config['name']}: {str(e)}")

return None

def create_benchmark_dataset() -> Dataset:

"""Create the complete benchmark dataset"""

all_examples = []

# Process each source dataset

for category, config in tqdm(SOURCES.items(), desc="Processing datasets"):

print(f"\nProcessing {category} dataset...")

# Load dataset

dataset = load_source_dataset(config)

if not dataset:

continue

# Select examples

try:

examples = select_challenging_examples(

dataset,

category,

config["samples"],

config["field_map"]

)

print(f"Selected {len(examples)} examples from {category}")

all_examples.extend(examples)

except Exception as e:

print(f"Error selecting examples from {category}: {str(e)}")

continue

# Shuffle final dataset

random.shuffle(all_examples)

# Create train/test splits

num_train = int(len(all_examples) * SPLIT_RATIO["train"])

train_examples = all_examples[:num_train]

test_examples = all_examples[num_train:]

# Convert to HuggingFace Dataset

dataset_dict = DatasetDict({

"train": Dataset.from_list(train_examples),

"test": Dataset.from_list(test_examples)

})

return dataset_dict

def push_to_hub(dataset: DatasetDict, repo_id: str):

"""Push the dataset to HuggingFace Hub"""

try:

# Create README content

readme_content = f"""# OptiLLMBench Dataset

A benchmark dataset for evaluating test-time optimization and scaling capabilities of language models.

## Dataset Description

OptiLLMBench contains {NUM_SAMPLES} carefully selected challenging problems across multiple domains:

- Mathematical reasoning (from competition_math)

- Code generation (from HumanEval)

- Word problems (from GSM8K)

- Multiple choice reasoning (from MMLU)

- Logical deduction (from BBH)

Each example is chosen to benefit from test-time optimization techniques like:

- Increased context length

- Chain-of-thought reasoning

- Self-consistency

- Multiple solution attempts

- And other scaling approaches

## Usage

```python

from datasets import load_dataset

dataset = load_dataset("codelion/optillmbench")

# Access examples

for example in dataset["train"]:

print(f"Category: {{example['category']}}")

print(f"Question: {{example['question']}}")

print(f"Answer: {{example['answer']}}")

print(f"Metadata: {{example['metadata']}}")

```

## Citation

If you use this dataset in your research, please cite:

```bibtex

@software{{optillm,

title = {{Optillm: Optimizing inference proxy for LLMs}},

author = {{Asankhaya Sharma}},

year = {{2024}},

publisher = {{GitHub}},

url = {{https://github.com/codelion/optillm}}

}}

```

"""

# Push to hub

dataset.push_to_hub(

repo_id,

private=False,

embed_external_files=True

)

# Update README

api = HfApi()

api.upload_file(

path_or_fileobj=readme_content.encode(),

path_in_repo="README.md",

repo_id=repo_id,

repo_type="dataset"

)

print(f"Successfully pushed dataset to {repo_id}")

except Exception as e:

print(f"Error pushing to hub: {str(e)}")

def main():

"""Main execution function"""

print("Starting OptILM Bench dataset generation...")

# Create dataset

dataset = create_benchmark_dataset()

# Print statistics

print("\nDataset Statistics:")

for split in dataset:

print(f"\n{split} split:")

print(f"Number of examples: {len(dataset[split])}")

categories = dataset[split].unique("category")

for category in categories:

count = len([ex for ex in dataset[split] if ex["category"] == category])

print(f"- {category}: {count} examples")

# Push to HuggingFace Hub

print("\nPushing dataset to HuggingFace Hub...")

push_to_hub(dataset, DATASET_NAME)

if __name__ == "__main__":

main()