olmo_tap.experiments.utils.model_builder

Functions to support loading models for inference and training.

Functions

build_base_model(config)
inject_lora(model, config[, head_idx])
load_and_merge_lora_weights(model, config, ...)

olmo_tap.experiments.utils.model_builder.build_base_model(config: HydraLoRAConfig) → HydraTransformer

Parameters:

config – Config file detailing architecture of model to be loaded

Returns HydraTransformer:

OLMo with base weights

olmo_tap.experiments.utils.model_builder.inject_lora(model: HydraTransformer, config: HydraLoRAConfig, head_idx: int = 0)

Parameters:

• model – HydraTransformer model to inject trainable LoRA weights into.

• config – Config file detailing LoRA params (rank, alpha, target_modules).

• head_idx – Which Hydra index to load trainable LoRA weights into (default 0).

olmo_tap.experiments.utils.model_builder.load_and_merge_lora_weights(model: HydraTransformer, config: HydraLoRAConfig, weights_path: Path | str, head_idx: int = 0)

Parameters:

• model – HydraTransformer model to add trained LoRA weights to.

• config – Config file detailing LoRA params (rank, alpha, target_modules).

• weights_path – Path of saved LoRA weights.

• head_idx – Which Hydra index to add trained LoRA weights to.