feat: embedding progress bar (#71)

* feat: embedding progress

* refactor: with auto __close__

* refactor: with __exit__ tqdm

README.md

@@ -2,7 +2,7 @@
 
 FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a Github issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
-The default embedding supports "query" and "passage" prefixes for the input text. The default model is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval%20with%20FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+The default embedding supports "query" and "passage" prefixes for the input text. The default model is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
 
 1. Light & Fast
     - Quantized model weights
@@ -27,6 +27,8 @@ pip install fastembed
 
 ```python
 from fastembed.embedding import FlagEmbedding as Embedding
+from typing import List
+import numpy as np
 
 documents: List[str] = [
     "passage: Hello, World!",

docs/Getting Started.ipynb

@@ -38,18 +38,11 @@
    "id": "b61c6552",
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
-     ]
-    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "torch.Size([384])\n"
+      "(384,)\n"
      ]
     }
    ],
@@ -64,9 +57,9 @@
     "    \"This is an example document.\",\n",
     "    \"fastembed is supported by and maintained by Qdrant.\",\n",
     "]\n",
-    "# Initialize the DefaultEmbedding class with the desired parameters\n",
-    "embedding_model = DefaultEmbedding(model_name=\"BAAI/bge-small-en\", max_length=512)\n",
-    "embeddings: List[np.ndarray] = embedding_model.embed(documents)\n",
+    "# Initialize the DefaultEmbedding class\n",
+    "embedding_model = DefaultEmbedding()\n",
+    "embeddings: List[np.ndarray] = list(embedding_model.embed(documents))\n",
     "print(embeddings[0].shape)"
    ]
   },
@@ -97,7 +90,7 @@
    "source": [
     "from typing import List\n",
     "import numpy as np\n",
-    "from fastembed.embedding import DefaultEmbedding as Embedding"
+    "from fastembed.embedding import DefaultEmbedding"
    ]
   },
   {
@@ -171,20 +164,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 8,
    "id": "8013eee9",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "embeddings: List[np.ndarray] = embedding_model.embed(documents)"
+    "embeddings: List[np.ndarray] = list(embedding_model.embed(documents))"
    ]
   },
   {
@@ -197,7 +182,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 9,
    "id": "0d8c8e08",
    "metadata": {},
    "outputs": [
@@ -205,7 +190,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "torch.Size([384])\n"
+      "(384,)\n"
      ]
     }
    ],

docs/examples/Supported_Models.ipynb

@@ -29,6 +29,7 @@
        "      <th>model</th>\n",
        "      <th>dim</th>\n",
        "      <th>description</th>\n",
+       "      <th>size_in_GB</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
@@ -36,25 +37,43 @@
        "      <th>0</th>\n",
        "      <td>BAAI/bge-small-en</td>\n",
        "      <td>384</td>\n",
-       "      <td>Fast and Default English model</td>\n",
+       "      <td>Fast English model</td>\n",
+       "      <td>0.20</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
-       "      <td>BAAI/bge-base-en</td>\n",
-       "      <td>768</td>\n",
-       "      <td>Base English model</td>\n",
+       "      <td>BAAI/bge-small-en-v1.5</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Fast and Default English model</td>\n",
+       "      <td>0.13</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
-       "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
-       "      <td>384</td>\n",
-       "      <td>Sentence Transformer model, MiniLM-L6-v2</td>\n",
+       "      <td>BAAI/bge-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Base English model</td>\n",
+       "      <td>0.50</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
+       "      <td>BAAI/bge-base-en-v1.5</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Base English model, v1.5</td>\n",
+       "      <td>0.44</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Sentence Transformer model, MiniLM-L6-v2</td>\n",
+       "      <td>0.09</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
        "      <td>intfloat/multilingual-e5-large</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Multilingual model, e5-large. Recommend using this model for non-English languages. Recommend using this via Torch implementation of FastEmbed</td>\n",
+       "      <td>Multilingual model, e5-large. Recommend using this model for non-English languages</td>\n",
+       "      <td>2.24</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -63,15 +82,27 @@
       "text/plain": [
        "                                    model   dim  \\\n",
        "0                       BAAI/bge-small-en   384   \n",
-       "1                        BAAI/bge-base-en   768   \n",
-       "2  sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "3          intfloat/multilingual-e5-large  1024   \n",
+       "1                  BAAI/bge-small-en-v1.5   384   \n",
+       "2                        BAAI/bge-base-en   768   \n",
+       "3                   BAAI/bge-base-en-v1.5   768   \n",
+       "4  sentence-transformers/all-MiniLM-L6-v2   384   \n",
+       "5          intfloat/multilingual-e5-large  1024   \n",
        "\n",
-       "                                                                                                                                      description  \n",
-       "0                                                                                                                  Fast and Default English model  \n",
-       "1                                                                                                                              Base English model  \n",
-       "2                                                                                                        Sentence Transformer model, MiniLM-L6-v2  \n",
-       "3  Multilingual model, e5-large. Recommend using this model for non-English languages. Recommend using this via Torch implementation of FastEmbed  "
+       "                                                                          description  \\\n",
+       "0                                                                  Fast English model   \n",
+       "1                                                      Fast and Default English model   \n",
+       "2                                                                  Base English model   \n",
+       "3                                                            Base English model, v1.5   \n",
+       "4                                            Sentence Transformer model, MiniLM-L6-v2   \n",
+       "5  Multilingual model, e5-large. Recommend using this model for non-English languages   \n",
+       "\n",
+       "   size_in_GB  \n",
+       "0        0.20  \n",
+       "1        0.13  \n",
+       "2        0.50  \n",
+       "3        0.44  \n",
+       "4        0.09  \n",
+       "5        2.24  "
       ]
      },
      "execution_count": 1,
@@ -88,6 +119,13 @@
     "pd.set_option('display.max_colwidth', None)\n",
     "pd.DataFrame(Embedding.list_supported_models())"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

docs/examples/Usage_With_Qdrant.ipynb

@@ -28,16 +28,7 @@
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\u001b[33mDEPRECATION: pytorch-lightning 1.6.5 has a non-standard dependency specifier torch>=1.8.*. pip 23.3 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
-      "\u001b[0m"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "!pip install 'qdrant-client[fastembed]' --quiet --upgrade"
    ]
@@ -113,26 +104,19 @@
    "execution_count": 4,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
-     ]
-    },
     {
      "data": {
       "text/plain": [
-       "['77e1e4724dd243b08608f57d5692f6aa',\n",
-       " '74841e5dc3594646bda2c6a6d2795dbd',\n",
-       " '6ef39a9445604d0da84d04f760cd7cf7',\n",
-       " 'e659503d3b3748ef90f23c778274835b',\n",
-       " 'b999675068cd413f93faa0cc890c3819',\n",
-       " '8e452f2935cf4e4b80d8eea68c2aad58',\n",
-       " '28ed4fd4592c48c9a0519618d51bb86e',\n",
-       " '59378c784c5f49109bef65fdc4061334',\n",
-       " 'a78c9b598f7942749156334283a6f24f',\n",
-       " 'f72bb24701c64fabb0182c9e757b581b']"
+       "['6e8fcf7e0ecc407b9b6bb011d169f629',\n",
+       " 'c9d26e7e0ea741b2b1082d097796b28b',\n",
+       " 'cf05747e7eb34d2490b1df1f8be94049',\n",
+       " '208c197266d547a880dfb65e46738b19',\n",
+       " '27bd985c5d6f49d68fc2cf73dac74199',\n",
+       " 'c5e929c8837f4370818c97f63996f8ef',\n",
+       " 'c12213c6cdac470aa2471f2d30dc4041',\n",
+       " '974e64a7d8624f6e9824fa7b9c94f99d',\n",
+       " '0129fae193c740eba092512d8e53ab4a',\n",
+       " '492cad6e741e4aeebb196bd818a97d17']"
       ]
      },
      "execution_count": 4,
@@ -194,6 +178,13 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Behind the scenes, Qdrant Client uses the FastEmbed library to make a passage embedding and then uses the Qdrant API to upsert the documents with metadata, put together as a Points into the collection."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 6,

docs/experimental/Binary Quantization with Qdrant.ipynb

@@ -46,7 +46,7 @@
    "source": [
     "## 2. Download and Slice Dataset\n",
     "\n",
-    "We will be using the [dbpedia-entitis-openai-1M](https://huggingface.co/datasets/KShivendu/dbpedia-entities-openai-1M) dataset from the [HuggingFace Datasets](https://huggingface.co/datasets) library. This contains 1M vectors of 1536 dimensions each. We will be using the first 10K vectors here."
+    "We will be using the [dbpedia-entitis-openai-1M](https://huggingface.co/datasets/KShivendu/dbpedia-entities-openai-1M) dataset from the [HuggingFace Datasets](https://huggingface.co/datasets) library. This contains 1M vectors of 1536 dimensions each. We will be using the first 100K vectors here."
    ]
   },
   {

fastembed/embedding.py

@@ -1,7 +1,10 @@
 import json
+import math
 import os
 import shutil
 import tarfile
+import tempfile
+
 from abc import ABC, abstractmethod
 from itertools import islice
 from multiprocessing import get_all_start_methods
@@ -38,7 +41,7 @@ def normalize(input_array, p=2, dim=1, eps=1e-12):
     return normalized_array
 
 
-class EmbeddingModel(ABC):
+class EmbeddingModel:
     @classmethod
     def load_tokenizer(cls, model_dir: Path, max_length: int = 512) -> Tokenizer:
         config_path = model_dir / "config.json"
@@ -110,7 +113,7 @@ class EmbeddingModel(ABC):
         self.tokenizer = self.load_tokenizer(self.path, max_length=max_length)
         self.model = ort.InferenceSession(str(model_path), providers=onnx_providers, sess_options=so)
 
-    def onnx_embed(self, documents: List[str]) -> np.ndarray:
+    def onnx_embed(self, documents: List[str]) -> Tuple[np.ndarray, np.ndarray]:
         encoded = self.tokenizer.encode_batch(documents)
         input_ids = np.array([e.ids for e in encoded])
         attention_mask = np.array([e.attention_mask for e in encoded])
@@ -126,9 +129,8 @@ class EmbeddingModel(ABC):
             )
 
         model_output = self.model.run(None, onnx_input)
-        last_hidden_state = model_output[0][:, 0]
-        embeddings = normalize(last_hidden_state).astype(np.float32)
-        return embeddings
+        embeddings = model_output[0]
+        return embeddings, attention_mask
 
 
 class EmbeddingWorker(Worker):
@@ -150,16 +152,16 @@ class EmbeddingWorker(Worker):
 
     def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
         for idx, batch in items:
-            embeddings = self.model.onnx_embed(batch)
-            yield idx, embeddings
+            embeddings, attn_mask = self.model.onnx_embed(batch)
+            yield idx, (embeddings, attn_mask)
 
 
 class Embedding(ABC):
     """
     Abstract class for embeddings.
 
-    Args:
-        ABC ():
+    Inherits:
+        ABC: Abstract base class
 
     Raises:
         NotImplementedError: Raised when you call an abstract method that has not been implemented.
@@ -196,6 +198,12 @@ class Embedding(ABC):
                 "description": "Fast and Default English model",
                 "size_in_GB": 0.13
             },
+            {
+                "model": "BAAI/bge-small-zh-v1.5",
+                "dim": 512,
+                "description": "Fast and recommended Chinese model",
+                "size_in_GB": 0.1
+            },
             {
                 "model": "BAAI/bge-base-en",
                 "dim": 768,
@@ -220,6 +228,18 @@ class Embedding(ABC):
                 "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
                 "size_in_GB": 2.24
             },
+            {
+                "model": "jinaai/jina-embeddings-v2-base-en",
+                "dim": 768,
+                "description": " English embedding model supporting 8192 sequence length",
+                "size_in_GB": 0.55
+            },
+            {
+                "model": "jinaai/jina-embeddings-v2-small-en",
+                "dim": 512,
+                "description": " English embedding model supporting 8192 sequence length",
+                "size_in_GB": 0.13
+            }
         ]
 
     @classmethod
@@ -276,6 +296,24 @@ class Embedding(ABC):
                 progress_bar.close()
         return output_path
 
+    @classmethod
+    def download_files_from_huggingface(cls, repod_id: str, cache_dir: Optional[str] = None) -> str:
+        """
+        Downloads a model from HuggingFace Hub.
+        Args:
+            repod_id (str): The HF hub id (name) of the model to retrieve.
+            cache_dir (Optional[str]): The path to the cache directory.
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. "jinaai/jina-embeddings-v2-small-en".
+        Returns:
+            Path: The path to the model directory.
+        """
+        from huggingface_hub import snapshot_download
+
+        return snapshot_download(
+            repo_id=repod_id, ignore_patterns=["model.safetensors", "pytorch_model.bin"], cache_dir=cache_dir
+        )
+
     @classmethod
     def decompress_to_cache(cls, targz_path: str, cache_dir: str):
         """
@@ -311,7 +349,7 @@ class Embedding(ABC):
 
         return cache_dir
 
-    def retrieve_model(self, model_name: str, cache_dir: str) -> Path:
+    def retrieve_model_gcs(self, model_name: str, cache_dir: str) -> Path:
         """
         Retrieves a model from Google Cloud Storage.
 
@@ -355,6 +393,24 @@ class Embedding(ABC):
 
         return model_dir
 
+    def retrieve_model_hf(self, model_name: str, cache_dir: str) -> Path:
+        """
+        Retrieves a model from HuggingFace Hub.
+        Args:
+            model_name (str): The name of the model to retrieve.
+            cache_dir (str): The path to the cache directory.
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        Returns:
+            Path: The path to the model directory.
+        """
+
+        assert (
+                "/" in model_name
+        ), "model_name must be in the format <org>/<model> e.g. jinaai/jina-embeddings-v2-small-en"
+
+        return Path(self.download_files_from_huggingface(repod_id=model_name, cache_dir=cache_dir))
+
     def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
         """
         Embeds a list of text passages into a list of embeddings.
@@ -397,7 +453,7 @@ class FlagEmbedding(Embedding):
 
     def __init__(
             self,
-            model_name: str = "BAAI/bge-small-en",
+            model_name: str = "BAAI/bge-small-en-v1.5",
             max_length: int = 512,
             cache_dir: str = None,
             threads: int = None,
@@ -406,7 +462,9 @@ class FlagEmbedding(Embedding):
         Args:
             model_name (str): The name of the model to use.
             max_length (int, optional): The maximum number of tokens. Defaults to 512. Unknown behavior for values > 512.
-            cache_dir (str, optional): The path to the cache directory. Defaults to `local_cache` in the current directory.
+            cache_dir (str, optional): The path to the cache directory.                             \
+                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.    \
+                                       Defaults to `fastembed_cache` in the system's temp directory.
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
 
         Raises:
@@ -415,18 +473,19 @@ class FlagEmbedding(Embedding):
         self.model_name = model_name
 
         if cache_dir is None:
-            cache_dir = Path(".").resolve() / "local_cache"
+            default_cache_dir = os.path.join(tempfile.gettempdir(), "fastembed_cache")
+            cache_dir = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
             cache_dir.mkdir(parents=True, exist_ok=True)
 
         self._cache_dir = cache_dir
-        self._model_dir = self.retrieve_model(model_name, cache_dir)
+        self._model_dir = self.retrieve_model_gcs(model_name, cache_dir)
         self._max_length = max_length
 
         self.model = EmbeddingModel(self._model_dir, self.model_name, max_length=max_length,
                                     max_threads=threads)
 
     def embed(
-            self, documents: Union[str, Iterable[str]], batch_size: int = 256, parallel: int = None
+            self, documents: Union[str, Iterable[str]], batch_size: int = 256, parallel: int = None, show_progress: bool = True
     ) -> Iterable[np.ndarray]:
         """
         Encode a list of documents into list of embeddings.
@@ -439,6 +498,7 @@ class FlagEmbedding(Embedding):
                 If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
                 If 0, use all available cores.
                 If None, don't use data-parallel processing, use default onnxruntime threading instead.
+            show_progress (bool, optional): Whether to show a progress bar. Defaults to True.
 
         Returns:
             List of embeddings, one per document
@@ -455,20 +515,34 @@ class FlagEmbedding(Embedding):
 
         if parallel == 0:
             parallel = os.cpu_count()
+        
+        with tqdm(total=len(documents), disable=not show_progress) as progress_bar:
+            batch_iterable = iter_batch(documents, batch_size)
+            if parallel is None or is_small:
+                for batch in batch_iterable:
+                    embeddings, _ = self.model.onnx_embed(batch)
+                    yield from normalize(embeddings[:, 0]).astype(np.float32)
+                    progress_bar.update(len(embeddings))
+            else:
+                start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+                params = {
+                    "path": self._model_dir,
+                    "model_name": self.model_name,
+                    "max_length": self._max_length,
+                }
+                pool = ParallelWorkerPool(parallel, EmbeddingWorker, start_method=start_method)
+                for batch in pool.ordered_map(batch_iterable, **params):
+                    embeddings, _ = batch
+                    yield from normalize(embeddings[:, 0]).astype(np.float32)
+                    progress_bar.update(len(embeddings))
 
-        if parallel is None or is_small:
-            for batch in iter_batch(documents, batch_size):
-                yield from self.model.onnx_embed(batch)
-        else:
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            params = {
-                "path": self._model_dir,
-                "model_name": self.model_name,
-                "max_length": self._max_length,
-            }
-            pool = ParallelWorkerPool(parallel, EmbeddingWorker, start_method=start_method)
-            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
-                yield from batch
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Union[str, Union[int, float]]]]:
+        """
+        Lists the supported models.
+        """
+        # jina models are not supported by this class
+        return [model for model in super().list_supported_models() if not model['model'].startswith('jinaai')]
 
 
 class DefaultEmbedding(FlagEmbedding):
@@ -499,3 +573,105 @@ class OpenAIEmbedding(Embedding):
         # Use your OpenAI model to embed the texts
         # return self.model.embed(texts)
         raise NotImplementedError
+
+
+class JinaEmbedding(Embedding):
+    def __init__(
+            self,
+            model_name: str = "jinaai/jina-embeddings-v2-base-en",
+            max_length: int = 512,
+            cache_dir: str = None,
+            threads: int = None,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            max_length (int, optional): The maximum number of tokens. Defaults to 512. Unknown behavior for values > 512.
+            cache_dir (str, optional): The path to the cache directory.                             \
+                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.    \
+                                       Defaults to `fastembed_cache` in the system's temp directory.
+            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+        self.model_name = model_name
+
+        if cache_dir is None:
+            default_cache_dir = os.path.join(tempfile.gettempdir(), "fastembed_cache")
+            cache_dir = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
+            cache_dir.mkdir(parents=True, exist_ok=True)
+
+        self._cache_dir = cache_dir
+        self._model_dir = self.retrieve_model_hf(model_name, cache_dir)
+        self._max_length = max_length
+
+        self.model = EmbeddingModel(self._model_dir, self.model_name, max_length=max_length,
+                                    max_threads=threads)
+
+    def embed(
+            self, documents: Union[str, Iterable[str]], batch_size: int = 256, parallel: int = None, show_progress: bool = True
+    ) -> Iterable[np.ndarray]:
+        """
+        Encode a list of documents into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+        Args:
+            documents: Iterator of documents or single document to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+            show_progress (bool, optional): Whether to show a progress bar. Defaults to True.
+        Returns:
+            List of embeddings, one per document
+        """
+        is_small = False
+
+        if isinstance(documents, str):
+            documents = [documents]
+            is_small = True
+
+        if isinstance(documents, list):
+            if len(documents) < batch_size:
+                is_small = True
+
+        if parallel == 0:
+            parallel = os.cpu_count()
+
+        with tqdm(total=len(documents), disable=not show_progress) as progress_bar:
+            batch_iterable = iter_batch(documents, batch_size)
+            if parallel is None or is_small:
+                for batch in batch_iterable:
+                    embeddings, attn_mask = self.model.onnx_embed(batch)
+                    yield from normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+                    progress_bar.update(len(embeddings))
+            else:
+                start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+                params = {
+                    "path": self._model_dir,
+                    "model_name": self.model_name,
+                    "max_length": self._max_length,
+                }
+                pool = ParallelWorkerPool(parallel, EmbeddingWorker, start_method=start_method)
+                for batch in pool.ordered_map(batch_iterable, **params):
+                    embeddings, attn_mask = batch
+                    yield from normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+                    progress_bar.update(len(embeddings))
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Union[str, Union[int, float]]]]:
+        """
+        Lists the supported models.
+        """
+        # only jina models are supported by this class
+        return [model for model in Embedding.list_supported_models() if model['model'].startswith('jinaai')]
+
+    @staticmethod
+    def mean_pooling(model_output, attention_mask):
+        token_embeddings = model_output
+        input_mask_expanded = (np.expand_dims(attention_mask, axis=-1)).astype(float)
+
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        mask_sum = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
+
+        return sum_embeddings / mask_sum

pyproject.toml

@@ -16,7 +16,8 @@ onnx = "^1.11"
 onnxruntime = "^1.15"
 tqdm = "^4.65"
 requests = "^2.31"
-tokenizers = "^0.13"
+tokenizers = "^0.15.0"
+huggingface-hub = "0.19.4"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.2"

tests/test_onnx_embeddings.py

@@ -1,29 +1,32 @@
 import os
-
+import pytest
 import numpy as np
-from tqdm import tqdm
 
-from fastembed.embedding import DefaultEmbedding, Embedding
+from fastembed.embedding import DefaultEmbedding, JinaEmbedding
 
 CANONICAL_VECTOR_VALUES = {
-    "BAAI/bge-small-en": np.array([-0.0232, -0.0255,  0.0174, -0.0639, -0.0006]),
-    "BAAI/bge-small-en-v1.5": np.array([0.01522374, -0.02271799,  0.00860278, -0.07424029,  0.00386434]),
-    "BAAI/bge-base-en": np.array([0.0115,  0.0372,  0.0295,  0.0121,  0.0346]),
+    "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
+    "BAAI/bge-small-en-v1.5": np.array([0.01522374, -0.02271799, 0.00860278, -0.07424029, 0.00386434]),
+    "BAAI/bge-small-zh-v1.5": np.array([-0.01023294, 0.07634465, 0.0691722, -0.04458365, -0.03160762]),
+    "BAAI/bge-base-en": np.array([0.0115, 0.0372, 0.0295, 0.0121, 0.0346]),
     "BAAI/bge-base-en-v1.5": np.array([0.01129394, 0.05493144, 0.02615099, 0.00328772, 0.02996045]),
-    "sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259,  0.0058,  0.0114,  0.0380, -0.0233]),
-    "intfloat/multilingual-e5-large": np.array([0.0098,  0.0045,  0.0066, -0.0354,  0.0070]),
+    "sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259, 0.0058, 0.0114, 0.0380, -0.0233]),
+    "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
+    "jinaai/jina-embeddings-v2-small-en": np.array([-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]),
+    "jinaai/jina-embeddings-v2-base-en": np.array([-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]),
 }
 
 
-def test_default_embedding():
+@pytest.mark.parametrize('embedding_class', [DefaultEmbedding, JinaEmbedding])
+def test_embedding(embedding_class):
     is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
 
-    for model_desc in Embedding.list_supported_models():
+    for model_desc in embedding_class.list_supported_models():
         if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
             continue
 
         dim = model_desc["dim"]
-        model = DefaultEmbedding(model_name=model_desc["model"])
+        model = embedding_class(model_name=model_desc["model"])
 
         docs = ["hello world", "flag embedding"]
         embeddings = list(model.embed(docs))
@@ -34,18 +37,20 @@ def test_default_embedding():
         assert np.allclose(embeddings[0, :canonical_vector.shape[0]], canonical_vector, atol=1e-3), model_desc["model"]
 
 
-def test_batch_embedding():
-    model = DefaultEmbedding()
+@pytest.mark.parametrize('n_dims,embedding_class', [(384, DefaultEmbedding), (768, JinaEmbedding)])
+def test_batch_embedding(n_dims, embedding_class):
+    model = embedding_class()
 
     docs = ["hello world", "flag embedding"] * 100
     embeddings = list(model.embed(docs, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
 
-    assert embeddings.shape == (200, 384)
+    assert embeddings.shape == (200, n_dims)
 
 
-def test_parallel_processing():
-    model = DefaultEmbedding()
+@pytest.mark.parametrize('n_dims,embedding_class', [(384, DefaultEmbedding), (768, JinaEmbedding)])
+def test_parallel_processing(n_dims, embedding_class):
+    model = embedding_class()
 
     docs = ["hello world", "flag embedding"] * 100
     embeddings = list(model.embed(docs, batch_size=10, parallel=2))
@@ -57,6 +62,6 @@ def test_parallel_processing():
     embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
     embeddings_3 = np.stack(embeddings_3, axis=0)
 
-    assert embeddings.shape == (200, 384)
+    assert embeddings.shape == (200, n_dims)
     assert np.allclose(embeddings, embeddings_2, atol=1e-3)
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)