Update app.py

app.py

@@ -8,6 +8,8 @@ from keras.models import load_model
 from sklearn.preprocessing import MinMaxScaler
 import time
 import os
+import torch.nn as nn
+import torch
 
 # --- Page Configuration ---
 st.set_page_config(layout="wide")
@@ -111,8 +113,6 @@ st.markdown("""
 
 # --- 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):
@@ -145,6 +145,13 @@ 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_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."""
@@ -161,8 +168,9 @@ 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, model_type: str) -> pd.DataFrame:
-
+def predict_with_model(data, n_days, model_path, model_type, model=None)-> pd.DataFrame:
+    if model is None:
+        model = load_model_from_disk(model_path, model_type=model_type)
     close_prices = data['Close'].values.reshape(-1, 1)
     scaler = MinMaxScaler(feature_range=(0, 1))
     scaled_prices = scaler.fit_transform(close_prices)
@@ -235,13 +243,9 @@ with st.sidebar:
         st.session_state.last_prediction_days = prediction_days
 # --- 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]
+        model = st.session_state.models[model_type]
+        
+        
         print(f"Inside main logic block. Current loading state: {st.session_state.loading}")
         try:
             if os.path.exists("AMZN_data.csv"):
@@ -254,7 +258,11 @@ 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, model_type)
+                st.session_state.predictions = predict_with_model(
+                st.session_state.data, prediction_days, model_path=None, model_type=model_type, model=model
+            )
+                print("model",model)
+                print("data", st.session_state.data)
                 st.session_state.error = None
 
         except FileNotFoundError as e: