finetune_models.py

"""
Fine-tuning Script for Medical AI Models
Trains models on real medical datasets for production use
"""

import os
import torch
import pandas as pd
import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from transformers import (
    ViTImageProcessor, 
    ViTForImageClassification,
    Trainer, 
    TrainingArguments,
    AutoTokenizer,
    AutoModelForSequenceClassification
)
from datasets import load_dataset
from sklearn.model_selection import train_test_split
import json


class SkinLesionDataset(Dataset):
    """Dataset for skin lesion images (HAM10000 format)"""
    
    def __init__(self, image_paths, labels, processor):
        self.image_paths = image_paths
        self.labels = labels
        self.processor = processor
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        image = Image.open(self.image_paths[idx]).convert('RGB')
        encoding = self.processor(images=image, return_tensors="pt")
        encoding = {key: val.squeeze() for key, val in encoding.items()}
        encoding['labels'] = torch.tensor(self.labels[idx])
        return encoding


class SymptomDataset(Dataset):
    """Dataset for symptom-to-disease classification"""
    
    def __init__(self, texts, labels, tokenizer, max_length=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'
        )
        encoding = {key: val.squeeze() for key, val in encoding.items()}
        encoding['labels'] = torch.tensor(self.labels[idx])
        return encoding


class MedicalModelTrainer:
    """Fine-tune models on medical datasets"""
    
    def __init__(self, output_dir="./trained_models"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")
    
    def finetune_skin_model(self, data_dir, num_epochs=10):
        """
        Fine-tune Vision Transformer on HAM10000 skin lesion dataset
        
        Dataset structure:
        data_dir/
        ├── images/
        │   ├── image1.jpg
        │   ├── image2.jpg
        └── labels.csv (columns: image_id, diagnosis)
        
        Download from: https://www.kaggle.com/datasets/kmader/skin-cancer-mnist-ham10000
        """
        print("🔬 Fine-tuning Skin Condition Model...")
        
        # Load dataset
        try:
            labels_df = pd.read_csv(os.path.join(data_dir, "HAM10000_metadata.csv"))
        except FileNotFoundError:
            print("❌ Dataset not found. Download HAM10000 from Kaggle:")
            print("   kaggle datasets download -d kmader/skin-cancer-mnist-ham10000")
            return None
        
        # Map diagnoses to indices
        diagnosis_map = {
            'akiec': 0,  # Actinic keratoses
            'bcc': 1,    # Basal cell carcinoma
            'bkl': 2,    # Benign keratosis
            'df': 3,     # Dermatofibroma
            'mel': 4,    # Melanoma
            'nv': 5,     # Melanocytic nevi
            'vasc': 6    # Vascular lesions
        }
        
        labels_df['label'] = labels_df['dx'].map(diagnosis_map)
        
        # Prepare image paths
        image_dir = os.path.join(data_dir, "images")
        labels_df['image_path'] = labels_df['image_id'].apply(
            lambda x: os.path.join(image_dir, f"{x}.jpg")
        )
        
        # Filter existing images
        labels_df = labels_df[labels_df['image_path'].apply(os.path.exists)]
        
        print(f"📊 Loaded {len(labels_df)} images")
        
        # Split dataset
        train_df, val_df = train_test_split(
            labels_df, 
            test_size=0.2, 
            stratify=labels_df['label'],
            random_state=42
        )
        
        # Load processor and model
        processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
        model = ViTForImageClassification.from_pretrained(
            'google/vit-base-patch16-224',
            num_labels=len(diagnosis_map),
            ignore_mismatched_sizes=True
        )
        
        # Create datasets
        train_dataset = SkinLesionDataset(
            train_df['image_path'].tolist(),
            train_df['label'].tolist(),
            processor
        )
        
        val_dataset = SkinLesionDataset(
            val_df['image_path'].tolist(),
            val_df['label'].tolist(),
            processor
        )
        
        # Training arguments
        training_args = TrainingArguments(
            output_dir=os.path.join(self.output_dir, "skin-condition-vit"),
            evaluation_strategy="epoch",
            save_strategy="epoch",
            learning_rate=2e-5,
            per_device_train_batch_size=16,
            per_device_eval_batch_size=16,
            num_train_epochs=num_epochs,
            weight_decay=0.01,
            load_best_model_at_end=True,
            metric_for_best_model="accuracy",
            logging_dir='./logs',
            logging_steps=100,
            save_total_limit=2
        )
        
        # Define metrics
        def compute_metrics(eval_pred):
            predictions, labels = eval_pred
            predictions = np.argmax(predictions, axis=1)
            accuracy = (predictions == labels).mean()
            return {"accuracy": accuracy}
        
        # Create trainer
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
            compute_metrics=compute_metrics
        )
        
        # Train
        print("🏋️ Training started...")
        trainer.train()
        
        # Save model
        model_path = os.path.join(self.output_dir, "skin-condition-vit-final")
        trainer.save_model(model_path)
        processor.save_pretrained(model_path)
        
        # Save label mapping
        with open(os.path.join(model_path, "label_map.json"), "w") as f:
            reverse_map = {v: k for k, v in diagnosis_map.items()}
            json.dump(reverse_map, f)
        
        print(f"✅ Model saved to {model_path}")
        return model_path
    
    def finetune_symptom_model(self, data_file, num_epochs=5):
        """
        Fine-tune BERT on symptom-to-disease dataset
        
        Dataset format (CSV):
        symptoms,disease
        "headache fever cough","Influenza"
        "chest pain shortness of breath","Heart Condition"
        
        Download from Kaggle: Disease Symptom Prediction Dataset
        """
        print("🔬 Fine-tuning Symptom Analysis Model...")
        
        try:
            # Load dataset
            df = pd.read_csv(data_file)
            
            # Create disease label mapping
            diseases = df['disease'].unique()
            disease_map = {disease: idx for idx, disease in enumerate(diseases)}
            df['label'] = df['disease'].map(disease_map)
            
            print(f"📊 Loaded {len(df)} examples with {len(diseases)} diseases")
            
            # Split dataset
            train_df, val_df = train_test_split(
                df, 
                test_size=0.2,
                stratify=df['label'],
                random_state=42
            )
            
            # Load tokenizer and model
            model_name = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract"
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForSequenceClassification.from_pretrained(
                model_name,
                num_labels=len(diseases)
            )
            
            # Create datasets
            train_dataset = SymptomDataset(
                train_df['symptoms'].tolist(),
                train_df['label'].tolist(),
                tokenizer
            )
            
            val_dataset = SymptomDataset(
                val_df['symptoms'].tolist(),
                val_df['label'].tolist(),
                tokenizer
            )
            
            # Training arguments
            training_args = TrainingArguments(
                output_dir=os.path.join(self.output_dir, "symptom-bert"),
                evaluation_strategy="epoch",
                save_strategy="epoch",
                learning_rate=2e-5,
                per_device_train_batch_size=16,
                per_device_eval_batch_size=16,
                num_train_epochs=num_epochs,
                weight_decay=0.01,
                load_best_model_at_end=True,
                metric_for_best_model="accuracy",
                logging_steps=50
            )
            
            # Define metrics
            def compute_metrics(eval_pred):
                predictions, labels = eval_pred
                predictions = np.argmax(predictions, axis=1)
                accuracy = (predictions == labels).mean()
                return {"accuracy": accuracy}
            
            # Create trainer
            trainer = Trainer(
                model=model,
                args=training_args,
                train_dataset=train_dataset,
                eval_dataset=val_dataset,
                compute_metrics=compute_metrics
            )
            
            # Train
            print("🏋️ Training started...")
            trainer.train()
            
            # Save model
            model_path = os.path.join(self.output_dir, "symptom-bert-final")
            trainer.save_model(model_path)
            tokenizer.save_pretrained(model_path)
            
            # Save label mapping
            with open(os.path.join(model_path, "disease_map.json"), "w") as f:
                reverse_map = {v: k for k, v in disease_map.items()}
                json.dump(reverse_map, f)
            
            print(f"✅ Model saved to {model_path}")
            return model_path
            
        except FileNotFoundError:
            print("❌ Dataset not found. Create or download symptom-disease dataset")
            print("   Format: CSV with columns 'symptoms' and 'disease'")
            return None
    
    def create_sample_symptom_dataset(self, output_file="symptom_dataset.csv"):
        """Create a sample symptom dataset for testing"""
        print("📝 Creating sample symptom dataset...")
        
        sample_data = [
            ("headache fever fatigue", "Influenza"),
            ("cough shortness of breath chest pain", "Pneumonia"),
            ("nausea vomiting diarrhea", "Gastroenteritis"),
            ("rash itching redness", "Allergic Reaction"),
            ("sore throat fever headache", "Strep Throat"),
            ("fatigue weakness pale skin", "Anemia"),
            ("headache sensitivity to light nausea", "Migraine"),
            ("chest pain shortness of breath", "Heart Condition"),
            ("fever cough body aches", "Common Cold"),
            ("abdominal pain nausea fever", "Appendicitis")
        ] * 50  # Duplicate for larger dataset
        
        df = pd.DataFrame(sample_data, columns=['symptoms', 'disease'])
        df.to_csv(output_file, index=False)
        
        print(f"✅ Sample dataset saved to {output_file}")
        return output_file


def main():
    """Main training pipeline"""
    trainer = MedicalModelTrainer()
    
    print("=" * 60)
    print("🏥 Medical AI Model Fine-tuning Pipeline")
    print("=" * 60)
    
    # Option 1: Fine-tune skin condition model
    print("\n1️⃣ Skin Condition Model")
    print("   Dataset: HAM10000 (download from Kaggle)")
    print("   Command: kaggle datasets download -d kmader/skin-cancer-mnist-ham10000")
    
    skin_data_dir = "./HAM10000"
    if os.path.exists(skin_data_dir):
        trainer.finetune_skin_model(skin_data_dir, num_epochs=3)
    else:
        print("   ⏭️  Skipping (dataset not found)")
    
    # Option 2: Fine-tune symptom model
    print("\n2️⃣ Symptom Analysis Model")
    
    symptom_dataset = "./symptom_dataset.csv"
    if not os.path.exists(symptom_dataset):
        symptom_dataset = trainer.create_sample_symptom_dataset()
    
    trainer.finetune_symptom_model(symptom_dataset, num_epochs=3)
    
    print("\n" + "=" * 60)
    print("✅ Training complete!")
    print("=" * 60)
    print("\n📦 Trained models saved in ./trained_models/")
    print("\n🚀 To use in production:")
    print("   1. Update ai_models.py to load from ./trained_models/")
    print("   2. Replace model_name with local path")
    print("   3. Test with test_api.py")


if __name__ == "__main__":
    main()