""" Training Service - Core training logic with database persistence """ import os import json import asyncio import logging from datetime import datetime from typing import Dict, Optional, Any, List import torch from transformers import ( AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoModelForTokenClassification, AutoModelForSequenceClassification, AutoModelForQuestionAnswering, AutoConfig, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, DataCollatorForTokenClassification, default_data_collator ) from datasets import load_dataset, Dataset from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training from trl import SFTTrainer # instead of from transformers import Trainer from app.config import settings from app.database import AsyncSessionLocal, TrainingJob, TrainingLog, JobStatus import importlib logger = logging.getLogger(__name__) class TrainingService: """Service for training models with progress tracking and DB persistence.""" def __init__(self): self.active_trainings: Dict[str, Dict] = {} async def _update_job_progress( self, job_id: str, progress: float = None, status: str = None, current_step: int = None, total_steps: int = None, train_loss: float = None, metrics: Dict = None, error_message: str = None ): """Update job progress in database.""" try: async with AsyncSessionLocal() as session: from sqlalchemy import select, update stmt = select(TrainingJob).where(TrainingJob.job_id == job_id) result = await session.execute(stmt) job = result.scalar_one_or_none() if job: if progress is not None: job.progress = progress if status is not None: job.status = status if current_step is not None: job.current_step = current_step if total_steps is not None: job.total_steps = total_steps if train_loss is not None: job.train_loss = train_loss if metrics is not None: job.metrics = metrics if error_message is not None: job.error_message = error_message job.updated_at = datetime.utcnow() if status == JobStatus.RUNNING.value and not job.started_at: job.started_at = datetime.utcnow() if status in [JobStatus.COMPLETED.value, JobStatus.FAILED.value, JobStatus.CANCELLED.value]: job.completed_at = datetime.utcnow() await session.commit() logger.info(f"Updated job {job_id}: progress={progress}, status={status}") except Exception as e: logger.error(f"Failed to update job progress: {e}") async def _add_training_log( self, job_id: str, message: str, level: str = "INFO", step: int = None, loss: float = None, metrics: Dict = None ): """Add a log entry for the training job.""" try: async with AsyncSessionLocal() as session: log = TrainingLog( job_id=(await session.execute( __import__('sqlalchemy').select(TrainingJob.id).where(TrainingJob.job_id == job_id) )).scalar_one_or_none(), level=level, message=message, step=step, loss=loss, metrics=metrics ) session.add(log) await session.commit() except Exception as e: logger.error(f"Failed to add training log: {e}") def _apply_prompt_template( self, example: Dict, column_mapping: Dict, prompt_template: Dict ) -> str: """Apply prompt template to create formatted text.""" # Get template preset or custom sections preset = prompt_template.get('preset', 'none') # Get columns from mapping text_col = column_mapping.get('text', column_mapping.get('input', '')) input_col = column_mapping.get('input', '') output_col = column_mapping.get('output', column_mapping.get('response', '')) instruction_col = column_mapping.get('instruction', '') question_col = column_mapping.get('question', '') answer_col = column_mapping.get('answer', '') context_col = column_mapping.get('context', '') reasoning_col = column_mapping.get('reasoning', '') # Extract values from example text = str(example.get(text_col, '')) if text_col else '' inp = str(example.get(input_col, '')) if input_col else '' output = str(example.get(output_col, '')) if output_col else '' instruction = str(example.get(instruction_col, '')) if instruction_col else '' question = str(example.get(question_col, '')) if question_col else '' answer = str(example.get(answer_col, '')) if answer_col else '' context = str(example.get(context_col, '')) if context_col else '' reasoning = str(example.get(reasoning_col, '')) if reasoning_col else '' # Apply preset templates if preset == 'none': # No template - just use the text column or raw text return text or instruction or question or str(example.get(list(example.keys())[0], '')) elif preset == 'alpaca': result = "" if instruction: result += f"### Instruction:\n{instruction}\n\n" if inp: result += f"### Input:\n{inp}\n\n" result += f"### Response:\n{output or answer}" return result elif preset == 'chatml': result = "" if instruction or context: result += f"<|im_start|>system\n{instruction or context}<|im_end|>\n" result += f"<|im_start|>user\n{question or inp or text}<|im_end|>\n" result += f"<|im_start|>assistant\n{output or answer}<|im_end|>" return result elif preset == 'llama3': result = "<|begin_of_text|>" if instruction or context: result += f"<|start_header_id|>system<|end_header_id|>\n\n{instruction or context}<|eot_id|>" result += f"<|start_header_id|>user<|end_header_id|>\n\n{question or inp or text}<|eot_id|>" result += f"<|start_header_id|>assistant<|end_header_id|>\n\n{output or answer}<|eot_id|>" return result elif preset == 'mistral': result = "" if instruction or context: result = f"[INST] {instruction or context}\n\n" result += f"[INST] {question or inp or text} [/INST] {output or answer}" return result elif preset == 'vicuna': result = "" if instruction or context: result += f"SYSTEM: {instruction or context}\n\n" result += f"USER: {question or inp or text}\nASSISTANT: {output or answer}" return result elif preset == 'reasoning': result = "" if context: result += f"Context:\n{context}\n\n" result += f"Question: {question or inp or text}\n\n" if reasoning: result += f"Reasoning:\n{reasoning}\n\n" result += f"Answer: {output or answer}" return result elif preset == 'phi3': result = "" if instruction or context: result += f"<|system|>\n{instruction or context}<|end|>\n" result += f"<|user|>\n{question or inp or text}<|end|>\n" result += f"<|assistant|>\n{output or answer}<|end|>" return result # Custom template - apply sections custom = prompt_template.get('custom', {}) result = "" if custom.get('system_enabled'): sys_template = custom.get('system_template', '') for col in [instruction, context]: if col: sys_template = sys_template.replace(f'{{{text_col}}}', col) result += custom.get('system_prefix', '') + sys_template + custom.get('system_suffix', '\n\n') if custom.get('user_enabled', True): user_template = custom.get('user_template', '{input}') user_template = user_template.replace('{input}', inp or question or text) user_template = user_template.replace('{question}', question) user_template = user_template.replace('{instruction}', instruction) result += custom.get('user_prefix', '') + user_template + custom.get('user_suffix', '') if custom.get('assistant_enabled', True): asst_template = custom.get('assistant_template', '{output}') asst_template = asst_template.replace('{output}', output or answer) asst_template = asst_template.replace('{answer}', answer) result += custom.get('assistant_prefix', '') + asst_template + custom.get('assistant_suffix', '') return result async def train(self, job_id: str, config: Dict) -> Dict: """Execute training job with progress tracking.""" self.active_trainings[job_id] = { "status": "initializing", "progress": 0.0, "config": config } try: await self._update_job_progress(job_id, status=JobStatus.RUNNING.value, progress=0.0) await self._add_training_log(job_id, "Initializing training job...") # Extract configuration task_type = config.get('task_type', 'causal-lm') base_model = config.get('base_model', config.get('model_name', 'gpt2')) dataset_config = config.get('dataset', {}) dataset_name = dataset_config.get('name', config.get('dataset_name', '')) training_args = config.get('training_args', {}) peft_config = config.get('peft_config', {}) column_mapping = dataset_config.get('column_mapping', {}) prompt_template = config.get('prompt_template', {'preset': 'none'}) # Training parameters epochs = training_args.get('epochs', 3) batch_size = training_args.get('batch_size', 1) learning_rate = training_args.get('learning_rate', 5e-5) max_length = dataset_config.get('max_length', 512) warmup_steps = training_args.get('warmup_steps', 100) # Output directory output_dir = os.path.join(settings.OUTPUT_DIR, job_id) os.makedirs(output_dir, exist_ok=True) logger.info(f"Loading model: {base_model}") await self._update_job_progress(job_id, progress=5.0) def load_tokenizer_for_model(model_name, token=None): """Dynamically load the correct tokenizer based on model config""" # Load config to get tokenizer_class config = AutoConfig.from_pretrained(model_name, token=token) # Get tokenizer class from config tokenizer_class_name = getattr(config, 'tokenizer_class', None) if tokenizer_class_name: try: # Dynamically import the correct tokenizer class # tokenizer_class is usually like 'GemmaTokenizer' or 'LlamaTokenizer' module_path = f"transformers.models.{config.model_type}.tokenization_{config.model_type}" # Try to import the specific module try: module = importlib.import_module(module_path) tokenizer_class = getattr(module, tokenizer_class_name) tokenizer = tokenizer_class.from_pretrained(model_name, token=token) print(f"Loaded {tokenizer_class_name} for {model_name}") return tokenizer except (ImportError, AttributeError): pass # Fallback: Try generic import from transformers tokenizer_class = getattr(importlib.import_module('transformers'), tokenizer_class_name) tokenizer = tokenizer_class.from_pretrained(model_name, token=token) return tokenizer except Exception as e: print(f"Failed to load {tokenizer_class_name}, falling back to AutoTokenizer: {e}") # Fallback to AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, token=token) return tokenizer # Usage in training_service.py: tokenizer = load_tokenizer_for_model( base_model, token=settings.HF_TOKEN if hasattr(settings, 'HF_TOKEN') else None ) # Ensure pad token exists if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token else '' logger.info(f"Successfully loaded tokenizer: {type(tokenizer).__name__}") # Load model based on task type if task_type == 'causal-lm' or task_type == 'reasoning': model = AutoModelForCausalLM.from_pretrained( base_model, torch_dtype=torch.float32, trust_remote_code=True ) elif task_type == 'seq2seq': model = AutoModelForSeq2SeqLM.from_pretrained( base_model, torch_dtype=torch.float32 ) elif task_type == 'token-classification': model = AutoModelForTokenClassification.from_pretrained( base_model, torch_dtype=torch.float32 ) elif task_type == 'text-classification': model = AutoModelForSequenceClassification.from_pretrained( base_model, torch_dtype=torch.float32 ) elif task_type == 'question-answering': model = AutoModelForQuestionAnswering.from_pretrained( base_model, torch_dtype=torch.float32 ) else: model = AutoModelForCausalLM.from_pretrained( base_model, torch_dtype=torch.float32 ) # Apply PEFT if enabled if peft_config and peft_config.get('enabled', False): logger.info("Applying PEFT/LoRA configuration") lora_config = LoraConfig( r=peft_config.get('r', 16), lora_alpha=peft_config.get('alpha', 32), lora_dropout=peft_config.get('dropout', 0.05), bias="none", task_type=TaskType.CAUSAL_LM if task_type in ['causal-lm', 'reasoning'] else TaskType.SEQ_2_SEQ_LM ) model = get_peft_model(model, lora_config) model.print_trainable_parameters() await self._update_job_progress(job_id, progress=10.0) # Load dataset logger.info(f"Loading dataset: {dataset_name}") dataset = load_dataset(dataset_name) train_split = dataset_config.get('train_split', 'train') val_split = dataset_config.get('validation_split', 'validation') train_data = dataset[train_split] if train_split in dataset else dataset['train'] val_data = dataset[val_split] if val_split in dataset else None await self._update_job_progress(job_id, progress=15.0) # Tokenization function with column mapping and prompt template def tokenize_function(example): # Apply prompt template to create text formatted_text = self._apply_prompt_template( example, column_mapping, prompt_template ) # Tokenize tokenized = tokenizer( formatted_text, truncation=True, max_length=max_length, padding='max_length', return_tensors=None ) # For causal LM, labels = input_ids if task_type in ['causal-lm', 'reasoning']: tokenized['labels'] = tokenized['input_ids'].copy() return tokenized logger.info("Tokenizing dataset...") tokenized_train = train_data.map(tokenize_function, batched=False) if val_data: tokenized_val = val_data.map(tokenize_function, batched=False) else: tokenized_val = None await self._update_job_progress(job_id, progress=20.0) # Data collator if task_type in ['causal-lm', 'reasoning']: data_collator = DataCollatorForLanguageModeling( tokenizer=tokenizer, mlm=False ) elif task_type == 'seq2seq': data_collator = DataCollatorForSeq2Seq( tokenizer=tokenizer, model=model ) else: data_collator = default_data_collator # Training arguments training_arguments = TrainingArguments( output_dir=output_dir, num_train_epochs=epochs, per_device_train_batch_size=batch_size, per_device_eval_batch_size=batch_size, learning_rate=learning_rate, warmup_steps=warmup_steps, logging_dir=os.path.join(output_dir, 'logs'), logging_steps=10, save_steps=500, save_total_limit=2, load_best_model_at_end=False, report_to='none', disable_tqdm=False, remove_unused_columns=False ) # Custom callback for progress updates progress_service = self progress_job_id = job_id class ProgressCallback: def __init__(self, total_steps): self.total_steps = total_steps self.current_step = 0 def on_step_end(self, args, state, control, **kwargs): self.current_step = state.global_step progress = 20.0 + (state.global_step / self.total_steps) * 75.0 # Can't call async from sync callback, so we update in-memory if progress_job_id in progress_service.active_trainings: progress_service.active_trainings[progress_job_id]['progress'] = progress progress_service.active_trainings[progress_job_id]['current_step'] = state.global_step return control # Calculate total steps total_steps = (len(tokenized_train) // batch_size) * epochs progress_callback = ProgressCallback(total_steps) # Trainer trainer = SFTTrainer( model=model, args=training_arguments, train_dataset=tokenized_train, eval_dataset=tokenized_val, tokenizer=tokenizer, data_collator=data_collator, callbacks=[progress_callback] ) logger.info("Starting training...") await self._add_training_log(job_id, f"Starting training for {epochs} epochs...") # Run training in executor to not block loop = asyncio.get_event_loop() train_result = await loop.run_in_executor(None, trainer.train) # Save model logger.info("Saving model...") trainer.save_model(output_dir) tokenizer.save_pretrained(output_dir) # Update job as completed await self._update_job_progress( job_id, status=JobStatus.COMPLETED.value, progress=100.0, current_step=total_steps, total_steps=total_steps, train_loss=train_result.training_loss, metrics={'train_loss': train_result.training_loss} ) await self._add_training_log( job_id, f"Training completed! Final loss: {train_result.training_loss:.4f}" ) self.active_trainings[job_id]['status'] = 'completed' return { "status": "completed", "job_id": job_id, "output_path": output_dir, "final_loss": train_result.training_loss } except Exception as e: logger.error(f"Training failed: {e}", exc_info=True) await self._update_job_progress( job_id, status=JobStatus.FAILED.value, error_message=str(e) ) await self._add_training_log(job_id, f"Training failed: {str(e)}", level="ERROR") self.active_trainings[job_id]['status'] = 'failed' return { "status": "failed", "job_id": job_id, "error": str(e) } async def get_job_status(self, job_id: str) -> Optional[Dict]: """Get current status of a training job.""" return self.active_trainings.get(job_id) async def cancel_training(self, job_id: str) -> bool: """Cancel an active training job.""" if job_id in self.active_trainings: self.active_trainings[job_id]['status'] = 'cancelled' await self._update_job_progress(job_id, status=JobStatus.CANCELLED.value) return True return False