Implementation:Run llama Llama index Adapter Training Utils

Overview

This module provides the training loop and loss function for fine-tuning embedding adapter models. It implements a custom multiple negatives ranking loss optimized for LlamaIndex's adapter embeddings, and a full training function with support for warmup scheduling, gradient clipping, checkpointing, and verbose logging. These utilities are used by the embedding fine-tuning engine to train adapter layers on top of pre-trained embedding models.

Source file: llama-index-finetuning/llama_index/finetuning/embeddings/adapter_utils.py (150 lines)

Dependencies

Dependency	Purpose
`pathlib.Path`	Directory creation for checkpoint paths
`torch`	Tensor operations, gradient clipping, device management
`torch.nn`	Base module class and loss functions
`torch.optim.Optimizer`	Optimizer type hint
`transformers`	Linear warmup scheduler via `get_linear_schedule_with_warmup`
`llama_index.core.utils.print_text`	Colored text output for verbose logging
`llama_index.embeddings.adapter.BaseAdapter`	The adapter model interface being trained
`sentence_transformers.util.cos_sim`	Default cosine similarity function for the loss
`tqdm.autonotebook.trange`	Progress bars for epoch and iteration loops

Class: MyMultipleNegativesRankingLoss

Inherits from: torch.nn.Module

A custom loss module that implements the multiple negatives ranking loss. This loss is similar to the one in the sentence_transformers library but is optimized for LlamaIndex's own embedding adapter models, where only the query embeddings are transformed through the adapter while context embeddings are used directly.

Constructor

def __init__(
    self,
    model: BaseAdapter,
    scale: float = 20.0,
    similarity_fct: Optional[Callable] = None,
)

Parameter	Type	Default	Description
`model`	`BaseAdapter`	required	The adapter model whose forward pass transforms query embeddings
`scale`	`float`	`20.0`	Scaling factor applied to similarity scores before cross-entropy loss
`similarity_fct`	`Optional[Callable]`	`None`	Similarity function; defaults to `cos_sim` from sentence-transformers

Internally stores a CrossEntropyLoss instance for computing the final loss value.

Method: forward

def forward(self, query_embeds: Tensor, context_embeds: Tensor) -> Tensor

Workflow:

Transforms query embeddings through the adapter: query_embeds_2 = self.model.forward(query_embeds).
Computes similarity scores between transformed queries and original context embeddings, scaled by self.scale.
Creates target labels as [0, 1, 2, ..., N-1] (each query at index i should match context at index i).
Computes and returns cross-entropy loss between the similarity scores and the target labels.

Design note: Only query embeddings pass through the adapter; context embeddings remain unchanged. This asymmetric design is intentional, as indicated by the commented-out line # context_embeds_2 = self.model.forward(context_embeds).

Function: train_model

def train_model(
    model: BaseAdapter,
    data_loader: torch.utils.data.DataLoader,
    device: torch.device,
    epochs: int = 1,
    steps_per_epoch: Optional[int] = None,
    warmup_steps: int = 10000,
    optimizer_class: Type[Optimizer] = torch.optim.AdamW,
    optimizer_params: Dict[str, Any] = {"lr": 2e-5},
    output_path: str = "model_output",
    max_grad_norm: float = 1,
    show_progress_bar: bool = True,
    verbose: bool = False,
    checkpoint_path: Optional[str] = None,
    checkpoint_save_steps: int = 500,
) -> None

Parameter	Type	Default	Description
`model`	`BaseAdapter`	required	The adapter model to train
`data_loader`	`DataLoader`	required	DataLoader yielding (query_embed, context_embed) tuples
`device`	`torch.device`	required	Target device (CPU or GPU)
`epochs`	`int`	`1`	Number of training epochs
`steps_per_epoch`	`Optional[int]`	`None`	Steps per epoch; defaults to `len(data_loader)`
`warmup_steps`	`int`	`10000`	Number of warmup steps for the linear scheduler
`optimizer_class`	`Type[Optimizer]`	`AdamW`	Optimizer class to use
`optimizer_params`	`Dict[str, Any]`	`{"lr": 2e-5}`	Parameters passed to the optimizer
`output_path`	`str`	`"model_output"`	Directory path where the final model is saved
`max_grad_norm`	`float`	`1`	Maximum gradient norm for gradient clipping
`show_progress_bar`	`bool`	`True`	Whether to display progress bars
`verbose`	`bool`	`False`	Whether to print loss values and status messages
`checkpoint_path`	`Optional[str]`	`None`	Directory for saving intermediate checkpoints; if `None`, no checkpoints are saved
`checkpoint_save_steps`	`int`	`500`	Save a checkpoint every N global steps

Training Loop Workflow

Setup phase:
- Moves the model to the target device.
- Creates a MyMultipleNegativesRankingLoss instance wrapping the model.
- Collects model parameters and creates the optimizer.
- Computes total training steps: steps_per_epoch * epochs.
- Creates a linear warmup scheduler via transformers.get_linear_schedule_with_warmup.
- If checkpoint_path is provided, creates the directory (with parents) if it does not exist.
Training phase: For each epoch:
- Zeros gradients and sets the loss model to training mode.
- For each step within the epoch:
  - Draws the next batch from the data iterator; if the iterator is exhausted, restarts it.
  - Unpacks the batch into query and context tensors and moves them to the device.
  - Computes loss via the loss model's forward pass.
  - Backpropagates the loss.
  - Clips gradients to max_grad_norm.
  - Steps the optimizer and the scheduler.
  - Increments the global step counter.
  - If checkpointing is enabled and the global step is a multiple of checkpoint_save_steps, saves the model to checkpoint_path/step_{global_step}.
Finalization: Saves the trained model to output_path.

Commented-Out Parameters

Several parameters are defined but commented out in the function signature, indicating planned future features:

callback -- Training callback function
scheduler -- Custom scheduler selection
weight_decay -- Weight decay for regularization
evaluation_steps -- Periodic evaluation
save_best_model -- Saving the best model based on evaluation
use_amp -- Automatic mixed precision training
checkpoint_save_total_limit -- Limit on total checkpoint files

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment