Skip to main content

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:
ParameterTypeDefaultDescription
modelstrRequiredBase model to fine-tune
training_datastrRequiredPath to training data file
validation_datastrNonePath to validation data
epochsint3Number of training epochs
batch_sizeint4Training 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}")