Update utils/model_trainer.py

utils/model_trainer.py

@@ -1,6 +1,6 @@
 """
 Embedded Model Training for HF Spaces
-Fixed version with proper data validation and cleaning
+Fixed version with dynamic column mapping for SAP SALT dataset
 """
 
 import pandas as pd
@@ -15,7 +15,7 @@ from pathlib import Path
 from datetime import datetime
 
 class EmbeddedChurnTrainer:
-    """Embedded trainer with proper data validation"""
+    """Embedded trainer with dynamic column mapping for real SAP SALT data"""
     
     def __init__(self):
         self.model_path = Path('models/churn_model_v1.pkl')
@@ -23,6 +23,7 @@ class EmbeddedChurnTrainer:
         self.model = None
         self.label_encoders = {}
         self.feature_columns = []
+        self.column_mapping = {}
         
     def model_exists(self):
         """Check if trained model exists"""
@@ -30,65 +31,159 @@ class EmbeddedChurnTrainer:
     
     @st.cache_data
     def load_sap_data(_self):
-        """Load real SAP SALT dataset using Hugging Face datasets library"""
+        """Load real SAP SALT dataset and inspect its structure"""
         try:
             from datasets import load_dataset
             
             st.info("🔄 Loading SAP SALT dataset from Hugging Face...")
             
-            # Load the dataset - this will fail gracefully if not accessible
+            # Load the dataset
             dataset = load_dataset("SAP/SALT", split="train")
             data_df = dataset.to_pandas()
             
-            # Add required aggregated fields
+            # Debug: Show actual columns
+            st.info(f"📋 Dataset columns: {list(data_df.columns)}")
+            st.info(f"📊 Dataset shape: {data_df.shape}")
+            
+            # Create column mapping based on available columns
+            _self.column_mapping = _self._create_column_mapping(data_df.columns)
+            st.info(f"🔗 Column mapping: {_self.column_mapping}")
+            
+            # Add aggregated fields
             data_df = _self._add_aggregated_fields(data_df)
             
             st.success(f"✅ Loaded {len(data_df)} records from SAP SALT dataset")
             return data_df
             
         except ImportError:
-            st.error("❌ Hugging Face datasets library not available. Install with: pip install datasets")
-            raise RuntimeError("datasets library required to load SAP SALT dataset")
+            st.error("❌ Hugging Face datasets library not available")
+            raise RuntimeError("datasets library required")
             
         except Exception as e:
-            if "gated" in str(e).lower() or "authentication" in str(e).lower() or "401" in str(e):
+            if "gated" in str(e).lower() or "authentication" in str(e).lower():
                 st.error("🔐 **SAP SALT Dataset Access Required**")
                 st.info("""
                 **To access SAP SALT dataset:**
                 1. Visit: https://huggingface.co/datasets/SAP/SALT
                 2. Click "Agree and access repository" 
-                3. Add your HF token to Spaces secrets:
-                   - Go to Space Settings → Variables and Secrets
-                   - Add secret: `HF_TOKEN` with your token value
+                3. Add HF token to Space secrets: `HF_TOKEN`
                 4. Restart the Space
                 """)
-                raise RuntimeError(f"SAP SALT dataset access denied: {str(e)}")
             else:
                 st.error(f"❌ Failed to load SAP SALT dataset: {str(e)}")
-                raise RuntimeError(f"Dataset loading failed: {str(e)}")
+            raise
+    
+    def _create_column_mapping(self, available_columns):
+        """Create mapping from expected columns to available columns"""
+        cols = [col.upper() for col in available_columns]  # Convert to uppercase for matching
+        available_upper = {col.upper(): col for col in available_columns}
+        
+        mapping = {}
+        
+        # Map customer identifier
+        customer_candidates = ['CUSTOMER', 'SOLDTOPARTY', 'CUSTOMERID', 'CUSTOMER_ID']
+        for candidate in customer_candidates:
+            if candidate in cols:
+                mapping['Customer'] = available_upper[candidate]
+                break
+        else:
+            mapping['Customer'] = available_columns[0] if available_columns else 'Customer'  # Fallback
+        
+        # Map customer name
+        name_candidates = ['CUSTOMERNAME', 'CUSTOMER_NAME', 'NAME', 'COMPANYNAME']
+        for candidate in name_candidates:
+            if candidate in cols:
+                mapping['CustomerName'] = available_upper[candidate]
+                break
+        else:
+            mapping['CustomerName'] = None
+        
+        # Map country
+        country_candidates = ['COUNTRY', 'COUNTRYKEY', 'COUNTRY_CODE', 'LAND1']
+        for candidate in country_candidates:
+            if candidate in cols:
+                mapping['Country'] = available_upper[candidate]
+                break
+        else:
+            mapping['Country'] = None
+        
+        # Map customer group
+        group_candidates = ['CUSTOMERGROUP', 'CUSTOMER_GROUP', 'CUSTOMERCLASSIFICATION', 'KTOKD']
+        for candidate in group_candidates:
+            if candidate in cols:
+                mapping['CustomerGroup'] = available_upper[candidate]
+                break
+        else:
+            mapping['CustomerGroup'] = None
+        
+        # Map sales document
+        doc_candidates = ['SALESDOCUMENT', 'SALES_DOCUMENT', 'VBELN', 'DOCUMENTNUMBER']
+        for candidate in doc_candidates:
+            if candidate in cols:
+                mapping['SalesDocument'] = available_upper[candidate]
+                break
+        else:
+            mapping['SalesDocument'] = None
+        
+        # Map creation date
+        date_candidates = ['CREATIONDATE', 'CREATION_DATE', 'ERDAT', 'REQUESTEDDELIVERYDATE', 'DATE']
+        for candidate in date_candidates:
+            if candidate in cols:
+                mapping['CreationDate'] = available_upper[candidate]
+                break
+        else:
+            mapping['CreationDate'] = None
+        
+        return mapping
     
     def _add_aggregated_fields(self, data):
-        """Add customer-level aggregations for churn modeling"""
-        # Identify key columns (adapt based on actual SAP SALT structure)
-        customer_col = next((col for col in ['CUSTOMER', 'Customer', 'SOLDTOPARTY', 'SoldToParty'] if col in data.columns), 'Customer')
-        date_col = next((col for col in ['CREATIONDATE', 'CreationDate', 'REQUESTEDDELIVERYDATE'] if col in data.columns), 'CreationDate')
+        """Add customer-level aggregations using dynamic column mapping"""
+        # Get actual column names
+        customer_col = self.column_mapping.get('Customer')
+        date_col = self.column_mapping.get('CreationDate')
+        sales_doc_col = self.column_mapping.get('SalesDocument')
+        
+        if not customer_col:
+            st.error("❌ No customer identifier column found")
+            raise ValueError("Cannot identify customer column")
         
         # Customer-level aggregations
-        customer_aggs = data.groupby(customer_col).agg({
-            date_col: ['count', 'min', 'max']
-        }).reset_index()
+        agg_dict = {}
+        
+        if sales_doc_col:
+            agg_dict[sales_doc_col] = 'count'
         
-        # Flatten column names
-        customer_aggs.columns = [customer_col, 'total_orders', 'first_order_date', 'last_order_date']
+        if date_col:
+            agg_dict[date_col] = ['min', 'max']
         
-        # Merge back to original data
-        data = data.merge(customer_aggs, on=customer_col, how='left')
+        if not agg_dict:
+            # If no aggregation columns available, create dummy data
+            data['total_orders'] = 1
+            data['first_order_date'] = '2024-01-01'
+            data['last_order_date'] = '2024-01-01'
+        else:
+            customer_aggs = data.groupby(customer_col).agg(agg_dict).reset_index()
+            
+            # Flatten column names
+            new_cols = [customer_col]
+            if sales_doc_col:
+                new_cols.append('total_orders')
+            if date_col:
+                new_cols.extend(['first_order_date', 'last_order_date'])
+            
+            customer_aggs.columns = new_cols
+            
+            # Merge back to original data
+            data = data.merge(customer_aggs, on=customer_col, how='left')
         
-        # Standardize column names
-        data = data.rename(columns={
-            customer_col: 'Customer',
-            date_col: 'CreationDate'
-        })
+        # Standardize column names for downstream processing
+        rename_dict = {}
+        for standard_name, actual_name in self.column_mapping.items():
+            if actual_name and actual_name in data.columns:
+                rename_dict[actual_name] = standard_name
+        
+        if rename_dict:
+            data = data.rename(columns=rename_dict)
         
         return data
     
@@ -132,42 +227,66 @@ class EmbeddedChurnTrainer:
             raise
     
     def engineer_features(self, data):
-        """Feature engineering with proper data validation and cleaning"""
+        """Feature engineering with dynamic column handling"""
         try:
-            # Customer-level aggregation
-            customer_features = data.groupby('Customer').agg({
-                'CustomerName': 'first',
-                'Country': 'first', 
-                'CustomerGroup': 'first',
-                'total_orders': 'first',
-                'last_order_date': 'first',
-                'first_order_date': 'first'
-            }).reset_index()
-            
-            # Handle dates
+            # Identify available columns for customer aggregation
+            agg_cols = ['Customer']  # Always need customer ID
+            
+            optional_cols = ['CustomerName', 'Country', 'CustomerGroup']
+            for col in optional_cols:
+                if col in data.columns and data[col].notna().any():
+                    agg_cols.append(col)
+            
+            # Customer-level aggregation with only available columns
+            agg_dict = {}
+            for col in agg_cols:
+                if col != 'Customer':
+                    agg_dict[col] = 'first'
+            
+            # Add order-related aggregations
+            if 'total_orders' in data.columns:
+                agg_dict['total_orders'] = 'first'
+            if 'first_order_date' in data.columns:
+                agg_dict['first_order_date'] = 'first'
+            if 'last_order_date' in data.columns:
+                agg_dict['last_order_date'] = 'first'
+            
+            customer_features = data.groupby('Customer').agg(agg_dict).reset_index()
+            
+            # Handle dates safely
             reference_date = pd.to_datetime('2024-12-31')
-            customer_features['last_order_date'] = pd.to_datetime(customer_features['last_order_date'], errors='coerce')
-            customer_features['first_order_date'] = pd.to_datetime(customer_features['first_order_date'], errors='coerce')
             
-            # RFM Features with proper handling of edge cases
-            customer_features['Recency'] = (reference_date - customer_features['last_order_date']).dt.days
-            customer_features['Recency'] = customer_features['Recency'].fillna(365).clip(0, 3650)  # Cap at 10 years
+            if 'last_order_date' in customer_features.columns:
+                customer_features['last_order_date'] = pd.to_datetime(customer_features['last_order_date'], errors='coerce')
+                customer_features['Recency'] = (reference_date - customer_features['last_order_date']).dt.days
+            else:
+                customer_features['Recency'] = 100  # Default recency
             
-            customer_features['Frequency'] = customer_features['total_orders'].fillna(0).clip(0, 1000)  # Cap at reasonable max
+            if 'first_order_date' in customer_features.columns:
+                customer_features['first_order_date'] = pd.to_datetime(customer_features['first_order_date'], errors='coerce')
+                customer_features['Tenure'] = (reference_date - customer_features['first_order_date']).dt.days
+            else:
+                customer_features['Tenure'] = 365  # Default tenure
             
-            # Monetary value (simplified calculation to avoid extreme values)
-            customer_features['Monetary'] = (customer_features['Frequency'] * 500).clip(0, 1000000)  # Cap at 1M
+            # RFM Features with safe handling
+            customer_features['Recency'] = customer_features['Recency'].fillna(365).clip(0, 3650)
             
-            # Customer lifecycle features with safe division
-            customer_features['Tenure'] = (reference_date - customer_features['first_order_date']).dt.days
-            customer_features['Tenure'] = customer_features['Tenure'].fillna(0).clip(0, 3650)  # Cap at 10 years
+            if 'total_orders' in customer_features.columns:
+                customer_features['Frequency'] = customer_features['total_orders'].fillna(1).clip(1, 1000)
+            else:
+                customer_features['Frequency'] = 1  # Default frequency
             
-            # OrderVelocity with safe division to prevent infinity
-            tenure_months = customer_features['Tenure'] / 30 + 1  # Add 1 to prevent division by zero
-            customer_features['OrderVelocity'] = (customer_features['Frequency'] / tenure_months).clip(0, 100)  # Cap at reasonable max
+            customer_features['Monetary'] = (customer_features['Frequency'] * 500).clip(100, 1000000)
+            customer_features['Tenure'] = customer_features['Tenure'].fillna(365).clip(1, 3650)
             
-            # Categorical encoding with error handling
+            # Safe OrderVelocity calculation
+            tenure_months = customer_features['Tenure'] / 30 + 1
+            customer_features['OrderVelocity'] = (customer_features['Frequency'] / tenure_months).clip(0, 50)
+            
+            # Categorical encoding only for available columns
             self.label_encoders = {}
+            categorical_features = []
+            
             for col in ['Country', 'CustomerGroup']:
                 if col in customer_features.columns and customer_features[col].notna().any():
                     try:
@@ -175,84 +294,94 @@ class EmbeddedChurnTrainer:
                         customer_features[f'{col}_encoded'] = self.label_encoders[col].fit_transform(
                             customer_features[col].fillna('Unknown')
                         )
-                    except:
-                        # If encoding fails, create dummy encoded column
-                        customer_features[f'{col}_encoded'] = 0
+                        categorical_features.append(f'{col}_encoded')
+                    except Exception as e:
+                        st.warning(f"⚠️ Could not encode {col}: {str(e)}")
             
             # Target variable (churn definition)
             customer_features['IsChurned'] = (
                 (customer_features['Recency'] > 90) & 
-                (customer_features['Frequency'] > 0)
+                (customer_features['Frequency'] > 1)
             ).astype(int)
             
-            # Select features for model
+            # Define feature columns
             self.feature_columns = ['Recency', 'Frequency', 'Monetary', 'Tenure', 'OrderVelocity']
-            
-            # Add encoded categorical features if they exist
-            for col in ['Country', 'CustomerGroup']:
-                if f'{col}_encoded' in customer_features.columns:
-                    self.feature_columns.append(f'{col}_encoded')
+            self.feature_columns.extend(categorical_features)
             
             # Prepare final dataset
-            final_data = customer_features[self.feature_columns + ['IsChurned', 'Customer', 'CustomerName']].copy()
+            required_cols = self.feature_columns + ['IsChurned', 'Customer']
+            
+            # Add CustomerName if available
+            if 'CustomerName' in customer_features.columns:
+                required_cols.append('CustomerName')
             
-            # **CRITICAL: Clean all infinite and NaN values**
+            # Filter to only existing columns
+            available_cols = [col for col in required_cols if col in customer_features.columns]
+            final_data = customer_features[available_cols].copy()
+            
+            # **CRITICAL: Clean all data**
             for col in self.feature_columns:
-                # Replace infinity with NaN, then fill with 0
-                final_data[col] = final_data[col].replace([np.inf, -np.inf], np.nan).fillna(0)
-                
-                # Clip extreme values to prevent float32 overflow
-                final_data[col] = final_data[col].clip(-1e9, 1e9)
-            
-            # Validate no infinite or NaN values remain
-            if not np.isfinite(final_data[self.feature_columns]).all().all():
-                st.warning("⚠️ Cleaning remaining non-finite values...")
-                final_data[self.feature_columns] = final_data[self.feature_columns].fillna(0)
-                final_data[self.feature_columns] = final_data[self.feature_columns].replace([np.inf, -np.inf], 0)
+                if col in final_data.columns:
+                    final_data[col] = final_data[col].replace([np.inf, -np.inf], np.nan).fillna(0)
+                    final_data[col] = final_data[col].clip(-1e9, 1e9)
+            
+            st.info(f"✅ Features engineered: {self.feature_columns}")
+            st.info(f"📊 Final dataset shape: {final_data.shape}")
             
             return final_data
             
         except Exception as e:
             st.error(f"Feature engineering failed: {str(e)}")
+            st.info(f"Available columns: {list(data.columns)}")
             raise
     
     def train_model(self, data):
-        """Train RandomForest model with additional data validation"""
+        """Train model with additional validation"""
         try:
+            # Ensure all feature columns exist
+            missing_features = [col for col in self.feature_columns if col not in data.columns]
+            if missing_features:
+                st.warning(f"⚠️ Missing features: {missing_features}")
+                # Use only available features
+                self.feature_columns = [col for col in self.feature_columns if col in data.columns]
+            
+            if not self.feature_columns:
+                raise ValueError("No valid features available for training")
+            
             X = data[self.feature_columns].copy()
             y = data['IsChurned'].copy()
             
-            # **FINAL VALIDATION: Ensure X contains only finite values**
+            # Final data cleaning
             if not np.isfinite(X).all().all():
-                st.warning("⚠️ Final data cleaning before training...")
                 X = X.replace([np.inf, -np.inf], np.nan).fillna(0)
             
-            # Check data sufficiency
+            # Check data quality
             if len(X) < 50:
-                raise ValueError("Insufficient training data (need at least 50 samples)")
+                raise ValueError(f"Insufficient training data: {len(X)} samples")
             
             if y.nunique() < 2:
-                st.warning("⚠️ All customers have same churn status - adjusting model...")
-                # Create some artificial variation for model training
-                y.iloc[:len(y)//4] = 1 - y.iloc[:len(y)//4]
+                st.warning("⚠️ Creating artificial target variation for training...")
+                # Create some variation for model training
+                variation_size = len(y) // 4
+                y.iloc[:variation_size] = 1 - y.iloc[:variation_size]
             
             # Train-test split
             X_train, X_test, y_train, y_test = train_test_split(
-                X, y, test_size=0.2, random_state=42, stratify=y if y.nunique() > 1 else None
+                X, y, test_size=0.2, random_state=42, 
+                stratify=y if y.nunique() > 1 else None
             )
             
-            # Train model with reduced complexity to prevent memory issues
+            # Train model
             self.model = RandomForestClassifier(
-                n_estimators=50,      # Reduced for HF Spaces
-                max_depth=8,          # Prevent overly deep trees
-                min_samples_split=20, # Require minimum samples for splits
-                min_samples_leaf=10,  # Minimum samples in leaf
+                n_estimators=50,
+                max_depth=8,
+                min_samples_split=20,
+                min_samples_leaf=10,
                 class_weight='balanced',
                 random_state=42,
-                n_jobs=1             # Single thread for HF Spaces
+                n_jobs=1
             )
             
-            # Fit model
             self.model.fit(X_train, y_train)
             
             # Evaluate
@@ -266,9 +395,12 @@ class EmbeddedChurnTrainer:
                 'training_samples': len(X_train),
                 'test_samples': len(X_test),
                 'churn_rate': float(y.mean()),
-                'feature_importance': dict(zip(self.feature_columns, self.model.feature_importances_))
+                'feature_importance': dict(zip(self.feature_columns, self.model.feature_importances_)),
+                'column_mapping': self.column_mapping
             }
             
+            st.success(f"✅ Model trained successfully! Accuracy: {test_score:.3f}")
+            
             return metrics
             
         except Exception as e:
@@ -283,6 +415,7 @@ class EmbeddedChurnTrainer:
             'model': self.model,
             'label_encoders': self.label_encoders,
             'feature_columns': self.feature_columns,
+            'column_mapping': self.column_mapping,
             'version': 'v1',
             'training_date': datetime.now().isoformat()
         }
@@ -295,7 +428,8 @@ class EmbeddedChurnTrainer:
             'training_date': datetime.now().isoformat(),
             'metrics': metrics,
             'status': 'trained',
-            'data_source': 'SAP/SALT dataset from Hugging Face'
+            'data_source': 'SAP/SALT dataset from Hugging Face',
+            'column_mapping': self.column_mapping
         }
         
         with open(self.metadata_path, 'w') as f: