Fix app and add predict endpoint

app.py

@@ -1,23 +1,39 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
+from huggingface_hub import hf_hub_download
 import joblib
-import numpy as np
+import pandas as pd
 
-# Load your trained model
-model = joblib.load("model.joblib")
+app = FastAPI()
 
-# Initialize FastAPI app
-app = FastAPI(title="Supermarket Sales Forecast API")
+# Load model from your HF model repo
+model_path = hf_hub_download(
+    repo_id="danialsiddiqui/nexus-task6-model",
+    filename="model.joblib"
+)
 
-# Define input schema (replace with actual model features)
-class SalesInput(BaseModel):
-    feature1: float
-    feature2: float
-    feature3: float
+model_data = joblib.load(model_path)
+model = model_data["model"]
+columns = model_data["columns"]
+
+class PredictInput(BaseModel):
+    gender: str
+    customer_type: str
+    product_line: str
+    unit_price: float
+    quantity: int
+    tax_5: float
 
 @app.post("/predict")
-def predict(data: SalesInput):
-    # Convert input to array for prediction
-    input_data = np.array([[data.feature1, data.feature2, data.feature3]])
-    prediction = model.predict(input_data)
+def predict(input_data: PredictInput):
+    df = pd.DataFrame([input_data.dict()])
+    df = pd.get_dummies(df)
+
+    # Align columns
+    for col in columns:
+        if col not in df.columns:
+            df[col] = 0
+
+    df = df[columns]
+    prediction = model.predict(df)
     return {"prediction": prediction.tolist()}

requirements.txt

@@ -5,3 +5,6 @@ joblib
 numpy
 requests
 streamlit
+huggingface_hub
+pandas
+scikit-learn