Skip to content

Harden API

Train-time defenses. Import from safetune.runner.harden. Each trainer takes model, tokenizer, **hyperparams and exposes train(train_dataset, out_dir=...).

from safetune.runner.harden import SafeGradTrainer

trainer = SafeGradTrainer(model, tokenizer)
path = trainer.train(train_dataset, out_dir="./hardened")

Available trainers

AntibodyTrainer, AsFTTrainer, BoosterTrainer, CSTTrainer, CTRAPTrainer, ConstrainedSFTTrainer, DOORTrainer, DeRTaTrainer, DeepRefusalTrainer, LisaTrainer, LoXHardenTrainer, LookAheadTrainer, MARTTrainer, PlainSFTTrainer, RepNoiseTrainer, SAPTrainer, SEALTrainer, SEAMTrainer, SPPFTTrainer, STARDSSTrainer, SaLoRATrainer, SafeGradTrainer, SurgeryTrainer, TARTrainer, TVaccineTrainer, VaccineTrainer.

See the Harden guide for method selection. A few methods (CSTTrainer, MARTTrainer, DeepRefusalTrainer, AntibodyTrainer) are programmatic-only — see the CLI Reference.

Reference

safetune.runner.harden.SafeGradTrainer

Bases: _HardenBase

SafeGrad: gradient surgery + KL alignment vs frozen reference.

Parameters:

Name Type Description Default
rho float

gradient surgery mixing weight. Default 1.0.

1.0
kl_temperature float

KL alignment temperature. Default 1.0.

1.0
reference_model_path str

HF path/ID for the reference model.

None

safetune.runner.harden.LisaTrainer

Bases: _HardenBase

Lisa: bi-state proximal optimization (alignment/finetune alternation).

Parameters:

Name Type Description Default
lisa_rho float

proximal constraint weight. Default 0.1.

0.1
lisa_warmup_steps int

warm-up steps before alternation starts. Default 10.

10
lisa_alignment_step int

steps per alignment phase. Default 20.

20
lisa_finetune_step int

steps per fine-tune phase. Default 20.

20

safetune.runner.harden.SAPTrainer

Bases: _HardenBase

SAP: contrastive alignment + perturbation.

Parameters:

Name Type Description Default
grad_rate float

gradient mixing rate. Default 0.1.

0.1
v_update_step float

virtual gradient update step size. Default 0.05.

0.05
contrastive_temperature float

contrastive loss temperature. Default 1.0.

1.0