Update app.py

app.py

@@ -1,119 +1,45 @@
-import pandas as pd
-import numpy as np
+import os
+import onnx
 import joblib
 
-from xgboost import XGBClassifier
-from sklearn.model_selection import train_test_split
-from sklearn.metrics import roc_auc_score, classification_report
-
-import onnxmltools
-from onnxmltools.convert.common.data_types import FloatTensorType
+ONNX_PATH = "collusion_xgb_model.onnx"
+FEATURES_PATH = "feature_order.joblib"
 
 
 def main():
 
-    # =============================
-    # 1. LOAD DATA
-    # =============================
-    df = pd.read_csv("synthetic_collusion_1M.csv")
-
-    # Robust timestamp parsing
-    df["timestamp"] = pd.to_datetime(df["timestamp"], errors="coerce")
-    df["hour"] = df["timestamp"].dt.hour
-    df["day_of_week"] = df["timestamp"].dt.dayofweek
-
-    # =============================
-    # 2. FEATURE ENGINEERING
-    # =============================
-    df["user_txn_count"] = df.groupby("user_id")["transaction_id"].transform("count")
-    df["driver_txn_count"] = df.groupby("driver_id")["transaction_id"].transform("count")
-    df["user_driver_pair_count"] = (
-        df.groupby(["user_id", "driver_id"])["transaction_id"]
-          .transform("count")
-    )
-
-    FEATURES = [
-        "amount",
-        "user_txn_count",
-        "driver_txn_count",
-        "user_driver_pair_count",
-        "hour",
-        "day_of_week"
-    ]
-
-    X = df[FEATURES].fillna(0)
-    y = df["is_collusion_fraud"]
-
-    # =============================
-    # 3. TRAIN / TEST SPLIT
-    # =============================
-    X_train, X_test, y_train, y_test = train_test_split(
-        X, y,
-        test_size=0.2,
-        stratify=y,
-        random_state=42
-    )
-
-    # =============================
-    # 4. TRAIN XGBOOST
-    # =============================
-    xgb_model = XGBClassifier(
-        n_estimators=300,
-        max_depth=6,
-        learning_rate=0.05,
-        subsample=0.8,
-        colsample_bytree=0.8,
-        scale_pos_weight=y_train.value_counts()[0] / y_train.value_counts()[1],
-        base_score=0.5,
-        objective="binary:logistic",
-        eval_metric="auc",
-        random_state=42,
-        n_jobs=-1
-    )
+    print("🔍 Hugging Face inference environment check")
 
-    xgb_model.fit(X_train, y_train)
+    # -----------------------------
+    # 1. Check ONNX model
+    # -----------------------------
+    if not os.path.exists(ONNX_PATH):
+        raise FileNotFoundError(
+            "ONNX model not found. Expected collusion_xgb_model.onnx"
+        )
 
-    # =============================
-    # 5. EVALUATION
-    # =============================
-    y_prob = xgb_model.predict_proba(X_test)[:, 1]
+    print("ONNX model found")
 
-    print("\n=== MODEL EVALUATION ===")
-    print("ROC-AUC:", roc_auc_score(y_test, y_prob))
-    print(classification_report(y_test, (y_prob > 0.7).astype(int)))
+    # -----------------------------
+    # 2. Load & verify ONNX
+    # -----------------------------
+    model = onnx.load(ONNX_PATH)
+    onnx.checker.check_model(model)
 
-    # =============================
-    # 6. SAVE MODEL ARTIFACTS
-    # =============================
-    joblib.dump(xgb_model, "collusion_xgb_model.joblib")
-    joblib.dump(FEATURES, "feature_order.joblib")
+    print("ONNX model is valid")
 
-    print("✅ Model and feature order saved")
+    # -----------------------------
+    # 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)")
 
-    # =============================
-    # 7. CONVERT TO ONNX
-    # =============================
-    booster = xgb_model.get_booster()
-
-    # Rename features to f0, f1, f2... (required by onnxmltools)
-    booster.feature_names = [f"f{i}" for i in range(len(FEATURES))]
-
-    initial_type = [
-        ("float_input", FloatTensorType([None, len(FEATURES)]))
-    ]
-
-    onnx_model = onnxmltools.convert_xgboost(
-        booster,
-        initial_types=initial_type,
-        target_opset=12
-    )
-
-    with open("collusion_xgb_model.onnx", "wb") as f:
-        f.write(onnx_model.SerializeToString())
-
-    print("✅ ONNX model exported successfully")
+    print("\n🚀 Environment ready for inference")
+    print("➡️ Use app.py to serve predictions")
 
 
 if __name__ == "__main__":
     main()
-