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 models module manages the lifecycle of pretrained sequence classification models: constructing them from the Hugging Face Hub, persisting trained checkpoints to disk, and reloading those checkpoints for evaluation or downstream fine-tuning. It is intentionally free of training, evaluation, and attack logic — its sole concern is the model object itself. Tokenizer saving is handled separately in src.train, not here, to keep responsibilities clearly separated.

​
Functions

​
build_model()

Construct a pretrained sequence classification model ready for fine-tuning. Calls AutoModelForSequenceClassification.from_pretrained() with the supplied model_name and num_labels. The classification head is randomly initialised on top of the pretrained backbone.
​
model_name
str
required
Hugging Face model identifier (e.g. "distilbert-base-uncased", "bert-base-cased"). Any model supported by AutoModelForSequenceClassification can be used.
​
num_labels
int
required
Number of output classes for the classification head. Use 2 for binary tasks such as SST-2 sentiment analysis.
​
return
PreTrainedModel
A PreTrainedModel instance with pretrained encoder weights and a freshly initialised classification head, ready for Trainer-based fine-tuning.
Building a model with num_labels different from the checkpoint’s original configuration reinitialises the classification head and will log a warning from the Transformers library. This is expected behaviour during task-specific fine-tuning.
Example: 
from src.models import build_model

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

print(model.config.num_labels)
# 2
print(type(model))
# <class 'transformers.models.distilbert.modeling_distilbert.DistilBertForSequenceClassification'>

​
save_model()

Persist a trained model checkpoint to disk. Calls model.save_pretrained(output_dir), which writes config.json and the model weights file (pytorch_model.bin or sharded equivalents) to output_dir. The tokenizer is not saved by this function — use tokenizer.save_pretrained() separately, or rely on train_model() in src.train which handles both in one call.
​
model
PreTrainedModel
required
The trained model to save.
​
output_dir
str | Path
required
Destination directory. The directory is created by Hugging Face if it does not already exist.
​
return
None
Returns None. Raises an OSError if the path is not writable.
save_model() does not save the tokenizer. If you later call load_model() and need to tokenize text, ensure you have saved the tokenizer separately. In practice, train_model() from src.train saves both atomically.
Example: 
from src.models import build_model, save_model

model = build_model("distilbert-base-uncased", num_labels=2)
# ... fine-tuning happens here ...
save_model(model, output_dir="results/baseline/checkpoints/final")

​
load_model()

Reload a previously saved model checkpoint from disk. Calls AutoModelForSequenceClassification.from_pretrained(checkpoint_dir). Both local paths and Hugging Face Hub identifiers are accepted, since the function delegates directly to the Transformers library.
​
checkpoint_dir
str | Path
required
Path to a local directory containing config.json and model weights, or a Hugging Face Hub repository identifier.
​
return
PreTrainedModel
A PreTrainedModel restored from the checkpoint, in evaluation mode by default.
Example — full save and reload round-trip: 
from src.models import build_model, save_model, load_model

# Build and (hypothetically) fine-tune
model = build_model("distilbert-base-uncased", num_labels=2)

# Save
save_model(model, output_dir="results/run_01/checkpoints/final")

# Reload later for evaluation
restored = load_model("results/run_01/checkpoints/final")
print(restored.config.num_labels)
# 2
Example — loading a backdoor-poisoned checkpoint for intervention experiments: 
from src.models import load_model
from src.interventions import apply_clean_intervention

poisoned_model = load_model("results/poisoned/checkpoints/final")

# Apply a clean fine-tuning intervention without touching the original
intervened = apply_clean_intervention(
    model=poisoned_model,
    clean_dataset=clean_train,
    tokenizer=tokenizer,
    ratio=0.25,
    output_dir="results/intervened",
    epochs=3,
    learning_rate=2e-5,
    batch_size=32,
    weight_decay=0.01,
    seed=42,
)

Build docs developers (and LLMs) love

Get started for freeTalk to us