Update app.py

app.py

@@ -126,24 +126,23 @@ class GRUModel(nn.Module):
         return self.fc(out[:, -1, :])
 
 class BiLSTMModel(nn.Module):
-    def __init__(self, input_dim=1, hidden_dim=100, output_dim=1):
+    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=input_dim,
             hidden_size=hidden_dim,
-            num_layers=1,  # <-- match this to the training model
+            num_layers=num_layers,
             batch_first=True,
-            dropout=0.0,   # <-- dropout does nothing with 1 layer, so remove
+            dropout=dropout_prob,
             bidirectional=True
         )
-        self.fc = nn.Linear(hidden_dim * 2, output_dim)
+        self.fc = nn.Linear(hidden_dim * 2, output_dim)  # because bidirectional
 
     def forward(self, x):
-        h0 = torch.zeros(2, x.size(0), 100).to(x.device)
-        c0 = torch.zeros(2, x.size(0), 100).to(x.device)
+        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()