Spaces:
Runtime error
Runtime error
File size: 2,830 Bytes
57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 57f7a07 599c0d8 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 | import os
import torch
import torch.nn as nn
from flask import Flask, request, jsonify
# --- 1. Define the exact same Model Architecture ---
# This is necessary to load the saved state_dict
class FraudClassifier(nn.Module):
def __init__(self):
super(FraudClassifier, self).__init__()
self.layer1 = nn.Linear(2, 16)
self.layer2 = nn.Linear(16, 8)
self.layer3 = nn.Linear(8, 1)
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 App and Load Assets ---
app = Flask(__name__)
model = None
scaler_params = None
try:
# Load scaler parameters
scaler_params = torch.load("scaler_params.pt")
mean = scaler_params['mean']
std = scaler_params['std']
print("✅ PyTorch scaler parameters loaded.")
# Load the trained model
model = FraudClassifier()
model.load_state_dict(torch.load("realistic_fraud_model.pth"))
model.eval() # Set to evaluation mode
print("✅ PyTorch model loaded successfully.")
except Exception as e:
print(f"❌ Error loading assets: {e}")
# --- 3. Define 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_params is not None
})
@app.route('/predict', methods=['POST'])
def predict():
"""Receives feature data and returns a fraud prediction."""
if not model or not scaler_params:
return jsonify({"error": "Model or scaler not loaded."}), 500
try:
data = request.get_json(force=True)
# Expected input: {"features": [value_ether, fee_ether]}
features = data['features']
# Convert to a PyTorch tensor
features_tensor = torch.tensor(features, dtype=torch.float32)
# Apply the scaling transformation
scaled_tensor = (features_tensor - mean) / std
# Get model prediction
with torch.no_grad():
probability = model(scaled_tensor)
prob_fraud = probability.item()
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) |