Update app.py

app.py

@@ -110,55 +110,41 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 # --- Python Backend Functions ---
+# Outside of any function
+import torch.nn as nn
+import torch
+
+class GRUModel(nn.Module):
+    def __init__(self, input_dim=1, hidden_dim=100, num_layers=2, output_dim=1, dropout_prob=0.2):
+        super(GRUModel, self).__init__()
+        self.gru = nn.GRU(input_dim, hidden_dim, num_layers, batch_first=True, dropout=dropout_prob)
+        self.fc = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, x):
+        h0 = torch.zeros(2, x.size(0), 100).to(x.device)
+        out, _ = self.gru(x, h0)
+        return self.fc(out[:, -1, :])
+
+class BiLSTMModel(nn.Module):
+    def __init__(self):
+        super(BiLSTMModel, self).__init__()
+        self.lstm = nn.LSTM(input_size=1, hidden_size=100, num_layers=1, batch_first=True, dropout=0.2, bidirectional=True)
+        self.fc = nn.Linear(200, 1)
+
+    def forward(self, x):
+        h0 = torch.zeros(2, x.size(0), 100)
+        c0 = torch.zeros(2, x.size(0), 100)
+        out, _ = self.lstm(x, (h0, c0))
+        return self.fc(out[:, -1, :])
 
 @st.cache_resource(ttl=3600)
-def load_pytorch_model(path, model_type='Bi-Directional LSTM', input_dim=1, hidden_dim=100, num_layers=2, output_dim=1, dropout_prob=0.2):
-    import torch.nn as nn
-    import torch
-
-    class GRUModel(nn.Module):
-        def __init__(self):
-            super(GRUModel, self).__init__()
-            self.gru = nn.GRU(input_dim, hidden_dim, num_layers, batch_first=True, dropout=dropout_prob)
-            self.fc = nn.Linear(hidden_dim, output_dim)
-
-        def forward(self, x):
-            h0 = torch.zeros(num_layers, x.size(0), hidden_dim).to(x.device)
-            out, _ = self.gru(x, h0)
-            return self.fc(out[:, -1, :])
-
-    class BiLSTMModel(nn.Module):
-        def __init__(self):
-            super(BiLSTMModel, self).__init__()
-            self.lstm = nn.LSTM(
-                input_size=1,
-                hidden_size=100,
-                num_layers=1,  # <- match saved model
-                batch_first=True,
-                dropout=0.2,
-                bidirectional=True
-            )
-            self.fc = nn.Linear(200, 1)  # 2 * hidden_size because of bidirectional
-
-        def forward(self, x):
-            h0 = torch.zeros(2 * 1, x.size(0), 100)
-            c0 = torch.zeros(2 * 1, x.size(0), 100)
-            out, _ = self.lstm(x, (h0, c0))
-            return self.fc(out[:, -1, :])
-    model_class = BiLSTMModel if model_type == 'Bi-Directional LSTM' else GRUModel
-    model = model_class()
-
-    checkpoint = torch.load(path, map_location=torch.device('cpu'))
-
-    # If full checkpoint was saved with keys like 'model_state_dict'
-    if 'model_state_dict' in checkpoint:
-        model.load_state_dict(checkpoint['model_state_dict'])
-    else:
-        model.load_state_dict(checkpoint)  # Just raw state_dict
-
+def load_model_from_disk(path, model_type):
+    model = BiLSTMModel() if model_type == "Bi-Directional LSTM" else GRUModel()
+    state = torch.load(path, map_location=torch.device("cpu"))
+    model.load_state_dict(state['model_state_dict'] if 'model_state_dict' in state else state)
     model.eval()
     return model
-
+   
 @st.cache_data(ttl=900) # Cache data for 15 minutes
 def get_stock_data(ticker):
     """Fetches historical stock data from Yahoo Finance for the last 4 years."""
@@ -175,14 +161,8 @@ def get_stock_data(ticker):
     print(f"Successfully fetched {len(data)} rows for {ticker}")
     return data
 
-def predict_with_model(data: pd.DataFrame, n_days: int, model_path: str, model_type: str) -> pd.DataFrame:
-    import torch
+def predict_with_model(data: pd.DataFrame, n_days: int, model, model_type: str) -> pd.DataFrame:
 
-    try:
-        model = load_pytorch_model(model_path, model_type=model_type)
-    except FileNotFoundError as e:
-        raise e
-    print("model:",model)
     close_prices = data['Close'].values.reshape(-1, 1)
     scaler = MinMaxScaler(feature_range=(0, 1))
     scaled_prices = scaler.fit_transform(close_prices)
@@ -253,9 +233,15 @@ with st.sidebar:
         st.session_state.last_ticker = ticker
         st.session_state.last_model_type = model_type
         st.session_state.last_prediction_days = prediction_days
-        st.rerun()
 # --- Main Application Logic ---
 if st.session_state.run_button_clicked:
+        model_key = "bilstm_model" if model_type == "Bi-Directional LSTM" else "gru_model"
+
+        if model_key not in st.session_state:
+            model_path = "best_bilstm_model.pth" if model_type == "Bi-Directional LSTM" else "best_gru_model.pth"
+            st.session_state[model_key] = load_model_from_disk(model_path, model_type)
+
+        model = st.session_state[model_key]
         print(f"Inside main logic block. Current loading state: {st.session_state.loading}")
         try:
             if os.path.exists("AMZN_data.csv"):
@@ -268,7 +254,7 @@ if st.session_state.run_button_clicked:
                 st.session_state.error = f"Could not fetch data for ticker '{ticker}'. It may be an invalid symbol or network issue."
             else:
                 model_path = "best_bilstm_model.pth" if model_type == "Bi-Directional LSTM" else "best_gru_model.pth"
-                st.session_state.predictions = predict_with_model(st.session_state.data, prediction_days, model_path,model_type)
+                st.session_state.predictions = predict_with_model(st.session_state.data, prediction_days, model, model_type)
                 st.session_state.error = None
 
         except FileNotFoundError as e: