docs: Updated README.md

* Opened image support

---------

Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>

README.md

@@ -6,7 +6,7 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 ## 📈 Why FastEmbed?
 
-1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda. 
+1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda.
 
 2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data parallelism for encoding large datasets.
 
@@ -48,6 +48,7 @@ len(embeddings_list[0]) # Vector of 384 dimensions
 
 Fastembed supports a variety of models for different tasks and modalities.
 The list of all the available models can be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/)
+
 ### 🎒 Dense text embeddings
 
 ```python
@@ -63,8 +64,6 @@ embeddings = list(model.embed(documents))
 
 ```
 
-
-
 ### 🔱 Sparse text embeddings
 
 * SPLADE++
@@ -81,10 +80,23 @@ embeddings = list(model.embed(documents))
 # ]
 ```
 
-<!--
-* BM42 - ([link](ToDo))
+* BM25
 
+```python
+from fastembed import SparseTextEmbedding
+
+model = SparseTextEmbedding(model_name="Qdrant/bm25")
+embeddings = list(model.embed(documents))
+
+# [
+#   SparseEmbedding(indices=[ 129793020, 1999429279,  819028769, ... ], values=[1.6477, 1.6327, 1.2377, ...]),
+#   SparseEmbedding(indices=[ 682147660, 1100855371,  339478471, ... ], values=[1.6741, 1.5432, 1.6741, ...])
+# ]
 ```
+
+* [BM42](https://qdrant.tech/articles/bm42/)
+
+```python
 from fastembed import SparseTextEmbedding
 
 model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
@@ -95,11 +107,15 @@ embeddings = list(model.embed(documents))
 #   SparseEmbedding(indices=[ 38,  12,  91, ... ], values=[0.11, 0.22, 0.39, ...])
 # ]
 ```
--->
+
+You can install [PyStemmer](https://pypi.org/project/PyStemmer/) to improve the stemming performance when using BM25, BM42.
+
+```shell
+pip install fastembed[pystemmer]
+```
 
 ### 🦥 Late interaction models (aka ColBERT)
 
-
 ```python
 from fastembed import LateInteractionTextEmbedding
 
@@ -137,7 +153,6 @@ embeddings = list(model.embed(images))
 # ]
 ```
 
-
 ## ⚡️ FastEmbed on a GPU
 
 FastEmbed supports running on GPU devices.
@@ -168,7 +183,7 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```
 
-or 
+or
 
 ```bash
 pip install qdrant-client[fastembed-gpu]

fastembed/image/onnx_embedding.py

@@ -51,6 +51,7 @@ supported_onnx_models = [
     },
 ]
 
+
 class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
     def __init__(
         self,
@@ -141,16 +142,10 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
 
         return onnx_input
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext
-    ) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         return normalize(output.model_output).astype(np.float32)
 
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
-    def init_embedding(
-        self, model_name: str, cache_dir: str, **kwargs
-    ) -> OnnxImageEmbedding:
-        return OnnxImageEmbedding(
-            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
-        )
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
+        return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

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, PilInput
+from fastembed.common import ImageInput, OnnxProvider
 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
@@ -54,9 +54,12 @@ class OnnxImageModel(OnnxModel[T]):
     def _build_onnx_input(self, encoded: np.ndarray) -> Dict[str, np.ndarray]:
         return {node.name: encoded for node in self.model.get_inputs()}
 
-    def onnx_embed(self, images: List[PathInput], **kwargs) -> OnnxOutputContext:
+    def onnx_embed(self, images: List[ImageInput], **kwargs) -> OnnxOutputContext:
         with contextlib.ExitStack():
-            image_files = [Image.open(image) if not isinstance(image, Image.Image) else 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 +78,11 @@ class OnnxImageModel(OnnxModel[T]):
     ) -> Iterable[T]:
         is_small = False
 
-        if isinstance(images, str) or isinstance(images, Path) or (isinstance(images, Image.Image)):
+        if (
+            isinstance(images, str)
+            or isinstance(images, Path)
+            or (isinstance(images, Image.Image))
+        ):
             images = [images]
             is_small = True
 
@@ -90,9 +97,7 @@ class OnnxImageModel(OnnxModel[T]):
             for batch in iter_batch(images, batch_size):
                 yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
-            start_method = (
-                "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            )
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
             params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
             pool = ParallelWorkerPool(
                 parallel, self._get_worker_class(), start_method=start_method

fastembed/sparse/bm25.py

@@ -17,38 +17,6 @@ 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",
@@ -58,14 +26,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": supported_languages,
+        "additional_files": ["stopwords.txt"],
         "requires_idf": True,
     },
 ]
 
-# MODEL_TO_LANGUAGE = {
-#     "Qdrant/bm25": "english",
-# }
+MODEL_TO_LANGUAGE = {
+    "Qdrant/bm25": "english",
+}
 
 
 class Bm25(SparseTextEmbeddingBase):
@@ -103,16 +71,10 @@ 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
@@ -126,7 +88,7 @@ class Bm25(SparseTextEmbeddingBase):
 
         self.punctuation = set(string.punctuation)
         self.stopwords = set(self._load_stopwords(model_dir))
-        self.stemmer = get_stemmer(language)
+        self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
         self.tokenizer = WordTokenizer
 
     @classmethod
@@ -138,8 +100,9 @@ class Bm25(SparseTextEmbeddingBase):
         """
         return supported_bm25_models
 
-    def _load_stopwords(self, model_dir: Path) -> List[str]:
-        stopwords_path = model_dir / self.language
+    @classmethod
+    def _load_stopwords(cls, model_dir: Path) -> List[str]:
+        stopwords_path = model_dir / "stopwords.txt"
         if not stopwords_path.exists():
             return []
 

pyproject.toml

@@ -25,8 +25,11 @@ numpy = [
 ]
 pillow = "^10.3.0"
 snowballstemmer = "^2.2.0"
-PyStemmer = "^2.2.0"
 mmh3 = "^4.0"
+PyStemmer = { version = "^2.2.0", optional = true }
+
+[tool.poetry.extras]
+pystemmer = ["PyStemmer"]
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.2"

tests/test_attention_embeddings.py

@@ -4,7 +4,9 @@ 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)
 
@@ -63,11 +65,13 @@ 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", "Mort aux vaches"] * 100
+    docs = ["hello world", "attention embedding"] * 100
     embeddings = list(model.embed(docs, batch_size=10, parallel=2))
 
     embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
@@ -81,23 +85,3 @@ 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

@@ -37,8 +37,8 @@ def test_embedding():
         images = [
             TEST_MISC_DIR / "image.jpeg",
             str(TEST_MISC_DIR / "small_image.jpeg"),
+            Image.open((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)
@@ -50,7 +50,7 @@ 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"]
+        assert np.allclose(embeddings[1], embeddings[2]), model_desc["model"]
 
 
 @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])