app/services/training_service.py

"""
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 '<pad>'
            
            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