Problem
Pooling::Cls in rust/src/embeddings/local/pooling.rs always takes the hidden state at sequence index 0:
fn cls(output: &ModelOutput) -> Result<PooledOutputType, anyhow::Error> {
match output {
ModelOutput::Tensor(tensor) => tensor
.get_on_dim(1, 0) // always position 0
...
}
}
For padded batches, CLS pooling should use the first non-padding token (via attention_mask), not a fixed index 0. This matters when sequences are left-padded.
Pooling::LastToken already has padding awareness in the same file; CLS does not.
Expected behavior
CLS pooling should:
- Take
attention_mask as input (like mean pooling)
- Use the index of the first real token per sequence
- Match the behavior used by common ModernBERT embedding checkpoints (e.g. Granite, gte-modernbert)
Impact
CLS-pooled models may produce incorrect embeddings when batching with padding, especially left padding.
Suggested fix
Update Pooling::Cls to accept attention_mask and gather the first real token per row.
Problem
Pooling::Clsinrust/src/embeddings/local/pooling.rsalways takes the hidden state at sequence index 0:For padded batches, CLS pooling should use the first non-padding token (via
attention_mask), not a fixed index 0. This matters when sequences are left-padded.Pooling::LastTokenalready has padding awareness in the same file; CLS does not.Expected behavior
CLS pooling should:
attention_maskas input (like mean pooling)Impact
CLS-pooled models may produce incorrect embeddings when batching with padding, especially left padding.
Suggested fix
Update
Pooling::Clsto acceptattention_maskand gather the first real token per row.