Update app.py

app.py

@@ -1,48 +1,3 @@
-import os
-import onnx
-import joblib
-
-ONNX_PATH = "collusion_xgb_model.onnx"
-FEATURES_PATH = "feature_order.joblib"
-
-
-def main():
-
-    print("🔍 Hugging Face inference environment check")
-
-    # -----------------------------
-    # 1. Check ONNX model
-    # -----------------------------
-    if not os.path.exists(ONNX_PATH):
-        raise FileNotFoundError(
-            "ONNX model not found. Expected collusion_xgb_model.onnx"
-        )
-
-    print("ONNX model found")
-
-    # -----------------------------
-    # 2. Load & verify ONNX
-    # -----------------------------
-    model = onnx.load(ONNX_PATH)
-    onnx.checker.check_model(model)
-
-    print("ONNX model is valid")
-
-    # -----------------------------
-    # 3. Feature order check (optional)
-    # -----------------------------
-    if os.path.exists(FEATURES_PATH):
-        features = joblib.load(FEATURES_PATH)
-        print("Feature order loaded:", features)
-    else:
-        print("feature_order.joblib not found (ok for inference-only)")
-
-    print("\n🚀 Environment ready for inference")
-    print("➡️ Use app.py to serve predictions")
-
-
-if __name__ == "__main__":
-    main()
 import onnxruntime as ort
 import gradio as gr
 import numpy as np
@@ -50,27 +5,27 @@ import numpy as np
 sess = ort.InferenceSession("collusion_xgb_model.onnx")
 
 def predict(amount, user_txn, driver_txn, pair_count, hour, day):
-    X = np.array(
-        [[amount, user_txn, driver_txn, pair_count, hour, day]],
-        dtype=np.float32
-    )
-
+    X = np.array([[amount, user_txn, driver_txn, pair_count, hour, day]], dtype=np.float32)
     outputs = sess.run(None, {sess.get_inputs()[0].name: X})
-
     scores = outputs[0]
 
-    # Case 1: class probabilities [prob_0, prob_1]
     if scores.ndim == 2 and scores.shape[1] == 2:
         fraud_prob = scores[0][1]
-
-    # Case 2: single probability
-    elif scores.ndim == 2 and scores.shape[1] == 1:
-        fraud_prob = scores[0][0]
-
-    # Case 3: raw score → sigmoid
     else:
-        raw_score = scores[0]
-        fraud_prob = 1 / (1 + np.exp(-raw_score))
+        fraud_prob = scores[0][0]
 
     return float(fraud_prob)
 
+gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Number(label="Amount"),
+        gr.Number(label="User Txn Count"),
+        gr.Number(label="Driver Txn Count"),
+        gr.Number(label="User-Driver Pair Count"),
+        gr.Number(label="Hour"),
+        gr.Number(label="Day of Week"),
+    ],
+    outputs=gr.Number(label="Fraud Probability"),
+    title="Collusion Fraud Detection (ONNX)"
+).launch()