train_models.py

"""
Train all models for the Triage Text Classifier
This script trains all models and saves them for use in the Streamlit app
"""

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
import joblib
import os
from preprocessing import TextPreprocessor
from baseline_models import BaselineModels
from bert_model import BertTriageClassifier

def create_models_directory():
    """Create models directory if it doesn't exist"""
    if not os.path.exists('models'):
        os.makedirs('models')
        print("✓ Created models directory")

def load_data():
    """Load the triage dataset"""
    print("\n📊 Loading Data...")
    try:
        # Try to load real data first
        df = pd.read_csv('real_triage_data.csv')
        print(f"✓ Loaded real_triage_data.csv: {len(df)} samples")
        
        # Convert real data format to standard format
        if 'triage_category' in df.columns:
            df['priority'] = df['triage_category']
            df['text'] = df['presenting_complaint']
    except FileNotFoundError:
        # Fall back to synthetic data
        df = pd.read_csv('synthetic_triage_data.csv')
        print(f"✓ Loaded synthetic_triage_data.csv: {len(df)} samples")
    
    return df

def preprocess_data(df):
    """Preprocess the text data"""
    print("\n🔧 Preprocessing Data...")
    preprocessor = TextPreprocessor()
    
    # Preprocess all texts
    df['processed_text'] = df['text'].apply(preprocessor.preprocess_text)
    print(f"✓ Preprocessed {len(df)} text samples")
    
    return df

def train_baseline_models(X_train, X_test, y_train, y_test):
    """Train baseline models"""
    print("\n🤖 Training Baseline Models...")
    
    # Create and fit vectorizer
    print("  → Creating TF-IDF vectorizer...")
    vectorizer = TfidfVectorizer(max_features=5000, ngram_range=(1, 2))
    X_train_tfidf = vectorizer.fit_transform(X_train)
    X_test_tfidf = vectorizer.transform(X_test)
    
    # Save vectorizer
    joblib.dump(vectorizer, 'models/tfidf_vectorizer.joblib')
    print("  ✓ Saved TF-IDF vectorizer")
    
    # Train baseline models
    baseline_models = BaselineModels()
    baseline_models.train_models(X_train_tfidf, y_train)
    print("  ✓ Trained Logistic Regression")
    print("  ✓ Trained Random Forest")
    print("  ✓ Trained Naive Bayes")
    
    # Evaluate models
    results = baseline_models.evaluate_models(X_test_tfidf, y_test)
    
    print("\n📈 Baseline Model Results:")
    for model_name, metrics in results.items():
        print(f"  {model_name}:")
        print(f"    Accuracy:  {metrics['accuracy']:.3f}")
        print(f"    F1 Score:  {metrics['f1_score']:.3f}")
    
    # Save models
    baseline_models.save_models('models')
    print("\n  ✓ Saved all baseline models")
    
    return vectorizer, baseline_models, results

def train_bert_model(X_train, X_test, y_train, y_test):
    """Train BERT model"""
    print("\n🤖 Training BERT Model...")
    print("  ⚠️  This may take several minutes depending on your hardware...")
    
    try:
        # Initialize BERT model
        bert_model = BertTriageClassifier(
            model_name='bert-base-uncased',
            num_labels=3,
            max_length=128
        )
        
        # Prepare datasets
        print("  → Preparing datasets...")
        train_dataset, test_dataset = bert_model.prepare_data(
            X_train.tolist(), 
            y_train.tolist(), 
            X_test.tolist(), 
            y_test.tolist()
        )
        
        # Train model
        print("  → Training BERT (this will take a while)...")
        bert_model.train(
            train_dataset, 
            test_dataset,
            num_epochs=3,
            batch_size=16
        )
        print("  ✓ Trained BERT model")
        
        # Evaluate model
        bert_results = bert_model.evaluate(test_dataset)
        print("\n📈 BERT Model Results:")
        print(f"    Accuracy:  {bert_results['accuracy']:.3f}")
        print(f"    F1 Score:  {bert_results['f1']:.3f}")
        
        # Save model
        bert_model.save_model('bert_model')
        print("  ✓ Saved BERT model")
        
        return bert_model, bert_results
        
    except Exception as e:
        print(f"\n  ⚠️  Warning: Could not train BERT model: {e}")
        print("  → The baseline models are still available for use")
        return None, None

def main():
    """Main training pipeline"""
    print("=" * 60)
    print("🏥 Triage Text Classifier - Model Training Pipeline")
    print("=" * 60)
    
    # Create models directory
    create_models_directory()
    
    # Load data
    df = load_data()
    print(f"\nDataset shape: {df.shape}")
    print(f"Priority distribution:\n{df['priority'].value_counts()}")
    
    # Preprocess data
    df = preprocess_data(df)
    
    # Split data
    print("\n✂️  Splitting Data...")
    X = df['processed_text']
    y = df['priority']
    
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, 
        test_size=0.2, 
        random_state=42, 
        stratify=y
    )
    print(f"  Training samples: {len(X_train)}")
    print(f"  Testing samples:  {len(X_test)}")
    
    # Train baseline models
    vectorizer, baseline_models, baseline_results = train_baseline_models(
        X_train, X_test, y_train, y_test
    )
    
    # Train BERT model (optional, may take a while)
    print("\n" + "=" * 60)
    user_input = input("Do you want to train the BERT model? (y/n): ").lower().strip()
    
    if user_input == 'y':
        bert_model, bert_results = train_bert_model(
            X_train, X_test, y_train, y_test
        )
    else:
        print("\n  → Skipping BERT model training")
        print("  → You can still use the baseline models in the app")
    
    # Summary
    print("\n" + "=" * 60)
    print("✅ Training Complete!")
    print("=" * 60)
    print("\nTrained models saved in:")
    print("  • models/logistic_regression.joblib")
    print("  • models/random_forest.joblib")
    print("  • models/naive_bayes.joblib")
    print("  • models/tfidf_vectorizer.joblib")
    if user_input == 'y':
        print("  • bert_model/ (directory)")
    
    print("\n🚀 You can now run the Streamlit app:")
    print("   streamlit run app.py")
    print("=" * 60)

if __name__ == "__main__":
    main()