model.py

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

# Set matplotlib backend to non-GUI 'Agg' at the beginning
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns

def preprocess_data(data_path, test_size=0.2, random_state=42):
    """
    Preprocess the salary dataset
    
    Args:
        data_path (str): Path to the dataset CSV file
        test_size (float): Proportion of data to use for testing
        random_state (int): Random state for reproducibility
        
    Returns:
        tuple: X_train, X_test, y_train, y_test, preprocessor
    """
    # Load the dataset
    df = pd.read_csv(data_path)
    
    # Print dataset info for analysis
    print("Dataset information:")
    print(df.info())
    print("\nData sample:")
    print(df.head())
    
    # Handle missing values if any
    if df.isnull().sum().sum() > 0:
        print("\nMissing values detected:")
        print(df.isnull().sum())
        # Fill numeric columns with median and categorical with mode
        for col in df.columns:
            if df[col].dtype == 'object':
                df[col] = df[col].fillna(df[col].mode()[0])
            else:
                df[col] = df[col].fillna(df[col].median())
    
    # Split features and target
    X = df.drop('Salary', axis=1)
    y = df['Salary']
    
    # Identify categorical and numerical columns
    categorical_cols = X.select_dtypes(include=['object']).columns.tolist()
    numerical_cols = X.select_dtypes(exclude=['object']).columns.tolist()
    
    # Create preprocessor
    preprocessor = ColumnTransformer(
        transformers=[
            ('num', StandardScaler(), numerical_cols),
            ('cat', OneHotEncoder(handle_unknown='ignore'), categorical_cols)
        ])
    
    # Split the data
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state
    )
    
    print(f"\nTraining set size: {X_train.shape[0]} samples")
    print(f"Testing set size: {X_test.shape[0]} samples")
    
    return X_train, X_test, y_train, y_test, preprocessor

def build_model(preprocessor):
    """
    Build the ML model pipeline
    
    Args:
        preprocessor (ColumnTransformer): The data preprocessor
        
    Returns:
        Pipeline: Model pipeline
    """
    # Create pipeline with preprocessing and model
    model = Pipeline(steps=[
        ('preprocessor', preprocessor),
        ('regressor', RandomForestRegressor(random_state=42))
    ])
    
    return model

def evaluate_model(model, X_test, y_test):
    """
    Evaluate the model performance
    
    Args:
        model: Trained model
        X_test: Test features
        y_test: Test target
        
    Returns:
        dict: Performance metrics
    """
    # Make predictions
    y_pred = model.predict(X_test)
    
    # Calculate metrics
    mse = mean_squared_error(y_test, y_pred)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    
    # Print metrics
    print("\nModel Performance:")
    print(f"Mean Squared Error: {mse:.2f}")
    print(f"Root Mean Squared Error: {rmse:.2f}")
    print(f"Mean Absolute Error: {mae:.2f}")
    print(f"R² Score: {r2:.4f}")
    
    try:
        # Visualize actual vs predicted values
        plt.figure(figsize=(10, 6))
        plt.scatter(y_test, y_pred, alpha=0.5)
        plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--')
        plt.xlabel('Actual Salary')
        plt.ylabel('Predicted Salary')
        plt.title('Actual vs Predicted Salary')
        plt.savefig('actual_vs_predicted.png')
        plt.close()
        
        # Plot feature importance if model supports it
        # Check if model is a Pipeline or GridSearchCV
        if hasattr(model, 'best_estimator_'):  # GridSearchCV case
            estimator = model.best_estimator_[-1]  # Get regressor from pipeline in best_estimator_
            if hasattr(estimator, 'feature_importances_'):
                # Get feature names from the preprocessor in the pipeline
                preprocessor = model.best_estimator_[0]
                feature_names = []
                for name, transformer, cols in preprocessor.transformers_:
                    if name == 'num':
                        feature_names.extend(cols)
                    else:
                        for col in cols:
                            for category in transformer.categories_[cols.index(col)]:
                                feature_names.append(f"{col}_{category}")
                
                feature_importances = estimator.feature_importances_
                if len(feature_names) == len(feature_importances):
                    importance_df = pd.DataFrame({
                        'Feature': feature_names,
                        'Importance': feature_importances
                    }).sort_values('Importance', ascending=False)
                    
                    plt.figure(figsize=(12, 8))
                    sns.barplot(x='Importance', y='Feature', data=importance_df.head(15))
                    plt.title('Top 15 Feature Importance')
                    plt.tight_layout()
                    plt.savefig('feature_importance.png')
                    plt.close()
        elif hasattr(model[-1], 'feature_importances_'):  # Pipeline case
            # Original code for Pipeline
            feature_names = []
            for name, transformer, cols in model[0].transformers_:
                if name == 'num':
                    feature_names.extend(cols)
                else:
                    for col in cols:
                        for category in transformer.categories_[cols.index(col)]:
                            feature_names.append(f"{col}_{category}")
            
            feature_importances = model[-1].feature_importances_
            if len(feature_names) == len(feature_importances):
                importance_df = pd.DataFrame({
                    'Feature': feature_names,
                    'Importance': feature_importances
                }).sort_values('Importance', ascending=False)
                
                plt.figure(figsize=(12, 8))
                sns.barplot(x='Importance', y='Feature', data=importance_df.head(15))
                plt.title('Top 15 Feature Importance')
                plt.tight_layout()
                plt.savefig('feature_importance.png')
                plt.close()
    except Exception as e:
        print(f"Visualization error: {str(e)}. Visualizations skipped, but model evaluation completed.")
    
    return {
        'MSE': mse,
        'RMSE': rmse,
        'MAE': mae,
        'R2': r2
    }

def tune_hyperparameters(X_train, y_train, preprocessor):
    """
    Perform hyperparameter tuning for the model
    
    Args:
        X_train: Training features
        y_train: Training target
        preprocessor: Data preprocessor
        
    Returns:
        GridSearchCV: Tuned model
    """
    # Create the pipeline
    model = Pipeline(steps=[
        ('preprocessor', preprocessor),
        ('regressor', RandomForestRegressor(random_state=42))
    ])
    
    # Define hyperparameter grid
    param_grid = {
        'regressor__n_estimators': [50, 100, 200],
        'regressor__max_depth': [None, 10, 20],
        'regressor__min_samples_split': [2, 5, 10],
        'regressor__min_samples_leaf': [1, 2, 4]
    }
    
    # Create grid search
    print("\nStarting hyperparameter tuning (this may take some time)...")
    grid_search = GridSearchCV(
        model, 
        param_grid=param_grid,
        cv=5,
        n_jobs=-1,
        scoring='neg_mean_squared_error',
        verbose=1
    )
    
    # Fit the grid search
    grid_search.fit(X_train, y_train)
    
    # Print best parameters
    print("\nBest parameters:", grid_search.best_params_)
    print("Best CV score (neg MSE):", grid_search.best_score_)
    
    return grid_search

def main(data_path='/home/Maanu/Documents/Salary_Prediction/Salary Data.csv'):
    """
    Main function to run the entire pipeline
    
    Args:
        data_path (str): Path to the dataset
    """
    # Step 1: Preprocess the data
    X_train, X_test, y_train, y_test, preprocessor = preprocess_data(data_path)
    
    # Step 2: Build the model
    model = build_model(preprocessor)
    
    # Step 3: Train the model
    print("\nTraining the base model...")
    model.fit(X_train, y_train)
    
    # Step 4: Evaluate the base model
    print("\nEvaluating the base model:")
    base_metrics = evaluate_model(model, X_test, y_test)
    
    # Step 5: Hyperparameter tuning
    tuned_model = tune_hyperparameters(X_train, y_train, preprocessor)
    
    # Step 6: Evaluate the tuned model
    print("\nEvaluating the tuned model:")
    tuned_metrics = evaluate_model(tuned_model, X_test, y_test)
    
    # Compare models
    print("\nModel Comparison:")
    print(f"Base Model R²: {base_metrics['R2']:.4f}")
    print(f"Tuned Model R²: {tuned_metrics['R2']:.4f}")
    
    # Save the tuned model if it's better
    if tuned_metrics['R2'] > base_metrics['R2']:
        best_model = tuned_model
        print("\nTuned model performs better! Saving the tuned model.")
    else:
        best_model = model
        print("\nBase model performs better! Saving the base model.")
    
    # Example of using the model for prediction
    print("\nExample prediction:")
    example = {
        'Age': [30, 45],
        'Gender': ['Male', 'Female'],
        'Education Level': ['Bachelor\'s', 'Master\'s'],
        'Job Title': ['Software Engineer', 'Data Scientist'],
        'Years of Experience': [5, 15]
    }
    example_df = pd.DataFrame(example)
    print(example_df)
    predictions = best_model.predict(example_df)
    print(f"Predicted Salaries: {predictions}")
    
    # Save the model for the frontend
    import pickle
    with open('salary_prediction_model.pkl', 'wb') as f:
        pickle.dump(best_model, f)
    print("\nModel saved to salary_prediction_model.pkl")
    
    return best_model

if __name__ == "__main__":
    model = main()