Implementation:Run llama Llama index MistralAIFinetuneEngine
Overview
The MistralAIFinetuneEngine provides a fine-tuning engine for Mistral AI language models. It handles training file upload, fine-tuning job creation and monitoring, and retrieval of the fine-tuned model as a LlamaIndex LLM. This module resides in the llama-index-finetuning package under the mistralai submodule.
Source file: llama-index-finetuning/llama_index/finetuning/mistralai/base.py (156 lines)
Dependencies
| Dependency | Purpose |
|---|---|
logging, sys |
Logging setup with stdout stream handlers |
os |
Reading the MISTRAL_API_KEY environment variable
|
time |
Sleep delay when waiting for file readiness |
mistralai.Mistral |
Mistral AI SDK client |
mistralai.models.JobsAPIRoutesFineTuningGetFineTuningJobResponse |
Type for fine-tuning job response objects |
mistralai.models.WandbIntegration |
Weights & Biases integration configuration |
mistralai.models.CompletionTrainingParametersIn |
Training hyperparameters container |
llama_index.core.llms.llm.LLM |
Return type interface for fine-tuned models |
llama_index.finetuning.callbacks.finetuning_handler.MistralAIFineTuningHandler |
Handler for capturing fine-tuning events |
llama_index.finetuning.mistralai.utils.reformat_jsonl |
Utility to validate and reformat JSONL training data |
llama_index.finetuning.types.BaseLLMFinetuneEngine |
Abstract base class defining the fine-tuning engine interface |
llama_index.llms.mistralai.MistralAI |
LlamaIndex Mistral AI LLM wrapper |
Module-Level Configuration
The module configures basic logging at the module level:
logging.basicConfig(stream=sys.stdout, level=logging.INFO) logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
This ensures all log messages (INFO and above) are directed to stdout.
Class: MistralAIFinetuneEngine
Inherits from: BaseLLMFinetuneEngine
Constructor
def __init__(
self,
base_model: str,
training_path: str,
validation_path: Optional[str] = None,
verbose: bool = False,
start_job_id: Optional[str] = None,
validate_json: bool = True,
training_steps: int = 10,
learning_rate: float = 0.0001,
wandb_integration_dict: Optional[Dict[str, str]] = None,
) -> None
| Parameter | Type | Default | Description |
|---|---|---|---|
base_model |
str |
required | Mistral AI model identifier to fine-tune |
training_path |
str |
required | Path to the JSONL training data file |
validation_path |
Optional[str] |
None |
Optional path to a JSONL validation data file |
verbose |
bool |
False |
Enable verbose output |
start_job_id |
Optional[str] |
None |
Resume tracking of an existing fine-tuning job |
validate_json |
bool |
True |
Whether to validate and reformat JSONL data before uploading |
training_steps |
int |
10 |
Number of training steps for the fine-tuning job |
learning_rate |
float |
0.0001 |
Learning rate for the fine-tuning job |
wandb_integration_dict |
Optional[Dict[str, str]] |
None |
Dictionary with keys "project", "run_name", and "api_key" for Weights & Biases logging
|
Initialization behavior:
- Stores all parameters as instance attributes.
- Creates a
Mistralclient using theMISTRAL_API_KEYenvironment variable. - If
start_job_idis provided, retrieves the corresponding job viaself._client.fine_tuning.jobs.get(start_job_id).
Class Method: from_finetuning_handler
@classmethod
def from_finetuning_handler(
cls,
finetuning_handler: MistralAIFineTuningHandler,
base_model: str,
training_path: str,
**kwargs: Any,
) -> "MistralAIFinetuneEngine"
Factory method that creates an engine instance from a MistralAIFineTuningHandler.
Workflow:
- Calls
finetuning_handler.save_finetuning_events(training_path)to persist captured events. - Returns a new
MistralAIFinetuneEngine. Note: the constructor is called withdata_path=training_pathrather thantraining_path=training_path, which appears to be a parameter name mismatch.
Method: finetune
def finetune(self) -> None
Executes the full fine-tuning workflow: validation, upload, and job creation.
Workflow:
- Validation phase: If
self._validate_jsonisTrue, callsreformat_jsonl()on both the training path and validation path (if provided). - Upload phase: Uploads the training file via
self._client.files.upload(). If a validation path is provided, uploads it as well. - Job creation phase: Enters a retry loop that attempts to create a fine-tuning job:
- Calls
self._client.fine_tuning.jobs.create()with:training_files-- list containing the uploaded training file IDvalidation_files-- list containing the uploaded validation file ID (if applicable), orNonemodel-- the base model identifierhyperparameters-- aCompletionTrainingParametersIninstance with the configuredtraining_stepsandlearning_rateintegrations-- optional Weights & Biases integration (aWandbIntegrationdumped to a dictionary)
- If the job creation fails (e.g., file not yet ready), waits 60 seconds and retries.
- On success, stores the job output in
self._start_joband breaks out of the loop.
- Calls
Method: get_current_job
def get_current_job(self) -> Optional[JobsAPIRoutesFineTuningGetFineTuningJobResponse]
Retrieves the current status of the fine-tuning job.
Workflow:
- Raises
ValueErroriffinetune()has not been called (no job exists). - Calls
self._client.fine_tuning.jobs.get(job_id)to fetch the latest job state.
Method: get_finetuned_model
def get_finetuned_model(self, **model_kwargs: Any) -> LLM
Returns a LlamaIndex MistralAI LLM instance configured with the fine-tuned model.
Workflow:
- Retrieves the current job via
self.get_current_job(). - Checks that
fine_tuned_modelis notNoneand status is"SUCCESS". RaisesValueErrorif either check fails. - Returns
MistralAI(model=model_id, **model_kwargs).
Note: Unlike the OpenAI and Azure engines which check for status "succeeded" (lowercase), the Mistral engine checks for "SUCCESS" (uppercase), matching the Mistral API's status convention.
Environment Variables
| Variable | Required | Default | Purpose |
|---|---|---|---|
MISTRAL_API_KEY |
Yes | None |
API key for Mistral AI authentication |
Weights and Biases Integration
The engine supports optional Weights & Biases (W&B) logging during fine-tuning. When wandb_integration_dict is provided, the dictionary must contain three keys:
| Key | Description |
|---|---|
"project" |
W&B project name |
"run_name" |
Name for the W&B run |
"api_key" |
W&B API key for authentication |
The integration is passed to the Mistral API as a serialized WandbIntegration model.
See Also
- Run_llama_Llama_index_AzureOpenAIFinetuneEngine -- Azure OpenAI fine-tuning engine
- Run_llama_Llama_index_CohereRerankerFinetuneEngine -- Cohere reranker fine-tuning engine