PyTorch Model

Model.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+
+
+def gather_last_relevant_hidden(hiddens, seq_lens):
+    """Extract and collect the last relevant
+    hidden state based on the sequence length."""
+    seq_lens = seq_lens.long().detach().cpu().numpy() - 1
+    out = []
+    for batch_index, column_index in enumerate(seq_lens):
+        out.append(hiddens[batch_index, column_index])
+    return torch.stack(out)
+
+
+class SkimlitModel(nn.Module):
+    def __init__(self, embedding_dim, vocab_size, hidden_dim, n_layers, linear_output, num_classes, pretrained_embeddings=None, padding_idx=0):
+        super(SkimlitModel, self).__init__()
+
+        # Initalizing embeddings
+        if pretrained_embeddings is None:
+            self.embeddings = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)
+        else:
+            pretrained_embeddings = torch.from_numpy(pretrained_embeddings).float()
+            self.embeddings = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim, _weight=pretrained_embeddings, padding_idx=padding_idx)
+
+        # LSTM layers
+        self.lstm1 = nn.LSTM(embedding_dim, hidden_dim, num_layers=n_layers, batch_first=True, bidirectional=True)
+
+        # FC layers
+        self.fc_text = nn.Linear(2*hidden_dim, linear_output)
+
+        self.fc_line_num = nn.Linear(20, 64)
+        self.fc_total_line = nn.Linear(24, 64)
+
+        self.fc_final = nn.Linear((64+64+linear_output), num_classes)
+        self.dropout = nn.Dropout(0.3)
+
+    def forward(self, inputs):
+        x_in, seq_lens, line_nums, total_lines = inputs
+        x_in = self.embeddings(x_in)
+
+        # RNN outputs 
+        out, b_n = self.lstm1(x_in)
+        x_1 = gather_last_relevant_hidden(hiddens=out, seq_lens=seq_lens)
+
+        # FC layers output
+        x_1 = F.relu(self.fc_text(x_1))
+        x_2 = F.relu(self.fc_line_num(line_nums))
+        x_3 = F.relu(self.fc_total_line(total_lines))
+
+        x = torch.cat((x_1, x_2, x_3), dim=1)
+        x = self.dropout(x)
+        x = self.fc_final(x)
+        return x