Module neuralops.data

Model‑agnostic data pipeline – dataset loading, tokenization, preprocessing, DataLoader creation, and augmentation for all common tasks and models.

Sub-modules

neuralops.data.augmentation

Simple data augmentation registry (random mask, back‑translation placeholder).

neuralops.data.base

Base data utilities: configuration, task‑aware preprocessing, DataLoader creation.

neuralops.data.benchmark_datasets

Unified Benchmark Dataset Loader & Registry for NLP / LLM Evaluation. Supports 12+ standard NLP/LLM benchmarks with streaming, local caching, …

neuralops.data.cached_tensors

Pre-tokenized Tensor Dataset Caching & Zero-Overhead In-Memory Loader …

neuralops.data.custom

Custom dataset loading utilities.

neuralops.data.custom_loader

Custom dataset loading from CSV, JSON, Parquet files.

neuralops.data.dataset_discovery

Utility to scan and list available datasets.

neuralops.data.dataset_loader

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

neuralops.data.glue

GLUE and SuperGLUE dataset loaders. Uses the HuggingFace datasets library and local files if available.

neuralops.data.glue_loader

Local GLUE Loader — loads pre-cached GLUE datasets from disk first, then HF Hub.

neuralops.data.multi_choice

Multi‑choice dataset (ARC, HellaSwag, etc.) handling.

neuralops.data.prepare_cache_once
neuralops.data.prepare_datasets

Dataset Preparation & Preprocessing Pipeline for GPU Server Benchmark & Training …

neuralops.data.preprocessing
neuralops.data.synthetic_generator

High-Fidelity Synthetic Dataset & Token Generator for NLP / LLM Evaluation …

neuralops.data.task_mapper

Task metadata: metrics, number of labels, text columns.

neuralops.data.validation

Simple data validation – missing values, infinite, zero variance, schema checks.

Functions

def convert_to_multichoice_format(df: pandas.DataFrame) ‑> pandas.DataFrame
Expand source code
def convert_to_multichoice_format(df: pd.DataFrame) -> pd.DataFrame:
    """Normalise a DataFrame into the standard multi‑choice format (question, choices, label)."""
    rows = []
    for _, row in df.iterrows():
        q = row.get("prompt", row.get("ctx", ""))
        choices = []
        if "choices_json" in row:
            try:
                choices = json.loads(str(row["choices_json"]))
            except Exception:
                pass
        elif "endings" in row:
            try:
                choices = json.loads(str(row["endings"]))
            except Exception:
                pass
        else:
            choice_cols = [c for c in df.columns if c.startswith("choice_")]
            choices = [str(row[c]) for c in sorted(choice_cols)]
        label = 0
        if "gold_answer_json" in row:
            try:
                ans = json.loads(str(row["gold_answer_json"]))
                label = ans.get("choice_index", 0)
            except Exception:
                pass
        elif "answer_json" in row:
            try:
                ans = json.loads(str(row["answer_json"]))
                label = ans.get("choice_index", 0)
            except Exception:
                pass
        elif "label" in row:
            label = int(row["label"])
        rows.append({"question": str(q), "choices": choices, "label": label})
    return pd.DataFrame(rows)

Normalise a DataFrame into the standard multi‑choice format (question, choices, label).

def create_dataloader(dataset,
tokenizer,
config: DataConfig,
**kwargs) ‑> torch.utils.data.dataloader.DataLoader
Expand source code
def create_dataloader(dataset, tokenizer, config: DataConfig, **kwargs) -> DataLoader:
    """
    Create a DataLoader from a HuggingFace Dataset, applying preprocessing if not already tokenized.
    Uses ProductionDataCollator to guarantee rock-solid batch dimensions and GPU Tensor Core alignment.
    """
    if "input_ids" in dataset.column_names and "attention_mask" in dataset.column_names:
        processed = dataset
    else:
        processed = dataset.map(
            lambda examples: preprocess_function(examples, tokenizer, config),
            batched=True,
            remove_columns=dataset.column_names,
        )
    columns = [
        col
        for col in [
            "input_ids",
            "attention_mask",
            "labels",
            "start_positions",
            "end_positions",
            "token_type_ids",
        ]
        if col in processed.column_names
    ]
    processed.set_format(type="torch", columns=columns)

    collate_fn = kwargs.pop("collate_fn", None)
    if collate_fn is None:
        collate_fn = ProductionDataCollator(
            pad_token_id=(getattr(tokenizer, "pad_token_id", 0) if tokenizer is not None else 0),
            pad_to_multiple_of=8,
            max_length=config.max_length,
        )

    return DataLoader(
        processed,
        batch_size=config.batch_size,
        shuffle=config.shuffle,
        num_workers=config.num_workers,
        collate_fn=collate_fn,
        **kwargs,
    )

Create a DataLoader from a HuggingFace Dataset, applying preprocessing if not already tokenized. Uses ProductionDataCollator to guarantee rock-solid batch dimensions and GPU Tensor Core alignment.

def generate_synthetic_data(output_path: str,
num_samples: int = 10000,
domain: str = 'movies',
seed: int = 42) ‑> pandas.DataFrame
Expand source code
def generate_synthetic_data(
    output_path: str, num_samples: int = 10000, domain: str = "movies", seed: int = 42
) -> pd.DataFrame:
    """Generate a single-domain synthetic classification dataset."""
    random.seed(seed)
    half = num_samples // 2
    pos = _generate_batch(domain, half, label=1, is_positive=True)
    neg = _generate_batch(domain, half, label=0, is_positive=False)
    df = pd.DataFrame(pos + neg).sample(frac=1, random_state=seed).reset_index(drop=True)
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(output_path, index=False)
    logger.info("Generated %d samples to %s", len(df), output_path)
    return df

Generate a single-domain synthetic classification dataset.

def get_calibration_dataloader(task_name: str = 'sst2',
tokenizer: Any = None,
batch_size: int = 8,
num_samples: int = 128,
seq_len: int = 128,
split: str = 'train') ‑> Any
Expand source code
def get_calibration_dataloader(
    task_name: str = "sst2",
    tokenizer: Any = None,
    batch_size: int = 8,
    num_samples: int = 128,
    seq_len: int = 128,
    split: str = "train",
) -> Any:
    """
    Builds a calibration DataLoader from local disk-cached GLUE data.
    Used by data-aware SVD compression algorithms (ASVD, A3, FRanDI, DynamicRank,
    SwiftSVD, SliceGPT, etc.) to collect activation covariance and Fisher information.
    """
    from torch.utils.data import DataLoader

    t_name = task_name.lower().replace("-", "_")
    ds_dict = load_glue_local(t_name)
    if split not in ds_dict:
        split = list(ds_dict.keys())[0]
    ds = ds_dict[split]

    if num_samples and len(ds) > num_samples:
        ds = ds.select(range(min(num_samples, len(ds))))

    if tokenizer is not None:
        key1, key2 = GLUE_TASK_KEYS.get(t_name, ("sentence", None))
        if key1 not in ds.column_names:
            cols = [c for c in ds.column_names if c not in ("label", "labels", "idx")]
            key1 = cols[0] if cols else "text"
            key2 = None

        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token or "[PAD]"

        def tokenize_fn(examples):
            args = (
                (examples[key1],)
                if key2 is None or key2 not in examples
                else (examples[key1], examples[key2])
            )
            return tokenizer(*args, padding="max_length", truncation=True, max_length=seq_len)

        remove_cols = [c for c in ds.column_names if c not in ("label", "labels")]
        ds = ds.map(tokenize_fn, batched=True, remove_columns=remove_cols)
        keep = [
            c for c in ["input_ids", "attention_mask", "token_type_ids"] if c in ds.column_names
        ]
        ds.set_format(type="torch", columns=keep)

    return DataLoader(ds, batch_size=batch_size, shuffle=False)

Builds a calibration DataLoader from local disk-cached GLUE data. Used by data-aware SVD compression algorithms (ASVD, A3, FRanDI, DynamicRank, SwiftSVD, SliceGPT, etc.) to collect activation covariance and Fisher information.

def get_metric_name(task: str) ‑> str
Expand source code
def get_metric_name(task: str) -> str:
    """Return the metric name."""
    return TASK_METRICS.get(task, "accuracy")

Return the metric name.

def get_num_labels(task: str, dataset=None) ‑> int
Expand source code
def get_num_labels(task: str, dataset=None) -> int:
    """Return the num labels."""
    if task == "stsb":
        return 1
    if task in TASK_NUM_LABELS:
        return TASK_NUM_LABELS[task]
    if dataset and "label" in dataset["train"].column_names:
        return len(dataset["train"].unique("label"))
    return 2

Return the num labels.

def get_num_labels_and_is_regression(task_name: str) ‑> tuple[int, bool]
Expand source code
def get_num_labels_and_is_regression(task_name: str) -> tuple[int, bool]:
    """Return (num_labels, is_regression) for GLUE tasks."""
    t = task_name.lower()
    if t == "stsb":
        return 1, True
    mapping = {
        "cola": 2,
        "sst2": 2,
        "mrpc": 2,
        "qqp": 2,
        "mnli": 3,
        "qnli": 2,
        "rte": 2,
        "wnli": 2,
    }
    return mapping.get(t, 2), False

Return (num_labels, is_regression) for GLUE tasks.

def get_text_columns(task: str) ‑> list
Expand source code
def get_text_columns(task: str) -> list:
    """Return the text columns."""
    return TASK_COLUMNS.get(task, ["sentence"])

Return the text columns.

def is_multichoice_dataset(df: pandas.DataFrame) ‑> tuple[bool, int]
Expand source code
def is_multichoice_dataset(df: pd.DataFrame) -> tuple[bool, int]:
    """Check if a DataFrame contains a multi‑choice format and return the number of choices."""
    if "choices_json" in df.columns:
        sample = df["choices_json"].iloc[0]
        if isinstance(sample, str):
            try:
                choices = json.loads(sample)
                if isinstance(choices, list) and len(choices) >= 2:
                    return True, len(choices)
            except Exception:
                pass
    if "endings" in df.columns:
        return True, 4
    choice_cols = [c for c in df.columns if c.startswith("choice_")]
    if len(choice_cols) >= 2:
        return True, len(choice_cols)
    return False, 0

Check if a DataFrame contains a multi‑choice format and return the number of choices.

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]].

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.

def load_glue_local(task_name: str, base_path: str = None) ‑> datasets.dataset_dict.DatasetDict
Expand source code
def load_glue_local(task_name: str, base_path: str = None) -> DatasetDict:
    """
    Load a GLUE task from local Parquet files, falling back to HuggingFace Datasets hub.
    Compatible with local, Kaggle, Colab, and remote GPU/HPC environments.
    """
    from datasets import load_from_disk

    candidates = [
        GLUE_LOCAL_DIR / task_name,
        Path.cwd() / "data" / "glue_data" / task_name,
    ]
    for cand in candidates:
        if cand.exists():
            return load_from_disk(str(cand))

    from datasets import load_dataset as hf_load_dataset

    try:
        return hf_load_dataset("glue", task_name)
    except Exception:
        return hf_load_dataset("nyu-mll/glue", task_name)

Load a GLUE task from local Parquet files, falling back to HuggingFace Datasets hub. Compatible with local, Kaggle, Colab, and remote GPU/HPC environments.

def preprocess_function(examples,
tokenizer,
config: DataConfig) ‑> Any
Expand source code
def preprocess_function(examples, tokenizer, config: DataConfig) -> Any:
    """
    Generic preprocessing function handling all TaskTypes.
    Returns a dictionary with 'input_ids', 'attention_mask', and appropriate labels.
    """
    if config.custom_preprocess is not None:
        return config.custom_preprocess(examples, tokenizer, config)

    text1 = examples.get(config.text_column_1) if config.text_column_1 else None
    text2 = examples.get(config.text_column_2) if config.text_column_2 else None

    # Fallback auto-detection for text columns if specified column is missing
    if text1 is None:
        for col in [
            "sentence",
            "text",
            "sentence1",
            "premise",
            "question",
            "inputs",
            "input",
        ]:
            if col in examples:
                text1 = examples[col]
                break
        if text1 is None and hasattr(examples, "keys"):
            for k in examples.keys():
                if k not in ("label", "labels", "idx", "id"):
                    text1 = examples[k]
                    break

    if text2 is None and config.text_column_2:
        for col in ["sentence2", "hypothesis", "context"]:
            if col in examples:
                text2 = examples[col]
                break

    if text1 is not None:
        if isinstance(text1, list):
            text1 = [str(x) if x is not None else "" for x in text1]
        elif not isinstance(text1, str):
            text1 = str(text1)

    if text2 is not None:
        if isinstance(text2, list):
            text2 = [str(x) if x is not None else "" for x in text2]
        elif not isinstance(text2, str):
            text2 = str(text2)

    tokenizer_kwargs = {
        "truncation": config.truncation,
        "padding": config.padding,
        "max_length": config.max_length,
        "return_tensors": config.return_tensors,
    }

    if config.task_type == TaskType.LANGUAGE_MODELING:
        result = (
            tokenizer(text1, **tokenizer_kwargs) if text1 else tokenizer(text2, **tokenizer_kwargs)
        )
        return result

    if config.task_type == TaskType.TOKEN_CLASSIFICATION:
        tokenized = tokenizer(
            examples[config.tokens_column],
            truncation=config.truncation,
            padding=config.padding,
            max_length=config.max_length,
            is_split_into_words=True,
            return_tensors=config.return_tensors,
        )
        labels = []
        for i, label_list in enumerate(examples[config.tags_column]):
            word_ids = tokenized.word_ids(batch_index=i)
            label_ids = [-100] * len(word_ids)
            previous_word_idx = None
            for j, word_idx in enumerate(word_ids):
                if word_idx is None:
                    continue
                if word_idx != previous_word_idx:
                    label_ids[j] = label_list[word_idx]
                    previous_word_idx = word_idx
                else:
                    label_ids[j] = -100
            labels.append(label_ids)
        tokenized["labels"] = labels
        return tokenized

    if config.task_type in (
        TaskType.SINGLE_CLASSIFICATION,
        TaskType.PAIR_CLASSIFICATION,
        TaskType.REGRESSION,
    ):
        if text1 is not None and text2 is not None:
            tokenized = tokenizer(text1, text2, **tokenizer_kwargs)
        else:
            text_input = text1 if text1 is not None else text2
            tokenized = tokenizer(text_input, **tokenizer_kwargs)
        label_col = config.label_column
        if label_col not in examples:
            for candidate in ["label", "labels", "target", "gold_label"]:
                if candidate in examples:
                    label_col = candidate
                    break
        if label_col in examples:
            if config.task_type == TaskType.REGRESSION:
                tokenized["labels"] = [float(l) for l in examples[label_col]]
            else:
                tokenized["labels"] = [int(l) for l in examples[label_col]]
        return tokenized

    if config.task_type == TaskType.QA:
        tokenized = tokenizer(
            examples[config.question_column],
            examples[config.context_column],
            truncation="only_second" if config.truncation else False,
            padding=config.padding,
            max_length=config.max_length,
            return_offsets_mapping=True,
            return_tensors=config.return_tensors,
        )
        start_positions, end_positions = [], []
        for i, (ans, context) in enumerate(
            zip(examples[config.answer_column], examples[config.context_column])
        ):
            if isinstance(ans, dict):
                start_char = ans["answer_start"]
                end_char = start_char + len(ans["text"])
            else:
                start_char = ans
                end_char = start_char + len(examples[config.answer_column][i])
            offsets = tokenized["offset_mapping"][i]
            sequence_ids = tokenized.sequence_ids(i)
            start_token = end_token = None
            for j, (start, end) in enumerate(offsets):
                if sequence_ids[j] != 1:
                    continue
                if start <= start_char < end:
                    start_token = j
                if start < end_char <= end:
                    end_token = j
                    break
            if start_token is None:
                start_token = 0
            if end_token is None:
                end_token = len(offsets) - 1
            start_positions.append(start_token)
            end_positions.append(end_token)
        tokenized["start_positions"] = start_positions
        tokenized["end_positions"] = end_positions
        tokenized.pop("offset_mapping")
        return tokenized

    raise ValueError(f"Unsupported task type: {config.task_type}")

Generic preprocessing function handling all TaskTypes. Returns a dictionary with 'input_ids', 'attention_mask', and appropriate labels.

def random_mask_augment(dataset, mask_prob=0.15) ‑> Any
Expand source code
@AUGMENTATION_REGISTRY.register("random_mask")
def random_mask_augment(dataset, mask_prob=0.15) -> Any:
    """Random mask augment."""

    def mask_tokens(example) -> Any:
        """Mask tokens."""
        input_ids = example["input_ids"]
        for i in range(len(input_ids)):
            if random.random() < mask_prob:
                input_ids[i] = 103  # [MASK] token id (BERT)
        example["input_ids"] = input_ids
        return example

    return dataset.map(mask_tokens)

Random mask augment.

def run_universal_generation(output_dir: str | Path,
total_samples_per_domain: int = 5000,
domains: dict[str, Any] | None = None,
seed: int = 42) ‑> pandas.DataFrame
Expand source code
def run_universal_generation(
    output_dir: str | Path,
    total_samples_per_domain: int = 5000,
    domains: dict[str, Any] | None = None,
    seed: int = 42,
) -> pd.DataFrame:
    """Generate synthetic dataset across multiple domains and combine them."""
    random.seed(seed)
    domains = domains or DEFAULT_DOMAINS
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    all_dfs = []
    for domain_key in domains:
        half = total_samples_per_domain // 2
        pos = _generate_batch(domain_key, half, label=1, is_positive=True)
        neg = _generate_batch(domain_key, half, label=0, is_positive=False)
        df = pd.DataFrame(pos + neg).sample(frac=1, random_state=seed).reset_index(drop=True)
        df.to_csv(out_dir / f"data_{domain_key}.csv", index=False)
        all_dfs.append(df)

    final_df = (
        pd.concat(all_dfs, ignore_index=True)
        .sample(frac=1, random_state=seed)
        .reset_index(drop=True)
    )
    final_df.to_csv(out_dir / "combined_master_dataset.csv", index=False)
    logger.info(
        "Universal generation complete: %d total samples across %d domains.",
        len(final_df),
        len(domains),
    )
    return final_df

Generate synthetic dataset across multiple domains and combine them.

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 DataConfig (task_type: TaskType = TaskType.SINGLE_CLASSIFICATION,
task_name: str | None = None,
max_length: int = 128,
padding: Any = 'max_length',
truncation: bool = True,
return_tensors: str = 'pt',
text_column: str | None = None,
text_column_1: str | None = None,
text_column_2: str | None = None,
label_column: str | None = None,
batch_size: int = 16,
shuffle: bool = True,
num_workers: int = 0,
mlm_probability: float = 0.15,
question_column: str | None = None,
context_column: str | None = None,
answer_column: str | None = None,
tokens_column: str | None = None,
tags_column: str | None = None,
custom_preprocess:  | None = None,
max_samples: int | None = None,
split: str | None = None)
Expand source code
@dataclass
class DataConfig:
    task_type: TaskType = TaskType.SINGLE_CLASSIFICATION
    task_name: str | None = None
    max_length: int = 128
    padding: Any = "max_length"
    truncation: bool = True
    return_tensors: str = "pt"
    text_column: str | None = None
    text_column_1: str | None = None
    text_column_2: str | None = None
    label_column: str | None = None
    batch_size: int = 16
    shuffle: bool = True
    num_workers: int = 0
    mlm_probability: float = 0.15
    question_column: str | None = None
    context_column: str | None = None
    answer_column: str | None = None
    tokens_column: str | None = None
    tags_column: str | None = None
    custom_preprocess: Callable | None = None
    max_samples: int | None = None
    split: str | None = None

DataConfig(task_type: neuralops.data.base.TaskType = , task_name: str | None = None, max_length: int = 128, padding: Any = 'max_length', truncation: bool = True, return_tensors: str = 'pt', text_column: str | None = None, text_column_1: str | None = None, text_column_2: str | None = None, label_column: str | None = None, batch_size: int = 16, shuffle: bool = True, num_workers: int = 0, mlm_probability: float = 0.15, question_column: str | None = None, context_column: str | None = None, answer_column: str | None = None, tokens_column: str | None = None, tags_column: str | None = None, custom_preprocess: collections.abc.Callable | None = None, max_samples: int | None = None, split: str | None = None)

Instance variables

var answer_column : str | None

The type of the None singleton.

var batch_size : int

The type of the None singleton.

var context_column : str | None

The type of the None singleton.

var custom_preprocess :  | None

The type of the None singleton.

var label_column : str | None

The type of the None singleton.

var max_length : int

The type of the None singleton.

var max_samples : int | None

The type of the None singleton.

var mlm_probability : float

The type of the None singleton.

var num_workers : int

The type of the None singleton.

var padding : Any

The type of the None singleton.

var question_column : str | None

The type of the None singleton.

var return_tensors : str

The type of the None singleton.

var shuffle : bool

The type of the None singleton.

var split : str | None

The type of the None singleton.

var tags_column : str | None

The type of the None singleton.

var task_name : str | None

The type of the None singleton.

var task_type : TaskType

The type of the None singleton.

var text_column : str | None

The type of the None singleton.

var text_column_1 : str | None

The type of the None singleton.

var text_column_2 : str | None

The type of the None singleton.

var tokens_column : str | None

The type of the None singleton.

var truncation : bool

The type of the None singleton.

class DataValidator (schema: dict[str, typing.Any] = None)
Expand source code
class DataValidator:
    def __init__(self, schema: dict[str, Any] = None) -> None:
        self.schema = schema or {}

    def validate(self, dataset: dict[str, np.ndarray]) -> ValidationReport:
        """Validate."""
        issues = []
        for col, values in dataset.items():
            v = np.array(values)
            # Missing values
            missing = (
                np.sum(np.isnan(v)) if np.issubdtype(v.dtype, np.floating) else np.sum(v is None)
            )
            if missing > 0:
                issues.append({"column": col, "issue": "missing_values", "count": int(missing)})
            # Infinite values
            if v.dtype in [np.float32, np.float64]:
                inf_count = np.sum(np.isinf(v))
                if inf_count > 0:
                    issues.append(
                        {
                            "column": col,
                            "issue": "infinite_values",
                            "count": int(inf_count),
                        }
                    )
            # Low variance
            if len(v) > 1 and v.dtype in [np.float32, np.float64]:
                if np.std(v) < 1e-8:
                    issues.append({"column": col, "issue": "zero_variance"})
            # Schema checks
            if col in self.schema:
                expected = self.schema[col]
                if "min" in expected and np.min(v) < expected["min"]:
                    issues.append(
                        {
                            "column": col,
                            "issue": "below_min",
                            "min": expected["min"],
                            "actual": float(np.min(v)),
                        }
                    )
                if "max" in expected and np.max(v) > expected["max"]:
                    issues.append(
                        {
                            "column": col,
                            "issue": "above_max",
                            "max": expected["max"],
                            "actual": float(np.max(v)),
                        }
                    )

        failed_cols = set(i["column"] for i in issues)
        return ValidationReport(
            checks_passed=len(dataset) - len(failed_cols),
            checks_failed=len(failed_cols),
            issues=issues,
            is_valid=len(issues) == 0,
        )

Methods

def validate(self, dataset: dict[str, numpy.ndarray]) ‑> ValidationReport
Expand source code
def validate(self, dataset: dict[str, np.ndarray]) -> ValidationReport:
    """Validate."""
    issues = []
    for col, values in dataset.items():
        v = np.array(values)
        # Missing values
        missing = (
            np.sum(np.isnan(v)) if np.issubdtype(v.dtype, np.floating) else np.sum(v is None)
        )
        if missing > 0:
            issues.append({"column": col, "issue": "missing_values", "count": int(missing)})
        # Infinite values
        if v.dtype in [np.float32, np.float64]:
            inf_count = np.sum(np.isinf(v))
            if inf_count > 0:
                issues.append(
                    {
                        "column": col,
                        "issue": "infinite_values",
                        "count": int(inf_count),
                    }
                )
        # Low variance
        if len(v) > 1 and v.dtype in [np.float32, np.float64]:
            if np.std(v) < 1e-8:
                issues.append({"column": col, "issue": "zero_variance"})
        # Schema checks
        if col in self.schema:
            expected = self.schema[col]
            if "min" in expected and np.min(v) < expected["min"]:
                issues.append(
                    {
                        "column": col,
                        "issue": "below_min",
                        "min": expected["min"],
                        "actual": float(np.min(v)),
                    }
                )
            if "max" in expected and np.max(v) > expected["max"]:
                issues.append(
                    {
                        "column": col,
                        "issue": "above_max",
                        "max": expected["max"],
                        "actual": float(np.max(v)),
                    }
                )

    failed_cols = set(i["column"] for i in issues)
    return ValidationReport(
        checks_passed=len(dataset) - len(failed_cols),
        checks_failed=len(failed_cols),
        issues=issues,
        is_valid=len(issues) == 0,
    )

Validate.

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.
class MultiChoiceDataset (df: pandas.DataFrame,
tokenizer,
max_length: int = 128,
question_column: str | None = None,
choices_column: str | None = None,
answer_column: str | None = None)
Expand source code
class MultiChoiceDataset(Dataset):
    """Dataset for multi‑choice questions (ARC, HellaSwag, etc.)."""

    def __init__(
        self,
        df: pd.DataFrame,
        tokenizer,
        max_length: int = 128,
        question_column: str | None = None,
        choices_column: str | None = None,
        answer_column: str | None = None,
    ) -> None:
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = []

        q_col = question_column or self._detect_question_column(df)
        c_col = choices_column or self._detect_choices_column(df)
        a_col = answer_column or self._detect_answer_column(df)

        valid = 0
        for _, row in df.iterrows():
            question = self._extract_question(row, q_col)
            choices = self._extract_choices(row, c_col)
            label = self._extract_label(row, a_col, len(choices))
            if len(choices) >= 2 and 0 <= label < len(choices):
                self.samples.append({"question": question, "choices": choices, "label": label})
                valid += 1
        print(f"Multi‑choice dataset: {valid} valid samples")

    def _detect_question_column(self, df) -> str:
        """detect question column."""
        for col in ["question", "prompt", "ctx", "ctx_a", "text", "sentence"]:
            if col in df.columns:
                return col
        return df.columns[0]

    def _detect_choices_column(self, df) -> str | None:
        """detect choices column."""
        for col in ["choices_json", "choices", "endings", "options"]:
            if col in df.columns:
                return col
        return None

    def _detect_answer_column(self, df) -> str | None:
        """detect answer column."""
        for col in ["gold_answer_json", "answer_json", "label", "answer", "gold"]:
            if col in df.columns:
                return col
        return None

    def _extract_question(self, row, q_col) -> str:
        """extract question."""
        return str(row[q_col]) if q_col else ""

    def _extract_choices(self, row, c_col) -> list[str]:
        """extract choices."""
        if c_col is None:
            choice_cols = [col for col in row.index if col.startswith("choice_")]
            if choice_cols:
                return [str(row[col]) for col in sorted(choice_cols)]
            return []
        val = row[c_col]
        if isinstance(val, str):
            try:
                parsed = json.loads(val)
                if isinstance(parsed, list):
                    return parsed
            except json.JSONDecodeError:
                pass
        if isinstance(val, list):
            return val
        return []

    def _extract_label(self, row, a_col, num_choices) -> int:
        """extract label."""
        if a_col is None:
            return row.get("label", 0)
        val = row[a_col]
        if isinstance(val, str):
            try:
                obj = json.loads(val)
                if "choice_index" in obj:
                    return int(obj["choice_index"])
                if "label" in obj:
                    return int(obj["label"])
            except json.JSONDecodeError:
                pass
        if isinstance(val, (int, float)):
            return int(val)
        return 0

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        sample = self.samples[idx]
        q = sample["question"]
        choices = sample["choices"]
        label = sample["label"]

        input_ids, attn_masks = [], []
        for choice in choices:
            text = f"{q} {choice}" if q else choice
            enc = self.tokenizer(
                text,
                truncation=True,
                padding="max_length",
                max_length=self.max_length,
                return_tensors="pt",
            )
            input_ids.append(enc["input_ids"].squeeze(0))
            attn_masks.append(enc["attention_mask"].squeeze(0))
        return {
            "input_ids": torch.stack(input_ids),
            "attention_mask": torch.stack(attn_masks),
            "labels": torch.tensor(label, dtype=torch.long),
        }

Dataset for multi‑choice questions (ARC, HellaSwag, etc.).

Ancestors

  • torch.utils.data.dataset.Dataset
  • typing.Generic
class ProductionDataCollator (pad_token_id: int = 0,
label_pad_token_id: int = -100,
pad_to_multiple_of: int | None = 8,
max_length: int | None = None)
Expand source code
class ProductionDataCollator:
    """
    Production-grade data collator for PyTorch DataLoaders.
    Guarantees standard batch dimensions and types for high-performance GPU servers:
    - input_ids: [B, L], torch.long (padded to max_len or multiple of 8 for Tensor Cores)
    - attention_mask: [B, L], torch.long
    - token_type_ids: [B, L], torch.long (if present)
    - labels: [B] (for 1D classification/regression) or [B, L] (for language modeling with -100 ignore index)
    """

    def __init__(
        self,
        pad_token_id: int = 0,
        label_pad_token_id: int = -100,
        pad_to_multiple_of: int | None = 8,
        max_length: int | None = None,
    ) -> None:
        self.pad_token_id = pad_token_id
        self.label_pad_token_id = label_pad_token_id
        self.pad_to_multiple_of = pad_to_multiple_of
        self.max_length = max_length

    def __call__(self, features: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
        if not features:
            return {}

        batch: dict[str, Any] = {}
        first = features[0]

        # Determine sequence length
        if "input_ids" in first:
            seq_lens = [
                (len(f["input_ids"]) if isinstance(f["input_ids"], (list, torch.Tensor)) else 1)
                for f in features
            ]
            max_len = max(seq_lens)
            if self.max_length is not None:
                max_len = min(max_len, self.max_length)
            if self.pad_to_multiple_of is not None and self.pad_to_multiple_of > 0:
                max_len = (
                    (max_len + self.pad_to_multiple_of - 1) // self.pad_to_multiple_of
                ) * self.pad_to_multiple_of
        else:
            max_len = None

        for key in first.keys():
            if key in ("input_ids", "attention_mask", "token_type_ids"):
                pad_val = 0 if key == "attention_mask" else self.pad_token_id
                tensors = []
                for f in features:
                    val = f[key]
                    if not isinstance(val, torch.Tensor):
                        val = torch.tensor(val, dtype=torch.long)
                    else:
                        val = val.long()
                    if max_len is not None:
                        if val.size(0) < max_len:
                            padd = torch.full((max_len - val.size(0),), pad_val, dtype=val.dtype)
                            val = torch.cat([val, padd], dim=0)
                        elif val.size(0) > max_len:
                            val = val[:max_len]
                    tensors.append(val)
                batch[key] = torch.stack(tensors, dim=0).contiguous()

            elif key in ("label", "labels"):
                # Handle 1D classification vs 2D LM labels
                val0 = first[key]
                is_2d = (isinstance(val0, list) and len(val0) > 1) or (
                    isinstance(val0, torch.Tensor) and val0.dim() > 0 and val0.numel() > 1
                )
                if is_2d and max_len is not None:
                    tensors = []
                    for f in features:
                        val = f[key]
                        if not isinstance(val, torch.Tensor):
                            val = torch.tensor(val, dtype=torch.long)
                        else:
                            val = val.long()
                        if val.size(0) < max_len:
                            padd = torch.full(
                                (max_len - val.size(0),),
                                self.label_pad_token_id,
                                dtype=val.dtype,
                            )
                            val = torch.cat([val, padd], dim=0)
                        elif val.size(0) > max_len:
                            val = val[:max_len]
                        tensors.append(val)
                    batch["labels"] = torch.stack(tensors, dim=0).contiguous()
                else:
                    vals = [f[key] for f in features]
                    if all(
                        isinstance(v, (int, bool, np.integer))
                        or (
                            isinstance(v, torch.Tensor)
                            and v.numel() == 1
                            and v.dtype in (torch.int64, torch.int32, torch.bool)
                        )
                        for v in vals
                    ):
                        batch["labels"] = torch.tensor([int(v) for v in vals], dtype=torch.long)
                    else:
                        batch["labels"] = torch.tensor(
                            [float(v) for v in vals], dtype=torch.float32
                        )
            else:
                # Other auxiliary tensors
                vals = [f[key] for f in features]
                if isinstance(vals[0], torch.Tensor):
                    batch[key] = torch.stack(vals, dim=0).contiguous()
                elif isinstance(vals[0], (int, float)):
                    batch[key] = torch.tensor(vals).contiguous()
                else:
                    batch[key] = vals

        return batch

Production-grade data collator for PyTorch DataLoaders. Guarantees standard batch dimensions and types for high-performance GPU servers: - input_ids: [B, L], torch.long (padded to max_len or multiple of 8 for Tensor Cores) - attention_mask: [B, L], torch.long - token_type_ids: [B, L], torch.long (if present) - labels: [B] (for 1D classification/regression) or [B, L] (for language modeling with -100 ignore index)

class TaskType (*values)
Expand source code
class TaskType(Enum):
    SINGLE_CLASSIFICATION = "single_classification"
    PAIR_CLASSIFICATION = "pair_classification"
    REGRESSION = "regression"
    QA = "qa"
    TOKEN_CLASSIFICATION = "token_classification"
    LANGUAGE_MODELING = "language_modeling"
    SEQ2SEQ = "seq2seq"

Create a collection of name/value pairs.

Example enumeration:

>>> class Color(Enum):
...     RED = 1
...     BLUE = 2
...     GREEN = 3

Access them by:

  • attribute access:

Color.RED

  • value lookup:

Color(1)

  • name lookup:

Color['RED']

Enumerations can be iterated over, and know how many members they have:

>>> len(Color)
3
>>> list(Color)
[<Color.RED: 1>, <Color.BLUE: 2>, <Color.GREEN: 3>]

Methods can be added to enumerations, and members can have their own attributes – see the documentation for details.

Ancestors

  • enum.Enum

Class variables

var LANGUAGE_MODELING

The type of the None singleton.

var PAIR_CLASSIFICATION

The type of the None singleton.

var QA

The type of the None singleton.

var REGRESSION

The type of the None singleton.

var SEQ2SEQ

The type of the None singleton.

var SINGLE_CLASSIFICATION

The type of the None singleton.

var TOKEN_CLASSIFICATION

The type of the None singleton.

class ValidationReport (checks_passed: int,
checks_failed: int,
issues: list[dict[str, typing.Any]],
is_valid: bool)
Expand source code
@dataclass
class ValidationReport:
    checks_passed: int
    checks_failed: int
    issues: list[dict[str, Any]]
    is_valid: bool

ValidationReport(checks_passed: int, checks_failed: int, issues: list[dict[str, typing.Any]], is_valid: bool)

Instance variables

var checks_failed : int

The type of the None singleton.

var checks_passed : int

The type of the None singleton.

var is_valid : bool

The type of the None singleton.

var issues : list[dict[str, typing.Any]]

The type of the None singleton.