From bbdad26b44c660ad3443eba0a8c737ba45a68220 Mon Sep 17 00:00:00 2001 From: Ramnath0521 Date: Wed, 2 Sep 2026 20:44:08 +0530 Subject: [PATCH] fix: look up custom model postprocessing by resolved name `CustomTextEmbedding.__init__` read `POSTPROCESSING_MAPPING` with the `model_name` argument, but the mapping is keyed by the name the model was registered under, and model lookup is case-insensitive. Registering `Org/Model` and instantiating `org/model` therefore raised: KeyError: 'org/model' `self.model_description` is already resolved by the base class through the case-insensitive matcher, so its `.model` is the canonical registered name. Use that. The autouse test fixture also now clears `POSTPROCESSING_MAPPING` alongside `SUPPORTED_MODELS`; it reset only the latter, so a registration leaked into subsequent tests. Co-Authored-By: Claude Opus 5 --- fastembed/text/custom_text_embedding.py | 8 +++-- tests/test_custom_models.py | 40 +++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index 060ecdc3e..d55d2984d 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -50,8 +50,12 @@ def __init__( specific_model_path=specific_model_path, **kwargs, ) - self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling - self._normalization = self.POSTPROCESSING_MAPPING[model_name].normalization + # POSTPROCESSING_MAPPING is keyed by the registered name, while model + # lookup is case-insensitive, so the caller's spelling need not match. + # Use the resolved description's name rather than the argument. + postprocessing = self.POSTPROCESSING_MAPPING[self.model_description.model] + self._pooling = postprocessing.pooling + self._normalization = postprocessing.normalization @classmethod def _list_supported_models(cls) -> list[DenseModelDescription]: diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index 2050a42cb..181376a7a 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -21,9 +21,11 @@ @pytest.fixture(autouse=True) def restore_custom_models_fixture(): CustomTextEmbedding.SUPPORTED_MODELS = [] + CustomTextEmbedding.POSTPROCESSING_MAPPING = {} CustomTextCrossEncoder.SUPPORTED_MODELS = [] yield CustomTextEmbedding.SUPPORTED_MODELS = [] + CustomTextEmbedding.POSTPROCESSING_MAPPING = {} CustomTextCrossEncoder.SUPPORTED_MODELS = [] @@ -250,3 +252,41 @@ def test_do_not_add_existing_cross_encoder(): ) CustomTextCrossEncoder.SUPPORTED_MODELS.clear() + + +def test_custom_model_postprocessing_lookup_is_case_insensitive(): + """A custom model may be registered and instantiated with different casing. + + `TextEmbedding` resolves model names case-insensitively, so + `CustomTextEmbedding` has to look its postprocessing config up by the resolved + canonical name rather than by whatever string the caller happened to type. + """ + TextEmbedding.add_custom_model( + "Org/Model", + pooling=PoolingType.MEAN, + normalization=True, + sources=ModelSource(hf="intfloat/multilingual-e5-small"), + dim=384, + ) + + model = TextEmbedding("org/model", lazy_load=True).model + + assert model.model_description.model == "Org/Model" + assert model._pooling == PoolingType.MEAN + assert model._normalization is True + + +def test_custom_model_postprocessing_lookup_with_matching_case_still_works(): + """Control: the exact-casing path worked before and must keep working.""" + TextEmbedding.add_custom_model( + "Org/Model", + pooling=PoolingType.CLS, + normalization=False, + sources=ModelSource(hf="intfloat/multilingual-e5-small"), + dim=384, + ) + + model = TextEmbedding("Org/Model", lazy_load=True).model + + assert model._pooling == PoolingType.CLS + assert model._normalization is False