Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/Antisource/heronaisec/llms.txt

Use this file to discover all available pages before exploring further.

The data module provides the complete dataset preparation pipeline for Heron AI Security experiments. It is intentionally model-agnostic and attack-agnostic: it knows nothing about backdoor injections, training loops, or evaluation metrics. Its responsibilities are strictly limited to loading datasets from the Hugging Face Hub, validating that the expected schema is present, stripping whitespace, building tokenizers, converting text to token IDs, and setting the PyTorch tensor format. Tokenization is deliberately kept separate from loading so that attacks can be applied to raw text before any token-level transformation occurs.

​
Constants

​
REQUIRED_COLUMNS

REQUIRED_COLUMNS = {"sentence", "label"}
The set of column names every dataset split must contain. If a custom dataset omits either column, validate_dataset() raises a ValueError before any processing begins. To support a dataset with different field names, rename the columns upstream (e.g. dataset.rename_column("text", "sentence")) before calling prepare_dataset().

​
Functions

​
load_text_dataset()

Load a text classification dataset from the Hugging Face Hub. A thin wrapper around datasets.load_dataset() that returns the full DatasetDict with all available splits.
​
dataset_name
str
required
Hugging Face dataset repository name (e.g. "glue", "imdb").
​
subset
str
required
Dataset configuration name within the repository (e.g. "sst2", "mnli").
​
return
DatasetDict
A DatasetDict containing all available splits ("train", "validation", "test", etc.).
Example: 
from src.data import load_text_dataset

dataset = load_text_dataset("glue", "sst2")
print(dataset)
# DatasetDict({
#     train: Dataset({features: ['sentence', 'label', 'idx'], num_rows: 67349}),
#     validation: Dataset({...}),
#     test: Dataset({...})
# })

​
validate_dataset()

Validate that every split in a DatasetDict contains the required columns. Iterates over every split and computes the set difference between REQUIRED_COLUMNS and the split’s column_names. Raises immediately on the first violation — it does not accumulate errors across splits.
​
dataset
DatasetDict
required
The dataset to validate. All splits are checked against REQUIRED_COLUMNS.
​
return
None
Returns None on success.
A ValueError is raised if any split is missing "sentence" or "label". The error message names the offending split and the missing column(s), e.g. "train split is missing columns: {'sentence'}".
Example: 
from datasets import DatasetDict, Dataset
from src.data import validate_dataset

# This will raise because "sentence" is absent
bad = DatasetDict({"train": Dataset.from_dict({"text": ["hello"], "label": [0]})})
try:
    validate_dataset(bad)
except ValueError as exc:
    print(exc)
# ValueError: train split is missing columns: {'sentence'}

​
preprocess_dataset()

Apply generic text-level preprocessing to a DatasetDict. Currently strips leading and trailing whitespace from every "sentence" field. This operation is idempotent and does not alter labels or any other metadata. The original dataset object is not mutated; a new mapped DatasetDict is returned.
​
dataset
DatasetDict
required
A validated DatasetDict containing a "sentence" column in every split.
​
return
DatasetDict
A new DatasetDict with whitespace stripped from all "sentence" values.
Example: 
from src.data import load_text_dataset, preprocess_dataset

dataset = load_text_dataset("glue", "sst2")
cleaned = preprocess_dataset(dataset)
print(cleaned["train"][0]["sentence"])
# "hide new secretions from the parental units"

​
build_tokenizer()

Load the tokenizer corresponding to a pretrained model checkpoint. A thin wrapper around AutoTokenizer.from_pretrained(). The returned tokenizer is fully configured for the target model and is shared across tokenize_dataset() and the Hugging Face Trainer.
​
model_name
str
required
Hugging Face model identifier (e.g. "distilbert-base-uncased", "bert-base-cased").
​
return
AutoTokenizer
A tokenizer instance ready to encode text sequences.
Example: 
from src.data import build_tokenizer

tokenizer = build_tokenizer("distilbert-base-uncased")
print(tokenizer.model_max_length)
# 512

​
tokenize_dataset()

Tokenize the "sentence" field of a dataset using a pretrained tokenizer. Applies the tokenizer in batched mode with truncation=True. Accepts either a Dataset (single split) or a DatasetDict (all splits are tokenized). The function adds "input_ids" and "attention_mask" columns; padding is handled dynamically at collation time in the training loop rather than here.
​
dataset
Dataset | DatasetDict
required
The dataset to tokenize. Must contain a "sentence" column.
​
tokenizer
AutoTokenizer
required
A tokenizer instance obtained from build_tokenizer() or AutoTokenizer.from_pretrained().
​
return
Dataset | DatasetDict
The dataset with "input_ids" and "attention_mask" columns appended. The return type mirrors the input type.
Call tokenize_dataset() after any attack transformations have been applied. Tokenizing before poisoning means the trigger token would need to be injected at the token-ID level instead of the text level.
Example: 
from src.data import load_text_dataset, build_tokenizer, tokenize_dataset

dataset = load_text_dataset("glue", "sst2")
tokenizer = build_tokenizer("distilbert-base-uncased")

tokenized = tokenize_dataset(dataset["train"], tokenizer)
print(tokenized.column_names)
# ['sentence', 'label', 'idx', 'input_ids', 'attention_mask']

​
prepare_for_training()

Convert a tokenized dataset into PyTorch tensor format. Calls dataset.with_format(type="torch", columns=["input_ids", "attention_mask", "label"]). Only the three columns required by the Hugging Face Trainer are surfaced; all other columns (e.g. "sentence", "idx") are dropped from the tensor view.
​
dataset
Dataset
required
A tokenized Dataset that contains "input_ids", "attention_mask", and "label" columns.
​
return
Dataset
The same dataset object formatted as PyTorch tensors, restricted to the three training columns.
Example: 
from src.data import (
    load_text_dataset,
    build_tokenizer,
    tokenize_dataset,
    prepare_for_training,
)

dataset = load_text_dataset("glue", "sst2")
tokenizer = build_tokenizer("distilbert-base-uncased")

tokenized = tokenize_dataset(dataset["train"], tokenizer)
ready = prepare_for_training(tokenized)

sample = ready[0]
print(type(sample["input_ids"]))
# <class 'torch.Tensor'>

​
prepare_dataset()

Execute the full load → validate → preprocess → tokenizer pipeline in a single call. This is the standard entry-point for experiment scripts. It returns the raw (untokenized) train and validation splits together with the tokenizer so that attack transformations can be applied to the text before tokenization occurs.
​
dataset_name
str
required
Hugging Face dataset repository name (e.g. "glue").
​
subset
str
required
Dataset configuration name (e.g. "sst2").
​
model_name
str
required
Hugging Face model identifier used to load the matching tokenizer.
​
return
tuple[Dataset, Dataset, AutoTokenizer]
A three-element tuple of (train_dataset, validation_dataset, tokenizer). Both dataset splits are preprocessed but not yet tokenized.
The typical experiment workflow is: call prepare_dataset() → apply attacks to train_dataset via apply_attack() → call tokenize_dataset() and prepare_for_training() on both splits → pass to build_trainer().
Example: 
from src.data import prepare_dataset, tokenize_dataset, prepare_for_training
from src.attacks import apply_attack

train, validation, tokenizer = prepare_dataset(
    dataset_name="glue",
    subset="sst2",
    model_name="distilbert-base-uncased",
)

# Apply attack to raw text before tokenization
poisoned_train = apply_attack(
    dataset=train,
    attack_type="badnet",
    poison_rate=0.10,
    trigger="cf",
    target_label=1,
    seed=42,
)

# Now tokenize and format
train_ready = prepare_for_training(tokenize_dataset(poisoned_train, tokenizer))
val_ready   = prepare_for_training(tokenize_dataset(validation, tokenizer))

Build docs developers (and LLMs) love

Get started for freeTalk to us