models/rnn.py

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