Home > Notebooks > ML Layer¶

import mlflow
from xgboost import XGBClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, precision_recall_curve

# Load feature data
df_features = spark.table("ml.player_churn_features").toPandas()

# Split data
X = df_features.drop(['player_id', 'churned'], axis=1)
y = df_features['churned']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y)

# Train with MLflow tracking
with mlflow.start_run(run_name="churn_model_v1"):
    model = XGBClassifier(
        n_estimators=100,
        max_depth=6,
        learning_rate=0.1,
        scale_pos_weight=len(y_train[y_train==0]) / len(y_train[y_train==1])
    )
    model.fit(X_train, y_train)

    # Log metrics
    y_pred_proba = model.predict_proba(X_test)[:, 1]
    auc = roc_auc_score(y_test, y_pred_proba)
    mlflow.log_metric("auc_roc", auc)

    # Log model
    mlflow.xgboost.log_model(model, "churn_model")

    # Register model
    mlflow.register_model(
        f"runs:/{mlflow.active_run().info.run_id}/churn_model",
        "player_churn_model"
    )

from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
import numpy as np

# Load transaction features
df_txn = spark.table("ml.transaction_features").toPandas()

# Feature engineering for fraud detection
features = [
    'transaction_amount',
    'transactions_per_hour',
    'unique_machines_per_session',
    'cash_to_chip_ratio',
    'time_between_transactions',
    'deviation_from_player_avg'
]

X = df_txn[features]
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Train Isolation Forest
with mlflow.start_run(run_name="fraud_model_v1"):
    model = IsolationForest(
        n_estimators=200,
        contamination=0.01,  # Expected fraud rate
        max_samples='auto',
        random_state=42
    )
    model.fit(X_scaled)

    # Score transactions
    anomaly_scores = -model.decision_function(X_scaled)  # Higher = more anomalous
    anomaly_scores_normalized = (anomaly_scores - anomaly_scores.min()) / (anomaly_scores.max() - anomaly_scores.min())

    df_txn['fraud_risk_score'] = anomaly_scores_normalized
    df_txn['is_anomaly'] = model.predict(X_scaled) == -1

    # Apply rule-based flags
    df_txn['structuring_flag'] = (
        (df_txn['transaction_amount'] >= 8000) & 
        (df_txn['transaction_amount'] < 10000) &
        (df_txn['transactions_per_day'] > 1)
    )

    mlflow.sklearn.log_model(model, "fraud_model")

# Verify churn predictions distribution
churn_dist = spark.sql("""
    SELECT 
        CASE 
            WHEN churn_probability < 0.3 THEN 'LOW'
            WHEN churn_probability < 0.7 THEN 'MEDIUM'
            ELSE 'HIGH'
        END as risk_tier,
        COUNT(*) as player_count,
        AVG(churn_probability) as avg_probability
    FROM ml.player_churn_scores
    GROUP BY 1
""")
display(churn_dist)

# Verify fraud scores
fraud_stats = spark.sql("""
    SELECT 
        COUNT(*) as total_transactions,
        SUM(CASE WHEN is_anomaly THEN 1 ELSE 0 END) as anomalies,
        SUM(CASE WHEN structuring_flag THEN 1 ELSE 0 END) as structuring_flags,
        AVG(fraud_risk_score) as avg_risk_score
    FROM ml.fraud_risk_scores
    WHERE score_timestamp > current_timestamp() - INTERVAL 24 HOURS
""")
display(fraud_stats)

# Model drift detection
from evidently import ColumnMapping
from evidently.report import Report
from evidently.metrics import DataDriftPreset

reference_data = spark.table("ml.churn_features_baseline").toPandas()
current_data = spark.table("ml.player_churn_features").toPandas()

report = Report(metrics=[DataDriftPreset()])
report.run(reference_data=reference_data, current_data=current_data)
report.save_html("drift_report.html")

# List registered models
from mlflow import MlflowClient
client = MlflowClient()

for model in client.search_registered_models():
    print(f"Model: {model.name}")
    for version in model.latest_versions:
        print(f"  Version {version.version}: {version.current_stage}")