Spaces:
Sleeping
Sleeping
File size: 883 Bytes
e7a44ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
import torch
import torch.nn as nn
class LSTMClassifier(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_dim, output_dim, padding_idx):
super(LSTMClassifier, self).__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx)
self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers=1, dropout=0.3, batch_first=True, bidirectional=True)
self.fc1 = nn.Linear(hidden_dim * 2, hidden_dim)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
embedded = self.embedding(x)
output, (hidden, _) = self.lstm(embedded)
hidden_cat = torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1) # concatenate last hidden states
x = self.fc1(hidden_cat)
x = self.relu(x)
out = self.fc2(x)
return out |