Update predict_impulse.py

predict_impulse.py

@@ -2,18 +2,17 @@ import joblib
 import pandas as pd
 import os
 from huggingface_hub import hf_hub_download
-from opik import Opik, track
+import opik
+from opik import track
 
 # -------------------------------------------------
-# API key handling (HF Spaces / Expo compatible)
+# Opik API key (HF Spaces safe)
 # -------------------------------------------------
 if "OPIK_API_KEY" not in os.environ:
-    os.environ["OPIK_API_KEY"] = os.environ.get(
-        "EXPO_PUBLIC_OPIK_API_KEY", ""
-    )
+    os.environ["OPIK_API_KEY"] = os.environ.get("EXPO_PUBLIC_OPIK_API_KEY", "")
 
 # -------------------------------------------------
-# Load model from Hugging Face Hub
+# Load model from HF Model Hub
 # -------------------------------------------------
 MODEL_REPO = "obx0x3/sensei-model"
 MODEL_FILE = "impulse_model.pkl"
@@ -26,61 +25,44 @@ model_path = hf_hub_download(
 impulse_model = joblib.load(model_path)
 
 # -------------------------------------------------
-# Opik client (NO init, just client)
-# -------------------------------------------------
-opik = None
-try:
-    opik = Opik(project_name="budgetbuddy-hackathon")
-except Exception as e:
-    print("Opik disabled:", e)
-
-
-# -------------------------------------------------
-# Tracked prediction function
+# TRACKED FUNCTION (this creates the TRACE)
 # -------------------------------------------------
 @track
-def predict_impulse(category: str, amount: float, payment_method: str, day: str):
-    """
-    This function creates ONE Opik trace per call.
-    All logs inside attach to that trace.
-    """
-
+def predict_impulse(category, amount, payment_method, day):
     input_data = {
         "category": category,
         "amount": float(amount),
         "payment_method": payment_method,
-        "day": day,
+        "day": day
     }
 
     df = pd.DataFrame([input_data])
 
-    prediction = impulse_model.predict(df)[0]
-    probability = impulse_model.predict_proba(df)[0].max()
+    pred = impulse_model.predict(df)[0]
+    prob = impulse_model.predict_proba(df)[0].max()
 
     result = {
-        "impulsive": bool(prediction),
-        "label": "Impulsive" if prediction else "Normal Spend",
-        "confidence": round(float(probability), 3),
+        "impulsive": bool(pred),
+        "label": "Impulsive" if pred else "Normal Spend",
+        "confidence": round(float(prob), 3)
     }
 
     # -------------------------------------------------
-    # Log structured data to the ACTIVE trace
+    # Attach metadata to the CURRENT TRACE
     # -------------------------------------------------
-    if opik:
-        opik.log_event(
-            name="impulse_prediction",
-            input=input_data,
-            output=result,
-            model="sensei-impulse-model",
-            metadata={
-                "feature": "budget-ai",
-                "task": "impulse-detection",
-                "confidence_band": (
-                    "high" if probability >= 0.75
-                    else "medium" if probability >= 0.5
-                    else "low"
-                ),
-            },
+    opik.log_metadata({
+        "source": "hf-space",
+        "model": "sensei-impulse-model",
+        "confidence_band": (
+            "high" if prob >= 0.75 else
+            "medium" if prob >= 0.5 else
+            "low"
         )
+    })
+
+    opik.log({
+        "input": input_data,
+        "output": result
+    })
 
     return result