Skip to content

Commit 4e5575f

Browse files
authored
fix: raise exception if pooling is incorrect in custom model (#522)
1 parent 4736f46 commit 4e5575f

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

fastembed/text/custom_text_embedding.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,11 @@ def _pool(
7777
if self._pooling == PoolingType.DISABLED:
7878
return embeddings
7979

80+
raise ValueError(
81+
f"Unsupported pooling type {self._pooling}. "
82+
f"Supported types are: {PoolingType.CLS}, {PoolingType.MEAN}, {PoolingType.DISABLED}."
83+
)
84+
8085
def _normalize(self, embeddings: NumpyArray) -> NumpyArray:
8186
return normalize(embeddings) if self._normalization else embeddings
8287

0 commit comments

Comments
 (0)