Update app.py

app.py

@@ -2,7 +2,23 @@ import onnxruntime as ort
 import gradio as gr
 import numpy as np
 
-sess = ort.InferenceSession("collusion_xgb_model.onnx")
+# Load ONNX model
+sess = ort.InferenceSession("collusion_xgb_model.onnx", providers=["CPUExecutionProvider"])
+
+# Inspect outputs ONCE (safe)
+output_info = sess.get_outputs()
+
+# Find probability output index
+PROB_INDEX = None
+for i, out in enumerate(output_info):
+    if len(out.shape) == 2 and out.shape[1] == 2:
+        PROB_INDEX = i
+        break
+
+if PROB_INDEX is None:
+    raise RuntimeError("Probability output not found in ONNX model")
+
+print(f" Using output index {PROB_INDEX} as probability tensor")
 
 def predict(amount, user_txn, driver_txn, pair_count, hour, day):
     X = np.array(
@@ -12,11 +28,22 @@ def predict(amount, user_txn, driver_txn, pair_count, hour, day):
 
     outputs = sess.run(None, {sess.get_inputs()[0].name: X})
 
-    # IMPORTANT:
-    # outputs[0] = label (0 or 1)
-    # outputs[1] = probability [p0, p1]
-    probs = outputs[1]
-
-    fraud_prob = probs[0][1]   # class-1 (fraud)
+    probs = outputs[PROB_INDEX]      # [p0, p1]
+    fraud_prob = probs[0][1]         # class-1 probability
 
     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()
+