File size: 4,499 Bytes
0d3096e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import os
import joblib # Use joblib to load files saved with joblib.dump
import pandas as pd
from typing import Dict, Any, List
# Define the models and the version based on the training script output
class EndpointHandler:
VERSION = "1.0"
# Maps user-friendly aliases to the EXACT filenames from your training script
MODEL_MAP = {
"decision_tree": "classifier_FULL_Decision_Tree.pkl",
"random_forest": "classifier_FULL_Random_Forest.pkl",
"xgboost": "classifier_FULL_XGBoost.pkl",
}
# List of all features the model expects (is_fraud is excluded as it's the target)
EXPECTED_FEATURES = [
"cc_num", "merchant", "category", "amt", "gender", "state", "zip",
"lat", "long", "city_pop", "job", "unix_time", "merch_lat",
"merch_long", "age", "trans_hour", "trans_day", "trans_month",
"trans_weekday", "distance"
]
def __init__(self, path="."):
"""Loads all three Pipeline objects using joblib."""
self.models = {}
print(f"Server starting up for version: {self.VERSION}")
for alias, filename in self.MODEL_MAP.items():
model_path = os.path.join(path, filename)
try:
# Use joblib.load for files saved with joblib.dump
self.models[alias] = joblib.load(model_path)
print(f"✅ Pipeline loaded for {alias}")
except Exception as e:
# Note: Errors here are often due to package version mismatch
print(f"❌ Error loading {filename}. Check scikit-learn/xgboost versions: {e}")
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""Handles the API request, selects the model, and performs inference."""
inputs = data.get("inputs", {})
target_model_alias = inputs.get("model_name")
requested_version = inputs.get("model_version")
features_list = inputs.get("features")
base_response = {"server_version": self.VERSION}
# 1. Validation Checks
if requested_version != self.VERSION:
return {
**base_response,
"error": f"Requested version '{requested_version}' does not match server version '{self.VERSION}'."
}
if not target_model_alias or target_model_alias not in self.models:
return {
**base_response,
"error": f"Model '{target_model_alias}' not found. Available: {list(self.MODEL_MAP.keys())}",
}
if not features_list:
return {
**base_response,
"error": "No transaction features provided in the 'features' list."
}
# 2. Prepare Data
model_pipeline = self.models[target_model_alias]
try:
# Convert list of dicts to DataFrame and ensure column order matches training data
df_features = pd.DataFrame(features_list)
# CRITICAL: Reindex to ensure the columns are in the exact order the pipeline expects
if set(df_features.columns) != set(self.EXPECTED_FEATURES):
return {
**base_response,
"error": "Input features do not match expected features. Check column names."
}
df_features = df_features[self.EXPECTED_FEATURES]
except Exception as e:
return {
**base_response,
"error": f"Data preparation failed. Ensure JSON fields match all expected features: {str(e)}"
}
# 3. Predict Probability
try:
# predict_proba runs the ColumnTransformer (preprocessing) and then the classifier
# We take the probability of the positive class (Fraud=1), which is column [:, 1]
probabilities = model_pipeline.predict_proba(df_features)[:, 1]
return {
**base_response,
"model_used": target_model_alias,
"model_version": requested_version,
"prediction_probabilities": probabilities.tolist(),
}
except Exception as e:
return {
**base_response,
"error": f"Prediction execution failed. This may indicate a data type mismatch: {str(e)}",
} |