Support of multilanguage for bm25

Tests for french

fastembed/common/__init__.py

@@ -1,3 +1,3 @@
-from fastembed.common.types import ImageInput, OnnxProvider, PathInput
+from fastembed.common.types import ImageInput, OnnxProvider, PathInput, PilInput
 
-__all__ = ["OnnxProvider", "ImageInput", "PathInput"]
+__all__ = ["OnnxProvider", "ImageInput", "PathInput", "PilInput"]

fastembed/common/types.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from PIL import Image
 from typing import Any, Dict, Iterable, Tuple, Union
 
 if sys.version_info >= (3, 10):
@@ -9,6 +10,7 @@ else:
 
 
 PathInput: TypeAlias = Union[str, os.PathLike]
-ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput]]
+PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
+ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
 
 OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]

fastembed/image/onnx_image_model.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type
 import numpy as np
 from PIL import Image
 
-from fastembed.common import ImageInput, OnnxProvider, PathInput
+from fastembed.common import ImageInput, OnnxProvider, PathInput, PilInput
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
 from fastembed.common.utils import iter_batch
@@ -56,7 +56,7 @@ class OnnxImageModel(OnnxModel[T]):
 
     def onnx_embed(self, images: List[PathInput], **kwargs) -> OnnxOutputContext:
         with contextlib.ExitStack():
-            image_files = [Image.open(image) for image in images]
+            image_files = [Image.open(image) if not isinstance(image, Image.Image) else image for image in images]
             encoded = self.processor(image_files)
         onnx_input = self._build_onnx_input(encoded)
         onnx_input = self._preprocess_onnx_input(onnx_input)
@@ -75,7 +75,7 @@ class OnnxImageModel(OnnxModel[T]):
     ) -> Iterable[T]:
         is_small = False
 
-        if isinstance(images, str) or isinstance(images, Path):
+        if isinstance(images, str) or isinstance(images, Path) or (isinstance(images, Image.Image)):
             images = [images]
             is_small = True
 

fastembed/sparse/bm25.py

@@ -17,6 +17,38 @@ from fastembed.sparse.sparse_embedding_base import (
 )
 from fastembed.sparse.utils.tokenizer import WordTokenizer
 
+supported_languages = [
+    "arabic",
+    "azerbaijani",
+    "basque",
+    "bengali",
+    "catalan",
+    "chinese",
+    "danish",
+    "dutch",
+    "english",
+    "finnish",
+    "french",
+    "german",
+    "greek",
+    "hebrew",
+    "hinglish",
+    "hungarian",
+    "indonesian",
+    "italian",
+    "kazakh",
+    "nepali",
+    "norwegian",
+    "portuguese",
+    "romanian",
+    "russian",
+    "slovene",
+    "spanish",
+    "swedish",
+    "tajik",
+    "turkish",
+]
+
 supported_bm25_models = [
     {
         "model": "Qdrant/bm25",
@@ -26,14 +58,14 @@ supported_bm25_models = [
             "hf": "Qdrant/bm25",
         },
         "model_file": "mock.file",  # bm25 does not require a model, so we just use a mock
-        "additional_files": ["stopwords.txt"],
+        "additional_files": supported_languages,
         "requires_idf": True,
     },
 ]
 
-MODEL_TO_LANGUAGE = {
-    "Qdrant/bm25": "english",
-}
+# MODEL_TO_LANGUAGE = {
+#     "Qdrant/bm25": "english",
+# }
 
 
 class Bm25(SparseTextEmbeddingBase):
@@ -71,10 +103,16 @@ class Bm25(SparseTextEmbeddingBase):
         k: float = 1.2,
         b: float = 0.75,
         avg_len: float = 256.0,
+        language: str = "english",
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, **kwargs)
 
+        if language not in supported_languages:
+            raise ValueError(f"{language} language is not supported")
+        else:
+            self.language = language
+
         self.k = k
         self.b = b
         self.avg_len = avg_len
@@ -88,7 +126,7 @@ class Bm25(SparseTextEmbeddingBase):
 
         self.punctuation = set(string.punctuation)
         self.stopwords = set(self._load_stopwords(model_dir))
-        self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
+        self.stemmer = get_stemmer(language)
         self.tokenizer = WordTokenizer
 
     @classmethod
@@ -100,9 +138,8 @@ class Bm25(SparseTextEmbeddingBase):
         """
         return supported_bm25_models
 
-    @classmethod
-    def _load_stopwords(cls, model_dir: Path) -> List[str]:
-        stopwords_path = model_dir / "stopwords.txt"
+    def _load_stopwords(self, model_dir: Path) -> List[str]:
+        stopwords_path = model_dir / self.language
         if not stopwords_path.exists():
             return []
 

tests/test_attention_embeddings.py

@@ -4,9 +4,7 @@ import pytest
 from fastembed import SparseTextEmbedding
 
 
-@pytest.mark.parametrize(
-    "model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]
-)
+@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
 def test_attention_embeddings(model_name):
     model = SparseTextEmbedding(model_name=model_name)
 
@@ -65,13 +63,11 @@ def test_attention_embeddings(model_name):
         assert len(result.indices) == 2
 
 
-@pytest.mark.parametrize(
-    "model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]
-)
+@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
 def test_parallel_processing(model_name):
     model = SparseTextEmbedding(model_name=model_name)
 
-    docs = ["hello world", "attention embedding"] * 100
+    docs = ["hello world", "attention embedding", "Mort aux vaches"] * 100
     embeddings = list(model.embed(docs, batch_size=10, parallel=2))
 
     embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
@@ -85,3 +81,23 @@ def test_parallel_processing(model_name):
         assert np.allclose(emb_1.indices, emb_3.indices)
         assert np.allclose(emb_1.values, emb_2.values)
         assert np.allclose(emb_1.values, emb_3.values)
+
+    print("Passed")
+    if model_name == "Qdrant/bm25":
+        docs = ["Mort aux vaches", "Je suis au lit"]
+
+        model = SparseTextEmbedding(model_name=model_name, language="french")
+        embeddings = list(model.embed(docs, parallel=2))[:2]
+        assert embeddings[0].values.shape == (2,)
+        assert embeddings[0].indices.shape == (2,)
+
+        assert embeddings[1].values.shape == (2,)
+        assert embeddings[1].indices.shape == (2,)
+
+        model = SparseTextEmbedding(model_name=model_name, language="english")
+        embeddings = list(model.embed(docs, parallel=2))[:2]
+        assert embeddings[0].values.shape == (3,)
+        assert embeddings[0].indices.shape == (3,)
+
+        assert embeddings[1].values.shape == (4,)
+        assert embeddings[1].indices.shape == (4,)

tests/test_image_onnx_embeddings.py

@@ -1,7 +1,10 @@
 import os
+from io import BytesIO
 
 import numpy as np
 import pytest
+import requests
+from PIL import Image
 
 from fastembed import ImageEmbedding
 from tests.config import TEST_MISC_DIR
@@ -12,12 +15,10 @@ CANONICAL_VECTOR_VALUES = {
         [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01046245, 0.01171397, 0.00705971, 0.0]
     ),
     "Qdrant/Unicom-ViT-B-16": np.array(
-        [ 0.0170, -0.0361,  0.0125, -0.0428, -0.0232,  0.0232, -0.0602, -0.0333,
-         0.0155,  0.0497]
+        [0.0170, -0.0361, 0.0125, -0.0428, -0.0232, 0.0232, -0.0602, -0.0333, 0.0155, 0.0497]
     ),
     "Qdrant/Unicom-ViT-B-32": np.array(
-        [0.0418,  0.0550,  0.0003,  0.0253, -0.0185,  0.0016, -0.0368, -0.0402,
-        -0.0891, -0.0186]
+        [0.0418, 0.0550, 0.0003, 0.0253, -0.0185, 0.0016, -0.0368, -0.0402, -0.0891, -0.0186]
     ),
 }
 
@@ -33,10 +34,15 @@ def test_embedding():
 
         model = ImageEmbedding(model_name=model_desc["model"])
 
-        images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")]
+        images = [
+            TEST_MISC_DIR / "image.jpeg",
+            str(TEST_MISC_DIR / "small_image.jpeg"),
+            Image.open(BytesIO(requests.get("https://qdrant.tech/img/logo.png").content)),
+            TEST_MISC_DIR / "logo.png",
+        ]
         embeddings = list(model.embed(images))
         embeddings = np.stack(embeddings, axis=0)
-        assert embeddings.shape == (2, dim)
+        assert embeddings.shape == (len(images), dim)
 
         canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
 
@@ -44,19 +50,24 @@ def test_embedding():
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
         ), model_desc["model"]
 
+        assert np.allclose(embeddings[3, :10], embeddings[2:10], atol=1e-3), model_desc["model"]
+
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
 def test_batch_embedding(n_dims, model_name):
     model = ImageEmbedding(model_name=model_name)
     n_images = 32
-    images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * (
-        n_images // 2
-    )
+    test_images = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+        Image.open(TEST_MISC_DIR / "small_image.jpeg"),
+    ]
+    images = test_images * n_images
 
     embeddings = list(model.embed(images, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
 
-    assert embeddings.shape == (n_images, n_dims)
+    assert embeddings.shape == (len(test_images) * n_images, n_dims)
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
@@ -64,9 +75,12 @@ def test_parallel_processing(n_dims, model_name):
     model = ImageEmbedding(model_name=model_name)
 
     n_images = 32
-    images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * (
-        n_images // 2
-    )
+    test_images = [
+        TEST_MISC_DIR / "image.jpeg",
+        str(TEST_MISC_DIR / "small_image.jpeg"),
+        Image.open(TEST_MISC_DIR / "small_image.jpeg"),
+    ]
+    images = test_images * n_images
     embeddings = list(model.embed(images, batch_size=10, parallel=2))
     embeddings = np.stack(embeddings, axis=0)
 
@@ -76,6 +90,6 @@ def test_parallel_processing(n_dims, model_name):
     embeddings_3 = list(model.embed(images, batch_size=10, parallel=0))
     embeddings_3 = np.stack(embeddings_3, axis=0)
 
-    assert embeddings.shape == (n_images, n_dims)
+    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)