diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index f33878e..c42ac41 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -3,8 +3,6 @@ name: Tests on: push: branches: [ master, main, gpu ] - schedule: - - cron: 0 0 * * * pull_request: env: @@ -42,13 +40,6 @@ jobs: poetry config virtualenvs.create false poetry install --no-interaction --no-ansi --without docs - - name: Install Test Dependencies - run: pip install pytest pytest-md pytest-emoji - - name: Run pytest - uses: pavelzw/pytest-action@v2 - with: - verbose: true - emoji: true - job-summary: true - report-title: 'FastEmbed Test Report' + run: | + poetry run pytest diff --git a/fastembed/image/onnx_embedding.py b/fastembed/image/onnx_embedding.py index c6a9cd4..a1357a6 100644 --- a/fastembed/image/onnx_embedding.py +++ b/fastembed/image/onnx_embedding.py @@ -77,12 +77,12 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]): model_description = self._get_model_description(model_name) self.cache_dir = define_cache_dir(cache_dir) - model_dir = self.download_model( + self._model_dir = self.download_model( model_description, self.cache_dir, local_files_only=self._local_files_only ) self.load_onnx_model( - model_dir=model_dir, + model_dir=self._model_dir, model_file=model_description["model_file"], threads=threads, providers=providers, diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index 3ff942e..0c55429 100644 --- a/fastembed/late_interaction/colbert.py +++ b/fastembed/late_interaction/colbert.py @@ -136,12 +136,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]): model_description = self._get_model_description(model_name) self.cache_dir = define_cache_dir(cache_dir) - model_dir = self.download_model( + self._model_dir = self.download_model( model_description, self.cache_dir, local_files_only=self._local_files_only ) self.load_onnx_model( - model_dir=model_dir, + model_dir=self._model_dir, model_file=model_description["model_file"], threads=threads, providers=providers, diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index 1dfaa39..eea3880 100644 --- a/fastembed/sparse/bm25.py +++ b/fastembed/sparse/bm25.py @@ -116,12 +116,12 @@ class Bm25(SparseTextEmbeddingBase): model_description = self._get_model_description(model_name) self.cache_dir = define_cache_dir(cache_dir) - model_dir = self.download_model( + self._model_dir = self.download_model( model_description, self.cache_dir, local_files_only=self._local_files_only ) self.punctuation = set(string.punctuation) - self.stopwords = set(self._load_stopwords(model_dir, self.language)) + self.stopwords = set(self._load_stopwords(self._model_dir, self.language)) self.stemmer = get_stemmer(language) self.tokenizer = WordTokenizer diff --git a/fastembed/sparse/bm42.py b/fastembed/sparse/bm42.py index 0abe119..ec8372d 100644 --- a/fastembed/sparse/bm42.py +++ b/fastembed/sparse/bm42.py @@ -84,12 +84,12 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): model_description = self._get_model_description(model_name) self.cache_dir = define_cache_dir(cache_dir) - model_dir = self.download_model( + self._model_dir = self.download_model( model_description, self.cache_dir, local_files_only=self._local_files_only ) self.load_onnx_model( - model_dir=model_dir, + model_dir=self._model_dir, model_file=model_description["model_file"], threads=threads, providers=providers, @@ -103,7 +103,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): self.special_tokens = set(self.special_token_to_id.keys()) self.special_tokens_ids = set(self.special_token_to_id.values()) self.punctuation = set(string.punctuation) - self.stopwords = set(self._load_stopwords(model_dir)) + self.stopwords = set(self._load_stopwords(self._model_dir)) self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name]) self.alpha = alpha diff --git a/fastembed/sparse/splade_pp.py b/fastembed/sparse/splade_pp.py index 3d31ba1..7b344ef 100644 --- a/fastembed/sparse/splade_pp.py +++ b/fastembed/sparse/splade_pp.py @@ -82,12 +82,12 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): model_description = self._get_model_description(model_name) self.cache_dir = define_cache_dir(cache_dir) - model_dir = self.download_model( + self._model_dir = self.download_model( model_description, self.cache_dir, local_files_only=self._local_files_only ) self.load_onnx_model( - model_dir=model_dir, + model_dir=self._model_dir, model_file=model_description["model_file"], threads=threads, providers=providers, diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index 7e1f35c..6cba8ec 100644 --- a/fastembed/text/onnx_embedding.py +++ b/fastembed/text/onnx_embedding.py @@ -190,12 +190,12 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]): model_description = self._get_model_description(model_name) self.cache_dir = define_cache_dir(cache_dir) - model_dir = self.download_model( + self._model_dir = self.download_model( model_description, self.cache_dir, local_files_only=self._local_files_only ) self.load_onnx_model( - model_dir=model_dir, + model_dir=self._model_dir, model_file=model_description["model_file"], threads=threads, providers=providers, diff --git a/tests/test_attention_embeddings.py b/tests/test_attention_embeddings.py index 31dfcf2..f198e7b 100644 --- a/tests/test_attention_embeddings.py +++ b/tests/test_attention_embeddings.py @@ -1,3 +1,6 @@ +import os +import shutil + import numpy as np import pytest @@ -6,6 +9,7 @@ from fastembed import SparseTextEmbedding @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]) def test_attention_embeddings(model_name): + is_ci = os.getenv("CI") model = SparseTextEmbedding(model_name=model_name) output = list( @@ -62,9 +66,14 @@ def test_attention_embeddings(model_name): assert len(result.indices) == len(result.values) assert len(result.indices) == 2 + if is_ci: + shutil.rmtree(model.model._model_dir) + @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]) def test_parallel_processing(model_name): + is_ci = os.getenv("CI") + model = SparseTextEmbedding(model_name=model_name) docs = ["hello world", "attention embedding", "Mangez-vous vraiment des grenouilles?"] * 100 @@ -82,9 +91,14 @@ def test_parallel_processing(model_name): assert np.allclose(emb_1.values, emb_2.values) assert np.allclose(emb_1.values, emb_3.values) + if is_ci: + shutil.rmtree(model.model._model_dir) + @pytest.mark.parametrize("model_name", ["Qdrant/bm25"]) def test_multilanguage(model_name): + is_ci = os.getenv("CI") + docs = ["Mangez-vous vraiment des grenouilles?", "Je suis au lit"] model = SparseTextEmbedding(model_name=model_name, language="french") @@ -102,3 +116,6 @@ def test_multilanguage(model_name): assert embeddings[1].values.shape == (4,) assert embeddings[1].indices.shape == (4,) + + if is_ci: + shutil.rmtree(model.model._model_dir) diff --git a/tests/test_image_onnx_embeddings.py b/tests/test_image_onnx_embeddings.py index be3703c..4810610 100644 --- a/tests/test_image_onnx_embeddings.py +++ b/tests/test_image_onnx_embeddings.py @@ -1,4 +1,5 @@ import os +import shutil from io import BytesIO import numpy as np @@ -52,9 +53,13 @@ def test_embedding(): assert np.allclose(embeddings[1], embeddings[2]), model_desc["model"] + if is_ci: + shutil.rmtree(model.model._model_dir) + @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")]) def test_batch_embedding(n_dims, model_name): + is_ci = os.getenv("CI") model = ImageEmbedding(model_name=model_name) n_images = 32 test_images = [ @@ -68,10 +73,13 @@ def test_batch_embedding(n_dims, model_name): embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (len(test_images) * n_images, n_dims) + if is_ci: + shutil.rmtree(model.model._model_dir) @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")]) def test_parallel_processing(n_dims, model_name): + is_ci = os.getenv("CI") model = ImageEmbedding(model_name=model_name) n_images = 32 @@ -93,3 +101,5 @@ def test_parallel_processing(n_dims, model_name): assert embeddings.shape == (n_images * len(test_images), n_dims) assert np.allclose(embeddings, embeddings_2, atol=1e-3) assert np.allclose(embeddings, embeddings_3, atol=1e-3) + if is_ci: + shutil.rmtree(model.model._model_dir) diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index 13ad612..f05db26 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/tests/test_late_interaction_embeddings.py @@ -1,3 +1,6 @@ +import os +import shutil + import numpy as np from fastembed.late_interaction.late_interaction_text_embedding import ( @@ -105,6 +108,7 @@ docs = ["Hello World"] def test_batch_embedding(): + is_ci = os.getenv("CI") docs_to_embed = docs * 10 for model_name, expected_result in CANONICAL_COLUMN_VALUES.items(): @@ -116,8 +120,12 @@ def test_batch_embedding(): token_num, abridged_dim = expected_result.shape assert np.allclose(value[:, :abridged_dim], expected_result, atol=10e-4) + if is_ci: + shutil.rmtree(model.model._model_dir) + def test_single_embedding(): + is_ci = os.getenv("CI") docs_to_embed = docs for model_name, expected_result in CANONICAL_COLUMN_VALUES.items(): @@ -127,8 +135,12 @@ def test_single_embedding(): token_num, abridged_dim = expected_result.shape assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4) + if is_ci: + shutil.rmtree(model.model._model_dir) + def test_single_embedding_query(): + is_ci = os.getenv("CI") queries_to_embed = docs for model_name, expected_result in CANONICAL_QUERY_VALUES.items(): @@ -138,8 +150,12 @@ def test_single_embedding_query(): token_num, abridged_dim = expected_result.shape assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4) + if is_ci: + shutil.rmtree(model.model._model_dir) + def test_parallel_processing(): + is_ci = os.getenv("CI") model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0") token_dim = 128 docs = ["hello world", "flag embedding"] * 100 @@ -155,3 +171,6 @@ def test_parallel_processing(): assert embeddings.shape[0] == len(docs) and embeddings.shape[-1] == token_dim assert np.allclose(embeddings, embeddings_2, atol=1e-3) assert np.allclose(embeddings, embeddings_3, atol=1e-3) + + if is_ci: + shutil.rmtree(model.model._model_dir) diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py index b99f51b..a6e982a 100644 --- a/tests/test_sparse_embeddings.py +++ b/tests/test_sparse_embeddings.py @@ -1,4 +1,8 @@ +import os +import shutil + import pytest +import numpy as np from fastembed.sparse.bm25 import Bm25 from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding @@ -46,6 +50,7 @@ docs = ["Hello World"] def test_batch_embedding(): + is_ci = os.getenv("CI") docs_to_embed = docs * 10 for model_name, expected_result in CANONICAL_COLUMN_VALUES.items(): @@ -55,9 +60,12 @@ def test_batch_embedding(): for i, value in enumerate(result.values): assert pytest.approx(value, abs=0.001) == expected_result["values"][i] + if is_ci: + shutil.rmtree(model.model._model_dir) def test_single_embedding(): + is_ci = os.getenv("CI") for model_name, expected_result in CANONICAL_COLUMN_VALUES.items(): model = SparseTextEmbedding(model_name=model_name) @@ -68,11 +76,12 @@ def test_single_embedding(): for i, value in enumerate(result.values): assert pytest.approx(value, abs=0.001) == expected_result["values"][i] + if is_ci: + shutil.rmtree(model.model._model_dir) def test_parallel_processing(): - import numpy as np - + is_ci = os.getenv("CI") model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1") docs = ["hello world", "flag embedding"] * 30 sparse_embeddings_duo = list(model.embed(docs, batch_size=10, parallel=2)) @@ -97,16 +106,23 @@ def test_parallel_processing(): assert np.allclose(sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3) assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3) + if is_ci: + shutil.rmtree(model.model._model_dir) + @pytest.fixture def bm25_instance(): - return Bm25("Qdrant/bm25", language="english") + ci = os.getenv("CI", True) + model = Bm25("Qdrant/bm25", language="english") + yield model + if ci: + shutil.rmtree(model._model_dir) def test_stem_with_stopwords_and_punctuation(bm25_instance): # Setup - bm25_instance.stopwords = set(["the", "is", "a"]) - bm25_instance.punctuation = set([".", ",", "!"]) + bm25_instance.stopwords = {"the", "is", "a"} + bm25_instance.punctuation = {".", ",", "!"} # Test data tokens = ["The", "quick", "brown", "fox", "is", "a", "test", "sentence", ".", "!"] @@ -121,8 +137,8 @@ def test_stem_with_stopwords_and_punctuation(bm25_instance): def test_stem_case_insensitive_stopwords(bm25_instance): # Setup - bm25_instance.stopwords = set(["the", "is", "a"]) - bm25_instance.punctuation = set([".", ",", "!"]) + bm25_instance.stopwords = {"the", "is", "a"} + bm25_instance.punctuation = {".", ",", "!"} # Test data tokens = ["THE", "Quick", "Brown", "Fox", "IS", "A", "Test", "Sentence", ".", "!"] diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 019a42f..7f58c6f 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -1,4 +1,5 @@ import os +import shutil import numpy as np import pytest @@ -83,6 +84,8 @@ def test_embedding(): assert np.allclose( embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3 ), model_desc["model"] + if is_ci: + shutil.rmtree(model.model._model_dir) @pytest.mark.parametrize( @@ -90,6 +93,7 @@ def test_embedding(): [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")], ) def test_batch_embedding(n_dims, model_name): + is_ci = os.getenv("CI") model = TextEmbedding(model_name=model_name) docs = ["hello world", "flag embedding"] * 100 @@ -97,6 +101,8 @@ def test_batch_embedding(n_dims, model_name): embeddings = np.stack(embeddings, axis=0) assert embeddings.shape == (200, n_dims) + if is_ci: + shutil.rmtree(model.model._model_dir) @pytest.mark.parametrize( @@ -104,6 +110,7 @@ def test_batch_embedding(n_dims, model_name): [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")], ) def test_parallel_processing(n_dims, model_name): + is_ci = os.getenv("CI") model = TextEmbedding(model_name=model_name) docs = ["hello world", "flag embedding"] * 100 @@ -119,3 +126,5 @@ def test_parallel_processing(n_dims, model_name): assert embeddings.shape == (200, n_dims) assert np.allclose(embeddings, embeddings_2, atol=1e-3) assert np.allclose(embeddings, embeddings_3, atol=1e-3) + if is_ci: + shutil.rmtree(model.model._model_dir)