From 41613b00fe037fe4c02016dd7f35960ec88e4e37 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Fri, 9 Oct 2026 14:38:23 +0700 Subject: [PATCH] fix: raise an error when DISABLED pooling gets per-token model output --- fastembed/text/custom_text_embedding.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index 118b7313..518370dc 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -110,6 +110,14 @@ def _pool( return last_token_pooling(embeddings, attention_mask) if self._pooling == PoolingType.DISABLED: + if embeddings.ndim != 2: + raise ValueError( + f"{PoolingType.DISABLED} pooling expects the model to output sentence " + f"embeddings of shape (batch_size, dim), got an output of shape " + f"{embeddings.shape}. Use {PoolingType.CLS}, {PoolingType.MEAN} or " + f"{PoolingType.LAST_TOKEN} pooling, or set `output_name` to a pooled output " + "of the model, e.g. `sentence_embedding`." + ) return embeddings raise ValueError(