Train Module

Training and fine-tuning capabilities for PraisonAI

Train Module

The Train module provides training and fine-tuning capabilities for PraisonAI agents and models.

Import

from praisonai.train import Trainer

Quick Example

from praisonai.train import Trainer

# Create trainer
trainer = Trainer(
    model="gpt-4o-mini",
    training_data="training_data.jsonl"
)

# Start training
result = trainer.train()
print(f"Fine-tuned model: {result.model_id}")

Features

• Fine-tune models on custom datasets
• Support for JSONL training data format
• Integration with OpenAI fine-tuning API
• Training progress monitoring
• Model validation and evaluation

Training Data Format

Training data should be in JSONL format:

{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}]}
{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is AI?"}, {"role": "assistant", "content": "AI stands for Artificial Intelligence..."}]}

Constructor

Trainer(model, training_data)

Creates a new Trainer instance.

Parameters:

Parameter	Type	Default	Description
model	str	Required	Base model to fine-tune
training_data	str	Required	Path to training data file
validation_data	str	None	Path to validation data
epochs	int	3	Number of training epochs
batch_size	int	4	Training batch size

Methods

train()

Start the training process.

Returns: TrainingResult - Contains model_id and metrics

validate()

Validate the training data format.

Returns: bool - True if valid

get_status(job_id)

Get the status of a training job.

Parameters:

• job_id (str): The training job ID

Returns: dict - Job status and progress

cancel(job_id)

Cancel a running training job.

Parameters:

• job_id (str): The training job ID

Example: Full Training Workflow

from praisonai.train import Trainer

# Prepare training data
training_data = [
    {
        "messages": [
            {"role": "system", "content": "You are a customer support agent."},
            {"role": "user", "content": "How do I reset my password?"},
            {"role": "assistant", "content": "To reset your password..."}
        ]
    }
]

# Save as JSONL
import json
with open("training.jsonl", "w") as f:
    for item in training_data:
        f.write(json.dumps(item) + "\n")

# Create trainer
trainer = Trainer(
    model="gpt-4o-mini",
    training_data="training.jsonl",
    epochs=3
)

# Validate data
if trainer.validate():
    # Start training
    result = trainer.train()
    
    # Monitor progress
    while True:
        status = trainer.get_status(result.job_id)
        print(f"Status: {status['status']}")
        if status['status'] in ['succeeded', 'failed']:
            break
    
    # Use fine-tuned model
    print(f"Fine-tuned model: {result.model_id}")

Related

• Agent Module - Use trained models with agents
• LLM Module - LLM configuration