olmo_tap.experiments.uncertainty.training

HydraTransformer Double-Head Uncertainty Finetuning Pipeline.

Uses the FrozenHeadHandler to cycle through frozen LLM heads, each loaded with OLMo base + prod security LoRA + robustness LoRA. Uncertainty head is trained to predict confidence score of the multiple choice answer.

Intended Usage (run from tap root)::

# cycles through randomly selected heads for 100 steps each pixi run python -m experiments.uncertainty.training –num-epochs 5 –swap-freq 100

Functions

compute_total_steps(num_shards, batch_size, ...)
main()
parse_args()

olmo_tap.experiments.uncertainty.training.compute_total_steps(num_shards: int, batch_size: int, num_epochs: int) → int

olmo_tap.experiments.uncertainty.training.main()

olmo_tap.experiments.uncertainty.training.parse_args() → Namespace