Create lstm_model.py

lstm_model.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+import torch.nn.utils.rnn as rnn_utils
+
+class LSTMClassifier(nn.Module):
+    def __init__(self, vocab_size, embed_dim=256, hidden_dim=256,
+                 n_layers=2, dropout=0.3, bidirectional=True):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
+        self.lstm = nn.LSTM(embed_dim,
+                            hidden_dim,
+                            n_layers,
+                            batch_first=True,
+                            dropout=dropout,
+                            bidirectional=bidirectional)
+        self.bi = 2 if bidirectional else 1
+        self.fc = nn.Linear(hidden_dim * self.bi, 1)
+
+    def forward(self, x, lengths):
+        x = self.embedding(x)
+        packed = rnn_utils.pack_padded_sequence(
+            x, lengths.cpu(), batch_first=True, enforce_sorted=False)
+        _, (h, _) = self.lstm(packed)
+        if self.bi == 2:
+            h = torch.cat((h[-2], h[-1]), dim=1)  # concat fwd+rev
+        else:
+            h = h[-1]
+        return self.fc(h).squeeze(1)