bug fix

app.py

@@ -85,21 +85,29 @@ def generate_plot(df, dates, preds):
     # Adjust spacing between subplots
     plt.tight_layout()
     
-    for ax in axs.flat:
-        ax.set_xlim(df['Date'].min(), df['Date'].max())
-        ax.set_ylim(df['sales'].min(), df['sales'].max())
+    #for ax in axs.flat:
+    #    ax.set_xlim(df['Date'].min(), df['Date'].max())
+    #    ax.set_ylim(df['sales'].min(), df['sales'].max())
     
     st.pyplot(fig)
-    
 
-def main():
-    ## Initiate Data
+@st.cache_data
+def load_data():
     with open('data/parameters.pkl', 'rb') as f:
         parameters = pickle.load(f)
-    model = TemporalFusionTransformer.load_from_checkpoint('model/tft_check.ckpt', map_location=torch.device('cpu'))
-    
     df = pd.read_pickle('data/test_data.pkl')
     df = df.loc[(df["Branch"] == "15") & (df["Group"].isin(["6","7","4","1"]))]
+    return parameters, df
+
+@st.cache_resource
+def init_model():
+    model = TemporalFusionTransformer.load_from_checkpoint('model/tft_check.ckpt', map_location=torch.device('cpu'))
+    return model
+    
+def main():
+    ## Initiate Data
+    parameters, df = load_data()
+    model = init_model()
     
     # Start App
     st.title("Temporal Fusion Transformers for Interpretable Multi-horizon Time Series Forecasting")