LLM Utils

Helpers for LLM training (FSDP gathers, model loading, eval sampling, chat templating). For Gymnasium-style LLM datasets, see agilerl.llm_envs.

agilerl.utils.llm_utils.apply_chat_template(conversation_template: list[dict[str, str]], question: str, answer: str, tokenizer: PreTrainedTokenizerBase) → BatchEncoding

Create and tokenize a chat template for a reasoning task.

Parameters:

  • conversation_template (list[dict[str, str]]) – The conversation template to be tokenized.

  • question (str) – The question to be tokenized.

  • answer (str) – The answer to be tokenized.

  • tokenizer (PreTrainedTokenizerBase) – The tokenizer to be used.

Returns:

The tokenized prompt.

Return type:

BatchEncoding

agilerl.utils.llm_utils.get_lora_params(model: Module) → list[Tensor]

Return adapter parameters for scoped gathers / export.

Pass the result to agilerl.distributed.gather_params() for adapter-only save, export, or copy so base-model shards stay unmaterialised.

Parameters:

model (nn.Module) – PEFT model (or a wrapper around one).

Returns:

Adapter parameters.

Return type:

list[torch.Tensor]

agilerl.utils.llm_utils.create_model_from_name_or_path(model_name_or_path: str, model_config: dict[str, Any] | None = None, add_value_head: bool = False, use_distributed: bool = False) → PreTrainedModel

Create a model from a name or path.

Parameters:

  • model_name_or_path (str) – The name or path of the model to create.

  • model_config (dict[str, Any ] | None) – Extra keyword arguments forwarded to from_pretrained (e.g. a quantization_config). torch_dtype fills if absent, then attn_implementation is resolved (env, family trainer default, explicit value, then flash/sdpa). Other family trainer keys fill via setdefault.

  • add_value_head (bool, optional) – Flag to indicate if a value head should be added to the model, defaults to False

  • use_distributed (bool, optional) – Whether the model is created for a distributed run, defaults to False

Returns:

The created model.

Return type:

PreTrainedModel

agilerl.utils.llm_utils.sample_eval_prompts(env: object, n: int = 5, seed: int = 0) → list[tuple[str, str | None, str | None]]

Randomly sample n (prompt, chosen, rejected) triples from env’s held-out test dataset.

Columns are resolved automatically per gym type:

  • SFTGym — chosen is env.response_column; rejected is None (SFT has no negative example).

  • PreferenceGym — chosen and rejected map to the dataset’s "chosen" / "rejected" columns.

  • Any other gym — both are None.

Parameters:

  • env (object) – AgileRL gym environment with a test_dataloader attribute.

  • n (int, optional) – Number of samples to draw, defaults to 5.

  • seed (int, optional) – Random seed for reproducible sampling, defaults to 0.

Returns:

List of (prompt, chosen, rejected) tuples; unused fields are None.

Return type:

list[tuple[str, str | None, str | None]]

agilerl.utils.llm_utils.compare_responses(agent: LLMAlgorithm, tokenizer: PreTrainedTokenizerBase, samples: list[tuple[str, str | None, str | None]], max_new_tokens: int = 200, temperature: float = 1.0, do_sample: bool = False, skip_special_tokens: bool = True, show_base_model: bool = True) → None

Run each prompt through the base model and the fine-tuned LoRA model, printing a formatted comparison to the terminal one sample at a time.

After each sample the user is prompted to press Enter to continue or q + Enter to quit early. Intended to be called at the end of a training script for a quick qualitative sanity-check.

Works with any LoRA-adapted LLMAlgorithm (SFT, DPO, …). When the model has no LoRA adapter the base-model column is omitted and only the current model’s output is shown.

Parameters:

  • agent (LLMAlgorithm) – Trained AgileRL LLM agent exposing agent.actor and agent.device.

  • tokenizer (object) – HuggingFace tokenizer matching the model.

  • samples (list[tuple[str, str | None, str | None]]) – (prompt, chosen, rejected) triples as returned by sample_eval_prompts(). None fields are silently skipped.

  • max_new_tokens (int, optional) – Maximum tokens to generate per response, defaults to 200.

  • temperature (float, optional) – Sampling temperature, defaults to 1.0.

  • do_sample (bool, optional) – Use sampling instead of greedy decoding, defaults to False. Set True together with a temperature != 1.0 for stochastic outputs.

  • skip_special_tokens (bool, optional) – Strip special tokens when decoding, defaults to True.

  • show_base_model (bool, optional) – If False, skip the base-model generation block (only the current model output is printed). Useful when the adapter is merged or base vs. adapter outputs are identical.