Sentence Embeddings: Masked Mean Pooling

easy
Unit 5 · Evaluation & Embeddings0/3 to milestone

A transformer encoder outputs one vector per token. To get a single sentence embedding for search, models like sentence-transformers average the token vectors, but only over real tokens, not padding.

Implement mean_pool(token_embeddings, attention_mask, normalize=True):

  • token_embeddings has shape (batch, seq_len, dim).
  • attention_mask has shape (batch, seq_len), with 1 for real tokens and 0 for padding.
  • Return shape (batch, dim): the average of each sequence's real token vectors.
  • If normalize is true, scale each result to unit length (L2), so cosine similarity becomes a dot product.
  • A sequence with no real tokens should produce a zero vector (not NaN).
#embeddings#rag#numpy2 visible + 2 hidden tests
solution.py · Python 3.13
Loading editor…
Run your code to see results. ⌘↵ to run · ⌘⇧↵ to submit
Sentence Embeddings: Masked Mean Pooling · TechGradify