nlp_proj / biLSTM1.py
IvaElen's picture
Update biLSTM1.py
326d46b
import torch
import torch.nn as nn
class biLSTM(nn.Module):
"""
The LSTM model that will be used to perform Sentiment analysis.
"""
def __init__(self,
# объем словаря, с которым мы работаем, размер входа для слоя Embedding
vocab_size: int,
# размер выходного эмбеддинга каждый элемент последовательности
# будет описан вектором такой размерности
embedding_dim: int,
# размерность hidden state LSTM слоя
hidden_dim: int,
# число слоев в LSTM
n_layers: int,
drop_prob=0.5,
seq_len = 128) -> None:
super().__init__()
self.hidden_dim = hidden_dim
self.n_layers = n_layers
self.seq_len = seq_len
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.lstm = nn.LSTM(embedding_dim,
hidden_dim,
n_layers,
dropout=drop_prob,
batch_first=True,
bidirectional=True
)
self.do = nn.Dropout()
self.fc1 = nn.Linear(2*hidden_dim * self.seq_len, 256)
self.fc2 = nn.Linear(256, 1)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
embeds = self.embedding(x)
lstm_out, _ = self.lstm(embeds)
out = self.fc2(torch.tanh(self.do(self.fc1(lstm_out.flatten(1)))))
sig_out = self.sigmoid(out)
return sig_out