115 lines
4.9 KiB
Python
Executable File
115 lines
4.9 KiB
Python
Executable File
#taken from https://github.com/prakashpandey9/Text-Classification-Pytorch/blob/master/models/LSTM.py
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.autograd import Variable
|
|
from models.helpers import *
|
|
|
|
|
|
class RNNMultilingualClassifier(nn.Module):
|
|
|
|
def __init__(self, output_size, hidden_size, lvocab_size, learnable_length, lpretrained=None,
|
|
drop_embedding_range=None, drop_embedding_prop=0, post_probabilities=True, only_post=False,
|
|
bert_embeddings=False):
|
|
|
|
super(RNNMultilingualClassifier, self).__init__()
|
|
self.output_size = output_size
|
|
self.hidden_size = hidden_size
|
|
self.drop_embedding_range = drop_embedding_range
|
|
self.drop_embedding_prop = drop_embedding_prop
|
|
self.post_probabilities = post_probabilities
|
|
self.bert_embeddings = bert_embeddings
|
|
assert 0 <= drop_embedding_prop <= 1, 'drop_embedding_prop: wrong range'
|
|
|
|
self.lpretrained_embeddings = nn.ModuleDict()
|
|
self.llearnable_embeddings = nn.ModuleDict()
|
|
self.embedding_length = None
|
|
self.langs = sorted(lvocab_size.keys())
|
|
self.only_post = only_post
|
|
|
|
self.n_layers = 1
|
|
self.n_directions = 1
|
|
|
|
self.dropout = nn.Dropout(0.6)
|
|
|
|
lstm_out = 256
|
|
ff1 = 512
|
|
ff2 = 256
|
|
|
|
lpretrained_embeddings = {}
|
|
llearnable_embeddings = {}
|
|
if only_post==False:
|
|
for l in self.langs:
|
|
pretrained = lpretrained[l] if lpretrained else None
|
|
pretrained_embeddings, learnable_embeddings, embedding_length = init_embeddings(
|
|
pretrained, lvocab_size[l], learnable_length
|
|
)
|
|
lpretrained_embeddings[l] = pretrained_embeddings
|
|
llearnable_embeddings[l] = learnable_embeddings
|
|
self.embedding_length = embedding_length
|
|
|
|
# self.lstm = nn.LSTM(self.embedding_length, hidden_size, dropout=0.2 if self.n_layers>1 else 0, num_layers=self.n_layers, bidirectional=(self.n_directions==2))
|
|
self.rnn = nn.GRU(self.embedding_length, hidden_size)
|
|
self.linear0 = nn.Linear(hidden_size * self.n_directions, lstm_out)
|
|
self.lpretrained_embeddings.update(lpretrained_embeddings)
|
|
self.llearnable_embeddings.update(llearnable_embeddings)
|
|
|
|
self.linear1 = nn.Linear(lstm_out, ff1)
|
|
self.linear2 = nn.Linear(ff1, ff2)
|
|
|
|
if only_post:
|
|
self.label = nn.Linear(output_size, output_size)
|
|
elif post_probabilities and not bert_embeddings:
|
|
self.label = nn.Linear(ff2 + output_size, output_size)
|
|
elif bert_embeddings and not post_probabilities:
|
|
self.label = nn.Linear(ff2 + 768, output_size)
|
|
elif post_probabilities and bert_embeddings:
|
|
self.label = nn.Linear(ff2 + output_size + 768, output_size)
|
|
else:
|
|
self.label = nn.Linear(ff2, output_size)
|
|
|
|
def forward(self, input, post, bert_embed, lang):
|
|
if self.only_post:
|
|
doc_embedding = post
|
|
else:
|
|
doc_embedding = self.transform(input, lang)
|
|
if self.post_probabilities:
|
|
doc_embedding = torch.cat([doc_embedding, post], dim=1)
|
|
if self.bert_embeddings:
|
|
doc_embedding = torch.cat([doc_embedding, bert_embed], dim=1)
|
|
|
|
logits = self.label(doc_embedding)
|
|
return logits
|
|
|
|
def transform(self, input, lang):
|
|
batch_size = input.shape[0]
|
|
input = embed(self, input, lang)
|
|
input = embedding_dropout(input, drop_range=self.drop_embedding_range, p_drop=self.drop_embedding_prop,
|
|
training=self.training)
|
|
input = input.permute(1, 0, 2)
|
|
h_0 = Variable(torch.zeros(self.n_layers*self.n_directions, batch_size, self.hidden_size).cuda())
|
|
# c_0 = Variable(torch.zeros(self.n_layers*self.n_directions, batch_size, self.hidden_size).cuda())
|
|
# output, (_, _) = self.lstm(input, (h_0, c_0))
|
|
output, _ = self.rnn(input, h_0)
|
|
output = output[-1, :, :]
|
|
output = F.relu(self.linear0(output))
|
|
output = self.dropout(F.relu(self.linear1(output)))
|
|
output = self.dropout(F.relu(self.linear2(output)))
|
|
return output
|
|
|
|
def finetune_pretrained(self):
|
|
for l in self.langs:
|
|
self.lpretrained_embeddings[l].requires_grad = True
|
|
self.lpretrained_embeddings[l].weight.requires_grad = True
|
|
|
|
def get_embeddings(self, input, lang):
|
|
batch_size = input.shape[0]
|
|
input = embed(self, input, lang)
|
|
input = embedding_dropout(input, drop_range=self.drop_embedding_range, p_drop=self.drop_embedding_prop,
|
|
training=self.training)
|
|
input = input.permute(1, 0, 2)
|
|
h_0 = Variable(torch.zeros(self.n_layers * self.n_directions, batch_size, self.hidden_size).cuda())
|
|
output, _ = self.rnn(input, h_0)
|
|
output = output[-1, :, :]
|
|
return output.cpu().detach().numpy()
|
|
|