bug fix

app.py

@@ -38,12 +38,7 @@ def raw_preds_to_df(raw,quantiles = None):
     preds_df['pred_idx'] = preds_df['time_idx'] + preds_df['h'] - 1
     return preds_df
 
-def prepare_dataset(parameters, df, rain, temperature, datepicker):
-    rain_mapping = {
-        "Yes" : 1,
-        "No" : 0
-    }
-    
+def prepare_dataset(parameters, df, rain, temperature, datepicker, mapping):
     if rain != "Default":
         df["MTXWTH_Day_precip"] = rain_mapping[rain]
     
@@ -59,7 +54,7 @@ def prepare_dataset(parameters, df, rain, temperature, datepicker):
     return df.to_dataloader(train=False, batch_size=256,num_workers = 0)
 
 @st.cache_data
-def predict(_model, dataloader):
+def predict(_model, _dataloader):
     out = model.predict(dataloader, mode="raw", return_x=True, return_index=True)
     preds = raw_preds_to_df(out, quantiles = None)
     
@@ -113,6 +108,10 @@ def main():
     ## Initiate Data
     parameters, df = load_data()
     model = init_model()
+    rain_mapping = {
+        "Yes" : 1,
+        "No" : 0
+    }
     
     # Start App
     st.title("Temporal Fusion Transformers for Interpretable Multi-horizon Time Series Forecasting")
@@ -146,7 +145,7 @@ def main():
     datepicker = st.date_input("Start of Forecast", datetime.date(2022, 12, 24), min_value=datetime.date(2022, 6, 26) + datetime.timedelta(days = 35), max_value=datetime.date(2023, 6, 26) - datetime.timedelta(days = 30))
 
     if st.button("Forecast Sales", type="primary"):
-        dataloader = prepare_dataset(parameters, df, rain, temperature, datepicker)
+        dataloader = prepare_dataset(parameters, df, rain, temperature, datepicker, rain_mapping)
         preds = predict(model, dataloader)
         generate_plot(df, preds)