Update app.py

app.py

@@ -8,48 +8,53 @@ from neuralforecast.losses.pytorch import HuberMQLoss
 from neuralforecast.utils import AirPassengersDF
 import time
 
-# Paths for saving models
-nhits_paths = {
-    'D': './M4/NHITS/daily',
-    'M': './M4/NHITS/monthly',
-    'H': './M4/NHITS/hourly',
-    'W': './M4/NHITS/weekly',
-    'Y': './M4/NHITS/yearly'
-}
 
-timesnet_paths = {
-    'D': './M4/TimesNet/daily',
-    'M': './M4/TimesNet/monthly',
-    'H': './M4/TimesNet/hourly',
-    'W': './M4/TimesNet/weekly',
-    'Y': './M4/TimesNet/yearly'
-}
-
-lstm_paths = {
-    'D': './M4/LSTM/daily',
-    'M': './M4/LSTM/monthly',
-    'H': './M4/LSTM/hourly',
-    'W': './M4/LSTM/weekly',
-    'Y': './M4/LSTM/yearly'
-}
-
-tft_paths = {
-    'D': './M4/TFT/daily',
-    'M': './M4/TFT/monthly',
-    'H': './M4/TFT/hourly',
-    'W': './M4/TFT/weekly',
-    'Y': './M4/TFT/yearly'
-}
 
 @st.cache_resource
 def load_model(path, freq):
     nf = NeuralForecast.load(path=path)
     return nf
 
-nhits_models = {freq: load_model(path, freq) for freq, path in nhits_paths.items()}
-timesnet_models = {freq: load_model(path, freq) for freq, path in timesnet_paths.items()}
-lstm_models = {freq: load_model(path, freq) for freq, path in lstm_paths.items()}
-tft_models = {freq: load_model(path, freq) for freq, path in tft_paths.items()}
+@st.cache_resource
+def load_all_models():
+        # Paths for saving models
+    nhits_paths = {
+        'D': './M4/NHITS/daily',
+        'M': './M4/NHITS/monthly',
+        'H': './M4/NHITS/hourly',
+        'W': './M4/NHITS/weekly',
+        'Y': './M4/NHITS/yearly'
+    }
+    
+    timesnet_paths = {
+        'D': './M4/TimesNet/daily',
+        'M': './M4/TimesNet/monthly',
+        'H': './M4/TimesNet/hourly',
+        'W': './M4/TimesNet/weekly',
+        'Y': './M4/TimesNet/yearly'
+    }
+    
+    lstm_paths = {
+        'D': './M4/LSTM/daily',
+        'M': './M4/LSTM/monthly',
+        'H': './M4/LSTM/hourly',
+        'W': './M4/LSTM/weekly',
+        'Y': './M4/LSTM/yearly'
+    }
+    
+    tft_paths = {
+        'D': './M4/TFT/daily',
+        'M': './M4/TFT/monthly',
+        'H': './M4/TFT/hourly',
+        'W': './M4/TFT/weekly',
+        'Y': './M4/TFT/yearly'
+    }
+    nhits_models = {freq: load_model(path, freq) for freq, path in nhits_paths.items()}
+    timesnet_models = {freq: load_model(path, freq) for freq, path in timesnet_paths.items()}
+    lstm_models = {freq: load_model(path, freq) for freq, path in lstm_paths.items()}
+    tft_models = {freq: load_model(path, freq) for freq, path in tft_paths.items()}
+
+    return nhits_models, timesnet_models, lstm_models, tft_models
 
 def generate_forecast(model, df):
     forecast_df = model.predict(df=df)
@@ -189,6 +194,8 @@ def transfer_learning_forecasting():
     else:
         df = load_default()  
 
+    nhits_models, timesnet_models, lstm_models, tft_models = load_all_models()
+    
     # Model selection and forecasting
     st.subheader("Model Selection and Forecasting")
     model_choice = st.selectbox("Select model", ["NHITS", "TimesNet", "LSTM", "TFT"])
@@ -251,6 +258,7 @@ pg = st.navigation({
     ]
 })
 
+
 try:
     pg.run()
 except Exception as e: