Merge pull request #146 from qdrant:v0.2.3

Publish to PyPi with SPLADE models

.github/workflows/ci.yml

@@ -1,8 +1,8 @@
-name: ci 
+name: ci
 on:
   push:
     branches:
-      - master 
+      - master
       - main
 permissions:
   contents: write
@@ -14,7 +14,7 @@ jobs:
       - uses: actions/setup-python@v4
         with:
           python-version: 3.x
-      - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV 
+      - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
       - uses: actions/cache@v3
         with:
           key: mkdocs-material-${{ env.cache_id }}

.github/workflows/python-tests.yml

@@ -18,8 +18,11 @@ jobs:
           - '3.9.x'
           - '3.10.x'
           - '3.11.x'
+          - '3.12.x'
         os:
           - ubuntu-latest
+          - macos-latest
+          - windows-latest
 
     runs-on: ${{ matrix.os }}
 
@@ -35,7 +38,9 @@ jobs:
         run: |
           python -m pip install poetry
           poetry config virtualenvs.create false
-          poetry install --no-interaction --no-ansi
+          poetry install --no-interaction --no-ansi --without docs
       - name: Run tests
-        run: pytest
+        run: |
+          export IS_UBUNTU_CI=$(test "${{ matrix.os }}" = "ubuntu-latest" && echo "true" || echo "false")
+          pytest
         shell: bash

.gitignore

@@ -168,11 +168,5 @@ local_cache/*/*
 docs/experimental/*.parquet
 docs/experimental/*.bin
 qdrant_storage/*
-fooling_around/fast-multilingual-e5-large/config.json
-fooling_around/fast-multilingual-e5-large/model_optimized.onnx
-fooling_around/fast-multilingual-e5-large/model_optimized.onnx.data
-fooling_around/fast-multilingual-e5-large/ort_config.json
-fooling_around/fast-multilingual-e5-large/sentencepiece.bpe.model
-fooling_around/fast-multilingual-e5-large/special_tokens_map.json
-fooling_around/fast-multilingual-e5-large/tokenizer_config.json
-fooling_around/fast-multilingual-e5-large/tokenizer.json
+fooling_around/*
+experiments/models/*

.pre-commit-config.yaml

@@ -1,12 +1,9 @@
 repos:
-- repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v3.2.0
-  hooks:
-  - id: trailing-whitespace
-  - id: end-of-file-fixer
-  - id: check-yaml
-  - id: check-added-large-files
-- repo: https://github.com/psf/black
-  rev: 23.7.0
-  hooks:
-    - id: black
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.1.13
+    hooks:
+      - id: ruff
+        types_or: [ python, pyi, jupyter ]
+        args: [ --fix ]
+      - id: ruff-format
+        types_or: [ python, pyi, jupyter ]

README.md

@@ -1,61 +1,93 @@
 # ⚡️ 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 text embedding (`TextEmbedding`) model is Flag Embedding, the top model in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. 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
-    - Quantized model weights
-    - ONNX Runtime for inference
-    - No hidden dependencies on PyTorch or TensorFlow via Huggingface Transformers
+## 📈 Why FastEmbed?
 
-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
+1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda. 
 
-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 
+2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data-parallelism for encoding large datasets.
+
+3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [supported](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
 
 ## 🚀 Installation
 
-To install the FastEmbed library, pip works: 
+To install the FastEmbed library, pip works:
 
 ```bash
 pip install fastembed
 ```
 
-## 📖 Usage
+## 📖 Quickstart
 
 ```python
-from fastembed.embedding import FlagEmbedding as Embedding
+import numpy as np
+from fastembed import TextEmbedding
+from typing import List
 
+# Example list of documents
 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." 
+    "This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.",
+    "fastembed is supported by and maintained by Qdrant.",
 ]
-embedding_model = Embedding(model_name="BAAI/bge-base-en", max_length=512) 
-embeddings: List[np.ndarray] = list(embedding_model.embed(documents))
+
+# This will trigger the model download and initialization
+embedding_model = TextEmbedding()
+print("The model BAAI/bge-small-en-v1.5 is ready to use.")
+
+embeddings_generator = embedding_model.embed(documents)  # reminder this is a generator
+embeddings_list = list(embedding_model.embed(documents))
+  # you can also convert the generator to a list, and that to a numpy array
+len(embeddings_list[0]) # Vector of 384 dimensions
 ```
 
-### 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
+You 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("localhost", port=6333) # For production
+# client = QdrantClient(":memory:") # For small experiments
+
+# Prepare your documents, metadata, and IDs
+docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
+metadata = [
+    {"source": "Langchain-docs"},
+    {"source": "Llama-index-docs"},
+]
+ids = [42, 2]
+
+# If you want to change the model:
+# client.set_model("sentence-transformers/all-MiniLM-L6-v2")
+# List of supported models: https://qdrant.github.io/fastembed/examples/Supported_Models
+
+# Use the new add() instead of upsert()
+# This internally calls embed() of the configured embedding model
+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

@@ -11,51 +11,137 @@
     "\n",
     "## Quick Start\n",
     "\n",
-    "The fastembed package is designed to be easy to use. The main class is the `Embedding` class. It takes a list of strings as input and returns a list of vectors as output. The `Embedding` class is initialized with a model file."
+    "The fastembed package is designed to be easy to use. We'll be using `TextEmbedding` class. It takes a list of strings as input and returns an generator of vectors. If you're seeing generators for the first time, don't worry, you can convert it to a list using `list()`.\n",
+    "\n",
+    "> 💡 You can learn more about generators from [Python Wiki](https://wiki.python.org/moin/Generators)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "ada95c6a",
    "metadata": {},
    "outputs": [],
    "source": [
-    "!pip install fastembed --upgrade # Install fastembed"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "ed81d725",
-   "metadata": {},
-   "source": [
-    "Make the necessary imports, initialize the `Embedding` class, and embed your data into vectors:"
+    "!pip install -Uqq fastembed # Install fastembed"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "id": "b61c6552",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "890cc3b969354eec8d149d143e301a7a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Fetching 9 files:   0%|          | 0/9 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The model BAAI/bge-small-en-v1.5 is ready to use.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "384"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
-    "from typing import List\n",
     "import numpy as np\n",
-    "from fastembed.embedding import DefaultEmbedding\n",
+    "from fastembed import TextEmbedding\n",
+    "from typing import List\n",
     "\n",
     "# Example list of documents\n",
     "documents: List[str] = [\n",
-    "    \"Hello, World!\",\n",
-    "    \"This is an example document.\",\n",
+    "    \"This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.\",\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",
-    "print(embeddings[0].shape)"
+    "# This will trigger the model download and initialization\n",
+    "embedding_model = TextEmbedding()\n",
+    "print(\"The model BAAI/bge-small-en-v1.5 is ready to use.\")\n",
+    "\n",
+    "embeddings_generator = embedding_model.embed(documents)  # reminder this is a generator\n",
+    "embeddings_list = list(embedding_model.embed(documents))\n",
+    "# you can also convert the generator to a list, and that to a numpy array\n",
+    "len(embeddings_list[0])  # Vector of 384 dimensions"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d772190b",
+   "metadata": {},
+   "source": [
+    "> 💡 **Why do we use generators?**\n",
+    "> \n",
+    "> We use them to save memory mostly. Instead of loading all the vectors into memory, we can load them one by one. This is useful when you have a large dataset and you don't want to load all the vectors at once."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "8a225cb8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Document: This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.\n",
+      "Vector of type: <class 'numpy.ndarray'> with shape: (384,)\n",
+      "Document: fastembed is supported by and maintained by Qdrant.\n",
+      "Vector of type: <class 'numpy.ndarray'> with shape: (384,)\n"
+     ]
+    }
+   ],
+   "source": [
+    "embeddings_generator = embedding_model.embed(documents)  # reminder this is a generator\n",
+    "\n",
+    "for doc, vector in zip(documents, embeddings_generator):\n",
+    "    print(\"Document:\", doc)\n",
+    "    print(f\"Vector of type: {type(vector)} with shape: {vector.shape}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "769a1be9",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(2, 384)"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "embeddings_list = np.array(\n",
+    "    list(embedding_model.embed(documents))\n",
+    ")  # you can also convert the generator to a list, and that to a numpy array\n",
+    "embeddings_list.shape"
    ]
   },
   {
@@ -63,128 +149,74 @@
    "id": "8c49ae50",
    "metadata": {},
    "source": [
-    "## Let's think step by step"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "92cf4b76",
-   "metadata": {},
-   "source": [
-    "### Setup\n",
-    "\n",
-    "Importing the required classes and modules:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "c0a6f634",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from typing import List\n",
-    "import numpy as np\n",
-    "from fastembed.embedding import DefaultEmbedding as Embedding"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "3fd03a71",
-   "metadata": {},
-   "source": [
-    "Notice that we are using the DefaultEmbedding -- which is a quantized, state of the Art Flag Embedding model which beats OpenAI's Embedding by a large margin. \n",
-    "\n",
-    "### Prepare your Documents\n",
-    "You can define a list of documents that you'd like to embed. These can be sentences, paragraphs, or even entire documents. \n",
+    "We're using [BAAI/bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5) a state of the art Flag Embedding model. The model does better than OpenAI text-embedding-ada-002. We've made it even faster by converting it to ONNX format and quantizing the model for you.\n",
     "\n",
     "#### Format of the Document List\n",
+    "\n",
     "1. List of Strings: Your documents must be in a list, and each document must be a string\n",
-    "2. For Retrieval Tasks: If you're working with queries and passages, you can add special labels to them:\n",
+    "2. For Retrieval Tasks with our default: If you're working with queries and passages, you can add special labels to them:\n",
     "- **Queries**: Add \"query:\" at the beginning of each query string\n",
-    "- **Passages**: Add \"passage:\" at the beginning of each passage string"
+    "- **Passages**: Add \"passage:\" at the beginning of each passage string\n",
+    "\n",
+    "## Beyond the default model\n",
+    "\n",
+    "The default model is built for speed and efficiency. If you need a more accurate model, you can use the `TextEmbedding` class to load any model from our list of available models. You can find the list of available models using `TextEmbedding.list_supported_models()`."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "id": "145a56ce",
+   "execution_count": 5,
+   "id": "2e9c8766",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9470ec542f3c4400a42452c2489a1abc",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Fetching 8 files:   0%|          | 0/8 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
-    "# Example list of documents\n",
-    "documents: List[str] = [\n",
-    "    \"passage: Hello, World!\",\n",
-    "    \"query: Hello, World!\",  # these are two different embedding\n",
-    "    \"passage: This is an example passage.\",\n",
-    "    # You can leave out the prefix but it's recommended\n",
-    "    \"fastembed is supported by and maintained by Qdrant.\",\n",
-    "]"
+    "multilingual_large_model = TextEmbedding(\"intfloat/multilingual-e5-large\")  # This can take a few minutes to download"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "a9e70f0e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(4, 1024)"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.array(\n",
+    "    list(multilingual_large_model.embed([\"Hello, world!\", \"你好世界\", \"¡Hola Mundo!\", \"नमस्ते!\"]))\n",
+    ").shape  # Vector of 1024 dimensions"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "1cb3cc87",
+   "id": "64fe20ed",
    "metadata": {},
    "source": [
-    "### Load the Embedding Model Weights\n",
-    "Next, initialize the Embedding class with the desired parameters. Here, \"BAAI/bge-small-en\" is the pre-trained model name, and max_length=512 is the maximum token length for each document.\n",
-    "\n",
-    "This will download the model weights, decompress to directory `local_cache` and load them into the Embedding class.\n",
-    "\n",
-    "#### Initialize DefaultEmbedding\n",
-    "\n",
-    "We will initialize Flag Embeddings with the model name and the maximum token length. That is the DefaultEmbedding class with the model name \"BAAI/bge-small-en\" and max_length=512."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "272c8915",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "embedding_model = DefaultEmbedding()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "5549d501",
-   "metadata": {},
-   "source": [
-    "### Embed your Documents\n",
-    "\n",
-    "Use the embed method of the embedding model to transform the documents into a List of np.array. The method returns a generator, so we cast it to a list to get the embeddings."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "8013eee9",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "embeddings: List[np.ndarray] = list(\n",
-    "    embedding_model.embed(documents)\n",
-    ")  # notice that we are casting the generator to a list"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e5b5a6ad",
-   "metadata": {},
-   "source": [
-    "You can print the shape of the embeddings to understand their dimensions. Typically, the shape will indicate the number of dimensions in the vector."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "0d8c8e08",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "print(embeddings[0].shape)  # (384,) or similar output"
+    "Next: Checkout how to use FastEmbed with Qdrant for similarity search: [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/)"
    ]
   }
  ],
@@ -204,7 +236,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

docs/examples/Retrieval_with_FastEmbed.ipynb

@@ -21,7 +21,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -37,13 +37,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 2,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2024-02-07 22:20:57.013\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36mfastembed.embedding\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m7\u001b[0m - \u001b[33m\u001b[1mDefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated. Use TextEmbedding instead.\u001b[0m\n"
+     ]
+    }
+   ],
    "source": [
     "from typing import List\n",
     "import numpy as np\n",
-    "from fastembed.embedding import FlagEmbedding as Embedding"
+    "from fastembed import TextEmbedding"
    ]
   },
   {
@@ -58,7 +66,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -84,7 +92,7 @@
     "    \"His life has been depicted in various films, TV shows, and books\",\n",
     "]\n",
     "# Initialize the DefaultEmbedding class with the desired parameters\n",
-    "embedding_model = Embedding(model_name=\"BAAI/bge-small-en\", max_length=512)\n",
+    "embedding_model = TextEmbedding(model_name=\"BAAI/bge-small-en\", max_length=512)\n",
     "\n",
     "# We'll use the passage_embed method to get the embeddings for the documents\n",
     "embeddings: List[np.ndarray] = list(\n",
@@ -105,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -124,65 +132,27 @@
     "        print(f\"Rank {i+1}: {documents[sorted_scores[i]]}\")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Running and Comparing Queries\n",
-    "Finally, we run our sample query using the `print_top_k` function.\n",
-    "\n",
-    "The differences between using query embeddings and plain embeddings can be observed in the retrieved ranks:\n",
-    "\n",
-    "Using query embeddings (from `query_embed` method):"
-   ]
-  },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Rank 1: Maharana Pratap was a Rajput warrior king from Mewar\n",
-      "Rank 2: Maharana Pratap is considered a symbol of Rajput resistance against foreign rule\n",
-      "Rank 3: His legacy is celebrated in Rajasthan through festivals and monuments\n",
-      "Rank 4: His capital was Chittorgarh, which he lost to the Mughals\n",
-      "Rank 5: He fought against the Mughal Empire led by Akbar\n"
-     ]
+     "data": {
+      "text/plain": [
+       "(array([-0.04393955,  0.04452892, -0.00760788, -0.03399807,  0.01951348],\n",
+       "       dtype=float32),\n",
+       " array([-0.06002192,  0.04322132, -0.00545516, -0.04419701, -0.00542277],\n",
+       "       dtype=float32))"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
    "source": [
-    "print_top_k(query_embedding, embeddings, documents)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Using plain embeddings (from `embed` method):"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Rank 1: He died in 1597 at the age of 57\n",
-      "Rank 2: His life has been depicted in various films, TV shows, and books\n",
-      "Rank 3: Maharana Pratap was a Rajput warrior king from Mewar\n",
-      "Rank 4: He had 11 wives and 17 sons, including Amar Singh I who succeeded him as ruler of Mewar\n",
-      "Rank 5: He fought against the Mughal Empire led by Akbar\n"
-     ]
-    }
-   ],
-   "source": [
-    "print_top_k(plain_query_embedding, embeddings, documents)"
+    "query_embedding[:5], plain_query_embedding[:5]"
    ]
   },
   {
@@ -213,7 +183,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.11.5"
   },
   "orig_nbformat": 4
  },

docs/examples/Supported_Models.ipynb

@@ -0,0 +1,289 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "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",
+       "      <th>sources</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>BAAI/bge-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Base English model</td>\n",
+       "      <td>0.50</td>\n",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</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",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz', 'hf': 'qdrant/bge-base-en-v1.5-onnx-q'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>BAAI/bge-large-en-v1.5-quantized</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Large English model, v1.5</td>\n",
+       "      <td>1.34</td>\n",
+       "      <td>{'hf': 'qdrant/bge-large-en-v1.5-onnx-q'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>BAAI/bge-large-en-v1.5</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Large English model, v1.5</td>\n",
+       "      <td>1.34</td>\n",
+       "      <td>{'hf': 'qdrant/bge-large-en-v1.5-onnx'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>BAAI/bge-small-en</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Fast English model</td>\n",
+       "      <td>0.20</td>\n",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</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",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz', 'hf': 'qdrant/bge-small-en-v1.5-onnx-q'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</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",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</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",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz', 'hf': 'qdrant/all-MiniLM-L6-v2-onnx'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>8</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
+       "      <td>768</td>\n",
+       "      <td>8192 context length english model</td>\n",
+       "      <td>0.54</td>\n",
+       "      <td>{'hf': 'nomic-ai/nomic-embed-text-v1'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>9</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
+       "      <td>768</td>\n",
+       "      <td>8192 context length english model</td>\n",
+       "      <td>0.54</td>\n",
+       "      <td>{'hf': 'nomic-ai/nomic-embed-text-v1.5'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>10</th>\n",
+       "      <td>thenlper/gte-large</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Large general text embeddings model</td>\n",
+       "      <td>1.34</td>\n",
+       "      <td>{'hf': 'qdrant/gte-large-onnx'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>11</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",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz', 'hf': 'qdrant/multilingual-e5-large-onnx'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>12</th>\n",
+       "      <td>sentence-transformers/paraphrase-multilingual-mpnet-base-v2</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Sentence-transformers model for tasks like clustering or semantic search</td>\n",
+       "      <td>1.11</td>\n",
+       "      <td>{'hf': 'xenova/paraphrase-multilingual-mpnet-base-v2'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>13</th>\n",
+       "      <td>sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2</td>\n",
+       "      <td>0.46</td>\n",
+       "      <td>{'hf': 'qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>14</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",
+       "      <td>{'hf': 'xenova/jina-embeddings-v2-base-en'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>15</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",
+       "      <td>{'hf': 'xenova/jina-embeddings-v2-small-en'}</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                                          model   dim  \\\n",
+       "0                                              BAAI/bge-base-en   768   \n",
+       "1                                         BAAI/bge-base-en-v1.5   768   \n",
+       "2                              BAAI/bge-large-en-v1.5-quantized  1024   \n",
+       "3                                        BAAI/bge-large-en-v1.5  1024   \n",
+       "4                                             BAAI/bge-small-en   384   \n",
+       "5                                        BAAI/bge-small-en-v1.5   384   \n",
+       "6                                        BAAI/bge-small-zh-v1.5   512   \n",
+       "7                        sentence-transformers/all-MiniLM-L6-v2   384   \n",
+       "8                                  nomic-ai/nomic-embed-text-v1   768   \n",
+       "9                                nomic-ai/nomic-embed-text-v1.5   768   \n",
+       "10                                           thenlper/gte-large  1024   \n",
+       "11                               intfloat/multilingual-e5-large  1024   \n",
+       "12  sentence-transformers/paraphrase-multilingual-mpnet-base-v2   768   \n",
+       "13  sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2   384   \n",
+       "14                            jinaai/jina-embeddings-v2-base-en   768   \n",
+       "15                           jinaai/jina-embeddings-v2-small-en   512   \n",
+       "\n",
+       "                                                                           description  \\\n",
+       "0                                                                   Base English model   \n",
+       "1                                                             Base English model, v1.5   \n",
+       "2                                                            Large English model, v1.5   \n",
+       "3                                                            Large English model, v1.5   \n",
+       "4                                                                   Fast English model   \n",
+       "5                                                       Fast and Default English model   \n",
+       "6                                                   Fast and recommended Chinese model   \n",
+       "7                                             Sentence Transformer model, MiniLM-L6-v2   \n",
+       "8                                                    8192 context length english model   \n",
+       "9                                                    8192 context length english model   \n",
+       "10                                                 Large general text embeddings model   \n",
+       "11  Multilingual model, e5-large. Recommend using this model for non-English languages   \n",
+       "12            Sentence-transformers model for tasks like clustering or semantic search   \n",
+       "13                   Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2   \n",
+       "14                             English embedding model supporting 8192 sequence length   \n",
+       "15                             English embedding model supporting 8192 sequence length   \n",
+       "\n",
+       "    size_in_GB  \\\n",
+       "0         0.50   \n",
+       "1         0.44   \n",
+       "2         1.34   \n",
+       "3         1.34   \n",
+       "4         0.20   \n",
+       "5         0.13   \n",
+       "6         0.10   \n",
+       "7         0.09   \n",
+       "8         0.54   \n",
+       "9         0.54   \n",
+       "10        1.34   \n",
+       "11        2.24   \n",
+       "12        1.11   \n",
+       "13        0.46   \n",
+       "14        0.55   \n",
+       "15        0.13   \n",
+       "\n",
+       "                                                                                                                                           sources  \n",
+       "0                                                               {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz'}  \n",
+       "1                  {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz', 'hf': 'qdrant/bge-base-en-v1.5-onnx-q'}  \n",
+       "2                                                                                                        {'hf': 'qdrant/bge-large-en-v1.5-onnx-q'}  \n",
+       "3                                                                                                          {'hf': 'qdrant/bge-large-en-v1.5-onnx'}  \n",
+       "4                                                              {'url': 'https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz'}  \n",
+       "5                {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz', 'hf': 'qdrant/bge-small-en-v1.5-onnx-q'}  \n",
+       "6                                                         {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz'}  \n",
+       "7   {'url': 'https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz', 'hf': 'qdrant/all-MiniLM-L6-v2-onnx'}  \n",
+       "8                                                                                                           {'hf': 'nomic-ai/nomic-embed-text-v1'}  \n",
+       "9                                                                                                         {'hf': 'nomic-ai/nomic-embed-text-v1.5'}  \n",
+       "10                                                                                                                 {'hf': 'qdrant/gte-large-onnx'}  \n",
+       "11         {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz', 'hf': 'qdrant/multilingual-e5-large-onnx'}  \n",
+       "12                                                                                          {'hf': 'xenova/paraphrase-multilingual-mpnet-base-v2'}  \n",
+       "13                                                                                   {'hf': 'qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q'}  \n",
+       "14                                                                                                     {'hf': 'xenova/jina-embeddings-v2-base-en'}  \n",
+       "15                                                                                                    {'hf': 'xenova/jina-embeddings-v2-small-en'}  "
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from fastembed import TextEmbedding\n",
+    "import pandas as pd\n",
+    "\n",
+    "pd.set_option(\"display.max_colwidth\", None)\n",
+    "pd.DataFrame(TextEmbedding.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.10.13"
+  },
+  "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,13 +42,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
     "from typing import List\n",
-    "import numpy as np\n",
-    "from fastembed.embedding import FlagEmbedding as Embedding\n",
     "from qdrant_client import QdrantClient"
    ]
   },
@@ -68,7 +62,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -105,25 +99,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 77.7M/77.7M [00:05<00:00, 14.6MiB/s]\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "['901ff9d7a90a4e56afe655d1de8e7c06',\n",
-       " '7bae5aa894164398a9be68d47d72ff7a',\n",
-       " 'd940bba124e24469ae3166c3110c62f8',\n",
-       " 'c8dcf956a4f444c6bfe69a00cbeb85ce',\n",
-       " 'd2aeb24c51f549c5b21851be05cb4048',\n",
-       " '6d4c672bafef4db68ae72c8986ee8a72',\n",
-       " 'ff198613768a4361a6d8230b40ee7f78',\n",
-       " 'a133ad72876b48d48e716533c2d78cf0',\n",
-       " 'bdf2b016136e4f57b7ac4b65534031b8',\n",
-       " '5f4f279f304c47839e1a74b0f28de136']"
+       "['4fa8b10c78da4b18ba0830ba8a57367a',\n",
+       " '2eae04b515ee4e9185a9a0e6be812bba',\n",
+       " 'c6039f88486f47f1835ae3b069c5823c',\n",
+       " 'c2c8c51e305144d1917b373125fb4d95',\n",
+       " '79fd23b9ec0648cdab38d1947c6b933e',\n",
+       " '036aa200d8c3492b8a438e4f825f5e7f',\n",
+       " 'c35c77f3ea37460a9a13723fb77b7367',\n",
+       " '6ebccbca571b40d0ab6e83e5e0f2f562',\n",
+       " '38048c2ccc1d4962a4f8f1bd89c8357a',\n",
+       " 'c6b09308360140c7b4f106af3658a31e']"
       ]
      },
-     "execution_count": 17,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -150,38 +151,56 @@
   },
   {
    "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(collection_name=\"demo_collection\", documents=docs, metadata=metadata, ids=ids)"
+   ]
+  },
+  {
+   "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.8276550115796268), 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.8265536935180283)]\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(collection_name=\"demo_collection\", query_text=\"This is a query document\")\n",
+    "print(search_result)"
    ]
   },
   {
@@ -214,7 +233,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.11.5"
   },
   "orig_nbformat": 4
  },

docs/experimental/Binary Quantization from Scratch.ipynb

@@ -46,7 +46,6 @@
    "source": [
     "import numpy as np\n",
     "import pandas as pd\n",
-    "import matplotlib.pyplot as plt\n",
     "from tqdm import tqdm"
    ]
   },
@@ -305,9 +304,11 @@
     "sampling_rate = [1, 2, 3, 5]\n",
     "results = []\n",
     "\n",
+    "\n",
     "def mean_accuracy(number_of_samples, limit, sampling_rate):\n",
     "    return np.mean([accuracy(i, limit=limit, oversampling=sampling_rate) for i in range(number_of_samples)])\n",
     "\n",
+    "\n",
     "for i in tqdm(sampling_rate):\n",
     "    for j in tqdm(limits):\n",
     "        result = {\"sampling_rate\": i, \"limit\": j, \"recall\": mean_accuracy(number_of_samples, j, i)}\n",

docs/experimental/Binary Quantization with Qdrant.ipynb

@@ -34,10 +34,8 @@
    "outputs": [],
    "source": [
     "import pandas as pd\n",
-    "import uuid\n",
     "from qdrant_client import QdrantClient\n",
-    "from qdrant_client.http import models\n",
-    "from qdrant_client.http.models import PointStruct"
+    "from qdrant_client.http import models"
    ]
   },
   {
@@ -46,7 +44,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."
    ]
   },
   {
@@ -71,6 +69,7 @@
    ],
    "source": [
     "import datasets\n",
+    "\n",
     "dataset = datasets.load_dataset(\"KShivendu/dbpedia-entities-openai-1M\", split=\"train[0:100000]\")"
    ]
   },
@@ -133,10 +132,8 @@
     }
    ],
    "source": [
-    "from qdrant_client import QdrantClient\n",
-    "\n",
     "# client = QdrantClient(\n",
-    "#     url=\"https://2aaa9439-b209-4ba6-8beb-d0b61dbd9388.us-east-1-0.aws.cloud.qdrant.io:6333\", \n",
+    "#     url=\"https://2aaa9439-b209-4ba6-8beb-d0b61dbd9388.us-east-1-0.aws.cloud.qdrant.io:6333\",\n",
     "#     api_key=\"FCF8_ADVuSRrtNGeg_rBJvAMJecEDgQhzuXMZGW8F7OzvaC9wYOPeQ\",\n",
     "#     prefer_grpc=True\n",
     "# )\n",
@@ -175,12 +172,10 @@
     "bs = 10000\n",
     "for i in range(0, len(dataset), bs):\n",
     "    client.upload_collection(\n",
-    "        collection_name=collection_name, \n",
-    "        ids=range(i, i+bs),\n",
-    "        vectors=dataset[i:i+bs][\"openai\"],\n",
-    "        payload=[\n",
-    "            {\"text\": x} for x in dataset[i:i+bs][\"text\"]\n",
-    "        ],\n",
+    "        collection_name=collection_name,\n",
+    "        ids=range(i, i + bs),\n",
+    "        vectors=dataset[i : i + bs][\"openai\"],\n",
+    "        payload=[{\"text\": x} for x in dataset[i : i + bs][\"text\"]],\n",
     "        parallel=10,\n",
     "    )"
    ]
@@ -203,10 +198,7 @@
    ],
    "source": [
     "client.update_collection(\n",
-    "    collection_name=f\"{collection_name}\",\n",
-    "    optimizer_config=models.OptimizersConfigDiff(\n",
-    "        indexing_threshold=20000\n",
-    "    )\n",
+    "    collection_name=f\"{collection_name}\", optimizer_config=models.OptimizersConfigDiff(indexing_threshold=20000)\n",
     ")"
    ]
   },
@@ -289,6 +281,7 @@
    "source": [
     "import random\n",
     "from random import randint\n",
+    "\n",
     "random.seed(37)\n",
     "\n",
     "query_indices = [randint(0, len(dataset)) for _ in range(100)]\n",
@@ -304,7 +297,10 @@
    "source": [
     "## Add Gaussian noise to any vector\n",
     "import numpy as np\n",
+    "\n",
     "np.random.seed(37)\n",
+    "\n",
+    "\n",
     "def add_noise(vector, noise=0.05):\n",
     "    return vector + noise * np.random.randn(*vector.shape)"
    ]
@@ -959,6 +955,8 @@
    ],
    "source": [
     "import time\n",
+    "\n",
+    "\n",
     "def correct(results, text):\n",
     "    result_texts = [x.payload[\"text\"] for x in results]\n",
     "    return text in result_texts\n",
@@ -977,7 +975,7 @@
     "                    rescore=rescore,\n",
     "                    oversampling=oversampling,\n",
     "                )\n",
-    "            )\n",
+    "            ),\n",
     "        )\n",
     "        correct_results += correct(results, text)\n",
     "    return correct_results\n",
@@ -996,14 +994,16 @@
     "            start = time.time()\n",
     "            correct_results = count_correct(query_dataset, limit=limit, oversampling=oversampling, rescore=rescore)\n",
     "            end = time.time()\n",
-    "            results.append({\n",
-    "                \"limit\": limit,\n",
-    "                \"oversampling\": oversampling,\n",
-    "                \"rescore\": rescore,\n",
-    "                \"correct\": correct_results,\n",
-    "                \"total queries\": len(query_dataset[\"text\"]),\n",
-    "                \"time\": end - start,\n",
-    "            })\n",
+    "            results.append(\n",
+    "                {\n",
+    "                    \"limit\": limit,\n",
+    "                    \"oversampling\": oversampling,\n",
+    "                    \"rescore\": rescore,\n",
+    "                    \"correct\": correct_results,\n",
+    "                    \"total queries\": len(query_dataset[\"text\"]),\n",
+    "                    \"time\": end - start,\n",
+    "                }\n",
+    "            )\n",
     "\n",
     "results_df = pd.DataFrame(results)\n",
     "results_df"

docs/index.md

@@ -1,12 +1,21 @@
 # ⚡️ 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](https://github.com/huggingface/optimum)
 
-To install the FastEmbed library, pip works: 
+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:
 
 ```bash
 pip install fastembed
@@ -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

@@ -5,7 +5,7 @@
     <a href="{{ page.nb_url }}" title="Download Notebook" class="md-content__button md-icon jp-DownloadNB">
         {% include ".icons/material/download.svg" %}
     </a>
-{% endif %} 
+{% endif %}
 
 {{ super() }}
 
@@ -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 %}
+{% endblock %}

experiments/01_ONNX_Port.ipynb

@@ -27,7 +27,6 @@
     "from transformers import AutoTokenizer, AutoModel\n",
     "\n",
     "from optimum.onnxruntime import AutoOptimizationConfig, ORTModelForFeatureExtraction, ORTOptimizer\n",
-    "from optimum.onnxruntime.configuration import OptimizationConfig\n",
     "from optimum.pipelines import pipeline\n",
     "import torch.nn.functional as F"
    ]
@@ -149,7 +148,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "onnx_quant_embed = pipeline(\"feature-extraction\", model=model, accelerator=\"ort\", tokenizer=tokenizer,return_tensors=True)"
+    "onnx_quant_embed = pipeline(\n",
+    "    \"feature-extraction\", model=model, accelerator=\"ort\", tokenizer=tokenizer, return_tensors=True\n",
+    ")"
    ]
   },
   {
@@ -159,9 +160,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import torch\n",
     "embeddings = onnx_quant_embed(inputs=english_texts)\n",
-    "F.normalize(embeddings[4])[:,0], english_texts[4], len(embeddings), len(english_texts)"
+    "F.normalize(embeddings[4])[:, 0], english_texts[4], len(embeddings), len(english_texts)"
    ]
   },
   {
@@ -171,7 +171,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
     "def measure_pipeline_time(pipeline, input_texts: List[str], num_runs=10, **kwargs: Any) -> Tuple[float, float]:\n",
     "    \"\"\"Measures the time it takes to run the pipeline on the input texts.\"\"\"\n",
     "    times = []\n",
@@ -256,6 +255,7 @@
     "\n",
     "save_dir = Path(\"../local_cache/fast-bge-small-en-v1.5\")\n",
     "\n",
+    "\n",
     "def compress(directory_path):\n",
     "    directory_path = Path(directory_path)\n",
     "    assert directory_path.exists(), f\"{directory_path} does not exist\"\n",
@@ -304,9 +304,9 @@
     }
    ],
    "source": [
-    "import os\n",
     "from google.cloud import storage\n",
     "\n",
+    "\n",
     "def upload(bucket_name, source_file_path):\n",
     "    storage_client = storage.Client(project=\"main\")\n",
     "    bucket = storage_client.bucket(bucket_name)\n",

experiments/02_SPLADE_to_ONNX.ipynb

@@ -0,0 +1,371 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import torch\n",
+    "from transformers import AutoModelForMaskedLM, AutoTokenizer"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Running the model with Transformers and Torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sentences = [\n",
+    "    \"Hello World\",\n",
+    "    \"Built by Nirant Kasliwal\",\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## PyTorch Code from the [SPLADERunner](https://github.com/PrithivirajDamodaran/SPLADERunner) library"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hf_token = \"<your_hf_token_here>\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Output Logits shape:  torch.Size([2, 10, 30522])\n",
+      "Output Attention mask shape:  torch.Size([2, 10])\n",
+      "Sparse Vector shape:  torch.Size([2, 30522])\n",
+      "SPLADE BOW rep for sentence:\tBuilt by Nirant Kasliwal\n",
+      "[('##rant', 2.02), ('built', 1.94), ('##wal', 1.79), ('##sl', 1.69), ('build', 1.57), ('ka', 1.4), ('ni', 1.26), ('made', 0.93), ('architect', 0.76), ('was', 0.69), ('who', 0.61), ('his', 0.5), ('wrote', 0.47), ('india', 0.45), ('company', 0.41), ('##i', 0.41), ('he', 0.37), ('manufacturer', 0.36), ('by', 0.35), ('engineer', 0.33), ('architecture', 0.33), ('ko', 0.23), ('him', 0.22), ('invented', 0.19), ('said', 0.14), ('k', 0.11), ('man', 0.11), ('statue', 0.11), ('bomb', 0.1), ('##wa', 0.1), ('builder', 0.09), ('.', 0.07), ('started', 0.06), (',', 0.04), ('ku', 0.03)]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Download the model and tokenizer\n",
+    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"prithivida/Splade_PP_en_v1\", token=hf_token)\n",
+    "reverse_voc = {v: k for k, v in tokenizer.vocab.items()}\n",
+    "model = AutoModelForMaskedLM.from_pretrained(\"prithivida/Splade_PP_en_v1\", token=hf_token)\n",
+    "model.to(device)\n",
+    "\n",
+    "# Tokenize the input\n",
+    "inputs = tokenizer(sentences, return_tensors=\"pt\", padding=True, truncation=True, max_length=512)\n",
+    "inputs = {key: val.to(device) for key, val in inputs.items()}\n",
+    "input_ids = inputs[\"input_ids\"]\n",
+    "attention_mask = inputs[\"attention_mask\"]\n",
+    "token_type_ids = inputs[\"token_type_ids\"]\n",
+    "\n",
+    "# Run model and prepare sparse vector\n",
+    "outputs = model(**inputs)\n",
+    "logits = outputs.logits\n",
+    "print(\"Output Logits shape: \", logits.shape)\n",
+    "print(\"Output Attention mask shape: \", attention_mask.shape)\n",
+    "relu_log = torch.log(1 + torch.relu(logits))\n",
+    "weighted_log = relu_log * attention_mask.unsqueeze(-1)\n",
+    "max_val, _ = torch.max(weighted_log, dim=1)\n",
+    "vector = max_val.squeeze()\n",
+    "print(\"Sparse Vector shape: \", vector.shape)\n",
+    "# print(\"Number of Actual Dimensions: \", len(cols))\n",
+    "cols = [vec.nonzero().squeeze().cpu().tolist() for vec in vector]\n",
+    "weights = [vec[col].cpu().tolist() for vec, col in zip(vector, cols)]\n",
+    "\n",
+    "idx = 1\n",
+    "cols, weights = cols[idx], weights[idx]\n",
+    "# Print the BOW representation\n",
+    "d = {k: v for k, v in zip(cols, weights)}\n",
+    "sorted_d = {k: v for k, v in sorted(d.items(), key=lambda item: item[1], reverse=True)}\n",
+    "bow_rep = []\n",
+    "for k, v in sorted_d.items():\n",
+    "    bow_rep.append((reverse_voc[k], round(v, 2)))\n",
+    "print(f\"SPLADE BOW rep for sentence:\\t{sentences[idx]}\\n{bow_rep}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Export with output_attentions and logits"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Exporting model to models/nirantk_SPLADE_PP_en_v1\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "('models/nirantk_SPLADE_PP_en_v1/tokenizer_config.json',\n",
+       " 'models/nirantk_SPLADE_PP_en_v1/special_tokens_map.json',\n",
+       " 'models/nirantk_SPLADE_PP_en_v1/vocab.txt',\n",
+       " 'models/nirantk_SPLADE_PP_en_v1/added_tokens.json',\n",
+       " 'models/nirantk_SPLADE_PP_en_v1/tokenizer.json')"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from transformers import AutoTokenizer\n",
+    "\n",
+    "model_id = \"nirantk/SPLADE_PP_en_v1\"\n",
+    "output_dir = f\"models/{model_id.replace('/', '_')}\"\n",
+    "model_kwargs = {\"output_attentions\": True, \"return_dict\": True}\n",
+    "\n",
+    "print(f\"Exporting model to {output_dir}\")\n",
+    "tokenizer.save_pretrained(output_dir)\n",
+    "# main_export(\n",
+    "#     model_id,\n",
+    "#     output=output_dir,\n",
+    "#     no_post_process=True,\n",
+    "#     model_kwargs=model_kwargs,\n",
+    "#     token=hf_token,\n",
+    "# )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Running the model with ONNX"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from optimum.onnxruntime import ORTModelForMaskedLM\n",
+    "\n",
+    "model = ORTModelForMaskedLM.from_pretrained(\"nirantk/SPLADE_PP_en_v1\")\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"nirantk/SPLADE_PP_en_v1\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inputs = tokenizer(sentences, return_tensors=\"pt\", padding=True, truncation=True, max_length=512)\n",
+    "inputs = {key: val.to(device) for key, val in inputs.items()}\n",
+    "input_ids = inputs[\"input_ids\"]\n",
+    "attention_mask = inputs[\"attention_mask\"]\n",
+    "token_type_ids = inputs[\"token_type_ids\"]\n",
+    "\n",
+    "onnx_input = {\n",
+    "    \"input_ids\": input_ids.cpu().numpy(),\n",
+    "    \"attention_mask\": attention_mask.cpu().numpy(),\n",
+    "    \"token_type_ids\": token_type_ids.cpu().numpy(),\n",
+    "}\n",
+    "\n",
+    "logits = model(**onnx_input).logits"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(2, 10, 30522)"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "logits.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Output Logits shape:  (2, 10, 30522)\n",
+      "Sparse Vector shape:  (2, 30522)\n",
+      "SPLADE BOW rep for sentence:\tBuilt by Nirant Kasliwal\n",
+      "[('##rant', 2.02), ('built', 1.94), ('##wal', 1.79), ('##sl', 1.69), ('build', 1.57), ('ka', 1.4), ('ni', 1.26), ('made', 0.93), ('architect', 0.76), ('was', 0.69), ('who', 0.61), ('his', 0.5), ('wrote', 0.47), ('india', 0.45), ('company', 0.41), ('##i', 0.41), ('he', 0.37), ('manufacturer', 0.36), ('by', 0.35), ('engineer', 0.33), ('architecture', 0.33), ('ko', 0.23), ('him', 0.22), ('invented', 0.19), ('said', 0.14), ('k', 0.11), ('man', 0.11), ('statue', 0.11), ('bomb', 0.1), ('##wa', 0.1), ('builder', 0.09), ('.', 0.07), ('started', 0.06), (',', 0.04), ('ku', 0.03)]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"Output Logits shape: \", logits.shape)\n",
+    "\n",
+    "relu_log = np.log(1 + np.maximum(logits, 0))\n",
+    "\n",
+    "# Equivalent to relu_log * attention_mask.unsqueeze(-1)\n",
+    "# For NumPy, you might need to explicitly expand dimensions if 'attention_mask' is not already 2D\n",
+    "weighted_log = relu_log * np.expand_dims(attention_mask, axis=-1)\n",
+    "\n",
+    "# Equivalent to torch.max(weighted_log, dim=1)\n",
+    "# NumPy's max function returns only the max values, not the indices, so we don't need to unpack two values\n",
+    "max_val = np.max(weighted_log, axis=1)\n",
+    "\n",
+    "# Equivalent to max_val.squeeze()\n",
+    "# This step may be unnecessary in NumPy if max_val doesn't have unnecessary dimensions\n",
+    "vector = np.squeeze(max_val)\n",
+    "print(\"Sparse Vector shape: \", vector.shape)\n",
+    "\n",
+    "# print(vector[0].nonzero())\n",
+    "\n",
+    "cols = [vec.nonzero()[0].squeeze().tolist() for vec in vector]\n",
+    "weights = [vec[col].tolist() for vec, col in zip(vector, cols)]\n",
+    "\n",
+    "idx = 1\n",
+    "cols, weights = cols[idx], weights[idx]\n",
+    "# Print the BOW representation\n",
+    "d = {k: v for k, v in zip(cols, weights)}\n",
+    "sorted_d = {k: v for k, v in sorted(d.items(), key=lambda item: item[1], reverse=True)}\n",
+    "bow_rep = []\n",
+    "for k, v in sorted_d.items():\n",
+    "    bow_rep.append((reverse_voc[k], round(v, 2)))\n",
+    "print(f\"SPLADE BOW rep for sentence:\\t{sentences[idx]}\\n{bow_rep}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "35"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "len(cols)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[1010,\n",
+       " 1012,\n",
+       " 1047,\n",
+       " 2001,\n",
+       " 2002,\n",
+       " 2010,\n",
+       " 2011,\n",
+       " 2032,\n",
+       " 2040,\n",
+       " 2056,\n",
+       " 2072,\n",
+       " 2081,\n",
+       " 2158,\n",
+       " 2194,\n",
+       " 2318,\n",
+       " 2328,\n",
+       " 2626,\n",
+       " 2634,\n",
+       " 3857,\n",
+       " 3992,\n",
+       " 4213,\n",
+       " 4294,\n",
+       " 4944,\n",
+       " 5968,\n",
+       " 6231,\n",
+       " 7751,\n",
+       " 8826,\n",
+       " 9152,\n",
+       " 10556,\n",
+       " 12508,\n",
+       " 12849,\n",
+       " 13476,\n",
+       " 13970,\n",
+       " 14540,\n",
+       " 17884]"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "cols"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

experiments/attention_export.py

@@ -0,0 +1,15 @@
+from optimum.exporters.onnx import main_export
+from transformers import AutoTokenizer
+
+model_id = "sentence-transformers/paraphrase-MiniLM-L6-v2"
+output_dir = f"models/{model_id.replace('/', '_')}"
+model_kwargs = {"output_attentions": True, "return_dict": True}
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+# export if the output model does not exist
+# try:
+#     sess = onnxruntime.InferenceSession(f"{output_dir}/model.onnx")
+#     print("Model already exported")
+# except FileNotFoundError:
+print(f"Exporting model to {output_dir}")
+main_export(model_id, output=output_dir, no_post_process=True, model_kwargs=model_kwargs)

experiments/try_attention_export.py

@@ -0,0 +1,29 @@
+import numpy as np
+import onnx
+import onnxruntime
+from transformers import AutoTokenizer
+
+model_id = "sentence-transformers/paraphrase-MiniLM-L6-v2"
+output_dir = f"models/{model_id.replace('/', '_')}"
+model_kwargs = {"output_attentions": True, "return_dict": True}
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+model_path = f"{output_dir}/model.onnx"
+onnx_model = onnx.load(model_path)
+ort_session = onnxruntime.InferenceSession(model_path)
+text = "This is a test sentence"
+tokenizer_output = tokenizer(text, return_tensors="np")
+input_ids = tokenizer_output["input_ids"]
+attention_mask = tokenizer_output["attention_mask"]
+print(attention_mask)
+# Prepare the input
+input_ids = np.array(input_ids).astype(np.int64)  # Replace your_input_ids with actual input data
+
+# Run the ONNX model
+outputs = ort_session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
+
+# Get the attention weights
+attentions = outputs[-1]
+
+# Print the attention weights for the first layer and first head
+print(attentions[0][0])

fastembed/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.text.text_embedding import TextEmbedding
+
+__all__ = ["TextEmbedding"]

fastembed/common/model_management.py

@@ -0,0 +1,226 @@
+import os
+import shutil
+import tarfile
+from pathlib import Path
+from typing import List, Optional, Dict, Any
+
+import requests
+from huggingface_hub import snapshot_download
+from huggingface_hub.utils import RepositoryNotFoundError
+from tqdm import tqdm
+from loguru import logger
+
+
+def locate_model_file(model_dir: Path, file_names: List[str]) -> Path:
+    """
+    Find model path for both TransformerJS style `onnx`  subdirectory structure and direct model weights structure used
+    by Optimum and Qdrant
+    """
+    if not model_dir.is_dir():
+        raise ValueError(f"Provided model path '{model_dir}' is not a directory.")
+
+    for file_name in file_names:
+        file_paths = [path for path in model_dir.rglob(file_name) if path.is_file()]
+
+        if file_paths:
+            return file_paths[0]
+
+    raise ValueError(f"Could not find either of {', '.join(file_names)} in {model_dir}")
+
+
+class ModelManagement:
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def _get_model_description(cls, model_name: str) -> Dict[str, Any]:
+        """
+        Gets the model description from the model_name.
+
+        Args:
+            model_name (str): The name of the model.
+
+        raises:
+            ValueError: If the model_name is not supported.
+
+        Returns:
+            Dict[str, Any]: The model description.
+        """
+        for model in cls.list_supported_models():
+            if model_name == model["model"]:
+                return model
+
+        raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
+
+    @classmethod
+    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
+        """
+        Downloads a file from Google Cloud Storage.
+
+        Args:
+            url (str): The URL to download the file from.
+            output_path (str): The path to save the downloaded file to.
+            show_progress (bool, optional): Whether to show a progress bar. Defaults to True.
+
+        Returns:
+            str: The path to the downloaded file.
+        """
+
+        if os.path.exists(output_path):
+            return output_path
+        response = requests.get(url, stream=True)
+
+        # Handle HTTP errors
+        if response.status_code == 403:
+            raise PermissionError(
+                "Authentication Error: You do not have permission to access this resource. "
+                "Please check your credentials."
+            )
+
+        # Get the total size of the file
+        total_size_in_bytes = int(response.headers.get("content-length", 0))
+
+        # Warn if the total size is zero
+        if total_size_in_bytes == 0:
+            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
+
+        show_progress = total_size_in_bytes and show_progress
+
+        with tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True, disable=not show_progress) as progress_bar:
+            with open(output_path, "wb") as file:
+                for chunk in response.iter_content(chunk_size=1024):
+                    if chunk:  # Filter out keep-alive new chunks
+                        progress_bar.update(len(chunk))
+                        file.write(chunk)
+        return output_path
+
+    @classmethod
+    def download_files_from_huggingface(cls, hf_source_repo: str, cache_dir: Optional[str] = None) -> str:
+        """
+        Downloads a model from HuggingFace Hub.
+        Args:
+            hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
+            cache_dir (Optional[str]): The path to the cache directory.
+        Returns:
+            Path: The path to the model directory.
+        """
+
+        return snapshot_download(
+            repo_id=hf_source_repo,
+            ignore_patterns=["model.safetensors", "pytorch_model.bin"],
+            cache_dir=cache_dir,
+        )
+
+    @classmethod
+    def decompress_to_cache(cls, targz_path: str, cache_dir: str):
+        """
+        Decompresses a .tar.gz file to a cache directory.
+
+        Args:
+            targz_path (str): Path to the .tar.gz file.
+            cache_dir (str): Path to the cache directory.
+
+        Returns:
+            cache_dir (str): Path to the cache directory.
+        """
+        # Check if targz_path exists and is a file
+        if not os.path.isfile(targz_path):
+            raise ValueError(f"{targz_path} does not exist or is not a file.")
+
+        # Check if targz_path is a .tar.gz file
+        if not targz_path.endswith(".tar.gz"):
+            raise ValueError(f"{targz_path} is not a .tar.gz file.")
+
+        try:
+            # Open the tar.gz file
+            with tarfile.open(targz_path, "r:gz") as tar:
+                # Extract all files into the cache directory
+                tar.extractall(path=cache_dir)
+        except tarfile.TarError as e:
+            # If any error occurs while opening or extracting the tar.gz file,
+            # delete the cache directory (if it was created in this function)
+            # and raise the error again
+            if "tmp" in cache_dir:
+                shutil.rmtree(cache_dir)
+            raise ValueError(f"An error occurred while decompressing {targz_path}: {e}")
+
+        return cache_dir
+
+    @classmethod
+    def retrieve_model_gcs(cls, model_name: str, source_url: str, cache_dir: str) -> Path:
+        fast_model_name = f"fast-{model_name.split('/')[-1]}"
+
+        cache_tmp_dir = Path(cache_dir) / "tmp"
+        model_tmp_dir = cache_tmp_dir / fast_model_name
+        model_dir = Path(cache_dir) / fast_model_name
+
+        # check if the model_dir and the model files are both present for macOS
+        if model_dir.exists() and len(list(model_dir.glob("*"))) > 0:
+            return model_dir
+
+        if model_tmp_dir.exists():
+            shutil.rmtree(model_tmp_dir)
+
+        cache_tmp_dir.mkdir(parents=True, exist_ok=True)
+
+        model_tar_gz = Path(cache_dir) / f"{fast_model_name}.tar.gz"
+
+        cls.download_file_from_gcs(
+            source_url,
+            output_path=str(model_tar_gz),
+        )
+
+        cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
+        assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
+
+        model_tar_gz.unlink()
+        # Rename from tmp to final name is atomic
+        model_tmp_dir.rename(model_dir)
+
+        return model_dir
+
+    @classmethod
+    def download_model(cls, model: Dict[str, Any], cache_dir: Path) -> Path:
+        """
+        Downloads a model from HuggingFace Hub or Google Cloud Storage.
+
+        Args:
+            model (Dict[str, Any]): The model description.
+                Example:
+                ```
+                {
+                    "model": "BAAI/bge-base-en-v1.5",
+                    "dim": 768,
+                    "description": "Base English model, v1.5",
+                    "size_in_GB": 0.44,
+                    "sources": {
+                        "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
+                        "hf": "qdrant/bge-base-en-v1.5-onnx-q",
+                    }
+                }
+                ```
+            cache_dir (str): The path to the cache directory.
+
+        Returns:
+            Path: The path to the downloaded model directory.
+        """
+
+        hf_source = model.get("sources", {}).get("hf")
+        url_source = model.get("sources", {}).get("url")
+
+        if hf_source:
+            try:
+                return Path(cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir)))
+            except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
+                logger.error(f"Could not download model from HuggingFace: {e}" "Falling back to other sources.")
+
+        if url_source:
+            return cls.retrieve_model_gcs(model["model"], url_source, str(cache_dir))
+
+        raise ValueError(f"Could not download model {model['model']} from any source.")

fastembed/common/models.py

@@ -0,0 +1,52 @@
+import json
+from pathlib import Path
+
+import numpy as np
+from tokenizers import Tokenizer, AddedToken
+
+
+def load_tokenizer(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}")
+
+    with open(str(config_path)) as config_file:
+        config = json.load(config_file)
+
+    with open(str(tokenizer_config_path)) as tokenizer_config_file:
+        tokenizer_config = json.load(tokenizer_config_file)
+
+    with open(str(tokens_map_path)) as tokens_map_file:
+        tokens_map = json.load(tokens_map_file)
+
+    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.get("pad_token_id", 0), 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 normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
+    # 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
+    normalized_array = input_array / norm
+    return normalized_array

fastembed/common/onnx_model.py

@@ -0,0 +1,130 @@
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, TypeVar, Union
+
+import numpy as np
+import onnxruntime as ort
+
+from fastembed.common.model_management import locate_model_file
+from fastembed.common.models import load_tokenizer
+from fastembed.common.utils import iter_batch
+from fastembed.parallel_processor import ParallelWorkerPool, Worker
+
+# Holds type of the embedding result
+T = TypeVar("T")
+
+
+class OnnxModel(Generic[T]):
+    @classmethod
+    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @classmethod
+    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[T]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def __init__(self) -> None:
+        self.model = None
+        self.tokenizer = None
+
+    def _preprocess_onnx_input(self, onnx_input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def load_onnx_model(self, model_dir: Path, threads: Optional[int], max_length: int) -> None:
+        model_path = locate_model_file(model_dir, ["model.onnx", "model_optimized.onnx"])
+
+        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+        onnx_providers = ["CPUExecutionProvider"]
+
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        if threads is not None:
+            so.intra_op_num_threads = threads
+            so.inter_op_num_threads = threads
+
+        self.tokenizer = load_tokenizer(model_dir=model_dir, 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),
+            "token_type_ids": np.array([np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64),
+        }
+
+        onnx_input = self._preprocess_onnx_input(onnx_input)
+
+        model_output = self.model.run(None, onnx_input)
+        embeddings = model_output[0]
+        return embeddings, attention_mask
+
+    def _embed_documents(
+        self,
+        model_name: str,
+        cache_dir: str,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+    ) -> Iterable[T]:
+        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):
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
+        else:
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+            }
+            pool = ParallelWorkerPool(parallel, self._get_worker_class(), start_method=start_method)
+            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
+                yield from self._post_process_onnx_output(batch)
+
+
+class EmbeddingWorker(Worker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> OnnxModel:
+        raise NotImplementedError()
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ):
+        self.model = self.init_embedding(model_name, cache_dir)
+
+    @classmethod
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
+        return cls(
+            model_name=model_name,
+            cache_dir=cache_dir,
+        )
+
+    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)

fastembed/common/utils.py

@@ -0,0 +1,33 @@
+import os
+import tempfile
+from itertools import islice
+from pathlib import Path
+from typing import Union, Iterable, Generator, Optional
+
+
+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 define_cache_dir(cache_dir: Optional[str] = None) -> Path:
+    """
+    Define the cache directory for fastembed
+    """
+    if cache_dir is None:
+        default_cache_dir = os.path.join(tempfile.gettempdir(), "fastembed_cache")
+        cache_path = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
+    else:
+        cache_path = Path(cache_dir)
+
+    cache_path.mkdir(parents=True, exist_ok=True)
+
+    return cache_path

fastembed/embedding.py

@@ -1,309 +1,23 @@
-import os
-import shutil
-import tarfile
-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Iterable, List
+from typing import Optional
 
-import numpy as np
-import requests
-from optimum.onnxruntime import ORTModelForFeatureExtraction
-from tqdm import tqdm
-from transformers import AutoTokenizer
+from loguru import logger
+
+from fastembed.text.text_embedding import TextEmbedding
+
+logger.warning(
+    "DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated." "Use from fastembed import TextEmbedding instead."
+)
+
+DefaultEmbedding = TextEmbedding
+FlagEmbedding = TextEmbedding
 
 
-def normalize(input_array, p=2.0, 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
-    normalized_array = input_array / norm
-    return normalized_array
-
-class Embedding(ABC):
-    """
-    Abstract class for embeddings.
-
-    Args:
-        ABC ():
-
-    Raises:
-        NotImplementedError: Raised when you call an abstract method that has not been implemented.
-        PermissionError: _description_
-        ValueError: Several possible reasons: 1) targz_path does not exist or is not a file, 2) targz_path is not a .tar.gz file, 3) An error occurred while decompressing targz_path, 4) Could not find model_dir in cache_dir, 5) Could not find tokenizer.json in model_dir, 6) Could not find model.onnx in model_dir.
-        NotImplementedError: _description_
-
-    Returns:
-        _type_: _description_
-
-    Yields:
-        _type_: _description_
-    """
-
-    @abstractmethod
-    def embed(self, texts: List[str]) -> List[np.ndarray]:
-        raise NotImplementedError
-
-    @classmethod
-    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
-        """
-        Downloads a file from Google Cloud Storage.
-
-        Args:
-            url (str): The URL to download the file from.
-            output_path (str): The path to save the downloaded file to.
-            show_progress (bool, optional): Whether to show a progress bar. Defaults to True.
-
-        Returns:
-            str: The path to the downloaded file.
-        """
-
-        if os.path.exists(output_path):
-            return output_path
-        response = requests.get(url, stream=True)
-
-        # Handle HTTP errors
-        if response.status_code == 403:
-            raise PermissionError(
-                "Authentication Error: You do not have permission to access this resource. Please check your credentials."
-            )
-
-        # Get the total size of the file
-        total_size_in_bytes = int(response.headers.get("content-length", 0))
-
-        # Warn if the total size is zero
-        if total_size_in_bytes == 0:
-            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
-
-        # Initialize the progress bar
-        progress_bar = (
-            tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
-            if total_size_in_bytes and show_progress
-            else None
-        )
-
-        # Attempt to download the file
-        try:
-            with open(output_path, "wb") as file:
-                for chunk in response.iter_content(chunk_size=1024):  # Adjust chunk size to your preference
-                    if chunk:  # Filter out keep-alive new chunks
-                        if progress_bar is not None:
-                            progress_bar.update(len(chunk))
-                        file.write(chunk)
-        except Exception as e:
-            print(f"An error occurred while trying to download the file: {str(e)}")
-            return
-        finally:
-            if progress_bar is not None:
-                progress_bar.close()
-        return output_path
-
-    @classmethod
-    def decompress_to_cache(cls, targz_path: str, cache_dir: str):
-        """
-        Decompresses a .tar.gz file to a cache directory.
-
-        Args:
-            targz_path (str): Path to the .tar.gz file.
-            cache_dir (str): Path to the cache directory.
-
-        Returns:
-            cache_dir (str): Path to the cache directory.
-        """
-        # Check if targz_path exists and is a file
-        if not os.path.isfile(targz_path):
-            raise ValueError(f"{targz_path} does not exist or is not a file.")
-
-        # Check if targz_path is a .tar.gz file
-        if not targz_path.endswith(".tar.gz"):
-            raise ValueError(f"{targz_path} is not a .tar.gz file.")
-
-        try:
-            # Open the tar.gz file
-            with tarfile.open(targz_path, "r:gz") as tar:
-                # Extract all files into the cache directory
-                tar.extractall(path=cache_dir)
-        except tarfile.TarError as e:
-            # If any error occurs while opening or extracting the tar.gz file,
-            # delete the cache directory (if it was created in this function)
-            # and raise the error again
-            if "tmp" in cache_dir:
-                shutil.rmtree(cache_dir)
-            raise ValueError(f"An error occurred while decompressing {targz_path}: {e}")
-
-        return cache_dir
-
-    def retrieve_model(self, model_name: str, cache_dir: str) -> Path:
-        """
-        Retrieves a model from Google Cloud Storage.
-
-        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. BAAI/bge-base-en"
-
-        fast_model_name = f"fast-{model_name.split('/')[-1]}"
-
-        model_dir = Path(cache_dir) / fast_model_name
-        if model_dir.exists():
-            return model_dir
-
-        model_tar_gz = Path(cache_dir) / f"{fast_model_name}.tar.gz"
-        try:
-            self.download_file_from_gcs(
-                f"https://storage.googleapis.com/qdrant-fastembed/{fast_model_name}.tar.gz",
-            output_path=str(model_tar_gz),
-            )
-        except PermissionError:
-            simple_model_name = model_name.replace("/", "-")
-            print(f"Was not able to download {fast_model_name}.tar.gz, trying {simple_model_name}.tar.gz")
-            self.download_file_from_gcs(
-                f"https://storage.googleapis.com/qdrant-fastembed/{simple_model_name}.tar.gz",
-                output_path=str(model_tar_gz),
-            )
-
-        self.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=cache_dir)
-        assert model_dir.exists(), f"Could not find {model_dir} in {cache_dir}"
-
-        model_tar_gz.unlink()
-
-        return model_dir
-
-    def passage_embed(self, texts: List[str], batch_size: int = 256) -> 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.
-
-        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]])
-
-    def query_embed(self, query: str) -> Iterable[np.ndarray]:
-        """
-        Embeds a query
-
-        Args:
-            query (str): The query to search for.
-
-        Returns:
-            Iterable[np.ndarray]: The embeddings.
-        """
-
-        # Prepend "query: " to the query
-        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
-
-
-class FlagEmbedding(Embedding):
-    """
-    Implementation of the Flag Embedding model.
-
-    Args:
-        Embedding (_type_): _description_
-    """
-
+class JinaEmbedding(TextEmbedding):
     def __init__(
         self,
-        model_name: str = "BAAI/bge-small-en",
-        max_length: int = 512,
-        cache_dir: str = None,
+        model_name: str = "jinaai/jina-embeddings-v2-base-en",
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
     ):
-        """
-        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.
-
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
-        """
-        if cache_dir is None:
-            cache_dir = Path(".").resolve() / "local_cache"
-            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))
-
-    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)
-
-    def embed(self, documents: List[str], batch_size: int = 256) -> 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
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-
-        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)
-        else:
-            return self.onnx_embed(documents)
-
-
-class DefaultEmbedding(FlagEmbedding):
-    """
-    Implementation of the default Flag Embedding model.
-
-    Args:
-        FlagEmbedding (_type_): _description_
-    """
-
-    def __init__(
-        self,
-        model_name: str = "BAAI/bge-small-en",
-        onnx_providers: List[str] = None,
-        max_length: int = 512,
-        cache_dir: str = None,
-    ):
-        # if onnx_providers is None:
-        #     onnx_providers = [ONNXProviders.CPU]
-        super().__init__(model_name, max_length=max_length, cache_dir=cache_dir)
-
-
-class OpenAIEmbedding(Embedding):
-    def __init__(self):
-        # Initialize your OpenAI model here
-        # self.model = ...
-        ...
-
-    def embed(self, texts):
-        # Use your OpenAI model to embed the texts
-        # return self.model.embed(texts)
-        raise NotImplementedError
-        raise NotImplementedError
+        super().__init__(model_name, cache_dir, threads, **kwargs)

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, Tuple, Type
+
+# 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()

fastembed/sparse/sparse_embedding_base.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+from typing import Dict, Iterable, Optional, Union
+
+import numpy as np
+
+from fastembed.common.model_management import ModelManagement
+
+
+@dataclass
+class SparseEmbedding:
+    values: np.ndarray
+    indices: np.ndarray
+
+    def as_object(self) -> Dict[str, np.ndarray]:
+        return {
+            "values": self.values,
+            "indices": self.indices,
+        }
+
+    def as_dict(self) -> Dict[int, float]:
+        return {i: v for i, v in zip(self.indices, self.values)}
+
+
+class SparseTextEmbeddingBase(ModelManagement):
+    def __init__(self, model_name: str, cache_dir: Optional[str] = None, threads: Optional[int] = None, **kwargs):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[SparseEmbedding]:
+        raise NotImplementedError()

fastembed/sparse/sparse_text_embedding.py

@@ -0,0 +1,82 @@
+from typing import List, Type, Dict, Any, Union, Iterable, Optional
+
+from fastembed.sparse.sparse_embedding_base import SparseTextEmbeddingBase, SparseEmbedding
+from fastembed.sparse.splade_pp import SpladePP
+
+
+class SparseTextEmbedding(SparseTextEmbeddingBase):
+    EMBEDDINGS_REGISTRY: List[Type[SparseTextEmbeddingBase]] = [
+        SpladePP,
+    ]
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """
+        Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+
+            Example:
+                ```
+                [
+                    {
+                        "model": "prithvida/SPLADE_PP_en_v1",
+                        "vocab_size": 30522,
+                        "description": "Independent Implementation of SPLADE++ Model for English",
+                        "size_in_GB": 0.532,
+                        "sources": {
+                            "hf": "qdrant/SPLADE_PP_en_v1",
+                        },
+                    }
+                ]
+                ```
+        """
+        result = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding.list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
+            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
+            if any(model_name == model["model"] for model in supported_models):
+                self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in SparseTextEmbedding."
+            "Please check the supported models using `SparseTextEmbedding.list_supported_models()`"
+        )
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[SparseEmbedding]:
+        """
+        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
+        """
+        yield from self.model.embed(documents, batch_size, parallel, **kwargs)

fastembed/sparse/splade_pp.py

@@ -0,0 +1,116 @@
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+
+import numpy as np
+
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel
+from fastembed.common.utils import define_cache_dir
+from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
+
+supported_splade_models = [
+    {
+        "model": "prithvida/SPLADE_PP_en_v1",
+        "vocab_size": 30522,
+        "description": "Independent Implementation of SPLADE++ Model for English",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "Qdrant/SPLADE_PP_en_v1",
+        },
+    },
+]
+
+
+class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
+    @classmethod
+    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[SparseEmbedding]:
+        logits, attention_mask = output
+        relu_log = np.log(1 + np.maximum(logits, 0))
+
+        weighted_log = relu_log * np.expand_dims(attention_mask, axis=-1)
+
+        max_val = np.max(weighted_log, axis=1)
+
+        # Score matrix of shape (batch_size, vocab_size)
+        # Most of the values are 0, only a few are non-zero
+        scores = np.squeeze(max_val)
+        for row_scores in scores:
+            indices = row_scores.nonzero()[0]
+            scores = row_scores[indices]
+            yield SparseEmbedding(values=scores, indices=indices)
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_splade_models
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            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.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        self.model_name = model_name
+        self._model_description = self._get_model_description(model_name)
+
+        self._cache_dir = define_cache_dir(cache_dir)
+        self._model_dir = self.download_model(self._model_description, self._cache_dir)
+        self._max_length = 512
+
+        self.load_onnx_model(self._model_dir, self.threads, self._max_length)
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[SparseEmbedding]:
+        """
+        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
+        """
+        yield from self._embed_documents(
+            model_name=self.model_name,
+            cache_dir=str(self._cache_dir),
+            documents=documents,
+            batch_size=batch_size,
+            parallel=parallel,
+        )
+
+
+class SpladePPEmbeddingWorker(EmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> SpladePP:
+        return SpladePP(model_name=model_name, cache_dir=cache_dir, threads=1)

fastembed/text/e5_onnx_embedding.py

@@ -0,0 +1,59 @@
+from typing import Type, List, Dict, Any
+
+import numpy as np
+
+from fastembed.common.onnx_model import EmbeddingWorker
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+
+supported_multilingual_e5_models = [
+    {
+        "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,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
+            "hf": "qdrant/multilingual-e5-large-onnx",
+        },
+    },
+    {
+        "model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+        "dim": 768,
+        "description": "Sentence-transformers model for tasks like clustering or semantic search",
+        "size_in_GB": 1.11,
+        "sources": {
+            "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
+        },
+    },
+]
+
+
+class E5OnnxEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
+        return E5OnnxEmbeddingWorker
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_multilingual_e5_models
+
+    def _preprocess_onnx_input(self, onnx_input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        onnx_input.pop("token_type_ids", None)
+        return onnx_input
+
+
+class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> E5OnnxEmbedding:
+        return E5OnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)

fastembed/text/jina_onnx_embedding.py

@@ -0,0 +1,63 @@
+from typing import Type, List, Dict, Any, Tuple, Iterable
+
+import numpy as np
+
+from fastembed.common.models import normalize
+from fastembed.common.onnx_model import EmbeddingWorker
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+
+supported_jina_models = [
+    {
+        "model": "jinaai/jina-embeddings-v2-base-en",
+        "dim": 768,
+        "description": "English embedding model supporting 8192 sequence length",
+        "size_in_GB": 0.55,
+        "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-small-en",
+        "dim": 512,
+        "description": "English embedding model supporting 8192 sequence length",
+        "size_in_GB": 0.13,
+        "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
+    },
+]
+
+
+class JinaOnnxEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[EmbeddingWorker]:
+        return JinaEmbeddingWorker
+
+    @classmethod
+    def mean_pooling(cls, model_output, attention_mask) -> np.ndarray:
+        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
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_jina_models
+
+    @classmethod
+    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[np.ndarray]:
+        embeddings, attn_mask = output
+        return normalize(cls.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+
+
+class JinaEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> OnnxTextEmbedding:
+        return JinaOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)

fastembed/text/onnx_embedding.py

@@ -0,0 +1,244 @@
+from typing import Dict, Optional, Tuple, Union, Iterable, Type, List, Any
+
+import numpy as np
+
+from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker
+from fastembed.common.models import normalize
+from fastembed.common.utils import define_cache_dir
+from fastembed.text.text_embedding_base import TextEmbeddingBase
+
+supported_onnx_models = [
+    {
+        "model": "BAAI/bge-base-en",
+        "dim": 768,
+        "description": "Base English model",
+        "size_in_GB": 0.5,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
+        },
+    },
+    {
+        "model": "BAAI/bge-base-en-v1.5",
+        "dim": 768,
+        "description": "Base English model, v1.5",
+        "size_in_GB": 0.44,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
+            "hf": "qdrant/bge-base-en-v1.5-onnx-q",
+        },
+    },
+    {
+        "model": "BAAI/bge-large-en-v1.5-quantized",
+        "dim": 1024,
+        "description": "Large English model, v1.5",
+        "size_in_GB": 1.34,
+        "sources": {
+            "hf": "qdrant/bge-large-en-v1.5-onnx-q",
+        },
+    },
+    {
+        "model": "BAAI/bge-large-en-v1.5",
+        "dim": 1024,
+        "description": "Large English model, v1.5",
+        "size_in_GB": 1.34,
+        "sources": {
+            "hf": "qdrant/bge-large-en-v1.5-onnx",
+        },
+    },
+    {
+        "model": "BAAI/bge-small-en",
+        "dim": 384,
+        "description": "Fast English model",
+        "size_in_GB": 0.2,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
+        },
+    },
+    # {
+    #     "model": "BAAI/bge-small-en",
+    #     "dim": 384,
+    #     "description": "Fast English model",
+    #     "size_in_GB": 0.2,
+    #     "hf_sources": [],
+    #     "compressed_url_sources": [
+    #         "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en.tar.gz",
+    #         "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz"
+    #     ]
+    # },
+    {
+        "model": "BAAI/bge-small-en-v1.5",
+        "dim": 384,
+        "description": "Fast and Default English model",
+        "size_in_GB": 0.13,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz",
+            "hf": "qdrant/bge-small-en-v1.5-onnx-q",
+        },
+    },
+    {
+        "model": "BAAI/bge-small-zh-v1.5",
+        "dim": 512,
+        "description": "Fast and recommended Chinese model",
+        "size_in_GB": 0.1,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
+        },
+    },
+    {  # todo: it is not a flag embedding
+        "model": "sentence-transformers/all-MiniLM-L6-v2",
+        "dim": 384,
+        "description": "Sentence Transformer model, MiniLM-L6-v2",
+        "size_in_GB": 0.09,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
+            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
+        },
+    },
+    {
+        "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+        "dim": 384,
+        "description": "Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2",
+        "size_in_GB": 0.46,
+        "sources": {
+            "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
+        },
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1",
+        },
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+    },
+    {
+        "model": "thenlper/gte-large",
+        "dim": 1024,
+        "description": "Large general text embeddings model",
+        "size_in_GB": 1.34,
+        "sources": {
+            "hf": "qdrant/gte-large-onnx",
+        },
+    },
+    # {
+    #     "model": "sentence-transformers/all-MiniLM-L6-v2",
+    #     "dim": 384,
+    #     "description": "Sentence Transformer model, MiniLM-L6-v2",
+    #     "size_in_GB": 0.09,
+    #     "hf_sources": [
+    #         "qdrant/all-MiniLM-L6-v2-onnx"
+    #     ],
+    #     "compressed_url_sources": [
+    #         "https://storage.googleapis.com/qdrant-fastembed/fast-all-MiniLM-L6-v2.tar.gz",
+    #         "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz"
+    #     ]
+    # }
+]
+
+
+class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
+    """Implementation of the Flag Embedding model."""
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """
+        Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_onnx_models
+
+    def __init__(
+        self,
+        model_name: str = "BAAI/bge-small-en-v1.5",
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            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.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        self.model_name = model_name
+        self._model_description = self._get_model_description(model_name)
+
+        self._cache_dir = define_cache_dir(cache_dir)
+        self._model_dir = self.download_model(self._model_description, self._cache_dir)
+        self._max_length = 512
+
+        self.load_onnx_model(self._model_dir, self.threads, self._max_length)
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> 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
+        """
+        yield from self._embed_documents(
+            model_name=self.model_name,
+            cache_dir=str(self._cache_dir),
+            documents=documents,
+            batch_size=batch_size,
+            parallel=parallel,
+        )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
+        return OnnxTextEmbeddingWorker
+
+    def _preprocess_onnx_input(self, onnx_input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    @classmethod
+    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[np.ndarray]:
+        embeddings, _ = output
+        return normalize(embeddings[:, 0]).astype(np.float32)
+
+
+class OnnxTextEmbeddingWorker(EmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> OnnxTextEmbedding:
+        return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)

fastembed/text/text_embedding.py

@@ -0,0 +1,89 @@
+from typing import Any, Dict, Iterable, List, Optional, Type, Union
+
+import numpy as np
+
+from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
+from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.text_embedding_base import TextEmbeddingBase
+
+
+class TextEmbedding(TextEmbeddingBase):
+    EMBEDDINGS_REGISTRY: List[Type[TextEmbeddingBase]] = [
+        OnnxTextEmbedding,
+        E5OnnxEmbedding,
+        JinaOnnxEmbedding,
+    ]
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """
+        Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+
+            Example:
+                ```
+                [
+                    {
+                        "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,
+                        "sources": {
+                            "gcp": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
+                            "hf": "qdrant/multilingual-e5-large-onnx",
+                        }
+                    }
+                ]
+                ```
+        """
+        result = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding.list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str = "BAAI/bge-small-en-v1.5",
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
+            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
+            if any(model_name == model["model"] for model in supported_models):
+                self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in TextEmbedding."
+            "Please check the supported models using `TextEmbedding.list_supported_models()`"
+        )
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> 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
+        """
+        yield from self.model.embed(documents, batch_size, parallel, **kwargs)

fastembed/text/text_embedding_base.py

@@ -0,0 +1,53 @@
+from typing import Iterable, Optional, Union
+
+import numpy as np
+
+from fastembed.common.model_management import ModelManagement
+
+
+class TextEmbeddingBase(ModelManagement):
+    def __init__(self, model_name: str, cache_dir: Optional[str] = None, threads: Optional[int] = None, **kwargs):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        raise NotImplementedError()
+
+    def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
+        """
+        Embeds a list of text passages into a list of embeddings.
+
+        Args:
+            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.
+        """
+
+        # This is model-specific, so that different models can have specialized implementations
+        yield from self.embed(texts, **kwargs)
+
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
+        """
+        Embeds queries
+
+        Args:
+            query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
+
+        Returns:
+            Iterable[np.ndarray]: The embeddings.
+        """
+
+        # This is model-specific, so that different models can have specialized implementations
+        if isinstance(query, str):
+            yield from self.embed([query], **kwargs)
+        if isinstance(query, Iterable):
+            yield from self.embed(query, **kwargs)

mkdocs.yml

@@ -2,7 +2,7 @@ site_name: FastEmbed
 site_url: https://qdrant.github.io/fastembed/
 site_author: Nirant Kasliwal
 repo_url: https://github.com/qdrant/fastembed/
-repo_name: qdrant/fastembed  
+repo_name: qdrant/fastembed
 
 remote_branch: gh-pages
 remote_name: origin
@@ -33,11 +33,11 @@ theme:
     # Text color for primary color
     text: "#ffffff"
 
-  palette: 
+  palette:
     # Palette toggle for light mode
     - scheme: default
       toggle:
-        icon: material/brightness-7 
+        icon: material/brightness-7
         name: Switch to dark mode
 
     # Palette toggle for dark mode

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.0.4"
+version = "0.2.3"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -11,40 +11,36 @@ repository = "https://github.com/qdrant/fastembed"
 keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-transformers"]
 
 [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"
+python = ">=3.8.0,<3.13"
+onnx = "^1.15.0"
+onnxruntime = "^1.17.0"
+tqdm = "^4.66"
+requests = "^2.31"
+tokenizers = "^0.15.1"
+huggingface-hub = "^0.20"
+loguru = "^0.7.2"
+numpy = [
+    { version = ">=1.21", python = "<3.12" },
+    { version = ">=1.26", python = ">=3.12" }
+]
 
-[tool.poetry.dev-dependencies]
-ruff = "^0.0.277"
-isort = "^5.12.0"
-black = "^23.7.0"
-onnx = "^1.11.0"
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.4.2"
+ruff = "^0.2.2"
 notebook = ">=7.0.2"
-mkdocs-material = "^9.1.21"
-mkdocstrings = "^0.22.0"
-pillow = "^10.0.0"
+pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
+
+[tool.poetry.group.docs.dependencies]
+mkdocs-material = "^9.5.10"
+mkdocstrings = "^0.24.0"
+pillow = "^10.2.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"]
 build-backend = "poetry.core.masonry.api"
 
-[tool.black]
-line-length = 120
-
-[tool.isort]
-profile = "black"
-
 [tool.ruff]
 line-length = 120

tests/profiling.py

@@ -1,27 +1,25 @@
 # %% [markdown]
 # # 🤗 Huggingface vs ⚡ FastEmbed️
-# 
+#
 # Comparing the performance of Huggingface's 🤗 Transformers and ⚡ FastEmbed️ on a simple task on the following machine: Apple M2 Max, 32 GB RAM
-# 
+#
 # ## 📦 Imports
-# 
+#
 # Importing the necessary libraries for this comparison.
 
 # %%
 import time
-from pathlib import Path
-from typing import Any, Callable, List, Tuple
+from typing import Callable, List, Tuple
 
-import numpy as np
+import matplotlib.pyplot as plt
 import torch.nn.functional as F
-from torch import Tensor
 from transformers import AutoModel, AutoTokenizer
 
 from fastembed.embedding import DefaultEmbedding
 
 # %% [markdown]
 # ## 📖 Data
-# 
+#
 # data is a list of strings, each string is a document.
 
 # %%
@@ -43,9 +41,10 @@ len(documents)
 
 # %% [markdown]
 # ## Setting up 🤗 Huggingface
-# 
+#
 # We'll be using the [Huggingface Transformers](https://huggingface.co/transformers/) with PyTorch library to generate embeddings. We'll be using the same model across both libraries for a fair(er?) comparison.
 
+
 # %%
 class HF:
     """
@@ -64,12 +63,13 @@ class HF:
         sentence_embeddings = F.normalize(sentence_embeddings)
         return sentence_embeddings
 
+
 hf = HF(model_id="BAAI/bge-small-en")
 hf.embed(documents).shape
 
 # %% [markdown]
 # ## Setting up ⚡️FastEmbed
-# 
+#
 # Sorry, don't have a lot to set up here. We'll be using the default model, which is Flag Embedding, same as the Huggingface model.
 
 # %%
@@ -77,18 +77,19 @@ embedding_model = DefaultEmbedding()
 
 # %% [markdown]
 # ## 📊 Comparison
-# 
+#
 # We'll be comparing the following metrics: Minimum, Maximum, Mean, across k runs. Let's write a function to do that:
-# 
+#
 # ### 🚀 Calculating Stats
 
+
 # %%
 def calculate_time_stats(embed_func: Callable, documents: list, k: int) -> Tuple[float, float, float]:
     times = []
     for _ in range(k):
         # Timing the embed_func call
         start_time = time.time()
-        embeddings = embed_func(documents)
+        embed_func(documents)
         end_time = time.time()
 
         times.append(end_time - start_time)
@@ -96,15 +97,13 @@ def calculate_time_stats(embed_func: Callable, documents: list, k: int) -> Tuple
     # Returning mean, max, and min time for the call
     return (sum(times) / k, max(times), min(times))
 
+
 # %%
 hf_stats = calculate_time_stats(hf.embed, documents, k=2)
 print(f"Huggingface Transformers (Average, Max, Min): {hf_stats}")
 fst_stats = calculate_time_stats(lambda x: list(embedding_model.embed(x)), documents, k=2)
 print(f"FastEmbed (Average, Max, Min): {fst_stats}")
 
-# %%
-import matplotlib.pyplot as plt
-
 
 # %%
 def plot_character_per_second_comparison(
@@ -141,5 +140,3 @@ def plot_character_per_second_comparison(
 
 
 plot_character_per_second_comparison(hf_stats, fst_stats, documents)
-
-

tests/test_onnx_embeddings.py

@@ -1,13 +0,0 @@
-import numpy as np
-
-from fastembed.embedding import DefaultEmbedding
-
-
-def test_default_embedding():
-    model = DefaultEmbedding()
-
-    docs = ["hello world", "flag embedding"]
-    embeddings = np.array(model.embed(docs))
-    assert embeddings.shape == (2, 384)
-
-test_default_embedding()

tests/test_sparse_embeddings.py

@@ -0,0 +1,42 @@
+import pytest
+from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
+
+CANONICAL_COLUMN_VALUES = {
+    "prithvida/SPLADE_PP_en_v1": {
+        "indices": [2040, 2047, 2088, 2299, 2748, 3011, 3376, 3795, 4774, 5304, 5798, 6160, 7592, 7632, 8484],
+        "values": [
+            0.4219532012939453,
+            0.4320072531700134,
+            2.766580104827881,
+            0.3314574658870697,
+            1.395172119140625,
+            0.021595917642116547,
+            0.43770670890808105,
+            0.0008370947907678783,
+            0.5187209844589233,
+            0.17124654352664948,
+            0.14742016792297363,
+            0.8142819404602051,
+            2.803262710571289,
+            2.1904349327087402,
+            1.0531445741653442,
+        ],
+    }
+}
+
+docs = ["Hello World"]
+
+
+def test_batch_embedding():
+    docs_to_embed = docs * 10
+
+    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
+        print("evaluating", model_name)
+        model = SparseTextEmbedding(model_name=model_name)
+        result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+        print(result.indices)
+
+        assert result.indices.tolist() == expected_result["indices"]
+
+        for i, value in enumerate(result.values):
+            assert pytest.approx(value, abs=0.001) == expected_result["values"][i]

tests/test_text_onnx_embeddings.py

@@ -0,0 +1,82 @@
+import os
+
+import numpy as np
+import pytest
+
+from fastembed.text.text_embedding import TextEmbedding
+
+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]),
+    "BAAI/bge-large-en-v1.5": np.array([0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]),
+    "BAAI/bge-large-en-v1.5-quantized": np.array([0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]),
+    "sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259, 0.0058, 0.0114, 0.0380, -0.0233]),
+    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array([0.0094, 0.0184, 0.0328, 0.0072, -0.0351]),
+    "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
+    "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array(
+        [-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
+    ),
+    "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]),
+    "nomic-ai/nomic-embed-text-v1": np.array([0.0061, 0.0103, -0.0296, -0.0242, -0.0170]),
+    "nomic-ai/nomic-embed-text-v1.5": np.array(
+        [-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
+    ),
+    "thenlper/gte-large": np.array([-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]),
+}
+
+
+def test_embedding():
+    is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
+
+    for model_desc in TextEmbedding.list_supported_models():
+        if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
+            continue
+
+        dim = model_desc["dim"]
+        model = TextEmbedding(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,model_name", [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")]
+)
+def test_batch_embedding(n_dims, model_name):
+    model = TextEmbedding(model_name=model_name)
+
+    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,model_name", [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")]
+)
+def test_parallel_processing(n_dims, model_name):
+    model = TextEmbedding(model_name=model_name)
+
+    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)