""" AI Model Training & Fine-tuning Service Allows clients to train custom models on their data """ from typing import List, Dict, Any, Optional from datetime import datetime from pathlib import Path import json import logging import httpx from enum import Enum logger = logging.getLogger(__name__) class TrainingDataFormat(str, Enum): """Supported training data formats""" CONVERSATIONS = "conversations" # Chat format QA_PAIRS = "qa_pairs" # Question-Answer pairs DOCUMENTS = "documents" # Document corpus JSON = "json" # Custom JSON format class TrainingStatus(str, Enum): """Training job status""" PENDING = "pending" PREPROCESSING = "preprocessing" TRAINING = "training" VALIDATING = "validating" COMPLETED = "completed" FAILED = "failed" CANCELLED = "cancelled" class TrainingJob: """Represents a model training job""" def __init__( self, job_id: str, client_id: int, model_name: str, base_model: str, training_data: List[Dict[str, Any]], validation_data: Optional[List[Dict[str, Any]]] = None ): self.job_id = job_id self.client_id = client_id self.model_name = model_name self.base_model = base_model self.training_data = training_data self.validation_data = validation_data or [] self.status = TrainingStatus.PENDING self.progress = 0.0 self.created_at = datetime.now() self.started_at: Optional[datetime] = None self.completed_at: Optional[datetime] = None self.error_message: Optional[str] = None self.metrics: Dict[str, float] = {} def to_dict(self) -> Dict[str, Any]: return { "job_id": self.job_id, "client_id": self.client_id, "model_name": self.model_name, "base_model": self.base_model, "status": self.status.value, "progress": self.progress, "created_at": self.created_at.isoformat(), "started_at": self.started_at.isoformat() if self.started_at else None, "completed_at": self.completed_at.isoformat() if self.completed_at else None, "error_message": self.error_message, "metrics": self.metrics, "training_samples": len(self.training_data), "validation_samples": len(self.validation_data) } class ModelTrainingService: """ Handles model training and fine-tuning Supports: - Ollama Modelfile creation and training - OpenAI fine-tuning (if API key provided) - Custom model training from conversation data """ def __init__(self, ollama_url: str = "http://ollama:11434"): self.ollama_url = ollama_url self.training_jobs: Dict[str, TrainingJob] = {} self.models_dir = Path("./custom_models") self.models_dir.mkdir(exist_ok=True) async def create_training_job( self, client_id: int, model_name: str, base_model: str, training_data: List[Dict[str, Any]], validation_data: Optional[List[Dict[str, Any]]] = None, data_format: TrainingDataFormat = TrainingDataFormat.CONVERSATIONS ) -> TrainingJob: """ Create a new training job Args: client_id: Client ID model_name: Name for the custom model base_model: Base model to fine-tune (e.g., "llama3.2:3b") training_data: Training examples validation_data: Validation examples (optional) data_format: Format of training data Returns: TrainingJob instance """ job_id = f"train_{client_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}" # Validate and preprocess data processed_data = self._preprocess_data(training_data, data_format) processed_validation = self._preprocess_data(validation_data or [], data_format) if validation_data else [] job = TrainingJob( job_id=job_id, client_id=client_id, model_name=f"{model_name}_{client_id}", base_model=base_model, training_data=processed_data, validation_data=processed_validation ) self.training_jobs[job_id] = job logger.info(f"Created training job: {job_id}") return job def _preprocess_data( self, data: List[Dict[str, Any]], data_format: TrainingDataFormat ) -> List[Dict[str, Any]]: """Preprocess training data based on format""" if data_format == TrainingDataFormat.CONVERSATIONS: # Format: [{"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}] return data elif data_format == TrainingDataFormat.QA_PAIRS: # Format: [{"question": "...", "answer": "..."}] # Convert to conversation format return [ { "messages": [ {"role": "user", "content": item["question"]}, {"role": "assistant", "content": item["answer"]} ] } for item in data ] elif data_format == TrainingDataFormat.DOCUMENTS: # Format: [{"content": "document text..."}] # Extract Q&A pairs from documents (simplified) conversations = [] for doc in data: # Split document into chunks and create Q&A content = doc["content"] chunks = self._split_into_chunks(content, chunk_size=500) for i, chunk in enumerate(chunks): conversations.append({ "messages": [ {"role": "user", "content": f"Tell me about: {doc.get('title', f'section {i}')}"}, {"role": "assistant", "content": chunk} ] }) return conversations else: # Assume already in correct format return data def _split_into_chunks(self, text: str, chunk_size: int = 500) -> List[str]: """Split text into chunks""" words = text.split() chunks = [] current_chunk = [] current_size = 0 for word in words: if current_size + len(word) > chunk_size and current_chunk: chunks.append(" ".join(current_chunk)) current_chunk = [word] current_size = len(word) else: current_chunk.append(word) current_size += len(word) + 1 if current_chunk: chunks.append(" ".join(current_chunk)) return chunks async def train_ollama_model(self, job_id: str, system_prompt: Optional[str] = None) -> bool: """ Train custom Ollama model using Modelfile This creates a fine-tuned version with custom system prompt and examples """ job = self.training_jobs.get(job_id) if not job: raise ValueError(f"Training job not found: {job_id}") try: job.status = TrainingStatus.PREPROCESSING job.started_at = datetime.now() # Create Modelfile modelfile_path = self.models_dir / f"{job.model_name}.modelfile" modelfile_content = self._create_ollama_modelfile( base_model=job.base_model, system_prompt=system_prompt, training_examples=job.training_data[:10] # Include some examples ) # Save Modelfile modelfile_path.write_text(modelfile_content) logger.info(f"Created Modelfile at: {modelfile_path}") job.status = TrainingStatus.TRAINING job.progress = 0.5 # Create model using Ollama async with httpx.AsyncClient(timeout=300.0) as client: # Read the modelfile with open(modelfile_path, 'r') as f: modelfile = f.read() response = await client.post( f"{self.ollama_url}/api/create", json={ "name": job.model_name, "modelfile": modelfile } ) if response.status_code != 200: raise Exception(f"Failed to create model: {response.text}") logger.info(f"Successfully created Ollama model: {job.model_name}") job.status = TrainingStatus.COMPLETED job.completed_at = datetime.now() job.progress = 1.0 job.metrics = { "training_samples": len(job.training_data), "model_size": "estimated_3b_params" } return True except Exception as e: logger.error(f"Training failed for job {job_id}: {e}") job.status = TrainingStatus.FAILED job.error_message = str(e) return False def _create_ollama_modelfile( self, base_model: str, system_prompt: Optional[str], training_examples: List[Dict[str, Any]] ) -> str: """ Create Ollama Modelfile with custom configuration """ lines = [ f"FROM {base_model}", "" ] # Add system prompt if system_prompt: lines.append(f'SYSTEM """{system_prompt}"""') lines.append("") # Add parameters for better performance lines.extend([ "PARAMETER temperature 0.7", "PARAMETER top_p 0.9", "PARAMETER top_k 40", "PARAMETER num_ctx 4096", "" ]) # Add few-shot examples as SYSTEM messages if training_examples: examples_text = "Here are some example interactions:\n\n" for i, example in enumerate(training_examples[:5], 1): messages = example.get("messages", []) if len(messages) >= 2: user_msg = next((m["content"] for m in messages if m["role"] == "user"), "") assistant_msg = next((m["content"] for m in messages if m["role"] == "assistant"), "") examples_text += f"Example {i}:\nUser: {user_msg}\nAssistant: {assistant_msg}\n\n" lines.append(f'SYSTEM """{examples_text}"""') return "\n".join(lines) async def get_training_status(self, job_id: str) -> Dict[str, Any]: """Get status of a training job""" job = self.training_jobs.get(job_id) if not job: raise ValueError(f"Training job not found: {job_id}") return job.to_dict() async def list_custom_models(self, client_id: Optional[int] = None) -> List[Dict[str, Any]]: """List all custom trained models""" try: async with httpx.AsyncClient() as client: response = await client.get(f"{self.ollama_url}/api/tags") if response.status_code != 200: return [] data = response.json() models = data.get("models", []) # Filter by client_id if provided if client_id: models = [m for m in models if f"_{client_id}" in m.get("name", "")] return [ { "name": m.get("name"), "size": m.get("size", 0), "modified_at": m.get("modified_at"), "client_id": self._extract_client_id(m.get("name", "")) } for m in models ] except Exception as e: logger.error(f"Failed to list custom models: {e}") return [] def _extract_client_id(self, model_name: str) -> Optional[int]: """Extract client ID from model name""" try: parts = model_name.split("_") return int(parts[-1]) except: return None async def delete_custom_model(self, model_name: str) -> bool: """Delete a custom model""" try: async with httpx.AsyncClient() as client: response = await client.delete( f"{self.ollama_url}/api/delete", json={"name": model_name} ) return response.status_code == 200 except Exception as e: logger.error(f"Failed to delete model {model_name}: {e}") return False def export_training_data(self, job_id: str, output_path: Optional[str] = None) -> str: """Export training data to file""" job = self.training_jobs.get(job_id) if not job: raise ValueError(f"Training job not found: {job_id}") if output_path is None: output_path = str(self.models_dir / f"{job.model_name}_training_data.jsonl") # Export as JSONL (one JSON per line) with open(output_path, 'w') as f: for example in job.training_data: f.write(json.dumps(example) + "\n") logger.info(f"Exported training data to: {output_path}") return output_path