File size: 1,917 Bytes
e35f247 |
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 |
from transformers import Pipeline
import numpy as np
import joblib
from typing import Dict, List, Union
class FraudDetectionPipeline(Pipeline):
def __init__(self):
super().__init__()
# Load the model and scaler
self.model = joblib.load("random_forest_model.joblib")
self.scaler = joblib.load("rf_scaler.joblib")
def preprocess(self, features: Dict[str, Union[int, float]]) -> np.ndarray:
"""
Preprocess the input features
Expected features:
- account_age: int (months)
- cred_changes_freq: float (per year)
- return_order_ratio: float
- vpn_usage: int (0 or 1)
- credit_score: int
"""
# Convert input to correct format
input_data = np.array([[
features['account_age'],
features['cred_changes_freq'],
features['return_order_ratio'],
features['vpn_usage'],
features['credit_score']
]])
# Scale the features
scaled_input = self.scaler.transform(input_data)
return scaled_input
def _forward(self, features: Dict[str, Union[int, float]]) -> Dict[str, Union[str, float]]:
"""
Make prediction using the model
"""
# Preprocess
scaled_input = self.preprocess(features)
# Get prediction and probability
prediction = self.model.predict(scaled_input)[0]
probabilities = self.model.predict_proba(scaled_input)[0]
# Return prediction and confidence
return {
"prediction": "Fraud" if prediction == 1 else "Not Fraud",
"confidence": float(probabilities[prediction]),
"fraud_probability": float(probabilities[1])
}
def postprocess(self, model_outputs):
return model_outputs
def load_pipeline():
return FraudDetectionPipeline() |