diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index 465ffd251..b3a7db147 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -50,8 +50,9 @@ 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_config = self.POSTPROCESSING_MAPPING[self.model_description.model] + self._pooling = postprocessing_config.pooling + self._normalization = postprocessing_config.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..f3a14e891 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 = [] @@ -74,9 +76,6 @@ def test_text_custom_model(): if is_ci: delete_model_cache(model.model._model_dir) - CustomTextEmbedding.SUPPORTED_MODELS.clear() - CustomTextEmbedding.POSTPROCESSING_MAPPING.clear() - def test_cross_encoder_custom_model(): is_ci = os.getenv("CI") @@ -114,8 +113,6 @@ def test_cross_encoder_custom_model(): if is_ci: delete_model_cache(model.model._model_dir) - CustomTextCrossEncoder.SUPPORTED_MODELS.clear() - def test_mock_add_custom_models(): dim = 5 @@ -175,8 +172,24 @@ def test_mock_add_custom_models(): ) assert np.allclose(post_processed_output, expected_output[model_name], atol=1e-3) - CustomTextEmbedding.SUPPORTED_MODELS.clear() - 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(): @@ -212,9 +225,6 @@ def test_do_not_add_existing_model(): size_in_gb=0.47, ) - CustomTextEmbedding.SUPPORTED_MODELS.clear() - CustomTextEmbedding.POSTPROCESSING_MAPPING.clear() - def test_do_not_add_existing_cross_encoder(): existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2" @@ -239,5 +249,3 @@ def test_do_not_add_existing_cross_encoder(): sources=ModelSource(hf=custom_model_name), size_in_gb=0.08, ) - - CustomTextCrossEncoder.SUPPORTED_MODELS.clear()