|
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__()
|
|
|
|
|
|
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)
|
|
|
|
|
|
self.lstm1 = nn.LSTM(embedding_dim, hidden_dim, num_layers=n_layers, batch_first=True, bidirectional=True)
|
|
|
|
|
|
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)
|
|
|
|
|
|
out, b_n = self.lstm1(x_in)
|
|
x_1 = gather_last_relevant_hidden(hiddens=out, seq_lens=seq_lens)
|
|
|
|
|
|
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 |
