Upload folder using huggingface_hub

app.py

@@ -1,147 +1,96 @@
 import pandas as pd
 from flask import Flask, request, jsonify
-from CatBoostWrapper import CatBoostWrapper
-from FeatureEngineering import FeatureEngineering
-# To serialize the model
 import joblib
-
-# Initialize Flask app with a name
+import numpy as np
+import logging
+
+# -----------------------
+# Setup logging
+# -----------------------
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# -----------------------
+# Initialize Flask app
+# -----------------------
 app = Flask("Store Capacity Predictor")
 
-# Load the trained pipeline (replace joblib.load with dill.load if you saved with dill)
-# Load the trained churn prediction model
+# Load the trained pipeline
 pipeline = joblib.load("catbooster_model_v1_0.joblib")
 
-# Define a route for the home page
+# -----------------------
+# Routes
+# -----------------------
 @app.get('/')
 def home():
     return "Welcome to the Store Capacity Prediction API"
 
-# Define an endpoint to predict churn for a single customer
+# ----------------------- SINGLE PREDICTION -----------------------
 @app.post('/v1/predict')
 def predict_capacity():
-    """
-    Expects JSON input with store and slot information.
-    Example JSON:
-    {
-        "Store": "BU1",
-        "QueueType": "Normal",
-        "StaffEmployed": 20,
-        "StaffOnline": 15,
-        "StaffLeave": 2,
-        "StoreArea": 1000,
-        "PickingArea": 200,
-        "TechAdvancement": 3,
-        "SpecialEvent": "",
-        "Date": "2025-10-06",
-        "Slot": "14:00"
-    }
-    """
     try:
-        # Get JSON data from request
         sales_data = request.get_json()
+        input_data = pd.DataFrame([sales_data])
+        input_data["Date"] = pd.to_datetime(input_data["Date"])
+
+        logger.info("Single prediction input:\n%s", input_data)
 
-        # Extract relevant features for the model
-        sample = {
-            'Store': sales_data['Store'],
-            'QueueType': sales_data['QueueType'],
-            'StaffEmployed': sales_data['StaffEmployed'],
-            'StaffOnline': sales_data['StaffOnline'],
-            'StaffLeave': sales_data['StaffLeave'],
-            'StoreArea': sales_data['StoreArea'],
-            'PickingArea': sales_data['PickingArea'],
-            'TechAdvancement': sales_data['TechAdvancement'],
-            'SpecialEvent': sales_data['SpecialEvent'],
-            'Date': sales_data['Date'],
-            'Slot': sales_data['Slot']
-        }
-
-        # Convert to DataFrame
-        input_data = pd.DataFrame([sample])
-        input_data["Date"] = pd.to_datetime(input_data["Date"])  # ⚡ convert to datetime
-        data_set = input_data.copy()  # optional backup
-
-        # Drop any IDs if your pipeline doesn’t need them
-        # input_data.drop(["Store_Id"], axis=1, inplace=True)  # example
-
-
-        # Predict using the trained pipeline
+        # Predict
         prediction = pipeline.predict(input_data).tolist()[0]
 
-        # Return prediction as JSON
+        # Sanitize and convert to integer
+        if pd.isna(prediction) or np.isinf(prediction) or not np.isfinite(prediction):
+            logger.warning("Single prediction invalid (%s), replacing with 0", prediction)
+            prediction = 0
+        else:
+            prediction = int(np.clip(prediction, 0, 10000))  # Adjust max if needed
+
+        logger.info("Single prediction output: %s", prediction)
+
         return jsonify({'Predicted_Capacity': prediction})
 
     except Exception as e:
+        logger.error("Error in single prediction: %s", e)
         return jsonify({'error': str(e)}), 400
 
-# Define an endpoint to predict churn for a batch of customers
+# ----------------------- BATCH PREDICTION -----------------------
 @app.post('/v1/predict_batch')
 def predict_capacity_batch():
-    """
-    Expects JSON input with a list of store capacity data.
-    Example JSON:
-    {
-        "data": [
-            {
-                "Store": "BU1",
-                "QueueType": "Normal",
-                "StaffEmployed": 20,
-                "StaffOnline": 15,
-                "StaffLeave": 2,
-                "StoreArea": 1000,
-                "PickingArea": 200,
-                "TechAdvancement": 3,
-                "SpecialEvent": "",
-                "Date": "2025-10-06",
-                "Slot": "14:00"
-            },
-            {
-                "Store": "BU2",
-                "QueueType": "Express",
-                "StaffEmployed": 25,
-                "StaffOnline": 20,
-                "StaffLeave": 1,
-                "StoreArea": 1200,
-                "PickingArea": 250,
-                "TechAdvancement": 4,
-                "SpecialEvent": "BlackFriday",
-                "Date": "2025-10-06",
-                "Slot": "10:00"
-            }
-        ]
-    }
-    """
     try:
-        # Get JSON data from request
         input_json = request.get_json()
         data_list = input_json.get("data", [])
-
-        # Convert list of dicts to DataFrame
         input_data = pd.DataFrame(data_list)
-        input_data["Date"] = pd.to_datetime(input_data["Date"])  # ⚡ convert to datetime
-        print("Batch prediction input shape:", input_data.shape)
-        print("Batch input preview:\n", input_data.head())
+        input_data["Date"] = pd.to_datetime(input_data["Date"])
 
+        logger.info("Batch input shape: %s", input_data.shape)
+        logger.info("Batch input preview:\n%s", input_data.head())
 
-        # Predict using pipeline
         predictions = pipeline.predict(input_data).tolist()
+        clean_predictions = []
+
+        for idx, p in enumerate(predictions):
+            if pd.isna(p) or np.isinf(p) or not np.isfinite(p):
+                logger.warning("Row %d prediction invalid (%s), replacing with 0", idx, p)
+                clean_predictions.append(0)
+            else:
+                clean_predictions.append(int(np.clip(p, 0, 10000)))
 
-        print("Batch predictions:", predictions)
+        logger.info("Batch predictions: %s", clean_predictions)
 
-        # Prepare output DataFrame with Date, Store, Predicted_Capacity
         output_df = pd.DataFrame({
             "Date": input_data["Date"],
             "Store": input_data["Store"],
-            "Predicted_Capacity": predictions
+            "Predicted_Capacity": clean_predictions
         })
 
-        # Return as HTML table
         return output_df.to_html(index=False)
 
-
     except Exception as e:
+        logger.error("Error in batch prediction: %s", e)
         return jsonify({"error": str(e)}), 400
 
-# Run the Flask app in debug mode
+# -----------------------
+# Run Flask app
+# -----------------------
 if __name__ == '__main__':
     app.run(debug=True)