Tameem7/Prompt-Injection-Classifier

1#!/usr/bin/env python32"""3Utility for loading Nvidia's Aegis AI Content Safety Dataset 2.0 with4the exact fields needed for prompt injection detection experiments.5 6Only the `prompt` text and the normalized `prompt_label` fields are kept.7Labels are mapped to integers: `safe -> 0`, `unsafe -> 1`.8"""9 10from __future__ import annotations11 12from typing import Dict, Optional13 14from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict, load_dataset15 16DATASET_NAME = "nvidia/Aegis-AI-Content-Safety-Dataset-2.0"17LABEL_MAP = {"safe": 0, "unsafe": 1}18SELECTED_COLUMNS = ["prompt", "prompt_label"]19 20 21def _map_labels(batch: Dict[str, list]) -> Dict[str, list]:22    """Batched mapping function that converts string labels to ints."""23    batch["prompt_label"] = [LABEL_MAP[label] for label in batch["prompt_label"]]24    return batch25 26 27def _prepare_split(ds: Dataset) -> Dataset:28    """29    Keep only the required columns and normalize labels for a single split.30    """31    subset = ds.select_columns(SELECTED_COLUMNS)32    return subset.map(_map_labels, batched=True)33 34 35def load_aegis_dataset(36    split: Optional[str] = None,37    streaming: bool = False,38) -> Dataset | DatasetDict | IterableDataset | IterableDatasetDict:39    """40    Load the Aegis dataset with normalized `prompt_label`.41 42    Args:43        split: Optional split name ("train", "validation", "test", etc.).44        streaming: Whether to stream the data instead of downloading it locally.45 46    Returns:47        A processed Dataset (if split is provided) or DatasetDict containing only48        `prompt` and integer `prompt_label` columns.49    """50    dataset = load_dataset(DATASET_NAME, split=split, streaming=streaming)51 52    if split is not None:53        if streaming:54            # IterableDataset does not support select_columns/map the same way.55            def generator():56                for row in dataset:57                    yield {58                        "prompt": row["prompt"],59                        "prompt_label": LABEL_MAP[row["prompt_label"]],60                    }61 62            return IterableDataset.from_generator(generator)63 64        return _prepare_split(dataset)65 66    # Multiple splits.67    if streaming:68        processed = {}69        for split_name, iterable in dataset.items():70            def make_iter(it):71                def generator():72                    for row in it:73                        yield {74                            "prompt": row["prompt"],75                            "prompt_label": LABEL_MAP[row["prompt_label"]],76                        }77 78                return IterableDataset.from_generator(generator)79 80            processed[split_name] = make_iter(iterable)81        return IterableDatasetDict(processed)82 83    return DatasetDict({split_name: _prepare_split(split_ds) for split_name, split_ds in dataset.items()})84 85 86if __name__ == "__main__":87    processed = load_aegis_dataset()88    for split_name, split_ds in processed.items():89        print(f"{split_name}: {len(split_ds)} samples")90        print(split_ds[0])91 92