Spaces:
Runtime error
Runtime error
File size: 3,231 Bytes
57f7a07 | 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 | import os
import joblib
import torch
import torch.nn as nn
import numpy as np
from flask import Flask, request, jsonify
from flask_cors import CORS
# --- 1. Define the same PyTorch Model Architecture ---
# This class must be identical to the one in your training script.
class FraudClassifier(nn.Module):
def __init__(self):
super(FraudClassifier, self).__init__()
self.layer1 = nn.Linear(2, 16) # 2 input features
self.layer2 = nn.Linear(16, 8)
self.layer3 = nn.Linear(8, 1) # 1 output for binary classification
self.relu = nn.ReLU()
self.sigmoid = nn.Sigmoid()
def forward(self, x):
x = self.relu(self.layer1(x))
x = self.relu(self.layer2(x))
x = self.sigmoid(self.layer3(x))
return x
# --- 2. Initialize Flask App and Load Assets ---
app = Flask(__name__)
CORS(app) # Enable Cross-Origin Resource Sharing
# The model and scaler are loaded once when the application starts.
model = None
scaler = None
try:
# Load the PyTorch model
model = FraudClassifier()
# Use map_location=torch.device('cpu') for compatibility if the server doesn't have a GPU
model.load_state_dict(torch.load("realistic_fraud_model.pth", map_location=torch.device('cpu')))
model.eval() # Set model to evaluation mode
print("✅ PyTorch model loaded successfully.")
# Load the scaler
scaler = joblib.load("scaler.pkl")
print("✅ Scaler loaded successfully.")
except Exception as e:
print(f"❌ An error occurred during asset loading: {e}")
# --- 3. API Endpoints ---
@app.route('/', methods=['GET'])
def home():
"""Root endpoint to check API status."""
return jsonify({
"status": "API is running",
"model_loaded": model is not None,
"scaler_loaded": scaler is not None
})
@app.route('/predict', methods=['POST'])
def predict():
"""Receives feature data and returns a fraud prediction."""
if not model or not scaler:
return jsonify({"error": "Model or scaler not loaded."}), 500
try:
data = request.get_json(force=True)
# Expected input: {"features": [value_ether, fee_ether]}
features = np.array(data['features']).reshape(1, -1)
# Apply the same scaling as in training
scaled_features = scaler.transform(features)
# Convert to PyTorch tensor
features_tensor = torch.tensor(scaled_features, dtype=torch.float32)
# Make prediction
with torch.no_grad(): # Disable gradient calculation for inference
probability = model(features_tensor)
prob_fraud = probability.item() # Get the float value from the tensor
prediction = 1 if prob_fraud > 0.5 else 0
# Return the result
return jsonify({
'prediction': prediction,
'is_fraud': bool(prediction == 1),
'prediction_probability': {
'not_fraud': 1.0 - prob_fraud,
'fraud': prob_fraud
}
})
except Exception as e:
return jsonify({"error": f"An error occurred during prediction: {str(e)}"}), 400
if __name__ == '__main__':
port = int(os.environ.get("PORT", 8080))
app.run(host='0.0.0.0', port=port) |