Update model/train.py

Adding LightGBM for Ensemble Model

model/train.py

@@ -1,4 +1,4 @@
-# Enhanced version with comprehensive cross-validation and advanced feature engineering
+# Enhanced version with LightGBM, ensemble voting, and statistical validation
 
 import seaborn as sns
 import matplotlib.pyplot as plt
@@ -14,9 +14,10 @@ from sklearn.model_selection import (
     train_test_split, cross_val_score, GridSearchCV,
     StratifiedKFold, validation_curve, cross_validate
 )
-from sklearn.ensemble import RandomForestClassifier
+from sklearn.ensemble import RandomForestClassifier, VotingClassifier
 from sklearn.linear_model import LogisticRegression
 from sklearn.feature_extraction.text import TfidfVectorizer
+import lightgbm as lgb
 import pandas as pd
 import numpy as np
 from pathlib import Path
@@ -31,6 +32,7 @@ from datetime import datetime, timedelta
 from typing import Dict, Tuple, Optional, Any, List
 import warnings
 import re
+from scipy import stats
 warnings.filterwarnings('ignore')
 
 # Import enhanced feature engineering components
@@ -161,7 +163,7 @@ class ProgressTracker:
 
 
 def estimate_training_time(dataset_size: int, enable_tuning: bool = True, cv_folds: int = 5, 
-                          use_enhanced_features: bool = False) -> Dict:
+                          use_enhanced_features: bool = False, enable_ensemble: bool = True) -> Dict:
     """Estimate training time based on dataset characteristics and feature complexity"""
     
     # Base time estimates (in seconds) based on empirical testing
@@ -180,12 +182,16 @@ def estimate_training_time(dataset_size: int, enable_tuning: bool = True, cv_fol
         base_times['feature_selection'] *= 2.0  # More features to select from
         base_times['enhanced_feature_extraction'] = max(2.0, dataset_size * 0.05)  # New step
     
-    # Hyperparameter tuning multipliers
+    # Hyperparameter tuning multipliers with LightGBM
     tuning_multipliers = {
         'logistic_regression': 8 if enable_tuning else 1,  # 8 param combinations
         'random_forest': 12 if enable_tuning else 1,       # 12 param combinations
+        'lightgbm': 6 if enable_tuning else 1,             # 6 param combinations (CPU optimized)
     }
     
+    # Ensemble multiplier
+    ensemble_multiplier = 1.3 if enable_ensemble else 1.0  # 30% overhead for ensemble
+    
     # Cross-validation multiplier
     cv_multiplier = cv_folds if dataset_size > 100 else 1
     
@@ -202,7 +208,7 @@ def estimate_training_time(dataset_size: int, enable_tuning: bool = True, cv_fol
     
     estimates['feature_selection'] = base_times['feature_selection']
     
-    # Model training (now includes CV)
+    # Model training (now includes CV and LightGBM)
     for model_name, multiplier in tuning_multipliers.items():
         model_time = base_times['simple_training'] * multiplier * cv_multiplier
         estimates[f'{model_name}_training'] = model_time
@@ -211,14 +217,19 @@ def estimate_training_time(dataset_size: int, enable_tuning: bool = True, cv_fol
     # Cross-validation overhead
     estimates['cross_validation'] = base_times['simple_training'] * cv_folds * 0.5
     
+    # Ensemble training and validation
+    if enable_ensemble:
+        estimates['ensemble_training'] = base_times['simple_training'] * 0.5
+        estimates['ensemble_validation'] = base_times['evaluation'] * 2
+    
     # Model saving
     estimates['model_saving'] = 1.0
     
     # Total estimate
-    total_estimate = sum(estimates.values())
+    total_estimate = sum(estimates.values()) * ensemble_multiplier
     
-    # Add buffer for overhead (more for enhanced features)
-    buffer_multiplier = 1.4 if use_enhanced_features else 1.2
+    # Add buffer for overhead (more for enhanced features and ensemble)
+    buffer_multiplier = 1.5 if (use_enhanced_features and enable_ensemble) else 1.4 if use_enhanced_features else 1.2
     total_estimate *= buffer_multiplier
     
     return {
@@ -228,7 +239,8 @@ def estimate_training_time(dataset_size: int, enable_tuning: bool = True, cv_fol
         'dataset_size': dataset_size,
         'enable_tuning': enable_tuning,
         'cv_folds': cv_folds,
-        'use_enhanced_features': use_enhanced_features
+        'use_enhanced_features': use_enhanced_features,
+        'enable_ensemble': enable_ensemble
     }
 
 
@@ -288,7 +300,7 @@ class CrossValidationManager:
                 cv=cv_strategy,
                 scoring=scoring_metrics,
                 return_train_score=True,
-                n_jobs=1,  # Use single job for stability
+                n_jobs=1,  # Use single job for stability on HFS
                 verbose=0
             )
             
@@ -381,7 +393,6 @@ class CrossValidationManager:
             scores2 = results2['test_scores'][metric]['scores']
             
             # Paired t-test
-            from scipy import stats
             t_stat, p_value = stats.ttest_rel(scores1, scores2)
             
             comparison = {
@@ -406,21 +417,135 @@ class CrossValidationManager:
             return {'error': str(e)}
 
 
+class EnsembleManager:
+    """Manage ensemble model creation and validation"""
+    
+    def __init__(self, random_state: int = 42):
+        self.random_state = random_state
+        
+    def create_ensemble(self, individual_models: Dict[str, Any], 
+                       voting: str = 'soft') -> VotingClassifier:
+        """Create ensemble from individual models"""
+        
+        estimators = [(name, model) for name, model in individual_models.items()]
+        
+        ensemble = VotingClassifier(
+            estimators=estimators,
+            voting=voting,
+            n_jobs=1  # CPU optimization for HFS
+        )
+        
+        logger.info(f"Created {voting} voting ensemble with {len(estimators)} models")
+        return ensemble
+    
+    def evaluate_ensemble_vs_individuals(self, ensemble, individual_models: Dict,
+                                       X_test, y_test) -> Dict:
+        """Compare ensemble performance against individual models"""
+        
+        results = {}
+        
+        # Evaluate individual models
+        for name, model in individual_models.items():
+            y_pred = model.predict(X_test)
+            y_pred_proba = model.predict_proba(X_test)[:, 1]
+            
+            results[name] = {
+                'accuracy': float(accuracy_score(y_test, y_pred)),
+                'precision': float(precision_score(y_test, y_pred, average='weighted')),
+                'recall': float(recall_score(y_test, y_pred, average='weighted')),
+                'f1': float(f1_score(y_test, y_pred, average='weighted')),
+                'roc_auc': float(roc_auc_score(y_test, y_pred_proba))
+            }
+        
+        # Evaluate ensemble
+        y_pred_ensemble = ensemble.predict(X_test)
+        y_pred_proba_ensemble = ensemble.predict_proba(X_test)[:, 1]
+        
+        results['ensemble'] = {
+            'accuracy': float(accuracy_score(y_test, y_pred_ensemble)),
+            'precision': float(precision_score(y_test, y_pred_ensemble, average='weighted')),
+            'recall': float(recall_score(y_test, y_pred_ensemble, average='weighted')),
+            'f1': float(f1_score(y_test, y_pred_ensemble, average='weighted')),
+            'roc_auc': float(roc_auc_score(y_test, y_pred_proba_ensemble))
+        }
+        
+        # Calculate improvement over best individual model
+        best_individual_f1 = max(results[name]['f1'] for name in individual_models.keys())
+        ensemble_f1 = results['ensemble']['f1']
+        improvement = ensemble_f1 - best_individual_f1
+        
+        results['ensemble_analysis'] = {
+            'best_individual_f1': best_individual_f1,
+            'ensemble_f1': ensemble_f1,
+            'improvement': improvement,
+            'improvement_percentage': (improvement / best_individual_f1) * 100 if best_individual_f1 > 0 else 0,
+            'is_better': improvement > 0
+        }
+        
+        return results
+    
+    def statistical_ensemble_comparison(self, ensemble, individual_models: Dict,
+                                      X, y, cv_manager: CrossValidationManager) -> Dict:
+        """Perform statistical comparison between ensemble and individual models"""
+        
+        cv_strategy = cv_manager.create_cv_strategy(X, y)
+        
+        results = {}
+        
+        # Get CV results for ensemble
+        ensemble_cv = cv_manager.perform_cross_validation(ensemble, X, y, cv_strategy)
+        results['ensemble'] = ensemble_cv
+        
+        # Get CV results for individual models
+        individual_cv_results = {}
+        for name, model in individual_models.items():
+            model_cv = cv_manager.perform_cross_validation(model, X, y, cv_strategy)
+            individual_cv_results[name] = model_cv
+            results[name] = model_cv
+        
+        # Compare ensemble with each individual model
+        comparisons = {}
+        for name, model_cv in individual_cv_results.items():
+            comparison = cv_manager.compare_cv_results(model_cv, ensemble_cv)
+            comparisons[f'ensemble_vs_{name}'] = comparison
+        
+        results['statistical_comparisons'] = comparisons
+        
+        # Determine if ensemble should be used
+        ensemble_f1_scores = ensemble_cv.get('test_scores', {}).get('f1', {}).get('scores', [])
+        
+        significantly_better_count = 0
+        for comparison in comparisons.values():
+            if comparison.get('paired_ttest', {}).get('significant', False) and comparison.get('difference', 0) > 0:
+                significantly_better_count += 1
+        
+        results['ensemble_recommendation'] = {
+            'use_ensemble': significantly_better_count > 0,
+            'significantly_better_than': significantly_better_count,
+            'total_comparisons': len(comparisons),
+            'confidence': significantly_better_count / len(comparisons) if comparisons else 0
+        }
+        
+        return results
+
+
 class EnhancedModelTrainer:
-    """Production-ready model trainer with enhanced feature engineering and comprehensive CV"""
+    """Production-ready model trainer with LightGBM, enhanced features, and ensemble voting"""
 
-    def __init__(self, use_enhanced_features: bool = None):
+    def __init__(self, use_enhanced_features: bool = None, enable_ensemble: bool = True):
         # Auto-detect enhanced features if not specified
         if use_enhanced_features is None:
             self.use_enhanced_features = ENHANCED_FEATURES_AVAILABLE
         else:
             self.use_enhanced_features = use_enhanced_features and ENHANCED_FEATURES_AVAILABLE
         
+        self.enable_ensemble = enable_ensemble
         self.setup_paths()
         self.setup_training_config()
         self.setup_models()
         self.progress_tracker = None
         self.cv_manager = CrossValidationManager()
+        self.ensemble_manager = EnsembleManager()
         
         # Enhanced feature tracking
         self.feature_engineer = None
@@ -480,14 +605,14 @@ class EnhancedModelTrainer:
         self.class_weight = 'balanced'
 
     def setup_models(self):
-        """Setup model configurations for comparison"""
+        """Setup model configurations including LightGBM for comparison"""
         self.models = {
             'logistic_regression': {
                 'model': LogisticRegression(
                     max_iter=self.max_iter,
                     class_weight=self.class_weight,
                     random_state=self.random_state,
-                    n_jobs=-1
+                    n_jobs=1  # CPU optimization
                 ),
                 'param_grid': {
                     'model__C': [0.1, 1, 10],
@@ -496,15 +621,34 @@ class EnhancedModelTrainer:
             },
             'random_forest': {
                 'model': RandomForestClassifier(
-                    n_estimators=50,
+                    n_estimators=50,  # Reduced for CPU efficiency
                     class_weight=self.class_weight,
                     random_state=self.random_state,
-                    n_jobs=-1
+                    n_jobs=1  # CPU optimization
                 ),
                 'param_grid': {
                     'model__n_estimators': [50, 100],
                     'model__max_depth': [10, None]
                 }
+            },
+            'lightgbm': {
+                'model': lgb.LGBMClassifier(
+                    objective='binary',
+                    boosting_type='gbdt',
+                    num_leaves=31,
+                    max_depth=10,
+                    learning_rate=0.1,
+                    n_estimators=100,
+                    class_weight=self.class_weight,
+                    random_state=self.random_state,
+                    n_jobs=1,  # CPU optimization
+                    verbose=-1  # Suppress LightGBM output
+                ),
+                'param_grid': {
+                    'model__n_estimators': [50, 100],
+                    'model__learning_rate': [0.05, 0.1],
+                    'model__num_leaves': [15, 31]
+                }
             }
         }
 
@@ -752,7 +896,7 @@ class EnhancedModelTrainer:
                 param_grid,
                 cv=cv_strategy,
                 scoring='f1_weighted',
-                n_jobs=1,  # Single job for stability
+                n_jobs=1,  # Single job for CPU optimization
                 verbose=0,  # Reduce verbosity for speed
                 return_train_score=True  # For overfitting analysis
             )
@@ -813,9 +957,10 @@ class EnhancedModelTrainer:
                 raise Exception(f"Both hyperparameter tuning and fallback training failed: {str(e)} | {str(e2)}")
 
     def train_and_evaluate_models(self, X_train, X_test, y_train, y_test) -> Dict:
-        """Train and evaluate multiple models with enhanced features and comprehensive CV"""
+        """Train and evaluate multiple models including LightGBM with enhanced features and comprehensive CV"""
         
         results = {}
+        individual_models = {}
 
         for model_name in self.models.keys():
             logger.info(f"Training {model_name} with {'enhanced' if self.use_enhanced_features else 'standard'} features...")
@@ -843,6 +988,9 @@ class EnhancedModelTrainer:
                     'feature_type': 'enhanced' if self.use_enhanced_features else 'standard'
                 }
 
+                # Store for ensemble creation
+                individual_models[model_name] = best_model
+
                 # Log results
                 test_f1 = evaluation_metrics['f1']
                 cv_results = evaluation_metrics.get('cross_validation', {})
@@ -857,10 +1005,68 @@ class EnhancedModelTrainer:
                 logger.error(f"Training failed for {model_name}: {str(e)}")
                 results[model_name] = {'error': str(e)}
 
+        # Create and evaluate ensemble if enabled and we have multiple successful models
+        if self.enable_ensemble and len(individual_models) >= 2:
+            logger.info("Creating ensemble model...")
+            
+            try:
+                # Create ensemble
+                ensemble = self.ensemble_manager.create_ensemble(individual_models, voting='soft')
+                
+                # Fit ensemble
+                X_full_train = np.concatenate([X_train, X_test])
+                y_full_train = np.concatenate([y_train, y_test])
+                
+                ensemble.fit(X_train, y_train)
+                
+                # Evaluate ensemble
+                ensemble_metrics = self.comprehensive_evaluation(
+                    ensemble, X_test, y_test, X_train, y_train
+                )
+                
+                # Compare ensemble with individual models
+                ensemble_comparison = self.ensemble_manager.evaluate_ensemble_vs_individuals(
+                    ensemble, individual_models, X_test, y_test
+                )
+                
+                # Statistical comparison
+                statistical_comparison = self.ensemble_manager.statistical_ensemble_comparison(
+                    ensemble, individual_models, X_full_train, y_full_train, self.cv_manager
+                )
+                
+                # Store ensemble results
+                results['ensemble'] = {
+                    'model': ensemble,
+                    'evaluation_metrics': ensemble_metrics,
+                    'ensemble_comparison': ensemble_comparison,
+                    'statistical_comparison': statistical_comparison,
+                    'training_time': datetime.now().isoformat(),
+                    'feature_type': 'enhanced' if self.use_enhanced_features else 'standard'
+                }
+                
+                # Add ensemble to individual models for selection
+                individual_models['ensemble'] = ensemble
+                
+                # Log ensemble results
+                ensemble_f1 = ensemble_metrics['f1']
+                ensemble_improvement = ensemble_comparison.get('ensemble_analysis', {}).get('improvement', 0)
+                logger.info(f"Ensemble F1: {ensemble_f1:.4f}, Improvement: {ensemble_improvement:.4f}")
+                
+                # Log recommendation
+                recommendation = statistical_comparison.get('ensemble_recommendation', {})
+                if recommendation.get('use_ensemble', False):
+                    logger.info(f"✅ Ensemble recommended (confidence: {recommendation.get('confidence', 0):.2f})")
+                else:
+                    logger.info(f"❌ Ensemble not recommended")
+                
+            except Exception as e:
+                logger.error(f"Ensemble creation failed: {str(e)}")
+                results['ensemble'] = {'error': str(e)}
+
         return results
 
     def select_best_model(self, results: Dict) -> Tuple[str, Any, Dict]:
-        """Select the best performing model based on CV results"""
+        """Select the best performing model based on CV results with ensemble consideration"""
         
         if self.progress_tracker:
             self.progress_tracker.update("Selecting best model")
@@ -870,28 +1076,50 @@ class EnhancedModelTrainer:
         best_score = -1
         best_metrics = None
 
-        for model_name, result in results.items():
-            if 'error' in result:
-                continue
-
-            # Prioritize CV F1 score if available, fallback to test F1
-            cv_results = result['evaluation_metrics'].get('cross_validation', {})
-            if 'test_scores' in cv_results and 'f1' in cv_results['test_scores']:
-                f1_score = cv_results['test_scores']['f1']['mean']
-                score_type = "CV F1"
-            else:
-                f1_score = result['evaluation_metrics']['f1']
-                score_type = "Test F1"
-
-            if f1_score > best_score:
-                best_score = f1_score
-                best_model_name = model_name
-                best_model = result['model']
-                best_metrics = result['evaluation_metrics']
+        # Consider ensemble first if it exists and is recommended
+        if 'ensemble' in results and 'error' not in results['ensemble']:
+            ensemble_result = results['ensemble']
+            statistical_comparison = ensemble_result.get('statistical_comparison', {})
+            recommendation = statistical_comparison.get('ensemble_recommendation', {})
+            
+            if recommendation.get('use_ensemble', False):
+                ensemble_metrics = ensemble_result['evaluation_metrics']
+                cv_results = ensemble_metrics.get('cross_validation', {})
+                
+                if 'test_scores' in cv_results and 'f1' in cv_results['test_scores']:
+                    f1_score = cv_results['test_scores']['f1']['mean']
+                    if f1_score > best_score:
+                        best_score = f1_score
+                        best_model_name = 'ensemble'
+                        best_model = ensemble_result['model']
+                        best_metrics = ensemble_metrics
+                        logger.info("✅ Ensemble selected as best model")
+
+        # If ensemble not selected, choose best individual model
+        if best_model_name is None:
+            for model_name, result in results.items():
+                if 'error' in result or model_name == 'ensemble':
+                    continue
+
+                # Prioritize CV F1 score if available, fallback to test F1
+                cv_results = result['evaluation_metrics'].get('cross_validation', {})
+                if 'test_scores' in cv_results and 'f1' in cv_results['test_scores']:
+                    f1_score = cv_results['test_scores']['f1']['mean']
+                    score_type = "CV F1"
+                else:
+                    f1_score = result['evaluation_metrics']['f1']
+                    score_type = "Test F1"
+
+                if f1_score > best_score:
+                    best_score = f1_score
+                    best_model_name = model_name
+                    best_model = result['model']
+                    best_metrics = result['evaluation_metrics']
 
         if best_model_name is None:
             raise ValueError("No models trained successfully")
 
+        score_type = "CV F1" if 'cross_validation' in best_metrics else "Test F1"
         logger.info(f"Best model: {best_model_name} with {score_type} score: {best_score:.4f}")
         return best_model_name, best_model, best_metrics
 
@@ -940,6 +1168,10 @@ class EnhancedModelTrainer:
                         }
                         joblib.dump(enhanced_ref, self.vectorizer_path)
                         logger.info(f"✅ Saved enhanced features reference to {self.vectorizer_path}")
+                elif model_name == 'ensemble':
+                    # Handle ensemble model saving
+                    joblib.dump(model, self.model_path)
+                    logger.info(f"✅ Saved ensemble model to {self.model_path}")
                         
             except Exception as e:
                 logger.warning(f"Could not save individual components: {e}")
@@ -954,6 +1186,7 @@ class EnhancedModelTrainer:
             metadata = {
                 'model_version': f"v1.0_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
                 'model_type': model_name,
+                'is_ensemble': model_name == 'ensemble',
                 'feature_engineering': {
                     'type': 'enhanced' if self.use_enhanced_features else 'standard',
                     'enhanced_features_available': ENHANCED_FEATURES_AVAILABLE,
@@ -973,7 +1206,8 @@ class EnhancedModelTrainer:
                     'max_features': self.max_features,
                     'ngram_range': self.ngram_range,
                     'feature_selection_k': self.feature_selection_k,
-                    'use_enhanced_features': self.use_enhanced_features
+                    'use_enhanced_features': self.use_enhanced_features,
+                    'enable_ensemble': self.enable_ensemble
                 }
             }
             
@@ -1030,6 +1264,18 @@ class EnhancedModelTrainer:
                     metadata['cv_accuracy_mean'] = cv_results['test_scores']['accuracy']['mean']
                     metadata['cv_accuracy_std'] = cv_results['test_scores']['accuracy']['std']
             
+            # Add ensemble information if applicable
+            if model_name == 'ensemble' and 'ensemble' in results:
+                ensemble_result = results['ensemble']
+                ensemble_comparison = ensemble_result.get('ensemble_comparison', {})
+                statistical_comparison = ensemble_result.get('statistical_comparison', {})
+                
+                metadata['ensemble_info'] = {
+                    'ensemble_analysis': ensemble_comparison.get('ensemble_analysis', {}),
+                    'statistical_recommendation': statistical_comparison.get('ensemble_recommendation', {}),
+                    'individual_models': list(ensemble_comparison.keys()) if ensemble_comparison else []
+                }
+            
             # Add model comparison results if available
             if len(results) > 1:
                 model_comparison = {}
@@ -1058,6 +1304,10 @@ class EnhancedModelTrainer:
                 for feature_type, count in feature_metadata.get('feature_types', {}).items():
                     logger.info(f"   {feature_type}: {count}")
 
+            # Log ensemble information
+            if model_name == 'ensemble':
+                logger.info(f"✅ Ensemble model selected and saved")
+
             logger.info(f"✅ Model artifacts saved successfully with {'enhanced' if self.use_enhanced_features else 'standard'} features")
             return True
 
@@ -1072,18 +1322,22 @@ class EnhancedModelTrainer:
                 logger.error(f"Failed to save backup pipeline: {str(e2)}")
                 return False
 
-    def train_model(self, data_path: str = None, force_enhanced: bool = None) -> Tuple[bool, str]:
-        """Main training function with enhanced feature engineering pipeline"""
+    def train_model(self, data_path: str = None, force_enhanced: bool = None, force_ensemble: bool = None) -> Tuple[bool, str]:
+        """Main training function with LightGBM, enhanced feature engineering, and ensemble voting"""
         try:
-            # Override enhanced features setting if specified
+            # Override settings if specified
             if force_enhanced is not None:
                 original_setting = self.use_enhanced_features
                 self.use_enhanced_features = force_enhanced and ENHANCED_FEATURES_AVAILABLE
                 if force_enhanced and not ENHANCED_FEATURES_AVAILABLE:
                     logger.warning("Enhanced features requested but not available, using standard features")
             
+            if force_ensemble is not None:
+                self.enable_ensemble = force_ensemble
+            
             feature_type = "enhanced" if self.use_enhanced_features else "standard"
-            logger.info(f"Starting {feature_type} model training with cross-validation...")
+            ensemble_info = "with ensemble" if self.enable_ensemble else "without ensemble"
+            logger.info(f"Starting {feature_type} model training {ensemble_info} including LightGBM...")
 
             # Override data path if provided
             if data_path:
@@ -1099,24 +1353,27 @@ class EnhancedModelTrainer:
                 len(df), 
                 enable_tuning=True, 
                 cv_folds=self.cv_folds,
-                use_enhanced_features=self.use_enhanced_features
+                use_enhanced_features=self.use_enhanced_features,
+                enable_ensemble=self.enable_ensemble
             )
             
             print(f"\n📊 Enhanced Training Configuration:")
             print(f"Dataset size: {len(df)} samples")
             print(f"Feature engineering: {feature_type.title()}")
             print(f"Cross-validation folds: {self.cv_folds}")
+            print(f"Models: Logistic Regression, Random Forest, LightGBM")
+            print(f"Ensemble voting: {'Enabled' if self.enable_ensemble else 'Disabled'}")
             print(f"Estimated time: {time_estimate['total_formatted']}")
-            print(f"Models to train: {len(self.models)}")
             print(f"Hyperparameter tuning: Enabled")
             if self.use_enhanced_features:
                 print(f"Enhanced features: Sentiment, Readability, Entities, Linguistic")
             print()
 
-            # Setup progress tracker (adjusted for enhanced features)
+            # Setup progress tracker (adjusted for LightGBM and ensemble)
             base_steps = 4 + (len(self.models) * 3) + 1  # Basic steps
             enhanced_steps = 2 if self.use_enhanced_features else 0  # Feature engineering steps
-            total_steps = base_steps + enhanced_steps
+            ensemble_steps = 3 if self.enable_ensemble else 0  # Ensemble creation and evaluation
+            total_steps = base_steps + enhanced_steps + ensemble_steps
             self.progress_tracker = ProgressTracker(total_steps, f"{feature_type.title()} Training Progress")
 
             # Prepare data
@@ -1152,10 +1409,10 @@ class EnhancedModelTrainer:
             if len(X_test) < 1:
                 return False, "Cannot create test set. Dataset too small."
 
-            # Train and evaluate models with enhanced features
+            # Train and evaluate models with LightGBM and enhanced features
             results = self.train_and_evaluate_models(X_train, X_test, y_train, y_test)
 
-            # Select best model
+            # Select best model (could be ensemble)
             best_model_name, best_model, best_metrics = self.select_best_model(results)
 
             # Save model artifacts with enhanced feature information
@@ -1165,7 +1422,7 @@ class EnhancedModelTrainer:
             # Finish progress tracking
             self.progress_tracker.finish()
 
-            # Create success message with enhanced feature information
+            # Create success message with comprehensive information
             cv_results = best_metrics.get('cross_validation', {})
             cv_info = ""
             if 'test_scores' in cv_results and 'f1' in cv_results['test_scores']:
@@ -1180,9 +1437,14 @@ class EnhancedModelTrainer:
                 if feature_metadata:
                     total_features = feature_metadata.get('total_features', 0)
                     feature_info = f", Enhanced Features: {total_features}"
+            
+            # Ensemble information
+            ensemble_info = ""
+            if best_model_name == 'ensemble':
+                ensemble_info = " (Ensemble Model Selected)"
 
             success_message = (
-                f"{feature_type.title()} model training completed successfully. "
+                f"{feature_type.title()} model training completed successfully{ensemble_info}. "
                 f"Best model: {best_model_name} "
                 f"(Test F1: {best_metrics['f1']:.4f}, Test Accuracy: {best_metrics['accuracy']:.4f}{cv_info}{feature_info})"
             )
@@ -1193,22 +1455,24 @@ class EnhancedModelTrainer:
         except Exception as e:
             if self.progress_tracker:
                 print()  # New line after progress bar
-            error_message = f"Enhanced model training failed: {str(e)}"
+            error_message = f"Enhanced model training with LightGBM failed: {str(e)}"
             logger.error(error_message)
             return False, error_message
 
 
 def main():
-    """Main execution function with enhanced feature engineering support"""
+    """Main execution function with LightGBM, enhanced features, and ensemble support"""
     import argparse
     
     # Parse command line arguments
-    parser = argparse.ArgumentParser(description='Train fake news detection model with enhanced features')
+    parser = argparse.ArgumentParser(description='Train fake news detection model with LightGBM and enhanced features')
     parser.add_argument('--data_path', type=str, help='Path to training data CSV file')
     parser.add_argument('--config_path', type=str, help='Path to training configuration JSON file')
     parser.add_argument('--cv_folds', type=int, default=5, help='Number of cross-validation folds')
     parser.add_argument('--enhanced_features', action='store_true', help='Force use of enhanced features')
     parser.add_argument('--standard_features', action='store_true', help='Force use of standard TF-IDF features only')
+    parser.add_argument('--enable_ensemble', action='store_true', help='Enable ensemble voting')
+    parser.add_argument('--disable_ensemble', action='store_true', help='Disable ensemble voting')
     args = parser.parse_args()
     
     # Determine feature engineering mode
@@ -1222,7 +1486,21 @@ def main():
         use_enhanced = False
         logger.info("Standard features explicitly requested")
     
-    trainer = EnhancedModelTrainer(use_enhanced_features=use_enhanced)
+    # Determine ensemble mode
+    enable_ensemble = None
+    if args.enable_ensemble and args.disable_ensemble:
+        logger.warning("Both --enable_ensemble and --disable_ensemble specified. Using default.")
+    elif args.enable_ensemble:
+        enable_ensemble = True
+        logger.info("Ensemble voting explicitly enabled")
+    elif args.disable_ensemble:
+        enable_ensemble = False
+        logger.info("Ensemble voting explicitly disabled")
+    
+    trainer = EnhancedModelTrainer(
+        use_enhanced_features=use_enhanced, 
+        enable_ensemble=enable_ensemble if enable_ensemble is not None else True
+    )
     
     # Apply CV folds from command line
     if args.cv_folds:
@@ -1246,6 +1524,10 @@ def main():
             if 'enhanced_features' in config and use_enhanced is None:
                 trainer.use_enhanced_features = config['enhanced_features'] and ENHANCED_FEATURES_AVAILABLE
             
+            # Ensemble configuration
+            if 'enable_ensemble' in config and enable_ensemble is None:
+                trainer.enable_ensemble = config['enable_ensemble']
+            
             # Filter models if specified
             selected_models = config.get('selected_models')
             if selected_models and len(selected_models) < len(trainer.models):
@@ -1258,6 +1540,8 @@ def main():
             logger.info(f"Applied custom configuration with {trainer.cv_folds} CV folds")
             if trainer.use_enhanced_features:
                 logger.info("Enhanced features enabled via configuration")
+            if trainer.enable_ensemble:
+                logger.info("Ensemble voting enabled via configuration")
             
         except Exception as e:
             logger.warning(f"Failed to load configuration: {e}, using defaults")
@@ -1277,6 +1561,14 @@ def main():
                     print(f"  {feature_type}: {count}")
             except Exception as e:
                 logger.warning(f"Could not display feature summary: {e}")
+        
+        # Print model information
+        print(f"\n🎯 Model Information:")
+        print(f"Models trained: {', '.join(trainer.models.keys())}")
+        if trainer.enable_ensemble:
+            print(f"Ensemble voting: Enabled")
+        else:
+            print(f"Ensemble voting: Disabled")
     else:
         print(f"❌ {message}")
         exit(1)