wip: remove model dir (#350)

* wip: remove model dir

* fix: update pytest run

* wip: disable some tests not used atm

* wip: disable some tests in ci

* wip: add debug print

* fix: fix ci, remove models after usage

* fix: fix bm25 deletion

* fix: remove redundant ci commands

Author: joein

.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

fastembed/image/onnx_embedding.py

@@ -77,12 +77,12 @@ def __init__(
 
         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,

fastembed/late_interaction/colbert.py

@@ -136,12 +136,12 @@ def __init__(
         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,

fastembed/sparse/bm25.py

@@ -116,12 +116,12 @@ def __init__(
         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
 

fastembed/sparse/bm42.py

@@ -84,12 +84,12 @@ def __init__(
         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 @@ def __init__(
         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
 

fastembed/sparse/splade_pp.py

@@ -82,12 +82,12 @@ def __init__(
         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,

fastembed/text/onnx_embedding.py

@@ -190,12 +190,12 @@ def __init__(
 
         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,

tests/test_attention_embeddings.py

@@ -1,3 +1,6 @@
+import os
+import shutil
+
 import numpy as np
 import pytest
 
@@ -6,6 +9,7 @@
 
 @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)

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)

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 @@
 
 
 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)

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 @@
 
 
 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", ".", "!"]