Module neuralops.data.dataset_loader

Unified dataset loader driven by an ExperimentConfig. Supports GLUE, custom CSV/JSON/Parquet, multi‑choice, synthetic, and HuggingFace hub datasets.

Functions

def load_dataset(config: ExperimentConfig) ‑> datasets.dataset_dict.DatasetDict

Expand source code

def load_dataset(config: ExperimentConfig) -> DatasetDict:
    """
    Load a dataset based on ExperimentConfig.
    Returns DatasetDict with at least 'train' and optionally 'validation'/'test' splits.
    """
    name = config.dataset_name.lower()

    # GLUE tasks (local Parquet)
    if name in {"cola", "sst2", "mrpc", "qqp", "stsb", "mnli", "qnli", "rte", "wnli"}:
        return _load_glue(config)

    # Synthetic data generation
    elif name == "synthetic":
        return _load_synthetic(config)

    # Custom file (CSV, JSON, Parquet)
    elif name.endswith((".csv", ".json", ".jsonl", ".parquet")):
        return _load_custom(config)

    # Multi‑choice dataset (ARC, HellaSwag, etc.)
    elif name == "multi_choice":
        return _load_multi_choice(config)

    # Fallback: HuggingFace dataset hub
    else:
        return _load_hub(config)

Load a dataset based on ExperimentConfig. Returns DatasetDict with at least 'train' and optionally 'validation'/'test' splits.

Classes

class ExperimentConfig (dataset_name: str,
dataset_path: str | None = None,
task_type: str = 'single',
text_columns: list[str] | None = None,
label_column: str | None = None,
num_labels: int = 2,
metric: str = 'accuracy',
max_seq_length: int = 128,
split_ratio: float = 0.1,
synthetic_samples: int = 5000,
glue_base_path: str | None = None,
**kwargs)

Expand source code

class ExperimentConfig:
    """
    Configuration for loading a dataset.

    Attributes:
        dataset_name: Name of the dataset (GLUE task, synthetic, file path, or HF identifier).
        dataset_path: Path to local dataset file (for custom or multi‑choice).
        task_type: "single", "pair", "multi_choice", "regression".
        text_columns: List of column names containing text (auto‑detected if None).
        label_column: Name of the label column (auto‑detected if None).
        num_labels: Number of classes (for classification).
        metric: Primary metric for evaluation (e.g., "accuracy").
        max_seq_length: Maximum sequence length for tokenization.
        split_ratio: Fraction of data to use for validation (for custom datasets).
        synthetic_samples: Number of synthetic samples to generate (if dataset_name="synthetic").
        glue_base_path: Base directory for local GLUE Parquet files.
        **kwargs: Additional arguments passed to underlying loaders.
    """

    def __init__(
        self,
        dataset_name: str,
        dataset_path: str | None = None,
        task_type: str = "single",
        text_columns: list[str] | None = None,
        label_column: str | None = None,
        num_labels: int = 2,
        metric: str = "accuracy",
        max_seq_length: int = 128,
        split_ratio: float = 0.1,
        synthetic_samples: int = 5000,
        glue_base_path: str | None = None,
        **kwargs,
    ) -> None:
        self.dataset_name = dataset_name
        self.dataset_path = dataset_path
        self.task_type = task_type
        self.text_columns = text_columns
        self.label_column = label_column
        self.num_labels = num_labels
        self.metric = metric
        self.max_seq_length = max_seq_length
        self.split_ratio = split_ratio
        self.synthetic_samples = synthetic_samples
        from neuralops.utils.paths import get_glue_dir

        self.glue_base_path = glue_base_path or str(get_glue_dir())
        self.kwargs = kwargs

Configuration for loading a dataset.

Attributes

-----=

dataset_name

Name of the dataset (GLUE task, synthetic, file path, or HF identifier).

dataset_path

Path to local dataset file (for custom or multi‑choice).

task_type

"single", "pair", "multi_choice", "regression".

text_columns

List of column names containing text (auto‑detected if None).

label_column

Name of the label column (auto‑detected if None).

num_labels

Number of classes (for classification).

metric

Primary metric for evaluation (e.g., "accuracy").

max_seq_length

Maximum sequence length for tokenization.

split_ratio

Fraction of data to use for validation (for custom datasets).

synthetic_samples

Number of synthetic samples to generate (if dataset_name="synthetic").

glue_base_path

Base directory for local GLUE Parquet files.

**kwargs

Additional arguments passed to underlying loaders.