Spaces:
Sleeping
Sleeping
import torch | |
import torch.nn as nn | |
class RNNClassifier(nn.Module): | |
def __init__(self, vocab_size, embed_dim, hidden_dim, output_dim, padding_idx): | |
super(RNNClassifier, self).__init__() | |
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx) | |
self.rnn = nn.RNN(embed_dim, hidden_dim, batch_first=True) | |
self.fc1 = nn.Linear(hidden_dim, hidden_dim // 2) # New hidden layer | |
self.relu = nn.ReLU() | |
self.fc2 = nn.Linear(hidden_dim // 2, output_dim) | |
def forward(self, x): | |
embedded = self.embedding(x) # [batch_size, seq_len, embed_dim] | |
output, hidden = self.rnn(embedded) # hidden: [1, batch_size, hidden_dim] | |
x = self.fc1(hidden.squeeze(0)) # [batch_size, hidden_dim//2] | |
x = self.relu(x) | |
out = self.fc2(x) # [batch_size, output_dim] | |
return out | |