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 train module encapsulates all Hugging Face Trainer configuration and fine-tuning execution. It is intentionally isolated from attack generation, evaluation, and experiment orchestration — its only job is to wire together a TrainingArguments object, a DataCollatorWithPadding, and a Trainer, then execute the training loop and persist the outputs. Both model weights and the tokenizer are saved atomically at the end of train_model() to ensure that every checkpoint directory is self-contained and immediately loadable.

​
Functions

​
build_trainer()

Construct a fully configured Hugging Face Trainer instance. Creates a TrainingArguments object with opinionated defaults for reproducible research (logging_strategy="epoch", save_strategy="epoch", report_to="none"), attaches a DataCollatorWithPadding for efficient dynamic padding, and returns a Trainer ready to call .train() on.
​
model
PreTrainedModel
required
The model to fine-tune. Typically produced by build_model() or load_model() from src.models.
​
train_dataset
Dataset
required
A PyTorch-formatted Dataset containing "input_ids", "attention_mask", and "label" columns. Prepare with prepare_for_training() from src.data.
​
tokenizer
AutoTokenizer
required
The tokenizer matching the model checkpoint. Passed to the Trainer as processing_class=tokenizer and to DataCollatorWithPadding for dynamic padding. Also saved alongside model weights by train_model().
​
output_dir
str
required
Directory where per-epoch checkpoints are written by the Trainer. Pass a path inside your run directory (e.g. "results/run_01/checkpoints").
​
epochs
int
required
Total number of training epochs (num_train_epochs).
​
learning_rate
float
required
Peak learning rate for the AdamW optimiser. Cast to float internally to tolerate YAML-loaded values.
​
batch_size
int
required
Per-device training batch size (per_device_train_batch_size).
​
weight_decay
float
required
L2 regularisation coefficient applied to all non-bias parameters. Cast to float internally.
​
seed
int
required
Random seed passed to TrainingArguments for reproducible data sampling and weight initialisation within the Trainer.
​
return
Trainer
A configured Trainer instance. Call .train() on the result, or pass directly to train_model().
report_to="none" disables all experiment-tracker integrations (Weights & Biases, TensorBoard, etc.) by default. To enable logging, override this in your config and pass the value explicitly when calling build_trainer().
Example: 
from src.train import build_trainer
from src.models import build_model
from src.data import prepare_dataset, tokenize_dataset, prepare_for_training
from src.attacks import apply_attack

train_raw, val_raw, tokenizer = prepare_dataset("glue", "sst2", "distilbert-base-uncased")

poisoned = apply_attack(train_raw, "badnet", poison_rate=0.1, trigger="cf", target_label=1, seed=42)
train_ready = prepare_for_training(tokenize_dataset(poisoned, tokenizer))

model = build_model("distilbert-base-uncased", num_labels=2)

trainer = build_trainer(
    model=model,
    train_dataset=train_ready,
    tokenizer=tokenizer,
    output_dir="results/run_01/checkpoints",
    epochs=3,
    learning_rate=2e-5,
    batch_size=32,
    weight_decay=0.01,
    seed=42,
)

​
train_model()

Execute the training loop and save the final model checkpoint and tokenizer. Calls trainer.train(), then saves the model via trainer.save_model(save_directory) and the tokenizer via tokenizer.save_pretrained(save_directory). This ensures the checkpoint directory is self-contained: it holds the model weights, the model config, and the tokenizer files needed to reload for inference or further fine-tuning.
​
trainer
Trainer
required
A configured Trainer instance, typically from build_trainer().
​
save_directory
str | Path
required
Directory to write the final checkpoint. Both model and tokenizer are saved here. This is distinct from the per-epoch output_dir in TrainingArguments, though they can be the same path.
​
tokenizer
AutoTokenizer
required
The tokenizer to save alongside the model weights. Should be the same tokenizer used to tokenize train_dataset.
​
return
PreTrainedModel
The trained trainer.model object, which remains on the same device it was trained on.
The returned model is live in memory and ready for immediate evaluation with predict() from src.evaluation. You do not need to reload it from disk unless starting a new Python process.
Example — full training pipeline: 
from src.train import build_trainer, train_model
from src.models import build_model
from src.data import prepare_dataset, tokenize_dataset, prepare_for_training
from src.attacks import apply_attack

train_raw, val_raw, tokenizer = prepare_dataset("glue", "sst2", "distilbert-base-uncased")
poisoned = apply_attack(train_raw, "badnet", poison_rate=0.1, trigger="cf", target_label=1, seed=42)
train_ready = prepare_for_training(tokenize_dataset(poisoned, tokenizer))

model = build_model("distilbert-base-uncased", num_labels=2)

trainer = build_trainer(
    model=model,
    train_dataset=train_ready,
    tokenizer=tokenizer,
    output_dir="results/run_01/checkpoints",
    epochs=3,
    learning_rate=2e-5,
    batch_size=32,
    weight_decay=0.01,
    seed=42,
)

trained_model = train_model(
    trainer=trainer,
    save_directory="results/run_01/checkpoints/final",
    tokenizer=tokenizer,
)

print(f"Training complete. Model saved to results/run_01/checkpoints/final")

Build docs developers (and LLMs) love

Get started for freeTalk to us