diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index 465ffd251..19133bf0a 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -50,8 +50,10 @@ 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 + self._pooling = self.POSTPROCESSING_MAPPING[self.model_description.model].pooling + self._normalization = self.POSTPROCESSING_MAPPING[ + self.model_description.model + ].normalization @classmethod def _list_supported_models(cls) -> list[DenseModelDescription]: diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index dcca7d89c..cb88be6e7 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -179,6 +179,25 @@ def test_mock_add_custom_models(): CustomTextEmbedding.POSTPROCESSING_MAPPING.clear() +def test_custom_text_model_lookup_is_case_insensitive(): + model_name = "Org/Model" + + TextEmbedding.add_custom_model( + model_name, + pooling=PoolingType.MEAN, + normalization=True, + sources=ModelSource(hf="artificial"), + dim=5, + size_in_gb=0.1, + ) + + model = TextEmbedding("org/model", lazy_load=True, specific_model_path="./") + + assert isinstance(model.model, CustomTextEmbedding) + assert model.model._pooling == PoolingType.MEAN + assert model.model._normalization is True + + def test_do_not_add_existing_model(): existing_base_model = "sentence-transformers/all-MiniLM-L6-v2" custom_model_name = "intfloat/multilingual-e5-small"