olmo_tap.experiments.utils.weights

Helpers for exporting slim LoRA-only state dicts from fat training checkpoints.

Functions

extract_lora_state(state_dict)	Return only the PEFT LoRA A/B tensors from a head state dict.
save_lora_export(ckpt_path, out_path)	Load a fat robustness/security checkpoint, slim to LoRA only, save, return bytes written.

olmo_tap.experiments.utils.weights.extract_lora_state(state_dict: Mapping[str, Tensor]) → dict[str, Tensor]

Return only the PEFT LoRA A/B tensors from a head state dict.

olmo_tap.experiments.utils.weights.save_lora_export(ckpt_path: Path | str, out_path: Path | str) → int

Load a fat robustness/security checkpoint, slim to LoRA only, save, return bytes written.