|
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 MachineModel(nn.Module): |
|
def __init__(self, embedding_dim, vocab_size, hidden_dim, n_layers, linear_output, num_classes, pretrained_embeddings=None, padding_idx=0): |
|
super(MachineModel, 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 |