import torch
import torch.nn as nn

class SimpleMCQModel(nn.Module):
    def __init__(self, vocab_size, embed_dim=128, hidden_dim=64):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.fc1 = nn.Linear(embed_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, 1)

    def forward(self, x):
        # x shape: (batch_size, 5, max_len)
        batch_size, num_opts, max_len = x.shape
        x = x.view(batch_size * num_opts, max_len)
        embedded = self.embedding(x)
        pooled = embedded.mean(dim=1)
        out = self.relu(self.fc1(pooled))
        scores = self.fc2(out)
        scores = scores.view(batch_size, num_opts)
        return scores