CoolFace
Apppublic

Yashraj64004/human-resource-data-analysis-streamlit

sourceHugging Faceupdated 2mo agoView on Hugging Face
0likes
example_train_model.py114 linesDownload Raw Back to root
1"""Example script demonstrating train_attrition_model function"""2 3import sys4import os5sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))6 7import pandas as pd8from ml_models import train_attrition_model, load_trained_model, predict_attrition9 10# Example 1: Load sample HR data and train model11if __name__ == "__main__":12    print("=" * 70)13    print("Attrition Model Training Example")14    print("=" * 70)15    16    # Load data17    try:18        df = pd.read_csv('input_data/raw_hr_data.csv')19        print(f"\n✓ Loaded data with {len(df)} records and {len(df.columns)} columns")20    except FileNotFoundError:21        print("\n✗ Error: Could not find 'input_data/raw_hr_data.csv'")22        sys.exit(1)23    24    print(f"  Columns: {list(df.columns[:5])}...")25    print(f"  Target column 'Attrition': {df['Attrition'].value_counts().to_dict()}")26    27    # Example 1: Train with default numeric features28    print("\n" + "-" * 70)29    print("Example 1: Train with automatic numeric feature selection")30    print("-" * 70)31    32    try:33        result = train_attrition_model(df, target='Attrition', test_size=0.2)34        35        print(f"\n✓ Model trained successfully!")36        print(f"  - Accuracy: {result['accuracy']:.4f} ({result['accuracy']*100:.2f}%)")37        print(f"  - ROC-AUC: {result['roc_auc']:.4f}")38        print(f"  - Precision: {result['precision']:.4f}")39        print(f"  - Recall: {result['recall']:.4f}")40        print(f"  - F1-Score: {result['f1']:.4f}")41        print(f"  - Model saved to: {result['model_path']}")42        print(f"  - Features used: {len(result['feature_names'])} columns")43        print(f"  - Training set size: {result['train_size']}")44        print(f"  - Test set size: {result['test_size']}")45        46        # Show top 5 important features47        print(f"\n  Top 5 Important Features:")48        sorted_features = sorted(result['feature_importance'].items(), 49                                key=lambda x: x[1], reverse=True)[:5]50        for i, (feature, importance) in enumerate(sorted_features, 1):51            print(f"    {i}. {feature}: {importance:.4f}")52    53    except Exception as e:54        print(f"\n✗ Error during training: {str(e)}")55        sys.exit(1)56    57    # Example 2: Train with specific features58    print("\n" + "-" * 70)59    print("Example 2: Train with manually selected features")60    print("-" * 70)61    62    selected_features = ['Age', 'MonthlyIncome', 'YearsAtCompany', 63                        'DistanceFromHome', 'JobSatisfaction', 'OverTime']64    65    # Check if all features exist66    available_features = [f for f in selected_features if f in df.columns]67    missing_features = [f for f in selected_features if f not in df.columns]68    69    if missing_features:70        print(f"\n  Note: Some features not available: {missing_features}")71        print(f"  Using available features: {available_features}")72        selected_features = available_features73    74    try:75        result2 = train_attrition_model(df, target='Attrition', 76                                       features=selected_features, test_size=0.2)77        78        print(f"\n✓ Model trained with selected features!")79        print(f"  - Accuracy: {result2['accuracy']:.4f} ({result2['accuracy']*100:.2f}%)")80        print(f"  - ROC-AUC: {result2['roc_auc']:.4f}")81        print(f"  - Features used: {result2['feature_names']}")82    83    except Exception as e:84        print(f"\n✗ Error: {str(e)}")85    86    # Example 3: Load model and make predictions87    print("\n" + "-" * 70)88    print("Example 3: Load model and make predictions")89    print("-" * 70)90    91    try:92        loaded_model = load_trained_model('model/attrition_model.pkl')93        print(f"\n✓ Model loaded successfully!")94        95        # Make predictions on first 10 records96        predictions = predict_attrition(loaded_model, df.iloc[:10], 97                                       features=result['feature_names'])98        proba = predict_attrition(loaded_model, df.iloc[:10], 99                                 features=result['feature_names'], return_proba=True)100        101        print(f"\n  First 10 Predictions:")102        print(f"  Index | Binary Pred | Probability")103        print(f"  ------|-------------|------------")104        for i in range(10):105            pred_label = "At Risk" if predictions[i] == 1 else "Stable"106            print(f"  {i:5d} | {pred_label:11s} | {proba[i]:10.2%}")107    108    except Exception as e:109        print(f"\n✗ Error: {str(e)}")110    111    print("\n" + "=" * 70)112    print("Example completed!")113    print("=" * 70)114