File size: 2,830 Bytes
57f7a07
 
 
 
 
529843b
 
57f7a07
 
 
529843b
57f7a07
529843b
57f7a07
 
 
 
 
 
 
 
 
529843b
57f7a07
 
529843b
57f7a07
 
529843b
 
 
 
 
 
 
57f7a07
529843b
 
57f7a07
 
 
529843b
57f7a07
529843b
57f7a07
 
 
 
 
 
529843b
57f7a07
 
 
 
 
529843b
57f7a07
 
 
 
 
529843b
 
 
 
57f7a07
529843b
 
57f7a07
529843b
 
 
57f7a07
529843b
57f7a07
 
 
 
 
 
 
 
 
 
 
529843b
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)