| 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 | |