Update app.py

app.py

@@ -1,19 +1,16 @@
 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.
+# --- 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) # 2 input features
+        self.layer1 = nn.Linear(2, 16)
         self.layer2 = nn.Linear(16, 8)
-        self.layer3 = nn.Linear(8, 1)  # 1 output for binary classification
+        self.layer3 = nn.Linear(8, 1)
         self.relu = nn.ReLU()
         self.sigmoid = nn.Sigmoid()
 
@@ -23,61 +20,59 @@ class FraudClassifier(nn.Module):
         x = self.sigmoid(self.layer3(x))
         return x
 
-# --- 2. Initialize Flask App and Load Assets ---
+# --- 2. Initialize 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
+scaler_params = None
 
 try:
-    # Load the PyTorch model
+    # 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()
-    # 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
+    model.load_state_dict(torch.load("realistic_fraud_model.pth"))
+    model.eval() # Set 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}")
+    print(f"❌ Error loading assets: {e}")
 
-# --- 3. API Endpoints ---
+# --- 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 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:
+    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 = np.array(data['features']).reshape(1, -1)
+        features = data['features']
+        
+        # Convert to a PyTorch tensor
+        features_tensor = torch.tensor(features, dtype=torch.float32)
 
-        # Apply the same scaling as in training
-        scaled_features = scaler.transform(features)
+        # Apply the scaling transformation
+        scaled_tensor = (features_tensor - mean) / std
 
-        # Convert to PyTorch tensor
-        features_tensor = torch.tensor(scaled_features, dtype=torch.float32)
+        # Get model prediction
+        with torch.no_grad():
+            probability = model(scaled_tensor)
 
-        # 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
+        prob_fraud = probability.item()
         prediction = 1 if prob_fraud > 0.5 else 0
 
         # Return the result
@@ -89,6 +84,7 @@ def predict():
                 'fraud': prob_fraud
             }
         })
+
     except Exception as e:
         return jsonify({"error": f"An error occurred during prediction: {str(e)}"}), 400