Upload folder using huggingface_hub

app.py

@@ -11,7 +11,7 @@ logger = logging.getLogger(__name__)
 #initialise flask app
 sales_forecast_api = Flask('Sales forecasting')
 
-# load the model
+# load the model and training columns
 try:
     # Log current working directory and files
     current_dir = os.getcwd()
@@ -19,14 +19,23 @@ try:
     files_in_dir = os.listdir(current_dir)
     logger.info(f"Files in current directory: {files_in_dir}")
 
-    model = joblib.load('sales_forecast_v1_0.joblib')
-    logger.info("Model loaded successfully.")
+    # Assuming model and columns are saved as a dictionary using joblib
+    model_and_columns_path = 'sales_forecast_model_and_columns.joblib' # Update path if needed
+    loaded_object = joblib.load(model_and_columns_path)
+
+    model = loaded_object['model']
+    training_columns = loaded_object['columns']
+
+
+    logger.info("Model and training columns loaded successfully.")
 except FileNotFoundError:
-    logger.error("Model file not found!")
-    model = None # Or handle the error as appropriate
+    logger.error(f"Model and training columns file not found at {model_and_columns_path}!")
+    model = None
+    training_columns = None # Ensure training_columns is also None if file not found
 except Exception as e:
-    logger.error(f"Error loading model: {e}")
-    model = None # Or handle the error as appropriate
+    logger.error(f"Error loading model or training columns: {e}")
+    model = None
+    training_columns = None
 
 
 #define home page
@@ -37,8 +46,8 @@ def home():
 #define an endpoint for prediction
 @sales_forecast_api.post('/v1/sales')
 def sales_predict():
-  if model is None:
-      return jsonify({"error": "Model not loaded"}), 500
+  if model is None or training_columns is None:
+      return jsonify({"error": "Model or training columns not loaded"}), 500
 
   #get data from json request
   sales_data = request.get_json()
@@ -62,8 +71,11 @@ def sales_predict():
   categorical_columns_for_dummies = ['Product_Sugar_Content','Product_Type','Store_Size','Store_Location_City_Type','Store_Type']
   input_df_dummies = pd.get_dummies(input_data, columns=categorical_columns_for_dummies, drop_first=True)
 
+  # Reindex input_df_dummies to match the columns of X_train used during training
+  input_df_aligned = input_df_dummies.reindex(columns=training_columns, fill_value=0)
+
   #make model to predict
-  prediction = model.predict(input_df_dummies)
+  prediction = model.predict(input_df_aligned)
 
   return jsonify({'Prediction':prediction[0]})
 
@@ -72,8 +84,8 @@ def sales_predict():
 @sales_forecast_api.post('/v1/salesbatch')
 
 def sales_batch_predict():
-  if model is None:
-      return jsonify({"error": "Model not loaded"}), 500
+  if model is None or training_columns is None:
+      return jsonify({"error": "Model or training columns not loaded"}), 500
 
   #get the file from the request
   file = request.files['file']
@@ -83,7 +95,9 @@ def sales_batch_predict():
 #convert the categorical to dummies
   categorical_columns_for_dummies = ['Product_Sugar_Content','Product_Type','Store_Size','Store_Location_City_Type','Store_Type']
   input_df_dummies = pd.get_dummies(input_data, columns=categorical_columns_for_dummies, drop_first=True)
-  
+
+  # Reindex input_df_dummies to match the columns of X_train used during training
+  input_df_aligned =input_df_dummies.reindex(columns=training_columns, fill_value=0)
 
 #predict
   predictions = model.predict(input_df_aligned).tolist() # Predict and convert to list