Update app.py

app.py

@@ -1,82 +1,97 @@
-import os
-import joblib
-import pandas as pd
-from flask import Flask, request, jsonify
-from flask_cors import CORS
-
-# Initialize the Flask application
-app = Flask(__name__)
-CORS(app) # Enable Cross-Origin Resource Sharing
-
-# --- Model Loading ---
-# The model is loaded once when the application starts.
-try:
-    # In Hugging Face Spaces, your files are all in the same root directory.
-    model_path = 'realistic_fraud_model.pkl'
-    model = joblib.load(model_path)
-    print("✅ Model loaded successfully.")
-except FileNotFoundError:
-    print(f"❌ Error: Model file not found at '{model_path}'")
-    model = None
-except Exception as e:
-    print(f"❌ An error occurred while loading the model: {e}")
-    model = None
-# --- End of Model Loading ---
-
-
-# --- API Endpoints ---
-
-# Root endpoint to check if the API is running
-@app.route('/', methods=['GET'])
-def home():
-    return jsonify({"status": "API is running", "model_loaded": model is not None})
-
-
-# Prediction endpoint
-@app.route('/predict', methods=['POST'])
-def predict():
-    """
-    Receives feature data in JSON format and returns a fraud prediction.
-    """
-    if model is None:
-        return jsonify({"error": "Model is not loaded."}), 500
-
-    try:
-        # Get data from the POST request
-        data = request.get_json(force=True)
-        
-        # IMPORTANT: You must know the order of features your model expects.
-        # The input JSON should be an object with a key like "features"
-        # which is a list of values in the correct order.
-        # Example: {"features": [0.1, -0.5, 1.2, ...]}
-        features = data['features']
-        
-        # Convert to a Pandas DataFrame for prediction
-        # The feature names here are placeholders, they don't affect a
-        # scikit-learn model prediction if the order is correct.
-        feature_df = pd.DataFrame([features])
-
-        # Make prediction
-        prediction = model.predict(feature_df)
-        prediction_proba = model.predict_proba(feature_df)
-
-        # Return the result as JSON
-        return jsonify({
-            'prediction': int(prediction[0]),
-            'is_fraud': bool(prediction[0] == 1),
-            'prediction_probability': {
-                'not_fraud': prediction_proba[0][0],
-                'fraud': prediction_proba[0][1]
-            }
-        })
-
-    except Exception as e:
-        return jsonify({"error": f"An error occurred during prediction: {str(e)}"}), 400
-
-# --- End of API Endpoints ---
-
-
-if __name__ == '__main__':
-    # The 'port' is set by Hugging Face Spaces, default to 8080 for local testing
-    port = int(os.environ.get("PORT", 8080))
-    app.run(host='0.0.0.0', port=port)
+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)