Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| import json | |
| import joblib | |
| import pandas as pd | |
| import numpy as np | |
| import xgboost as xgb | |
| from sklearn.pipeline import Pipeline | |
| from sklearn.model_selection import train_test_split, RandomizedSearchCV, cross_val_predict | |
| from sklearn.metrics import precision_recall_curve, classification_report, accuracy_score, f1_score, precision_score, recall_score | |
| try: | |
| from imblearn.over_sampling import SMOTE | |
| except ImportError: | |
| SMOTE = None | |
| from supabase import create_client, Client | |
| from dotenv import load_dotenv | |
| # Import local modules | |
| # Assuming feature_engineering.py is in the same directory | |
| try: | |
| from feature_engineering import FraudFeatureEngineer | |
| except ImportError: | |
| import sys | |
| sys.path.append(os.path.dirname(os.path.abspath(__file__))) | |
| from feature_engineering import FraudFeatureEngineer | |
| load_dotenv() | |
| # Initialize Supabase client | |
| url: str = os.environ.get("SUPABASE_URL") | |
| key: str = os.environ.get("SUPABASE_SERVICE_ROLE_KEY") # Use Service Role Key for backend updates | |
| supabase: Client = create_client(url, key) if url and key else None | |
| def update_job_status(job_id: str, status: str, metrics: dict = None, file_path: str = None): | |
| """Update job status in Supabase""" | |
| if not supabase: | |
| print(f"⚠️ Supabase not configured. Job {job_id} status: {status}") | |
| return | |
| data = {"status": status} | |
| if metrics: | |
| data["metrics"] = metrics | |
| if file_path: | |
| data["file_path"] = file_path | |
| try: | |
| supabase.table("model_registry").update(data).eq("id", job_id).execute() | |
| except Exception as e: | |
| print(f"❌ Failed to update Supabase for job {job_id}: {str(e)}") | |
| def make_splits(df, test_frac=0.05, time_col='step', min_test_fraud=100, random_state=42): | |
| """ | |
| Returns: X_train, X_test, y_train, y_test | |
| """ | |
| # Normalize column names | |
| if 'is_fraud' in df.columns: | |
| df = df.rename(columns={'is_fraud': 'isFraud'}) | |
| if 'fraud' in df.columns: | |
| df = df.rename(columns={'fraud': 'isFraud'}) | |
| # Temporal test split | |
| if time_col in df.columns: | |
| cutoff = df[time_col].quantile(1 - test_frac) | |
| train_df = df[df[time_col] <= cutoff].reset_index(drop=True) | |
| test_df = df[df[time_col] > cutoff].reset_index(drop=True) | |
| else: | |
| # Random split if no time column | |
| train_df, test_df = train_test_split(df, test_size=test_frac, stratify=df.get('isFraud'), random_state=42) | |
| train_df = train_df.reset_index(drop=True) | |
| test_df = test_df.reset_index(drop=True) | |
| # Fallback to stratified split if test set is poor | |
| if (test_df.empty) or (test_df['isFraud'].sum() < min_test_fraud): | |
| X_temp = df.drop(columns=['isFraud']) | |
| y_temp = df['isFraud'] | |
| X_train, X_test, y_train, y_test = train_test_split( | |
| X_temp, y_temp, test_size=test_frac, stratify=y_temp, random_state=random_state | |
| ) | |
| return X_train.reset_index(drop=True), X_test.reset_index(drop=True), y_train.reset_index(drop=True), y_test.reset_index(drop=True) | |
| else: | |
| X_train = train_df.drop(columns=['isFraud']) | |
| y_train = train_df['isFraud'] | |
| X_test = test_df.drop(columns=['isFraud']) | |
| y_test = test_df['isFraud'] | |
| return X_train.reset_index(drop=True), X_test.reset_index(drop=True), y_train.reset_index(drop=True), y_test.reset_index(drop=True) | |
| async def train_model_async(job_id: str, file_path: str, params: dict): | |
| """ | |
| Background task to train model | |
| """ | |
| print(f"🚀 Starting training job {job_id} with params: {params}") | |
| update_job_status(job_id, "training") | |
| try: | |
| # 1. Load Data | |
| if not os.path.exists(file_path): | |
| raise FileNotFoundError(f"Dataset not found at {file_path}") | |
| df = pd.read_csv(file_path) | |
| # Basic Preprocessing | |
| if 'oldbalanceOrg' in df.columns: | |
| df = df.rename(columns={'oldbalanceOrg':'oldBalanceOrig', 'newbalanceOrig':'newBalanceOrig', | |
| 'oldbalanceDest':'oldBalanceDest', 'newbalanceDest':'newBalanceDest'}) | |
| # Filter types if needed (Notebook filters to TRANSFER/CASH_OUT) | |
| # We'll stick to what the notebook did for consistency | |
| if 'type' in df.columns: | |
| df = df[df['type'].isin(['TRANSFER', 'CASH_OUT'])].reset_index(drop=True) | |
| # 2. Split Data | |
| test_frac = params.get('test_split', 0.05) | |
| X_train, X_test, y_train, y_test = make_splits(df, test_frac=test_frac) | |
| print(f"📊 Data split: Train={len(X_train)}, Test={len(X_test)}") | |
| # 3. Feature Engineering | |
| pagerank_limit = params.get('pagerank_limit', 10000) | |
| advanced_features = params.get('advanced_feature_engineering', False) | |
| fe = FraudFeatureEngineer(pagerank_limit=pagerank_limit, advanced_features=advanced_features) | |
| print("⚙️ Fitting feature engineer...") | |
| fe.fit(X_train, y_train) | |
| X_train_trans = fe.transform(X_train) | |
| X_test_trans = fe.transform(X_test) | |
| # Apply Advanced Preprocessing (SMOTE) if requested | |
| if params.get('advanced_preprocessing'): | |
| if SMOTE: | |
| print("⚖️ Applying SMOTE for class imbalance...") | |
| try: | |
| smote = SMOTE(random_state=42) | |
| X_train_trans, y_train = smote.fit_resample(X_train_trans, y_train) | |
| print(f" New training shape: {X_train_trans.shape}") | |
| except Exception as e: | |
| print(f"⚠️ SMOTE failed: {str(e)}. Proceeding with original data.") | |
| else: | |
| print("⚠️ SMOTE requested but imbalanced-learn not installed. Proceeding without it.") | |
| # 4. Configure XGBoost | |
| # Calculate scale_pos_weight | |
| neg = (y_train == 0).sum() | |
| pos = (y_train == 1).sum() | |
| spw = int(max(1, neg / max(1, pos))) | |
| xgb_params = { | |
| 'objective': 'binary:logistic', | |
| 'tree_method': 'hist', | |
| 'random_state': 42, | |
| 'n_jobs': 1, | |
| 'n_estimators': int(params.get('n_estimators', 300)), | |
| 'max_depth': int(params.get('max_depth', 6)), | |
| 'learning_rate': float(params.get('learning_rate', 0.05)), | |
| 'scale_pos_weight': spw | |
| } | |
| clf = xgb.XGBClassifier(**xgb_params) | |
| # 5. Train | |
| print("🏋️ Training classifier...") | |
| clf.fit(X_train_trans, y_train) | |
| # 6. Evaluate | |
| print("📝 Evaluating...") | |
| y_pred = clf.predict(X_test_trans) | |
| metrics = { | |
| "accuracy": float(accuracy_score(y_test, y_pred)), | |
| "precision": float(precision_score(y_test, y_pred)), | |
| "recall": float(recall_score(y_test, y_pred)), | |
| "f1": float(f1_score(y_test, y_pred)) | |
| } | |
| print(f"✅ Metrics: {metrics}") | |
| # 7. Save Model | |
| # Ensure Models directory exists | |
| models_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "Models") | |
| os.makedirs(models_dir, exist_ok=True) | |
| model_filename = f"model_{job_id}.pkl" | |
| model_save_path = os.path.join(models_dir, model_filename) | |
| # Save ONLY the classifier (as per notebook strategy) | |
| joblib.dump(clf, model_save_path) | |
| print(f"💾 Model saved to {model_save_path}") | |
| # 8. Update Registry | |
| update_job_status(job_id, "ready", metrics=metrics, file_path=f"Models/{model_filename}") | |
| except Exception as e: | |
| print(f"❌ Training failed: {str(e)}") | |
| update_job_status(job_id, "failed", metrics={"error": str(e)}) | |