Revert "feat: embedding progress bar (#71)" (#77)

This reverts commit 2c7fee3b95.

.github/workflows/python-tests.yml

@@ -20,6 +20,8 @@ jobs:
           - '3.11.x'
         os:
           - ubuntu-latest
+          - macos-latest
+          - windows-latest
 
     runs-on: ${{ matrix.os }}
 
@@ -37,5 +39,7 @@ jobs:
           poetry config virtualenvs.create false
           poetry install --no-interaction --no-ansi
       - name: Run tests
-        run: pytest
+        run: |
+          export IS_UBUNTU_CI=$(test "${{ matrix.os }}" = "ubuntu-latest" && echo "true" || echo "false")
+          pytest
         shell: bash

README.md

@@ -1,22 +1,19 @@
 # ⚡️ What is FastEmbed?
 
-FastEmbed is an easy to use -- lightweight, fast, Python library built for retrieval embedding generation. 
+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/)
+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
+1. Light & Fast
     - Quantized model weights
-    - ONNX Runtime for inference
-    - No hidden dependencies on PyTorch or TensorFlow via Huggingface Transformers
+    - ONNX Runtime, no PyTorch dependency
+    - CPU-first design
+    - Data-parallelism for encoding of large datasets
 
 2. Accuracy/Recall
     - Better than OpenAI Ada-002
     - Default is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard
-
-3. Fast
-    - About 2x faster than Huggingface (PyTorch) transformers on single queries
-    - Lot faster for batches!
-    - ONNX Runtime allows you to use dedicated runtimes for even higher throughput and lower latency 
+    - List of [supported models](https://qdrant.github.io/fastembed/examples/Supported_Models/) - including multilingual models
 
 ## 🚀 Installation
 
@@ -30,32 +27,58 @@ 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!",
     "query: Hello, World!", # these are two different embedding
     "passage: This is an example passage.",
-    # You can leave out the prefix but it's recommended
-    "fastembed is supported by and maintained by Qdrant." 
+    "fastembed is supported by and maintained by Qdrant." # You can leave out the prefix but it's recommended
 ]
 embedding_model = Embedding(model_name="BAAI/bge-base-en", max_length=512) 
-embeddings: List[np.ndarray] = list(embedding_model.embed(documents))
+embeddings: List[np.ndarray] = list(embedding_model.embed(documents)) # Note the list() call - this is a generator 
 ```
 
-### Why fast?
+## Usage with Qdrant
 
-It's important we justify the "fast" in FastEmbed. FastEmbed is fast because:
+Installation with Qdrant Client in Python:
 
-1. Quantized model weights
-2. ONNX Runtime which allows for fast inference on CPU and other dedicated runtimes
+```bash
+pip install qdrant-client[fastembed]
+```
 
-### Why light?
-1. No hidden dependencies on PyTorch or TensorFlow via Huggingface Transformers
-2. We do use the tokenizer from Huggingface Transformers, but it's a light dependency
+Might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh. 
 
-### Why accurate?
-1. Better than OpenAI Ada-002
-2. Top of the Embedding leaderboards e.g. [MTEB](https://huggingface.co/spaces/mteb/leaderboard)
+```python
+from qdrant_client import QdrantClient
+
+# Initialize the client
+client = QdrantClient(":memory:")  # or QdrantClient(path="path/to/db")
+
+# Prepare your documents, metadata, and IDs
+docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
+metadata = [
+    {"source": "Langchain-docs"},
+    {"source": "Linkedin-docs"},
+]
+ids = [42, 2]
+
+# Use the new add method
+client.add(
+    collection_name="demo_collection",
+    documents=docs,
+    metadata=metadata,
+    ids=ids
+)
+
+search_result = client.query(
+    collection_name="demo_collection",
+    query_text="This is a query document"
+)
+print(search_result)
+```
 
 #### Similar Work
-Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.
+
+Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.

docs/Getting Started.ipynb

@@ -16,12 +16,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "ada95c6a",
    "metadata": {},
    "outputs": [],
    "source": [
-    "!pip install fastembed --upgrade # Install fastembed"
+    "!pip install fastembed --upgrade --quiet # Install fastembed "
    ]
   },
   {
@@ -34,10 +34,33 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "id": "b61c6552",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 76.7M/76.7M [00:05<00:00, 15.0MiB/s]\n",
+      "100%|██████████| 3/3 [00:00<00:00, 455.37it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(384,)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
    "source": [
     "from typing import List\n",
     "import numpy as np\n",
@@ -49,12 +72,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] = list(\n",
-    "    embedding_model.embed(documents)\n",
-    ")  # notice that we are casting the generator to a list\n",
-    "\n",
+    "# Initialize the DefaultEmbedding class\n",
+    "embedding_model = DefaultEmbedding()\n",
+    "embeddings: List[np.ndarray] = list(embedding_model.embed(documents))\n",
     "print(embeddings[0].shape)"
    ]
   },
@@ -78,14 +98,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "c0a6f634",
    "metadata": {},
    "outputs": [],
    "source": [
     "from typing import List\n",
     "import numpy as np\n",
-    "from fastembed.embedding import DefaultEmbedding as Embedding"
+    "from fastembed.embedding import DefaultEmbedding"
    ]
   },
   {
@@ -107,7 +127,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "145a56ce",
    "metadata": {},
    "outputs": [],
@@ -139,7 +159,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "id": "272c8915",
    "metadata": {},
    "outputs": [],
@@ -159,14 +179,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "id": "8013eee9",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 4/4 [00:00<00:00, 361.82it/s]\n"
+     ]
+    }
+   ],
    "source": [
-    "embeddings: List[np.ndarray] = list(\n",
-    "    embedding_model.embed(documents)\n",
-    ")  # notice that we are casting the generator to a list"
+    "embeddings: List[np.ndarray] = list(embedding_model.embed(documents))"
    ]
   },
   {
@@ -179,10 +205,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "id": "0d8c8e08",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(384,)\n"
+     ]
+    }
+   ],
    "source": [
     "print(embeddings[0].shape)  # (384,) or similar output"
    ]

docs/examples/Supported_Models.ipynb

@@ -0,0 +1,176 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>model</th>\n",
+       "      <th>dim</th>\n",
+       "      <th>description</th>\n",
+       "      <th>size_in_GB</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>BAAI/bge-small-en</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Fast English model</td>\n",
+       "      <td>0.20</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\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>BAAI/bge-small-zh-v1.5</td>\n",
+       "      <td>512</td>\n",
+       "      <td>Fast and recommended Chinese model</td>\n",
+       "      <td>0.10</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\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>4</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>5</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>6</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</td>\n",
+       "      <td>2.24</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</th>\n",
+       "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>English embedding model supporting 8192 sequence length</td>\n",
+       "      <td>0.55</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>8</th>\n",
+       "      <td>jinaai/jina-embeddings-v2-small-en</td>\n",
+       "      <td>512</td>\n",
+       "      <td>English embedding model supporting 8192 sequence length</td>\n",
+       "      <td>0.13</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                    model   dim  \\\n",
+       "0                       BAAI/bge-small-en   384   \n",
+       "1                  BAAI/bge-small-en-v1.5   384   \n",
+       "2                  BAAI/bge-small-zh-v1.5   512   \n",
+       "3                        BAAI/bge-base-en   768   \n",
+       "4                   BAAI/bge-base-en-v1.5   768   \n",
+       "5  sentence-transformers/all-MiniLM-L6-v2   384   \n",
+       "6          intfloat/multilingual-e5-large  1024   \n",
+       "7       jinaai/jina-embeddings-v2-base-en   768   \n",
+       "8      jinaai/jina-embeddings-v2-small-en   512   \n",
+       "\n",
+       "                                                                          description  \\\n",
+       "0                                                                  Fast English model   \n",
+       "1                                                      Fast and Default English model   \n",
+       "2                                                  Fast and recommended Chinese model   \n",
+       "3                                                                  Base English model   \n",
+       "4                                                            Base English model, v1.5   \n",
+       "5                                            Sentence Transformer model, MiniLM-L6-v2   \n",
+       "6  Multilingual model, e5-large. Recommend using this model for non-English languages   \n",
+       "7                             English embedding model supporting 8192 sequence length   \n",
+       "8                             English embedding model supporting 8192 sequence length   \n",
+       "\n",
+       "   size_in_GB  \n",
+       "0        0.20  \n",
+       "1        0.13  \n",
+       "2        0.10  \n",
+       "3        0.50  \n",
+       "4        0.44  \n",
+       "5        0.09  \n",
+       "6        2.24  \n",
+       "7        0.55  \n",
+       "8        0.13  "
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "from fastembed.embedding import Embedding\n",
+    "import pandas as pd\n",
+    "pd.set_option('display.max_colwidth', None)\n",
+    "pd.DataFrame(Embedding.list_supported_models())"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "fst",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.17"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/examples/Usage_With_Qdrant.ipynb

@@ -5,9 +5,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# [Experimental] Usage With Qdrant\n",
-    "\n",
-    "> **Note:** This notebook is experimental and is subject to change. For working with this, use the dev branch of QdrantClient.\n",
+    "# Usage With Qdrant\n",
     "\n",
     "This notebook demonstrates how to use FastEmbed and Qdrant to perform vector search and retrieval. Qdrant is an open-source vector similarity search engine that is used to store, organize, and query collections of high-dimensional vectors. \n",
     "\n",
@@ -28,13 +26,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# !pip install fastembed --quiet --upgrade\n",
-    "\n",
-    "# !pip install git+https://github.com/qdrant/qdrant_client.git@dev"
+    "!pip install 'qdrant-client[fastembed]' --quiet --upgrade"
    ]
   },
   {
@@ -46,7 +42,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -68,7 +64,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -105,25 +101,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "['901ff9d7a90a4e56afe655d1de8e7c06',\n",
-       " '7bae5aa894164398a9be68d47d72ff7a',\n",
-       " 'd940bba124e24469ae3166c3110c62f8',\n",
-       " 'c8dcf956a4f444c6bfe69a00cbeb85ce',\n",
-       " 'd2aeb24c51f549c5b21851be05cb4048',\n",
-       " '6d4c672bafef4db68ae72c8986ee8a72',\n",
-       " 'ff198613768a4361a6d8230b40ee7f78',\n",
-       " 'a133ad72876b48d48e716533c2d78cf0',\n",
-       " 'bdf2b016136e4f57b7ac4b65534031b8',\n",
-       " '5f4f279f304c47839e1a74b0f28de136']"
+       "['6e8fcf7e0ecc407b9b6bb011d169f629',\n",
+       " 'c9d26e7e0ea741b2b1082d097796b28b',\n",
+       " 'cf05747e7eb34d2490b1df1f8be94049',\n",
+       " '208c197266d547a880dfb65e46738b19',\n",
+       " '27bd985c5d6f49d68fc2cf73dac74199',\n",
+       " 'c5e929c8837f4370818c97f63996f8ef',\n",
+       " 'c12213c6cdac470aa2471f2d30dc4041',\n",
+       " '974e64a7d8624f6e9824fa7b9c94f99d',\n",
+       " '0129fae193c740eba092512d8e53ab4a',\n",
+       " '492cad6e741e4aeebb196bd818a97d17']"
       ]
      },
-     "execution_count": 17,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -150,38 +146,64 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[42, 2]"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Prepare your documents, metadata, and IDs\n",
+    "docs = [\"Qdrant has Langchain integrations\", \"Qdrant also has Llama Index integrations\"]\n",
+    "metadata = [\n",
+    "    {\"source\": \"Langchain-docs\"},\n",
+    "    {\"source\": \"Linkedin-docs\"},\n",
+    "]\n",
+    "ids = [42, 2]\n",
+    "\n",
+    "# Use the new add method\n",
+    "client.add(\n",
+    "    collection_name=\"demo_collection\",\n",
+    "    documents=docs,\n",
+    "    metadata=metadata,\n",
+    "    ids=ids\n",
+    ")"
+   ]
+  },
+  {
+   "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,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Top 5 results:\n",
-      "Rank 1: Maharana Pratap was a Rajput warrior king from Mewar. Score: 0.77\n",
-      "Rank 2: Maharana Pratap is considered a symbol of Rajput resistance against foreign rule. Score: 0.77\n",
-      "Rank 3: His legacy is celebrated in Rajasthan through festivals and monuments. Score: 0.69\n",
-      "Rank 4: He had 11 wives and 17 sons, including Amar Singh I who succeeded him as ruler of Mewar. Score: 0.68\n",
-      "Rank 5: He fought against the Mughal Empire led by Akbar. Score: 0.67\n"
+      "[QueryResponse(id='42', embedding=None, metadata={'document': 'Qdrant has Langchain integrations', 'source': 'Langchain-docs'}, document='Qdrant has Langchain integrations', score=0.8496814051311954), QueryResponse(id='2', embedding=None, metadata={'document': 'Qdrant also has Llama Index integrations', 'source': 'Linkedin-docs'}, document='Qdrant also has Llama Index integrations', score=0.8478494193031256)]\n"
      ]
     }
    ],
    "source": [
-    "from qdrant_client.qdrant_fastembed import QueryResponse\n",
-    "\n",
-    "\n",
-    "def print_top_k_results(results: List[QueryResponse], k: int = 5):\n",
-    "    print(f\"Top {k} results:\")\n",
-    "    for i, result in enumerate(results[:k]):\n",
-    "        print(f\"Rank {i + 1}: {result.document}. Score: {result.score:.2f}\")\n",
-    "\n",
-    "\n",
-    "query_text = \"Who is Maharana Pratap?\"\n",
-    "results = client.query(\n",
-    "    collection_name=\"test_collection\", query_text=query_text, limit=7\n",
-    ")  # Returns limit most relevant documents\n",
-    "\n",
-    "print_top_k_results(results)"
+    "search_result = client.query(\n",
+    "    collection_name=\"demo_collection\",\n",
+    "    query_text=[\"This is a query document\"]\n",
+    ")\n",
+    "print(search_result)"
    ]
   },
   {

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."
    ]
   },
   {

docs/index.md

@@ -1,10 +1,19 @@
 # ⚡️ What is FastEmbed?
 
-FastEmbed is an easy to use -- lightweight, fast, Python library built for retrieval embedding generation. 
+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](https://github.com/FlagOpen/FlagEmbedding), which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. 
+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/).
 
-Advanced user? Skip ahead to [Retrieval with FastEmbed](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/)
+1. Light & Fast
+    - Quantized model weights
+    - ONNX Runtime for inference via [Optimum](github.com/huggingface/optimum)
+
+2. Accuracy/Recall
+    - Better than OpenAI Ada-002
+    - Default is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard
+    - List of [supported models](https://qdrant.github.io/fastembed/examples/Supported_Models/) - including multilingual models
+
+## 🚀 Installation
 
 To install the FastEmbed library, pip works: 
 
@@ -15,31 +24,53 @@ pip install fastembed
 ## 📖 Usage
 
 ```python
-from fastembed.embedding import DefaultEmbedding
+from fastembed.embedding import FlagEmbedding as Embedding
 
 documents: List[str] = [
     "passage: Hello, World!",
     "query: Hello, World!", # these are two different embedding
     "passage: This is an example passage.",
-    # You can leave out the prefix but it's recommended
-    "fastembed is supported by and maintained by Qdrant." 
+    "fastembed is supported by and maintained by Qdrant." # You can leave out the prefix but it's recommended
 ]
-embedding_model = DefaultEmbedding() 
-embeddings: List[np.ndarray] = list(embedding_model.embed(documents))
+embedding_model = Embedding(model_name="BAAI/bge-base-en", max_length=512) 
+embeddings: List[np.ndarray] = embedding_model.embed(documents) # If you use  
 ```
 
-## 🚒 Under the hood
+## Usage with Qdrant
 
-### Why fast?
+Installation with Qdrant Client in Python:
 
-It's important we justify the "fast" in FastEmbed. FastEmbed is fast because:
+```bash
+pip install qdrant-client[fastembed]
+```
 
-1. Quantized model weights
-2. ONNX Runtime which allows for inference on CPU, GPU, and other dedicated runtimes
+Might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh. 
 
-### Why light?
-1. No hidden dependencies on PyTorch or TensorFlow via Huggingface Transformers
+```python
+from qdrant_client import QdrantClient
 
-### Why accurate?
-1. Better than OpenAI Ada-002
-2. Top of the Embedding leaderboards e.g. [MTEB](https://huggingface.co/spaces/mteb/leaderboard)
+# Initialize the client
+client = QdrantClient(":memory:")  # or QdrantClient(path="path/to/db")
+
+# Prepare your documents, metadata, and IDs
+docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
+metadata = [
+    {"source": "Langchain-docs"},
+    {"source": "Linkedin-docs"},
+]
+ids = [42, 2]
+
+# Use the new add method
+client.add(
+    collection_name="demo_collection",
+    documents=docs,
+    metadata=metadata,
+    ids=ids
+)
+
+search_result = client.query(
+    collection_name="demo_collection",
+    query_text="This is a query document"
+)
+print(search_result)
+```

docs/overrides/main.html

@@ -20,7 +20,8 @@
         </span>
         <strong>Qdrant Discord server</strong>
     </a>
-    to get help and share your work! Or check out <a rel="me" href="https://login.cloud.qdrant.io/">Qdrant Cloud</a> to
+    to get help and share your work! Or check out <a rel="me"
+        href="https://cloud.qdrant.io?utm_source=twitter&utm_medium=website&utm_campaign=fastembed">Qdrant Cloud</a> to
     get started with vector search!
 </div>
 {% endblock %}

fastembed/embedding.py

@@ -1,30 +1,166 @@
+import json
 import os
 import shutil
 import tarfile
+import tempfile
+
 from abc import ABC, abstractmethod
+from itertools import islice
+from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Iterable, List
+from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
+import onnxruntime as ort
 import requests
-from optimum.onnxruntime import ORTModelForFeatureExtraction
+from tokenizers import AddedToken, Tokenizer
 from tqdm import tqdm
-from transformers import AutoTokenizer
+
+from fastembed.parallel_processor import ParallelWorkerPool, Worker
 
 
-def normalize(input_array, p=2.0, dim=1, eps=1e-12):
+def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
+    """
+    >>> list(iter_batch([1,2,3,4,5], 3))
+    [[1, 2, 3], [4, 5]]
+    """
+    source_iter = iter(iterable)
+    while source_iter:
+        b = list(islice(source_iter, size))
+        if len(b) == 0:
+            break
+        yield b
+
+
+def normalize(input_array, p=2, dim=1, eps=1e-12):
     # Calculate the Lp norm along the specified dimension
     norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
-    norm = np.maximum(norm, eps) # Avoid division by zero
+    norm = np.maximum(norm, eps)  # Avoid division by zero
     normalized_array = input_array / norm
     return normalized_array
 
+
+class EmbeddingModel:
+    @classmethod
+    def load_tokenizer(cls, model_dir: Path, max_length: int = 512) -> Tokenizer:
+        config_path = model_dir / "config.json"
+        if not config_path.exists():
+            raise ValueError(f"Could not find config.json in {model_dir}")
+
+        tokenizer_path = model_dir / "tokenizer.json"
+        if not tokenizer_path.exists():
+            raise ValueError(f"Could not find tokenizer.json in {model_dir}")
+
+        tokenizer_config_path = model_dir / "tokenizer_config.json"
+        if not tokenizer_config_path.exists():
+            raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
+
+        tokens_map_path = model_dir / "special_tokens_map.json"
+        if not tokens_map_path.exists():
+            raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
+
+        config = json.load(open(str(config_path)))
+        tokenizer_config = json.load(open(str(tokenizer_config_path)))
+        tokens_map = json.load(open(str(tokens_map_path)))
+
+        tokenizer = Tokenizer.from_file(str(tokenizer_path))
+        tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
+        tokenizer.enable_padding(pad_id=config["pad_token_id"], pad_token=tokenizer_config["pad_token"])
+
+        for token in tokens_map.values():
+            if isinstance(token, str):
+                tokenizer.add_special_tokens([token])
+            elif isinstance(token, dict):
+                tokenizer.add_special_tokens([AddedToken(**token)])
+
+        return tokenizer
+
+    def __init__(
+            self,
+            path: Path,
+            model_name: str,
+            max_length: int = 512,
+            max_threads: int = None,
+    ):
+        self.path = path
+        self.model_name = model_name
+        model_path = self.path / "model.onnx"
+        optimized_model_path = self.path / "model_optimized.onnx"
+
+        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+        onnx_providers = ["CPUExecutionProvider"]
+
+        if not model_path.exists():
+            # Rename file model_optimized.onnx to model.onnx if it exists
+            if optimized_model_path.exists():
+                optimized_model_path.rename(model_path)
+            else:
+                raise ValueError(f"Could not find model.onnx in {self.path}")
+
+        # Hacky support for multilingual model
+        self.exclude_token_type_ids = False
+        if model_name == "intfloat/multilingual-e5-large":
+            self.exclude_token_type_ids = True
+
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        if max_threads is not None:
+            so.intra_op_num_threads = max_threads
+            so.inter_op_num_threads = max_threads
+
+        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]) -> 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])
+
+        onnx_input = {
+            "input_ids": np.array(input_ids, dtype=np.int64),
+            "attention_mask": np.array(attention_mask, dtype=np.int64),
+        }
+
+        if not self.exclude_token_type_ids:
+            onnx_input["token_type_ids"] = np.array(
+                [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
+            )
+
+        model_output = self.model.run(None, onnx_input)
+        embeddings = model_output[0]
+        return embeddings, attention_mask
+
+
+class EmbeddingWorker(Worker):
+    def __init__(
+            self,
+            path: Path,
+            model_name: str,
+            max_length: int = 512,
+    ):
+        self.model = EmbeddingModel(path=path, model_name=model_name, max_length=max_length, max_threads=1)
+
+    @classmethod
+    def start(cls, path: Path, model_name: str, max_length: int = 512, **kwargs: Any) -> "EmbeddingWorker":
+        return cls(
+            path=path,
+            model_name=model_name,
+            max_length=max_length,
+        )
+
+    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+        for idx, batch in items:
+            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.
@@ -40,9 +176,71 @@ class Embedding(ABC):
     """
 
     @abstractmethod
-    def embed(self, texts: List[str]) -> List[np.ndarray]:
+    def embed(self, texts: Iterable[str], batch_size: int = 256, parallel: int = None) -> List[np.ndarray]:
         raise NotImplementedError
 
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Union[str, Union[int, float]]]]:
+        """
+        Lists the supported models.
+        """
+        return [
+            {
+                "model": "BAAI/bge-small-en",
+                "dim": 384,
+                "description": "Fast English model",
+                "size_in_GB": 0.2
+            },
+            {
+                "model": "BAAI/bge-small-en-v1.5",
+                "dim": 384,
+                "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,
+                "description": "Base English model",
+                "size_in_GB": 0.5
+            },
+            {
+                "model": "BAAI/bge-base-en-v1.5",
+                "dim": 768,
+                "description": "Base English model, v1.5",
+                "size_in_GB": 0.44
+            },
+            {
+                "model": "sentence-transformers/all-MiniLM-L6-v2",
+                "dim": 384,
+                "description": "Sentence Transformer model, MiniLM-L6-v2",
+                "size_in_GB": 0.09
+            },
+            {
+                "model": "intfloat/multilingual-e5-large",
+                "dim": 1024,
+                "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
     def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
         """
@@ -97,6 +295,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):
         """
@@ -132,7 +348,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.
 
@@ -159,7 +375,7 @@ class Embedding(ABC):
         try:
             self.download_file_from_gcs(
                 f"https://storage.googleapis.com/qdrant-fastembed/{fast_model_name}.tar.gz",
-            output_path=str(model_tar_gz),
+                output_path=str(model_tar_gz),
             )
         except PermissionError:
             simple_model_name = model_name.replace("/", "-")
@@ -176,21 +392,37 @@ class Embedding(ABC):
 
         return model_dir
 
-    def passage_embed(self, texts: List[str], batch_size: int = 256) -> Iterable[np.ndarray]:
+    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.
 
         Args:
-            texts (List[str]): The list of texts to embed.
-            batch_size (int, optional): The batch size. Defaults to 256.
+            texts (Iterable[str]): The list of texts to embed.
+            **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
             Iterable[np.ndarray]: The embeddings.
         """
 
-        for i in range(0, len(texts), batch_size):
-            # Prepend "passage: " to each text
-            yield from self.embed([f"passage: {t}" for t in texts[i : i + batch_size]])
+        yield from self.embed((f"passage: {t}" for t in texts), **kwargs)
 
     def query_embed(self, query: str) -> Iterable[np.ndarray]:
         """
@@ -207,7 +439,6 @@ class Embedding(ABC):
         query = f"query: {query}"
         # Embed the query
         query_embedding = self.embed([query])
-        # Compute the cosine similarity between the query embedding and the document embeddings
         return query_embedding
 
 
@@ -220,60 +451,92 @@ class FlagEmbedding(Embedding):
     """
 
     def __init__(
-        self,
-        model_name: str = "BAAI/bge-small-en",
-        max_length: int = 512,
-        cache_dir: str = None,
+            self,
+            model_name: str = "BAAI/bge-small-en-v1.5",
+            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:
-            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)
-        model_dir = self.retrieve_model(model_name, cache_dir)
-        if not (model_dir / "tokenizer.json").exists():
-            raise ValueError(f"Could not find tokenizer.json in {model_dir}")
-        if not (model_dir / "model.onnx").exists():
-            # Rename file model_optimized.onnx to model.onnx if it exists
-            if (model_dir / "model_optimized.onnx").exists():
-                (model_dir / "model_optimized.onnx").rename(model_dir / "model.onnx")
-            else:
-                raise ValueError(f"Could not find model.onnx in {model_dir}")
 
-        self.tokenizer = AutoTokenizer.from_pretrained(str(model_dir))
-        self.model = ORTModelForFeatureExtraction.from_pretrained(str(model_dir))
+        self._cache_dir = cache_dir
+        self._model_dir = self.retrieve_model_gcs(model_name, cache_dir)
+        self._max_length = max_length
 
-    def onnx_embed(self, documents: List[str]) -> Iterable[np.ndarray]:
-        encoded_input = self.tokenizer(documents, padding=True, truncation=True, return_tensors='pt')
-        model_output = self.model(**encoded_input)
-        embeddings = model_output[0][:, 0]
-        return normalize(embeddings, p=2, dim=1)
+        self.model = EmbeddingModel(self._model_dir, self.model_name, max_length=max_length,
+                                    max_threads=threads)
 
-    def embed(self, documents: List[str], batch_size: int = 256) -> Iterable[np.ndarray]:
+    def embed(
+            self, documents: Union[str, Iterable[str]], batch_size: int = 256, parallel: int = None
+    ) -> 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: List of documents to embed
+            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.
 
         Returns:
             List of embeddings, one per document
         """
-        # TODO: Replace loop with parallelized batching
-        if len(documents) >= batch_size:
-            for i in range(0, len(documents), batch_size):
-                batch = documents[i : i + batch_size]
-                return self.onnx_embed(batch)
+        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()
+
+        if parallel is None or is_small:
+            for batch in iter_batch(documents, batch_size):
+                embeddings, _ = self.model.onnx_embed(batch)
+                yield from normalize(embeddings[:, 0]).astype(np.float32)
         else:
-            return self.onnx_embed(documents)
+            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):
+                embeddings, _ = batch
+                yield from normalize(embeddings[:, 0]).astype(np.float32)
+
+    @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):
@@ -285,15 +548,13 @@ class DefaultEmbedding(FlagEmbedding):
     """
 
     def __init__(
-        self,
-        model_name: str = "BAAI/bge-small-en",
-        onnx_providers: List[str] = None,
-        max_length: int = 512,
-        cache_dir: str = None,
+            self,
+            model_name: str = "BAAI/bge-small-en-v1.5",
+            max_length: int = 512,
+            cache_dir: Optional[str] = None,
+            threads: Optional[int] = None,
     ):
-        # if onnx_providers is None:
-        #     onnx_providers = [ONNXProviders.CPU]
-        super().__init__(model_name, max_length=max_length, cache_dir=cache_dir)
+        super().__init__(model_name, max_length=max_length, cache_dir=cache_dir, threads=threads)
 
 
 class OpenAIEmbedding(Embedding):
@@ -302,8 +563,104 @@ class OpenAIEmbedding(Embedding):
         # self.model = ...
         ...
 
-    def embed(self, texts):
+    def embed(self, texts, batch_size: int = 256, parallel: int = None):
         # Use your OpenAI model to embed the texts
         # return self.model.embed(texts)
         raise NotImplementedError
-        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
+    ) -> 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.
+        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()
+
+        if parallel is None or is_small:
+            for batch in iter_batch(documents, batch_size):
+                embeddings, attn_mask = self.model.onnx_embed(batch)
+                yield from normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+        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):
+                embeddings, attn_mask = batch
+                yield from normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+
+    @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

fastembed/parallel_processor.py

@@ -0,0 +1,207 @@
+import logging
+import os
+from collections import defaultdict
+from enum import Enum
+from multiprocessing import Queue, get_context
+from multiprocessing.context import BaseContext
+from multiprocessing.process import BaseProcess
+from multiprocessing.sharedctypes import Synchronized as BaseValue
+from queue import Empty
+from typing import Any, Dict, Iterable, List, Optional, Type, Tuple
+
+# Single item should be processed in less than:
+processing_timeout = 10 * 60  # seconds
+
+max_internal_batch_size = 200
+
+
+class QueueSignals(str, Enum):
+    stop = "stop"
+    confirm = "confirm"
+    error = "error"
+
+
+class Worker:
+    @classmethod
+    def start(cls, **kwargs: Any) -> "Worker":
+        raise NotImplementedError()
+
+    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+        raise NotImplementedError()
+
+
+def _worker(
+    worker_class: Type[Worker],
+    input_queue: Queue,
+    output_queue: Queue,
+    num_active_workers: BaseValue,
+    worker_id: int,
+    kwargs: Optional[Dict[str, Any]] = None,
+) -> None:
+    """
+    A worker that pulls data pints off the input queue, and places the execution result on the output queue.
+    When there are no data pints left on the input queue, it decrements
+    num_active_workers to signal completion.
+    """
+
+    if kwargs is None:
+        kwargs = {}
+
+    logging.info(f"Reader worker: {worker_id} PID: {os.getpid()}")
+    try:
+        worker = worker_class.start(**kwargs)
+
+        # Keep going until you get an item that's None.
+        def input_queue_iterable() -> Iterable[Any]:
+            while True:
+                item = input_queue.get()
+                if item == QueueSignals.stop:
+                    break
+                yield item
+
+        for processed_item in worker.process(input_queue_iterable()):
+            output_queue.put(processed_item)
+    except Exception as e:  # pylint: disable=broad-except
+        logging.exception(e)
+        output_queue.put(QueueSignals.error)
+    finally:
+        # It's important that we close and join the queue here before
+        # decrementing num_active_workers. Otherwise our parent may join us
+        # before the queue's feeder thread has passed all buffered items to
+        # the underlying pipe resulting in a deadlock.
+        #
+        # See:
+        # https://docs.python.org/3.6/library/multiprocessing.html?highlight=process#pipes-and-queues
+        # https://docs.python.org/3.6/library/multiprocessing.html?highlight=process#programming-guidelines
+        output_queue.close()
+        output_queue.join_thread()
+
+        with num_active_workers.get_lock():
+            num_active_workers.value -= 1
+
+        logging.info(f"Reader worker {worker_id} finished")
+
+
+class ParallelWorkerPool:
+    def __init__(self, num_workers: int, worker: Type[Worker], start_method: Optional[str] = None):
+        self.worker_class = worker
+        self.num_workers = num_workers
+        self.input_queue: Optional[Queue] = None
+        self.output_queue: Optional[Queue] = None
+        self.ctx: BaseContext = get_context(start_method)
+        self.processes: List[BaseProcess] = []
+        self.queue_size = self.num_workers * max_internal_batch_size
+
+        self.num_active_workers: Optional[BaseValue] = None
+
+    def start(self, **kwargs: Any) -> None:
+        self.input_queue = self.ctx.Queue(self.queue_size)
+        self.output_queue = self.ctx.Queue(self.queue_size)
+
+        ctx_value = self.ctx.Value("i", self.num_workers)
+        assert isinstance(ctx_value, BaseValue)
+        self.num_active_workers = ctx_value
+
+        for worker_id in range(0, self.num_workers):
+            assert hasattr(self.ctx, "Process")
+            process = self.ctx.Process(
+                target=_worker,
+                args=(
+                    self.worker_class,
+                    self.input_queue,
+                    self.output_queue,
+                    self.num_active_workers,
+                    worker_id,
+                    kwargs.copy(),
+                ),
+            )
+            process.start()
+            self.processes.append(process)
+
+    def ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Any]:
+        buffer = defaultdict(Any)
+        next_expected = 0
+
+        for idx, item in self.semi_ordered_map(stream, *args, **kwargs):
+            buffer[idx] = item
+            while next_expected in buffer:
+                yield buffer.pop(next_expected)
+                next_expected += 1
+
+    def semi_ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Tuple[int, Any]]:
+        try:
+            self.start(**kwargs)
+
+            assert self.input_queue is not None, "Input queue was not initialized"
+            assert self.output_queue is not None, "Output queue was not initialized"
+
+            pushed = 0
+            read = 0
+            for idx, item in enumerate(stream):
+                if pushed - read < self.queue_size:
+                    try:
+                        out_item = self.output_queue.get_nowait()
+                    except Empty:
+                        out_item = None
+                else:
+                    try:
+                        out_item = self.output_queue.get(timeout=processing_timeout)
+                    except Empty as e:
+                        self.join_or_terminate()
+                        raise e
+
+                if out_item is not None:
+                    if out_item == QueueSignals.error:
+                        self.join_or_terminate()
+                        raise RuntimeError("Thread unexpectedly terminated")
+                    yield out_item
+                    read += 1
+
+                self.input_queue.put((idx, item))
+                pushed += 1
+
+            for _ in range(self.num_workers):
+                self.input_queue.put(QueueSignals.stop)
+
+            while read < pushed:
+                out_item = self.output_queue.get(timeout=processing_timeout)
+                if out_item == QueueSignals.error:
+                    self.join_or_terminate()
+                    raise RuntimeError("Thread unexpectedly terminated")
+                yield out_item
+                read += 1
+        finally:
+            assert self.input_queue is not None, "Input queue is None"
+            assert self.output_queue is not None, "Output queue is None"
+            self.input_queue.close()
+            self.output_queue.close()
+
+    def join_or_terminate(self, timeout: Optional[int] = 1) -> None:
+        """
+        Emergency shutdown
+        @param timeout:
+        @return:
+        """
+        for process in self.processes:
+            process.join(timeout=timeout)
+            if process.is_alive():
+                process.terminate()
+        self.processes.clear()
+
+    def join(self) -> None:
+        for process in self.processes:
+            process.join()
+        self.processes.clear()
+
+    def __del__(self) -> None:
+        """
+        Terminate processes if the user hasn't joined. This is necessary as
+        leaving stray processes running can corrupt shared state. In brief,
+        we've observed shared memory counters being reused (when the memory was
+        free from the perspective of the parent process) while the stray
+        workers still held a reference to them.
+        For a discussion of using destructors in Python in this manner, see
+        https://eli.thegreenplace.net/2009/06/12/safely-using-destructors-in-python/.
+        """
+        for process in self.processes:
+            process.terminate()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.0.4"
+version = "0.1.2"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -12,29 +12,25 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
 
 [tool.poetry.dependencies]
 python = ">=3.8.0,<3.12"
-onnxruntime = "^1.15.1"
-torch = ">=2.0.0, !=2.0.1"
-optimum = ">1.12.0"
-tqdm = "^4.65.0"
-requests = "^2.31.0"
-tokenizers = "^0.13.3"
+onnx = "^1.11"
+onnxruntime = "^1.15"
+tqdm = "^4.65"
+requests = "^2.31"
+tokenizers = "^0.15.0"
+huggingface-hub = "0.19.4"
 
-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.4.2"
 ruff = "^0.0.277"
 isort = "^5.12.0"
 black = "^23.7.0"
-onnx = "^1.11.0"
 notebook = ">=7.0.2"
 mkdocs-material = "^9.1.21"
 mkdocstrings = "^0.22.0"
 pillow = "^10.0.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
-pytest = "^7.4.0"
 
-[tool.poetry.dependencies.onnxruntime-silicon]
-version = "^1.15.0"
-markers = "sys_platform == 'darwin'"  # This makes it macOS specific
 
 [build-system]
 requires = ["poetry-core"]

tests/test_onnx_embeddings.py

@@ -1,13 +1,67 @@
+import os
+import pytest
 import numpy as np
 
-from fastembed.embedding import DefaultEmbedding
+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-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]),
+    "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():
-    model = DefaultEmbedding()
+@pytest.mark.parametrize('embedding_class', [DefaultEmbedding, JinaEmbedding])
+def test_embedding(embedding_class):
+    is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
 
-    docs = ["hello world", "flag embedding"]
-    embeddings = np.array(model.embed(docs))
-    assert embeddings.shape == (2, 384)
+    for model_desc in embedding_class.list_supported_models():
+        if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
+            continue
 
-test_default_embedding()
+        dim = model_desc["dim"]
+        model = embedding_class(model_name=model_desc["model"])
+
+        docs = ["hello world", "flag embedding"]
+        embeddings = list(model.embed(docs))
+        embeddings = np.stack(embeddings, axis=0)
+        assert embeddings.shape == (2, dim)
+
+        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
+        assert np.allclose(embeddings[0, :canonical_vector.shape[0]], canonical_vector, atol=1e-3), model_desc["model"]
+
+
+@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, n_dims)
+
+
+@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))
+    embeddings = np.stack(embeddings, axis=0)
+
+    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+    embeddings_2 = np.stack(embeddings_2, axis=0)
+
+    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+    embeddings_3 = np.stack(embeddings_3, axis=0)
+
+    assert embeddings.shape == (200, n_dims)
+    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
+    assert np.allclose(embeddings, embeddings_3, atol=1e-3)