From b18630119d72dab5da996e232cd36cc8f6a0c6e8 Mon Sep 17 00:00:00 2001 From: "d.rudenko" Date: Sat, 8 Feb 2025 14:40:29 +0100 Subject: [PATCH] Clear supported_models list to keep only supported models --- tests/test_add_custom_model.py | 53 ++++++++++++++++++---------------- 1 file changed, 28 insertions(+), 25 deletions(-) diff --git a/tests/test_add_custom_model.py b/tests/test_add_custom_model.py index 2a4174a..e493daf 100644 --- a/tests/test_add_custom_model.py +++ b/tests/test_add_custom_model.py @@ -45,12 +45,6 @@ SOURCES = { @pytest.mark.parametrize("scenario", canonical_vectors) def test_add_custom_model_variations(scenario): - """ - Tests add_custom_model for different base models and different - (mean_pooling, normalization) configs. Checks the first 5 dims - of "hello world" match the scenario's canonical vector. - """ - is_ci = bool(os.getenv("CI", False)) base_model_name = scenario["model"] @@ -58,6 +52,10 @@ def test_add_custom_model_variations(scenario): normalization = scenario["normalization"] cv = np.array(scenario["canonical_vector"], dtype=np.float32) + backup_supported_models = {} + for embedding_cls in TextEmbedding.EMBEDDINGS_REGISTRY: + backup_supported_models[embedding_cls] = embedding_cls.list_supported_models().copy() + suffixes = [] suffixes.append("mean" if mean_pooling else "no-mean") suffixes.append("norm" if normalization else "no-norm") @@ -81,32 +79,37 @@ def test_add_custom_model_variations(scenario): "additional_files": [], } - if is_ci and model_info["size_in_GB"] > 1: + if is_ci and model_info["size_in_GB"] > 1.0: pytest.skip( f"Skipping {custom_model_name} on CI due to size_in_GB={model_info['size_in_GB']}" ) - TextEmbedding.add_custom_model( - model_info=model_info, mean_pooling=mean_pooling, normalization=normalization - ) + try: + TextEmbedding.add_custom_model( + model_info=model_info, mean_pooling=mean_pooling, normalization=normalization + ) - model = TextEmbedding(model_name=custom_model_name) + model = TextEmbedding(model_name=custom_model_name) - docs = ["hello world", "flag embedding"] - embeddings = list(model.embed(docs)) - embeddings = np.stack(embeddings, axis=0) + docs = ["hello world", "flag embedding"] + embeddings = list(model.embed(docs)) + embeddings = np.stack(embeddings, axis=0) - assert embeddings.shape == ( - 2, - dim, - ), f"Expected shape (2, {dim}) for {custom_model_name}, but got {embeddings.shape}" + assert embeddings.shape == ( + 2, + dim, + ), f"Expected shape (2, {dim}) for {custom_model_name}, but got {embeddings.shape}" - num_compare_dims = cv.shape[0] - assert np.allclose( - embeddings[0, :num_compare_dims], cv, atol=1e-3 - ), f"Embedding mismatch for {custom_model_name} (first {num_compare_dims} dims)." + num_compare_dims = cv.shape[0] + assert np.allclose( + embeddings[0, :num_compare_dims], cv, atol=1e-3 + ), f"Embedding mismatch for {custom_model_name} (first {num_compare_dims} dims)." - assert not np.allclose(embeddings[0, :], 0.0), "Embedding should not be all zeros." + assert not np.allclose(embeddings[0, :], 0.0), "Embedding should not be all zeros." - if is_ci: - delete_model_cache(model.model._model_dir) + if is_ci: + delete_model_cache(model.model._model_dir) + + finally: + for embedding_cls, old_list in backup_supported_models.items(): + embedding_cls.supported_models = old_list