Module neuralops.data.benchmark_datasets

Unified Benchmark Dataset Loader & Registry for NLP / LLM Evaluation. Supports 12+ standard NLP/LLM benchmarks with streaming, local caching, tokenized PyTorch DataLoader creation, and offline synthetic fallbacks.

Functions

def get_synthetic_benchmark_sample(dataset_name: str, num_samples: int = 4, seed: int = 42) ‑> list[str]

Expand source code

def get_synthetic_benchmark_sample(
    dataset_name: str, num_samples: int = 4, seed: int = 42
) -> list[str]:
    """Retrieve raw text strings of synthetic benchmark samples."""
    raw = load_benchmark_dataset(
        dataset_name,
        max_samples=num_samples,
        use_synthetic=True,
        seed=seed,
        return_dataloader=False,
    )
    results = []
    for item in raw:
        if isinstance(item, dict):
            text = item.get("text") or item.get("prompt") or item.get("sentence") or str(item)
            results.append(text)
        else:
            results.append(str(item))
    return results

Retrieve raw text strings of synthetic benchmark samples.

def load_benchmark_dataset(dataset_name: str,
tokenizer: Any = None,
seq_len: int = 2048,
batch_size: int = 8,
split: str | None = None,
max_samples: int = 1000,
streaming: bool = True,
cache_dir: str | None = None,
shuffle: bool = False,
num_workers: int = 0,
use_synthetic: bool = False,
seed: int = 42,
return_dataloader: bool = True,
**kwargs: Any) ‑> torch.utils.data.dataloader.DataLoader | list[dict[str, typing.Any]]

Expand source code

def load_benchmark_dataset(
    dataset_name: str,
    tokenizer: Any = None,
    seq_len: int = 2048,
    batch_size: int = 8,
    split: str | None = None,
    max_samples: int = 1000,
    streaming: bool = True,
    cache_dir: str | None = None,
    shuffle: bool = False,
    num_workers: int = 0,
    use_synthetic: bool = False,
    seed: int = 42,
    return_dataloader: bool = True,
    **kwargs: Any,
) -> DataLoader | list[dict[str, Any]]:
    """
    Unified loader for 12+ NLP / LLM benchmark datasets.

    Supports:
      - wikitext2 / wikitext103
      - c4 / c4_en
      - ptb
      - pg19
      - gsm8k
      - humaneval
      - mmlu
      - pile / pile_val
      - arc (ARC-Easy & ARC-Challenge)
      - hellaswag
      - lambada
      - glue (sst2, mrpc, qqp, mnli, qnli, rte, wnli, cola, stsb)

    Args:
        dataset_name: Benchmark name or alias (e.g. 'wikitext2', 'c4', 'gsm8k', 'arc', 'glue/sst2').
        tokenizer: HuggingFace tokenizer instance (optional).
        seq_len: Maximum sequence length for tokenization.
        batch_size: DataLoader batch size.
        split: Dataset split ('train', 'validation', 'test'). Defaults to benchmark standard.
        max_samples: Maximum number of samples to load/stream.
        streaming: Whether to stream from HuggingFace Hub to avoid large downloads.
        cache_dir: Local directory to cache datasets.
        shuffle: Whether to shuffle samples.
        num_workers: Number of workers for DataLoader.
        use_synthetic: Force synthetic data generation for offline testing.
        seed: Random seed for reproducibility.
        return_dataloader: If True, returns PyTorch DataLoader; otherwise returns list of dicts.
        **kwargs: Additional arguments passed to `datasets.load_dataset`.

    Returns:
        torch.utils.data.DataLoader or List[Dict[str, Any]].
    """
    clean_name = dataset_name.lower().replace("-", "_")
    info = BenchmarkDatasetRegistry.get(clean_name)

    # Fallback to general lookup if not matched directly
    if info is None and "/" in dataset_name:
        prefix, subtask = dataset_name.split("/", 1)
        if prefix.lower() == "glue":
            info = BenchmarkDatasetRegistry.get(f"glue_{subtask}")

    if info is None:
        logger.warning(
            "Dataset '%s' not in official benchmark registry. Creating generic fallback.",
            dataset_name,
        )
        info = BenchmarkDatasetInfo(
            name=clean_name,
            aliases=[],
            category="custom",
            default_hf_path=dataset_name,
            default_split=split or "validation",
            text_formatter=lambda r: {"text": r.get("text", str(r))},
            synthetic_generator=lambda n, s: [
                {
                    "text": f"Synthetic custom text sample #{i} for {dataset_name}.",
                    "id": i,
                }
                for i in range(n)
            ],
        )

    target_split = split or info.default_split
    raw_samples: list[dict[str, Any]] = []

    # Attempt online loading if not forced synthetic
    if not use_synthetic:
        try:
            from datasets import load_dataset as hf_load_dataset

            hf_kwargs = dict(kwargs)
            if info.default_subset:
                args = (info.default_hf_path, info.default_subset)
            else:
                args = (info.default_hf_path,)

            if cache_dir:
                hf_kwargs["cache_dir"] = cache_dir

            logger.info("Attempting to load benchmark '%s' from HF (%s)...", info.name, args)

            try:
                # First try streaming if requested
                if streaming:
                    hf_ds = hf_load_dataset(*args, split=target_split, streaming=True, **hf_kwargs)
                    count = 0
                    for row in hf_ds:
                        formatted = info.text_formatter(row) if info.text_formatter else dict(row)
                        # Filter out empty text rows (e.g. WikiText empty headers)
                        if (
                            formatted.get("text", "").strip()
                            or "prompt" in formatted
                            or "question" in formatted
                        ):
                            raw_samples.append(formatted)
                            count += 1
                            if count >= max_samples:
                                break
                else:
                    hf_ds = hf_load_dataset(*args, split=target_split, **hf_kwargs)
                    sub_ds = hf_ds.select(range(min(len(hf_ds), max_samples)))
                    for row in sub_ds:
                        formatted = info.text_formatter(row) if info.text_formatter else dict(row)
                        raw_samples.append(formatted)

            except Exception as inner_e:
                logger.warning(
                    "Streaming/split load failed for %s (%s). Retrying non-streaming or default split...",
                    info.name,
                    inner_e,
                )
                hf_ds = hf_load_dataset(*args, **hf_kwargs)
                actual_split = target_split if target_split in hf_ds else list(hf_ds.keys())[0]
                ds_split = hf_ds[actual_split]
                sub_ds = ds_split.select(range(min(len(ds_split), max_samples)))
                for row in sub_ds:
                    formatted = info.text_formatter(row) if info.text_formatter else dict(row)
                    raw_samples.append(formatted)

        except Exception as e:
            logger.warning(
                "Failed to load '%s' from HuggingFace Hub: %s. Falling back to synthetic generator.",
                dataset_name,
                str(e),
            )
            raw_samples = []

    # If empty or failed, run synthetic generator
    if not raw_samples:
        logger.info(
            "Generating %d synthetic samples for benchmark '%s' (seed=%d)...",
            max_samples,
            info.name,
            seed,
        )
        if info.synthetic_generator:
            raw_samples = info.synthetic_generator(max_samples, seed)
        else:
            raw_samples = [
                {"text": f"Synthetic benchmark text #{i} for {info.name}", "id": i}
                for i in range(max_samples)
            ]

    if shuffle:
        rng = random.Random(seed)
        rng.shuffle(raw_samples)

    if not return_dataloader:
        return raw_samples

    # Build PyTorch Dataset and DataLoader
    is_lm = info.category in ("language_modeling", "code")
    torch_ds = BenchmarkTorchDataset(
        samples=raw_samples,
        tokenizer=tokenizer,
        seq_len=seq_len,
        is_causal_lm=is_lm,
    )

    loader = DataLoader(
        torch_ds,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
    )
    return loader

Unified loader for 12+ NLP / LLM benchmark datasets.

Supports -----= - wikitext2 / wikitext103 - c4 / c4_en - ptb - pg19 - gsm8k - humaneval - mmlu - pile / pile_val - arc (ARC-Easy & ARC-Challenge) - hellaswag - lambada - glue (sst2, mrpc, qqp, mnli, qnli, rte, wnli, cola, stsb)

Args

-----=

dataset_name

Benchmark name or alias (e.g. 'wikitext2', 'c4', 'gsm8k', 'arc', 'glue/sst2').

tokenizer

HuggingFace tokenizer instance (optional).

seq_len

Maximum sequence length for tokenization.

batch_size

DataLoader batch size.

split

Dataset split ('train', 'validation', 'test'). Defaults to benchmark standard.

max_samples

Maximum number of samples to load/stream.

streaming

Whether to stream from HuggingFace Hub to avoid large downloads.

cache_dir

Local directory to cache datasets.

shuffle

Whether to shuffle samples.

num_workers

Number of workers for DataLoader.

use_synthetic

Force synthetic data generation for offline testing.

seed

Random seed for reproducibility.

return_dataloader

If True, returns PyTorch DataLoader; otherwise returns list of dicts.

**kwargs

Additional arguments passed to datasets.load_dataset.

Returns -----= torch.utils.data.DataLoader or List[Dict[str, Any]].

Classes

class BenchmarkDatasetInfo (name: str,
aliases: list[str],
category: str,
default_hf_path: str,
default_subset: str | None = None,
default_split: str = 'validation',
text_formatter: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
synthetic_generator: Callable[[int, int], list[dict[str, Any]]] | None = None,
description: str = '')

Expand source code

@dataclass
class BenchmarkDatasetInfo:
    """Metadata describing a benchmark dataset and its default loading strategy."""

    name: str
    aliases: list[str]
    category: str  # e.g., "language_modeling", "reasoning", "code", "multi_choice", "glue"
    default_hf_path: str
    default_subset: str | None = None
    default_split: str = "validation"
    text_formatter: Callable[[dict[str, Any]], dict[str, Any]] | None = None
    synthetic_generator: Callable[[int, int], list[dict[str, Any]]] | None = None
    description: str = ""

Metadata describing a benchmark dataset and its default loading strategy.

Instance variables

var aliases : list[str]

The type of the None singleton.

var category : str

The type of the None singleton.

var default_hf_path : str

The type of the None singleton.

var default_split : str

The type of the None singleton.

var default_subset : str | None

The type of the None singleton.

var description : str

The type of the None singleton.

var name : str

The type of the None singleton.

var synthetic_generator : Callable[[int, int], list[dict[str, typing.Any]]] | None

The type of the None singleton.

var text_formatter : Callable[[dict[str, typing.Any]], dict[str, typing.Any]] | None

The type of the None singleton.

class BenchmarkDatasetRegistry

Expand source code

class BenchmarkDatasetRegistry:
    """Registry maintaining metadata and formatters for NLP/LLM benchmarks."""

    _registry: dict[str, BenchmarkDatasetInfo] = {}

    @classmethod
    def register(cls, info: BenchmarkDatasetInfo) -> None:
        """Register a new benchmark dataset info."""
        key = info.name.lower()
        cls._registry[key] = info
        for alias in info.aliases:
            cls._registry[alias.lower()] = info

    @classmethod
    def get(cls, name: str) -> BenchmarkDatasetInfo | None:
        """Retrieve benchmark dataset info by name or alias."""
        return cls._registry.get(
            name.lower().replace("-", "").replace("_", "")
        ) or cls._registry.get(name.lower())

    @classmethod
    def is_registered(cls, name: str) -> bool:
        """Check if a benchmark dataset name is registered."""
        return cls.get(name) is not None

    @classmethod
    def list_benchmarks(cls) -> list[str]:
        """List canonical names of all registered benchmarks."""
        canonical = set()
        for info in cls._registry.values():
            canonical.add(info.name)
        return sorted(list(canonical))

    @classmethod
    def list_datasets(cls) -> list[str]:
        """Alias for list_benchmarks."""
        return cls.list_benchmarks()

Registry maintaining metadata and formatters for NLP/LLM benchmarks.

Static methods

def get(name: str) ‑> BenchmarkDatasetInfo | None

Retrieve benchmark dataset info by name or alias.

def is_registered(name: str) ‑> bool

Check if a benchmark dataset name is registered.

def list_benchmarks() ‑> list[str]

List canonical names of all registered benchmarks.

def list_datasets() ‑> list[str]

Alias for list_benchmarks.

def register(info: BenchmarkDatasetInfo) ‑> None

Register a new benchmark dataset info.

class BenchmarkTorchDataset (samples: list[dict[str, Any]],
tokenizer: Any = None,
seq_len: int = 2048,
is_causal_lm: bool = True)

Expand source code

class BenchmarkTorchDataset(Dataset):
    """PyTorch Dataset wrapper for processed benchmark records."""

    def __init__(
        self,
        samples: list[dict[str, Any]],
        tokenizer: Any = None,
        seq_len: int = 2048,
        is_causal_lm: bool = True,
    ):
        self.samples = samples
        self.tokenizer = tokenizer
        self.seq_len = seq_len
        self.is_causal_lm = is_causal_lm

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> dict[str, Any]:
        item = self.samples[idx]
        if self.tokenizer is None:
            return item

        # Extract primary text string
        text = item.get("text", "")
        if not text:
            for k in ["prompt", "sentence", "question", "sentence1"]:
                if k in item and item[k]:
                    text = str(item[k])
                    break

        enc = self.tokenizer(
            text,
            truncation=True,
            max_length=self.seq_len,
            padding="max_length",
            return_tensors="pt",
        )

        input_ids = enc["input_ids"].squeeze(0)
        attention_mask = enc["attention_mask"].squeeze(0)

        data = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
        }

        # Handle causal LM labels or classification labels
        if "label" in item:
            data["labels"] = torch.tensor(
                item["label"],
                dtype=torch.long if isinstance(item["label"], int) else torch.float,
            )
        elif self.is_causal_lm:
            # Shifted causal LM targets (replace padding token ids with -100 in labels)
            labels = input_ids.clone()
            if self.tokenizer.pad_token_id is not None:
                labels[labels == self.tokenizer.pad_token_id] = -100
            data["labels"] = labels

        return data

PyTorch Dataset wrapper for processed benchmark records.

Ancestors

• torch.utils.data.dataset.Dataset
• typing.Generic