Classify Padded Text with Masked Mean Pooling

Build a small text classifier that gives the same answer however much padding a sentence carries.

Sentences have different lengths, so a batch is a padded (B, T) tensor of token ids. Id 0 is the padding token. Write MeanPoolClassifier, an nn.Module that:

1. looks up an embedding for each token with nn.Embedding(vocab_size, dim, padding_idx=0), stored as self.embedding,
2. averages the embeddings of the real tokens only (a masked mean),
3. maps that average to class logits with nn.Linear(dim, n_classes), stored as self.head.

mask   = tokens != 0                      # (B, T)
pooled = sum of embeddings where mask / number of real tokens
logits = head(pooled)                     # (B, n_classes)

Then write train_text_classifier, which trains a new model on the data it gets and returns it in eval mode.

Function signature

class MeanPoolClassifier(torch.nn.Module):
    def __init__(self, vocab_size, dim, n_classes): ...
    def forward(self, tokens): ...


def train_text_classifier(tokens, labels, vocab_size, n_classes):

Example

Input

model = MeanPoolClassifier(vocab_size=30, dim=8, n_classes=3)
short  = torch.tensor([[5, 7, 0]])
padded = torch.tensor([[5, 7, 0, 0, 0, 0, 0]])

Output

torch.allclose(model(short), model(padded))  ->  True

Reasoning

Both rows hold the same two real tokens, 5 and 7. The mask removes every padding position, and the mean divides by 2 in both cases. A plain .mean(dim=1) would divide by 3 for one row and by 7 for the other, so the logits would change with the padding.

Constraints and notes

• tokens is int64 with shape (B, T). Padding is always 0 and comes after the real tokens.
• A row with no real tokens must give finite logits. Clamp the count to at least 1.
• The tests train your model on synthetic three-class text and need at least 95% accuracy on held-out sentences, also when those sentences are padded to a much greater length.
• Keep the training to a few seconds on a CPU.