Update app.py

app.py

@@ -126,17 +126,23 @@ class GRUModel(nn.Module):
         return self.fc(out[:, -1, :])
 
 class BiLSTMModel(nn.Module):
-    def __init__(self):
+    def __init__(self, input_dim=1, hidden_dim=100, num_layers=2, output_dim=1, dropout_prob=0.2):
         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)
+        self.lstm = nn.LSTM(
+            input_size=input_dim,
+            hidden_size=hidden_dim,
+            num_layers=num_layers,
+            batch_first=True,
+            dropout=dropout_prob,
+            bidirectional=True
+        )
+        self.fc = nn.Linear(hidden_dim * 2, output_dim)  # because bidirectional
 
     def forward(self, x):
-        h0 = torch.zeros(2, x.size(0), 100)
-        c0 = torch.zeros(2, x.size(0), 100)
+        h0 = torch.zeros(self.lstm.num_layers * 2, x.size(0), self.lstm.hidden_size).to(x.device)
+        c0 = torch.zeros(self.lstm.num_layers * 2, x.size(0), self.lstm.hidden_size).to(x.device)
         out, _ = self.lstm(x, (h0, c0))
         return self.fc(out[:, -1, :])
-
 @st.cache_resource(ttl=3600)
 def load_model_from_disk(path, model_type):
     model = BiLSTMModel() if model_type == "Bi-Directional LSTM" else GRUModel()
@@ -145,13 +151,14 @@ def load_model_from_disk(path, model_type):
     model.eval()
     return model
    
-
-if "models" not in st.session_state:
-    st.session_state.models = {
-        "Bi-Directional LSTM": load_model_from_disk("best_bilstm_model.pth", "Bi-Directional LSTM"),
-        "Gated Recurrent Unit (GRU)": load_model_from_disk("best_gru_model.pth", "GRU")
+@st.cache_resource
+def preload_models():
+    return {
+        "Bi-Directional LSTM": load_model_from_disk("best_bilstm_model.pth", model_type="Bi-Directional LSTM"),
+        "Gated Recurrent Unit (GRU)": load_model_from_disk("best_gru_model.pth", model_type="GRU")
     }
 
+MODELS = preload_models()
 @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."""
@@ -243,7 +250,7 @@ with st.sidebar:
         st.session_state.last_prediction_days = prediction_days
 # --- Main Application Logic ---
 if st.session_state.run_button_clicked:
-        model = st.session_state.models[model_type]
+        model = MODELS[model_type]
         
         
         print(f"Inside main logic block. Current loading state: {st.session_state.loading}")