Update model_folder/modeling_bilstm.py

model_folder/modeling_bilstm.py

@@ -2,23 +2,34 @@ import torch.nn as nn
 import torch
 
 class BiLSTM(nn.Module):
-    def __init__(self, embedding_matrix, embed_size, num_labels, hidden_size=64, dropout=0.3):
+    def __init__(self, embedding_matrix, embed_size, num_labels, hidden_size=64, dropout=0.1):
         super(BiLSTM, self).__init__()
+        self.hidden_size = hidden_size
         self.embedding = nn.Embedding.from_pretrained(
             torch.tensor(embedding_matrix, dtype=torch.float32), freeze=True)
-        self.lstm = nn.LSTM(embed_size, hidden_size, bidirectional=True, batch_first=True)
-        self.fc = nn.Sequential(
-            nn.Linear(hidden_size * 4, 128),
-            nn.ReLU(),
-            nn.Dropout(dropout),
-            nn.Linear(128, num_labels)
-        )
+        self.lstm = nn.LSTM(embed_size, self.hidden_size, bidirectional=True, batch_first=True)
+        self.linear = nn.Linear(self.hidden_size * 4, 64)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(dropout)
+        self.out = nn.Linear(64, num_labels)
 
     def forward(self, x):
-        x = self.embedding(x)
-        h_lstm, _ = self.lstm(x)
+        # Embedding layer
+        h_embedding = self.embedding(x)
+        
+        # LSTM layer
+        h_lstm, _ = self.lstm(h_embedding)
+        
+        # Pooling layers
         avg_pool = torch.mean(h_lstm, 1)
         max_pool, _ = torch.max(h_lstm, 1)
-        x = torch.cat((avg_pool, max_pool), dim=1)
-        x = self.fc(x)
-        return x
+        
+        # Concatenate pooling outputs
+        conc = torch.cat((avg_pool, max_pool), 1)
+        
+        # Fully connected layers
+        conc = self.relu(self.linear(conc))
+        conc = self.dropout(conc)
+        out = self.out(conc)
+        
+        return out