debug: try out with markers

* feat: Quantized models

* refactor: use model_file for GCS

* refactoring: refactor model downloading (#209)

* refactoring: refactor model downloading

* refactor: update docstring

Co-authored-by: Anush <anushshetty90@gmail.com>

* Update fastembed/common/model_management.py

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

* fix: model_file for Snowflake models

---------

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

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -0,0 +1,56 @@
+name: Bug/New Model Request
+description: File a bug report/Request a new Model
+title: "[Bug/Model Request]: "
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Thanks for taking the time to fill out this bug report!
+  - type: textarea
+    id: what-happened
+    attributes:
+      label: What happened?
+      description: Also tell us, what did you expect to happen?
+      placeholder: Tell us what you see!
+      value: "A bug happened!"
+    validations:
+      required: true
+  - type: textarea
+    id: Python version
+    attributes:
+      label: What Python version are you on? e.g. python --version
+      description: Also tell us, what package manager are you using e.g. conda, pip, poetry?
+      placeholder: Python3.10
+    validations:
+      required: true
+  - type: dropdown
+    id: version
+    attributes:
+      label: Version
+      description: What version of FastEmbed are you running? python -c "import fastembed; print(fastembed.__version__)". If you're not on the latest, please upgrade and see if the problem persists.
+      options:
+        - 0.2.6 (Latest)
+        - 0.2.5
+        - 0.2.4 
+        - 0.2.3 
+        - 0.2.2
+        - 0.2.1
+        - 0.1.x 
+      default: 0
+    validations:
+      required: true
+  - type: dropdown
+    id: os
+    attributes:
+      label: What os are you seeing the problem on?
+      multiple: true
+      options:
+        - Linux
+        - MacOS
+        - Windows
+  - type: textarea
+    id: logs
+    attributes:
+      label: Relevant stack traces and/or logs
+      description: Please copy and paste any relevant raised exceptions. This will be automatically formatted into code, so no need for backticks.
+      render: shell

.github/ISSUE_TEMPLATE/config.yaml

@@ -0,0 +1,5 @@
+blank_issues_enabled: false
+contact_links:
+  - name: GitHub Community Support
+    url: https://github.com/qdrant/fastembed/discussions
+    about: Please ask and answer questions here.

.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

@@ -3,6 +3,8 @@ name: Tests
 on:
   push:
     branches: [ master, main ]
+  schedule:
+    - cron: 0 0 * * *
   pull_request:
 
 env:
@@ -18,6 +20,7 @@ jobs:
           - '3.9.x'
           - '3.10.x'
           - '3.11.x'
+          - '3.12.x'
         os:
           - ubuntu-latest
           - macos-latest
@@ -28,16 +31,24 @@ jobs:
     name: Python ${{ matrix.python-version }} on ${{ matrix.os }} test
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
           python -m pip install poetry
           poetry config virtualenvs.create false
-          poetry install --no-interaction --no-ansi
-      - name: Run tests
-        run: pytest
-        shell: bash
+          poetry install --no-interaction --no-ansi --without docs
+
+      - name: Install Test Dependencies
+        run: pip install pytest pytest-md pytest-emoji
+
+      - name: Run pytest
+        uses: pavelzw/pytest-action@v2
+        with:
+            verbose: true
+            emoji: true
+            job-summary: true
+            report-title: 'FastEmbed Test Report'

.gitignore

@@ -85,28 +85,8 @@ ipython_config.py
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
-# .python-version
+.python-version
 
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/#use-with-ide
 .pdm.toml
 
 # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
@@ -152,27 +132,11 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/
 .DS_Store
-nbs/*.tar.gz
 *.tar.gz
-nbs/fast-*/*
-local_cache/*/*
-*/local_cache/*/*
-*/*/local_cache/*/*
+**/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
+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.3.4
+    hooks:
+      - id: ruff
+        types_or: [ python, pyi, jupyter ]
+        args: [ --fix ]
+      - id: ruff-format
+        types_or: [ python, pyi, jupyter ]

CONTRIBUTING.md

@@ -0,0 +1,78 @@
+# Contributing to FastEmbed!
+
+:+1::tada: First off, thanks for taking the time to contribute! :tada::+1:
+
+The following is a set of guidelines for contributing to FastEmbed. These are mostly guidelines, not rules. Use your best judgment, and feel free to propose changes to this document in a pull request.
+
+## Table Of Contents
+
+[I don't want to read this whole thing, I just have a question!!!](#i-dont-want-to-read-this-whole-thing-i-just-have-a-question)
+
+[How Can I Contribute?](#how-can-i-contribute)
+  * [Your First Code Contribution](#your-first-code-contribution)
+  * [Adding New Models](#adding-new-models)
+
+[Styleguides](#styleguides)
+  * [Code Lint](#code-lint)
+  * [Pre-Commit Hooks](#pre-commit-hooks)
+ 
+## I don't want to read this whole thing I just have a question!!!
+
+> **Note:** Please don't file an issue to ask a question. You'll get faster results by using the resources below:
+
+* [FastEmbed Docs](https://qdrant.github.io/fastembed/)
+* [Qdrant Discord](https://discord.gg/Qy6HCJK9Dc)
+
+## How Can I Contribute?
+
+## How Do I Submit A (Good) Bug Report?
+
+Bugs are tracked as [GitHub issues](https://guides.github.com/features/issues/). 
+
+Explain the problem and include additional details to help maintainers reproduce the problem:
+
+* **Use a clear and descriptive title** for the issue to identify the problem.
+* **Describe the exact steps which reproduce the problem** in as many details as possible. For example, start by explaining how you are using FastEmbed, e.g. with Langchain, Qdrant Client, Llama Index and which command exactly you used. When listing steps, **don't just say what you did, but explain how you did it**.
+* **Provide specific examples to demonstrate the steps**. Include links to files or GitHub projects, or copy/pasteable snippets, which you use in those examples. If you're providing snippets in the issue, use [Markdown code blocks](https://help.github.com/articles/markdown-basics/#multiple-lines).
+* **Describe the behavior you observed after following the steps** and point out what exactly is the problem with that behavior.
+* **Explain which behavior you expected to see instead and why.**
+* **If the problem is related to performance or memory**, include a [call stack profile capture](https://github.com/joerick/pyinstrument) and your observations. 
+
+Include details about your configuration and environment:
+
+* **Which version of FastEmbed are you using?** You can get the exact version by running `python -c "import fastembed; print(fastembed.__version__)"`.
+* **What's the name and version of the OS you're using**?
+* **Which packages do you have installed?** You can get that list by running `pip freeze`
+
+### Your First Code Contribution
+
+Unsure where to begin contributing to FastEmbed? You can start by looking through these `good-first-issue`issues:
+
+* [Good First Issue](https://github.com/qdrant/fastembed/labels/good%20first%20issue) - issues which should only require a few lines of code, and a test or two. These are a great way to get started with FastEmbed. This includes adding new models which are already tested and ready on Huggingface Hub. 
+
+## Pull Requests
+
+The best way to learn about the mechanics of FastEmbed is to start working on it. 
+
+### Your First Code Contribution
+Your first code contribution can be small bug fixes:
+1. This PR adds a small bug fix for a single input: https://github.com/qdrant/fastembed/pull/148
+2. This PR adds a check for the right file location and extension, specific to an OS: https://github.com/qdrant/fastembed/pull/128
+
+Even documentation improvements and tests are most welcome:
+1. This PR fixes a README link: https://github.com/qdrant/fastembed/pull/143
+
+### Adding New Models
+1. Open Requests for New Models are [here](https://github.com/qdrant/fastembed/labels/model%20request). 
+2. There are quite a few pull requests that were merged for this purpose and you can use them as a reference. Here is an example: https://github.com/qdrant/fastembed/pull/129
+3. Make sure to add tests for the new model
+    - The CANONICAL_VECTOR values must come from a reference implementation usually from Huggingface Transformers or Sentence Transformers
+    - Here is a reference [Colab Notebook](https://colab.research.google.com/drive/1tNdV3DsiwsJzu2AXnUnoeF5av1Hp8HF1?usp=sharing) for how we will evaluate whether your VECTOR values in the test are correct or not.
+
+## Styleguides
+
+### Code Lint
+We use ruff for code linting. It should be installed with poetry since it's a dev dependency.
+
+### Pre-Commit Hooks
+We use pre-commit hooks to ensure that the code is linted before it's committed. You can install pre-commit hooks by running `pre-commit install` in the root directory of the project.

README.md

@@ -1,40 +1,45 @@
 # ⚡️ What is FastEmbed?
 
-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.
+FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
-The default embedding supports "query" and "passage" prefixes for the input text. The default model is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval%20with%20FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+The default text embedding (`TextEmbedding`) model is Flag Embedding, presented 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 & Fast
-    - Quantized model weights
-    - ONNX Runtime, no PyTorch dependency
-    - CPU-first design
+## 📈 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
-    - List of [supported models](https://qdrant.github.io/fastembed/examples/Supported_Models/) - including multilingual models
+1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda. 
+
+2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data-parallelism for encoding large datasets.
+
+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
+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.",
-    "fastembed is supported by and maintained by Qdrant." # You can leave out the prefix but it's recommended
+    "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)) # If you use  
+
+# 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
 ```
 
 ## Usage with Qdrant
@@ -45,23 +50,29 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```
 
-Might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh. 
+You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
 from qdrant_client import QdrantClient
 
 # Initialize the client
-client = QdrantClient(":memory:")  # or QdrantClient(path="path/to/db")
+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": "Linkedin-docs"},
+    {"source": "Llama-index-docs"},
 ]
 ids = [42, 2]
 
-# Use the new add method
+# 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,

docs/Getting Started.ipynb

@@ -11,7 +11,9 @@
     "\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)"
    ]
   },
   {
@@ -21,15 +23,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!pip install fastembed --upgrade --quiet # 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"
    ]
   },
   {
@@ -39,35 +33,115 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
-     ]
+     "data": {
+      "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": [
-      "torch.Size([384])\n"
+      "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": [
+    "import numpy as np\n",
+    "from fastembed import TextEmbedding\n",
+    "from typing import List\n",
+    "\n",
+    "# Example list of documents\n",
+    "documents: List[str] = [\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",
+    "\n",
+    "# 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(embeddings_generator)\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": [
-    "from typing import List\n",
-    "import numpy as np\n",
-    "from fastembed.embedding import DefaultEmbedding\n",
+    "embeddings_generator = embedding_model.embed(documents)  # reminder this is a generator\n",
     "\n",
-    "# Example list of documents\n",
-    "documents: List[str] = [\n",
-    "    \"Hello, World!\",\n",
-    "    \"This is an example document.\",\n",
-    "    \"fastembed is supported by and maintained by Qdrant.\",\n",
-    "]\n",
-    "# Initialize the DefaultEmbedding class with the desired parameters\n",
-    "embedding_model = DefaultEmbedding(model_name=\"BAAI/bge-small-en\", max_length=512)\n",
-    "embeddings: List[np.ndarray] = embedding_model.embed(documents)\n",
-    "print(embeddings[0].shape)"
+    "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"
    ]
   },
   {
@@ -75,142 +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": 3,
-   "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"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "145a56ce",
-   "metadata": {},
-   "outputs": [],
-   "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",
-    "]"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "1cb3cc87",
-   "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",
+    "- **Passages**: Add \"passage:\" at the beginning of each passage string\n",
     "\n",
-    "This will download the model weights, decompress to directory `local_cache` and load them into the Embedding class.\n",
+    "## Beyond the default model\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."
+    "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": 5,
-   "id": "272c8915",
+   "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": [
-    "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."
+    "multilingual_large_model = TextEmbedding(\"intfloat/multilingual-e5-large\")  # This can take a few minutes to download"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "8013eee9",
+   "id": "a9e70f0e",
    "metadata": {},
    "outputs": [
     {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
-     ]
+     "data": {
+      "text/plain": [
+       "(4, 1024)"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
    "source": [
-    "embeddings: List[np.ndarray] = embedding_model.embed(documents)"
+    "np.array(\n",
+    "    list(multilingual_large_model.embed([\"Hello, world!\", \"你好世界\", \"¡Hola Mundo!\", \"नमस्ते!\"]))\n",
+    ").shape  # Vector of 1024 dimensions"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "e5b5a6ad",
+   "id": "64fe20ed",
    "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": 7,
-   "id": "0d8c8e08",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "torch.Size([384])\n"
-     ]
-    }
-   ],
-   "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/)"
    ]
   }
  ],
@@ -230,7 +236,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

docs/examples/SPLADE_with_FastEmbed.ipynb

@@ -0,0 +1,384 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Introduction to SPLADE with FastEmbed\n",
+    "\n",
+    "In this notebook, we will explore how to generate Sparse Vectors -- in particular a variant of the [SPLADE](https://arxiv.org/abs/2107.05720).\n",
+    "\n",
+    "> 💡 The original [naver/SPLADE](https://github.com/naver/splade) models were licensed CC BY-NC-SA 4.0 -- Not for Commercial Use. This [SPLADE++](https://huggingface.co/prithivida/Splade_PP_en_v1) model is Apache License and hence, licensed for commercial use. \n",
+    "\n",
+    "## Outline:\n",
+    "1. [What is SPLADE?](#What-is-SPLADE?)\n",
+    "2. [Setting up the environment](#Setting-up-the-environment)\n",
+    "3. [Generating SPLADE vectors with FastEmbed](#Generating-SPLADE-vectors-with-FastEmbed)\n",
+    "4. [Understanding SPLADE vectors](#Understanding-SPLADE-vectors)\n",
+    "5. [Observations and Design Choices](#Observations-and-Model-Design-Choices)\n",
+    "\n",
+    "\n",
+    "## What is SPLADE?\n",
+    "\n",
+    "SPLADE was a novel method for _learning_ sparse vectors for text representation. This model beats BM25 -- the underlying approach for the Elastic/Lucene family of implementations. Thus making it highly effective for tasks such as information retrieval, document classification, and more. \n",
+    "\n",
+    "The key advantage of SPLADE is its ability to generate sparse vectors, which are more efficient and interpretable than dense vectors. This makes SPLADE a powerful tool for handling large-scale text data.\n",
+    "\n",
+    "## Setting up the environment\n",
+    "\n",
+    "This notebook uses few dependencies, which are installed below: "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# !pip install -q fastembed"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's get started! 🚀"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T00:49:20.516644Z",
+     "start_time": "2024-03-30T00:49:20.188543Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from fastembed import SparseTextEmbedding, SparseEmbedding\n",
+    "from typing import List"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "> You can find the list of all supported Sparse Embedding models by calling this API: `SparseTextEmbedding.list_supported_models()`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T00:49:22.366294Z",
+     "start_time": "2024-03-30T00:49:22.362384Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'model': 'prithvida/Splade_PP_en_v1',\n",
+       "  'vocab_size': 30522,\n",
+       "  'description': 'Misspelled version of the model. Retained for backward compatibility. Independent Implementation of SPLADE++ Model for English',\n",
+       "  'size_in_GB': 0.532,\n",
+       "  'sources': {'hf': 'Qdrant/SPLADE_PP_en_v1'}},\n",
+       " {'model': 'prithivida/Splade_PP_en_v1',\n",
+       "  'vocab_size': 30522,\n",
+       "  'description': 'Independent Implementation of SPLADE++ Model for English',\n",
+       "  'size_in_GB': 0.532,\n",
+       "  'sources': {'hf': 'Qdrant/SPLADE_PP_en_v1'}}]"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "SparseTextEmbedding.list_supported_models()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T00:49:27.193530Z",
+     "start_time": "2024-03-30T00:49:26.139248Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2aa47b26ab01475e8d3577433037f685",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Fetching 9 files:   0%|          | 0/9 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "model_name = \"prithvida/Splade_PP_en_v1\"\n",
+    "# This triggers the model download\n",
+    "model = SparseTextEmbedding(model_name=model_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T00:49:28.624109Z",
+     "start_time": "2024-03-30T00:49:28.399960Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "documents: List[str] = [\n",
+    "    \"Chandrayaan-3 is India's third lunar mission\",\n",
+    "    \"It aimed to land a rover on the Moon's surface - joining the US, China and Russia\",\n",
+    "    \"The mission is a follow-up to Chandrayaan-2, which had partial success\",\n",
+    "    \"Chandrayaan-3 will be launched by the Indian Space Research Organisation (ISRO)\",\n",
+    "    \"The estimated cost of the mission is around $35 million\",\n",
+    "    \"It will carry instruments to study the lunar surface and atmosphere\",\n",
+    "    \"Chandrayaan-3 landed on the Moon's surface on 23rd August 2023\",\n",
+    "    \"It consists of a lander named Vikram and a rover named Pragyan similar to Chandrayaan-2. Its propulsion module would act like an orbiter.\",\n",
+    "    \"The propulsion module carries the lander and rover configuration until the spacecraft is in a 100-kilometre (62 mi) lunar orbit\",\n",
+    "    \"The mission used GSLV Mk III rocket for its launch\",\n",
+    "    \"Chandrayaan-3 was launched from the Satish Dhawan Space Centre in Sriharikota\",\n",
+    "    \"Chandrayaan-3 was launched earlier in the year 2023\",\n",
+    "]\n",
+    "sparse_embeddings_list: List[SparseEmbedding] = list(\n",
+    "    model.embed(documents, batch_size=6)\n",
+    ")  # batch_size is optional, notice the generator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T00:49:29.646340Z",
+     "start_time": "2024-03-30T00:49:29.643411Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "SparseEmbedding(values=array([0.05297208, 0.01963477, 0.36459631, 1.38508618, 0.71776593,\n",
+       "       0.12667948, 0.46230844, 0.446771  , 0.26897505, 1.01519883,\n",
+       "       1.5655334 , 0.29412213, 1.53102326, 0.59785569, 1.1001817 ,\n",
+       "       0.02079751, 0.09955651, 0.44249091, 0.09747757, 1.53519952,\n",
+       "       1.36765671, 0.15740395, 0.49882549, 0.38629025, 0.76612782,\n",
+       "       1.25805044, 0.39058095, 0.27236196, 0.45152301, 0.48262018,\n",
+       "       0.26085234, 1.35912788, 0.70710695, 1.71639752]), indices=array([ 1010,  1011,  1016,  1017,  2001,  2018,  2034,  2093,  2117,\n",
+       "        2319,  2353,  2509,  2634,  2686,  2796,  2817,  2922,  2959,\n",
+       "        3003,  3148,  3260,  3390,  3462,  3523,  3822,  4231,  4316,\n",
+       "        4774,  5590,  5871,  6416, 11926, 12076, 16469]))"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "index = 0\n",
+    "sparse_embeddings_list[index]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The previous output is a SparseEmbedding object for the first document in our list.\n",
+    "\n",
+    "It contains two arrays: values and indices. \n",
+    "- The 'values' array represents the weights of the features (tokens) in the document.\n",
+    "- The 'indices' array represents the indices of these features in the model's vocabulary.\n",
+    "\n",
+    "Each pair of corresponding values and indices represents a token and its weight in the document."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T00:49:31.549533Z",
+     "start_time": "2024-03-30T00:49:31.546398Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Token at index 1010 has weight 0.05297207832336426\n",
+      "Token at index 1011 has weight 0.01963476650416851\n",
+      "Token at index 1016 has weight 0.36459630727767944\n",
+      "Token at index 1017 has weight 1.385086178779602\n",
+      "Token at index 2001 has weight 0.7177659273147583\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Let's print the first 5 features and their weights for better understanding.\n",
+    "for i in range(5):\n",
+    "    print(f\"Token at index {sparse_embeddings_list[0].indices[i]} has weight {sparse_embeddings_list[0].values[i]}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Understanding SPLADE vectors\n",
+    "\n",
+    "This is still a little abstract, so let's use the tokenizer vocab to make sense of these indices."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T00:49:36.203640Z",
+     "start_time": "2024-03-30T00:49:34.889654Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from transformers import AutoTokenizer\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(SparseTextEmbedding.list_supported_models()[0][\"sources\"][\"hf\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T00:49:36.210049Z",
+     "start_time": "2024-03-30T00:49:36.206825Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{\n",
+      "    \"chandra\": 1.7163975238800049,\n",
+      "    \"third\": 1.5655333995819092,\n",
+      "    \"##ya\": 1.535199522972107,\n",
+      "    \"india\": 1.5310232639312744,\n",
+      "    \"3\": 1.385086178779602,\n",
+      "    \"mission\": 1.3676567077636719,\n",
+      "    \"lunar\": 1.3591278791427612,\n",
+      "    \"moon\": 1.2580504417419434,\n",
+      "    \"indian\": 1.1001816987991333,\n",
+      "    \"##an\": 1.015198826789856,\n",
+      "    \"3rd\": 0.7661278247833252,\n",
+      "    \"was\": 0.7177659273147583,\n",
+      "    \"spacecraft\": 0.7071069478988647,\n",
+      "    \"space\": 0.5978556871414185,\n",
+      "    \"flight\": 0.4988254904747009,\n",
+      "    \"satellite\": 0.4826201796531677,\n",
+      "    \"first\": 0.46230843663215637,\n",
+      "    \"expedition\": 0.4515230059623718,\n",
+      "    \"three\": 0.4467709958553314,\n",
+      "    \"fourth\": 0.44249090552330017,\n",
+      "    \"vehicle\": 0.390580952167511,\n",
+      "    \"iii\": 0.3862902522087097,\n",
+      "    \"2\": 0.36459630727767944,\n",
+      "    \"##3\": 0.2941221296787262,\n",
+      "    \"planet\": 0.27236196398735046,\n",
+      "    \"second\": 0.26897504925727844,\n",
+      "    \"missions\": 0.2608523368835449,\n",
+      "    \"launched\": 0.15740394592285156,\n",
+      "    \"had\": 0.12667948007583618,\n",
+      "    \"largest\": 0.09955651313066483,\n",
+      "    \"leader\": 0.09747757017612457,\n",
+      "    \",\": 0.05297207832336426,\n",
+      "    \"study\": 0.02079751156270504,\n",
+      "    \"-\": 0.01963476650416851\n",
+      "}\n"
+     ]
+    }
+   ],
+   "source": [
+    "def get_tokens_and_weights(sparse_embedding, tokenizer):\n",
+    "    token_weight_dict = {}\n",
+    "    for i in range(len(sparse_embedding.indices)):\n",
+    "        token = tokenizer.decode([sparse_embedding.indices[i]])\n",
+    "        weight = sparse_embedding.values[i]\n",
+    "        token_weight_dict[token] = weight\n",
+    "\n",
+    "    # Sort the dictionary by weights\n",
+    "    token_weight_dict = dict(sorted(token_weight_dict.items(), key=lambda item: item[1], reverse=True))\n",
+    "    return token_weight_dict\n",
+    "\n",
+    "\n",
+    "# Test the function with the first SparseEmbedding\n",
+    "print(json.dumps(get_tokens_and_weights(sparse_embeddings_list[index], tokenizer), indent=4))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Observations and Model Design Choices\n",
+    "\n",
+    "1. The relative order of importance is quite useful. The most important tokens in the sentence have the highest weights.\n",
+    "1. **Term Expansion**: The model can expand the terms in the document. This means that the model can generate weights for tokens that are not present in the document but are related to the tokens in the document. This is a powerful feature that allows the model to capture the context of the document. Here, you'll see that the model has added the tokens '3' from 'third' and 'moon' from 'lunar' to the sparse vector.\n",
+    "\n",
+    "### Design Choices\n",
+    "\n",
+    "1. The weights are not normalized. This means that the sum of the weights is not 1 or 100. This is a common practice in sparse embeddings, as it allows the model to capture the importance of each token in the document.\n",
+    "1. Tokens are included in the sparse vector only if they are present in the model's vocabulary. This means that the model will not generate a weight for tokens that it has not seen during training.\n",
+    "1. Tokens do not map to words directly -- allowing you to gracefully handle typo errors and out-of-vocabulary tokens."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "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
+}

docs/examples/Supported_Models.ipynb

@@ -3,6 +3,39 @@
   {
    "cell_type": "code",
    "execution_count": 1,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T11:18:52.052764Z",
+     "start_time": "2024-03-30T11:18:52.039616Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "\n",
+    "from fastembed import SparseTextEmbedding, TextEmbedding"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Supported Text Embedding Models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -29,70 +62,317 @@
        "      <th>model</th>\n",
        "      <th>dim</th>\n",
        "      <th>description</th>\n",
+       "      <th>size_in_GB</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
-       "      <td>BAAI/bge-small-en</td>\n",
+       "      <td>BAAI/bge-small-en-v1.5</td>\n",
        "      <td>384</td>\n",
        "      <td>Fast and Default English model</td>\n",
+       "      <td>0.067</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
-       "      <td>BAAI/bge-base-en</td>\n",
-       "      <td>768</td>\n",
-       "      <td>Base English model</td>\n",
+       "      <td>BAAI/bge-small-zh-v1.5</td>\n",
+       "      <td>512</td>\n",
+       "      <td>Fast and recommended Chinese model</td>\n",
+       "      <td>0.090</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
        "      <td>384</td>\n",
        "      <td>Sentence Transformer model, MiniLM-L6-v2</td>\n",
+       "      <td>0.090</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-xs</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Based on all-MiniLM-L6-v2 model with only 22m ...</td>\n",
+       "      <td>0.090</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>jinaai/jina-embeddings-v2-small-en</td>\n",
+       "      <td>512</td>\n",
+       "      <td>English embedding model supporting 8192 sequen...</td>\n",
+       "      <td>0.120</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-s</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Based on infloat/e5-small-unsupervised, does n...</td>\n",
+       "      <td>0.130</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
+       "      <td>BAAI/bge-small-en</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Fast English model</td>\n",
+       "      <td>0.130</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</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.210</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>8</th>\n",
+       "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Sentence Transformer model, paraphrase-multili...</td>\n",
+       "      <td>0.220</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>9</th>\n",
+       "      <td>BAAI/bge-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Base English model</td>\n",
+       "      <td>0.420</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>10</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-m</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Based on intfloat/e5-base-unsupervised model, ...</td>\n",
+       "      <td>0.430</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>11</th>\n",
+       "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>English embedding model supporting 8192 sequen...</td>\n",
+       "      <td>0.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>12</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.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>13</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.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>14</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-m-long</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Based on nomic-ai/nomic-embed-text-v1-unsuperv...</td>\n",
+       "      <td>0.540</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>15</th>\n",
+       "      <td>mixedbread-ai/mxbai-embed-large-v1</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>MixedBread Base sentence embedding model, does...</td>\n",
+       "      <td>0.640</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>16</th>\n",
+       "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Sentence-transformers model for tasks like clu...</td>\n",
+       "      <td>1.000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>17</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-l</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Based on intfloat/e5-large-unsupervised, large...</td>\n",
+       "      <td>1.020</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>18</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.200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>19</th>\n",
+       "      <td>thenlper/gte-large</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Large general text embeddings model</td>\n",
+       "      <td>1.200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>20</th>\n",
        "      <td>intfloat/multilingual-e5-large</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Multilingual model, e5-large. Recommend using this model for non-English languages. Recommend using this via Torch implementation of FastEmbed</td>\n",
+       "      <td>Multilingual model, e5-large. Recommend using ...</td>\n",
+       "      <td>2.240</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
-       "                                    model   dim  \\\n",
-       "0                       BAAI/bge-small-en   384   \n",
-       "1                        BAAI/bge-base-en   768   \n",
-       "2  sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "3          intfloat/multilingual-e5-large  1024   \n",
+       "                                                model   dim  \\\n",
+       "0                              BAAI/bge-small-en-v1.5   384   \n",
+       "1                              BAAI/bge-small-zh-v1.5   512   \n",
+       "2              sentence-transformers/all-MiniLM-L6-v2   384   \n",
+       "3                 snowflake/snowflake-arctic-embed-xs   384   \n",
+       "4                  jinaai/jina-embeddings-v2-small-en   512   \n",
+       "5                  snowflake/snowflake-arctic-embed-s   384   \n",
+       "6                                   BAAI/bge-small-en   384   \n",
+       "7                               BAAI/bge-base-en-v1.5   768   \n",
+       "8   sentence-transformers/paraphrase-multilingual-...   384   \n",
+       "9                                    BAAI/bge-base-en   768   \n",
+       "10                 snowflake/snowflake-arctic-embed-m   768   \n",
+       "11                  jinaai/jina-embeddings-v2-base-en   768   \n",
+       "12                       nomic-ai/nomic-embed-text-v1   768   \n",
+       "13                     nomic-ai/nomic-embed-text-v1.5   768   \n",
+       "14            snowflake/snowflake-arctic-embed-m-long   768   \n",
+       "15                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
+       "16  sentence-transformers/paraphrase-multilingual-...   768   \n",
+       "17                 snowflake/snowflake-arctic-embed-l  1024   \n",
+       "18                             BAAI/bge-large-en-v1.5  1024   \n",
+       "19                                 thenlper/gte-large  1024   \n",
+       "20                     intfloat/multilingual-e5-large  1024   \n",
        "\n",
-       "                                                                                                                                      description  \n",
-       "0                                                                                                                  Fast and Default English model  \n",
-       "1                                                                                                                              Base English model  \n",
-       "2                                                                                                        Sentence Transformer model, MiniLM-L6-v2  \n",
-       "3  Multilingual model, e5-large. Recommend using this model for non-English languages. Recommend using this via Torch implementation of FastEmbed  "
+       "                                          description  size_in_GB  \n",
+       "0                      Fast and Default English model       0.067  \n",
+       "1                  Fast and recommended Chinese model       0.090  \n",
+       "2            Sentence Transformer model, MiniLM-L6-v2       0.090  \n",
+       "3   Based on all-MiniLM-L6-v2 model with only 22m ...       0.090  \n",
+       "4   English embedding model supporting 8192 sequen...       0.120  \n",
+       "5   Based on infloat/e5-small-unsupervised, does n...       0.130  \n",
+       "6                                  Fast English model       0.130  \n",
+       "7                            Base English model, v1.5       0.210  \n",
+       "8   Sentence Transformer model, paraphrase-multili...       0.220  \n",
+       "9                                  Base English model       0.420  \n",
+       "10  Based on intfloat/e5-base-unsupervised model, ...       0.430  \n",
+       "11  English embedding model supporting 8192 sequen...       0.520  \n",
+       "12                  8192 context length english model       0.520  \n",
+       "13                  8192 context length english model       0.520  \n",
+       "14  Based on nomic-ai/nomic-embed-text-v1-unsuperv...       0.540  \n",
+       "15  MixedBread Base sentence embedding model, does...       0.640  \n",
+       "16  Sentence-transformers model for tasks like clu...       1.000  \n",
+       "17  Based on intfloat/e5-large-unsupervised, large...       1.020  \n",
+       "18                          Large English model, v1.5       1.200  \n",
+       "19                Large general text embeddings model       1.200  \n",
+       "20  Multilingual model, e5-large. Recommend using ...       2.240  "
       ]
      },
-     "execution_count": 1,
+     "execution_count": 2,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2\n",
-    "\n",
-    "from fastembed.embedding import Embedding\n",
-    "import pandas as pd\n",
-    "pd.set_option('display.max_colwidth', None)\n",
-    "pd.DataFrame(Embedding.list_supported_models())"
+    "supported_models = (\n",
+    "    pd.DataFrame(TextEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=\"sources\")\n",
+    "    .reset_index(drop=True)\n",
+    ")\n",
+    "supported_models"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Supported Sparse Text Embedding Models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-03-30T11:19:01.564291Z",
+     "start_time": "2024-03-30T11:19:01.538768Z"
+    }
+   },
+   "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>vocab_size</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>prithvida/Splade_PP_en_v1</td>\n",
+       "      <td>30522</td>\n",
+       "      <td>Misspelled version of the model. Retained for ...</td>\n",
+       "      <td>0.532</td>\n",
+       "      <td>{'hf': 'Qdrant/SPLADE_PP_en_v1'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>prithivida/Splade_PP_en_v1</td>\n",
+       "      <td>30522</td>\n",
+       "      <td>Independent Implementation of SPLADE++ Model f...</td>\n",
+       "      <td>0.532</td>\n",
+       "      <td>{'hf': 'Qdrant/SPLADE_PP_en_v1'}</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                        model  vocab_size  \\\n",
+       "0   prithvida/Splade_PP_en_v1       30522   \n",
+       "1  prithivida/Splade_PP_en_v1       30522   \n",
+       "\n",
+       "                                         description  size_in_GB  \\\n",
+       "0  Misspelled version of the model. Retained for ...       0.532   \n",
+       "1  Independent Implementation of SPLADE++ Model f...       0.532   \n",
+       "\n",
+       "                            sources  \n",
+       "0  {'hf': 'Qdrant/SPLADE_PP_en_v1'}  \n",
+       "1  {'hf': 'Qdrant/SPLADE_PP_en_v1'}  "
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pd.DataFrame(SparseTextEmbedding.list_supported_models())"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "fst",
+   "display_name": "Python 3.8.18 ('base')",
    "language": "python",
    "name": "python3"
   },
@@ -106,9 +386,14 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.10.13"
   },
-  "orig_nbformat": 4
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "c4a27af61e455bc18dcf16f5867a2ff0402fa12b01dd0f6ce3a79ae73ad15e91"
+   }
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2

docs/experimental/Binary Quantization from Scratch.ipynb

@@ -3,22 +3,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": [
-    "# Binary Quantization of OpenAI Embedding\n",
-    "---\n",
-    "\n",
-    "In the world of large-scale data retrieval and processing, efficiency is crucial. With the exponential growth of data, the ability to retrieve information quickly and accurately can significantly affect system performance. This blog post explores a technique known as binary quantization applied to OpenAI embeddings, demonstrating how it can enhance **retrieval latency by 20x** or more.\n",
-    "\n",
-    "## What Are OpenAI Embeddings?\n",
-    "OpenAI embeddings are numerical representations of textual information. They transform text into a vector space where semantically similar texts are mapped close together. This mathematical representation enables computers to understand and process human language more effectively.\n",
-    "\n",
-    "## Binary Quantization\n",
-    "Binary quantization is a method which converts continuous numerical values into binary values (0 or 1). It simplifies the data structure, allowing faster computations. Here's a brief overview of the binary quantization process applied to OpenAI embeddings:\n",
-    "\n",
-    "1. **Load Embeddings**: OpenAI embeddings are loaded from parquet files.\n",
-    "2. **Binary Transformation**: The continuous valued vectors are converted into binary form. Here, values greater than 0 are set to 1, and others remain 0.\n",
-    "3. **Comparison & Retrieval**: Binary vectors are used for comparison using logical XOR operations and other efficient algorithms."
-   ]
+   "source": []
   },
   {
    "cell_type": "markdown",
@@ -46,7 +31,6 @@
    "source": [
     "import numpy as np\n",
     "import pandas as pd\n",
-    "import matplotlib.pyplot as plt\n",
     "from tqdm import tqdm"
    ]
   },
@@ -305,9 +289,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/index.md

@@ -6,7 +6,7 @@ The default embedding supports "query" and "passage" prefixes for the input text
 
 1. Light & Fast
     - Quantized model weights
-    - ONNX Runtime for inference via [Optimum](github.com/huggingface/optimum)
+    - ONNX Runtime for inference via [Optimum](https://github.com/huggingface/optimum)
 
 2. Accuracy/Recall
     - Better than OpenAI Ada-002
@@ -15,7 +15,7 @@ The default embedding supports "query" and "passage" prefixes for the input text
 
 ## 🚀 Installation
 
-To install the FastEmbed library, pip works: 
+To install the FastEmbed library, pip works:
 
 ```bash
 pip install fastembed
@@ -32,8 +32,8 @@ documents: List[str] = [
     "passage: This is an example passage.",
     "fastembed is supported by and maintained by Qdrant." # You can leave out the prefix but it's recommended
 ]
-embedding_model = Embedding(model_name="BAAI/bge-base-en", max_length=512) 
-embeddings: List[np.ndarray] = embedding_model.embed(documents) # If you use  
+embedding_model = Embedding(model_name="BAAI/bge-base-en", max_length=512)
+embeddings: List[np.ndarray] = embedding_model.embed(documents) # If you use
 ```
 
 ## Usage with Qdrant
@@ -44,7 +44,7 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```
 
-Might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh. 
+Might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
 from qdrant_client import QdrantClient
@@ -73,4 +73,4 @@ search_result = client.query(
     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 %}

docs/qdrant/Binary_Quantization_with_Qdrant.ipynb

@@ -0,0 +1,377 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Binary Quantization with Qdrant & OpenAI Embedding\n",
+    "\n",
+    "---\n",
+    "In the world of large-scale data retrieval and processing, efficiency is crucial. With the exponential growth of data, the ability to retrieve information quickly and accurately can significantly affect system performance. This blog post explores a technique known as binary quantization applied to OpenAI embeddings, demonstrating how it can enhance **retrieval latency by 20x** or more.\n",
+    "\n",
+    "## What Are OpenAI Embeddings?\n",
+    "OpenAI embeddings are numerical representations of textual information. They transform text into a vector space where semantically similar texts are mapped close together. This mathematical representation enables computers to understand and process human language more effectively.\n",
+    "\n",
+    "## Binary Quantization\n",
+    "Binary quantization is a method which converts continuous numerical values into binary values (0 or 1). It simplifies the data structure, allowing faster computations. Here's a brief overview of the binary quantization process applied to OpenAI embeddings:\n",
+    "\n",
+    "1. **Load Embeddings**: OpenAI embeddings are loaded from parquet files.\n",
+    "2. **Binary Transformation**: The continuous valued vectors are converted into binary form. Here, values greater than 0 are set to 1, and others remain 0.\n",
+    "3. **Comparison & Retrieval**: Binary vectors are used for comparison using logical XOR operations and other efficient algorithms.\n",
+    "\n",
+    "Binary Quantization is a promising approach to improve retrieval speeds and reduce memory footprint of vector search engines. In this notebook we will show how to use Qdrant to perform binary quantization of vectors and perform fast similarity search on the resulting index.\n",
+    "\n",
+    "## Table of Contents\n",
+    "1. Imports\n",
+    "2. Download and Slice Dataset\n",
+    "3. Create Qdrant Collection\n",
+    "4. Indexing\n",
+    "5. Search\n",
+    "\n",
+    "## 1. Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:33:03.013948Z",
+     "start_time": "2024-04-01T16:33:01.019043Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "!pip install qdrant-client pandas dataset --quiet --upgrade"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:33:03.914729Z",
+     "start_time": "2024-04-01T16:33:03.015394Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/joein/work/qdrant/fastembed/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import random\n",
+    "import time\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "from qdrant_client import QdrantClient, models\n",
+    "\n",
+    "random.seed(37)\n",
+    "np.random.seed(37)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Download and Slice Dataset\n",
+    "\n",
+    "We will be using the [dbpedia-entities](https://huggingface.co/datasets/Qdrant/dbpedia-entities-openai3-text-embedding-3-small-1536-100K) dataset from the [HuggingFace Datasets](https://huggingface.co/datasets) library. This contains 100K vectors of 1536 dimensions each"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:33:09.085853Z",
+     "start_time": "2024-04-01T16:33:03.912688Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "100000"
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import datasets\n",
+    "\n",
+    "dataset = datasets.load_dataset(\n",
+    "    \"Qdrant/dbpedia-entities-openai3-text-embedding-3-small-1536-100K\", split=\"train\"\n",
+    ")\n",
+    "len(dataset)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:33:09.176212Z",
+     "start_time": "2024-04-01T16:33:09.084550Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "True"
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "client = QdrantClient(\n",
+    "    prefer_grpc=True,\n",
+    ")\n",
+    "\n",
+    "collection_name = \"binary-quantization\"\n",
+    "client.recreate_collection(\n",
+    "    collection_name=collection_name,\n",
+    "    vectors_config=models.VectorParams(\n",
+    "        size=1536,\n",
+    "        distance=models.Distance.DOT,\n",
+    "        on_disk=True,\n",
+    "    ),\n",
+    "    quantization_config=models.BinaryQuantization(\n",
+    "        binary=models.BinaryQuantizationConfig(always_ram=True),\n",
+    "    ),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.383986Z",
+     "start_time": "2024-04-01T16:33:09.175725Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "def iter_dataset(dataset):\n",
+    "    for point in dataset:\n",
+    "        yield point[\"openai\"], {\"text\": point[\"text\"]}\n",
+    "\n",
+    "\n",
+    "vectors, payload = zip(*iter_dataset(dataset))\n",
+    "client.upload_collection(\n",
+    "    collection_name=collection_name,\n",
+    "    vectors=vectors,\n",
+    "    payload=payload,\n",
+    "    parallel=max(1, (os.cpu_count() // 2)),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.390886Z",
+     "start_time": "2024-04-01T16:34:13.385961Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "{'status': <CollectionStatus.YELLOW: 'yellow'>,\n 'optimizer_status': <OptimizersStatusOneOf.OK: 'ok'>,\n 'vectors_count': 116640,\n 'indexed_vectors_count': 43520,\n 'points_count': 116640,\n 'segments_count': 6,\n 'config': {'params': {'vectors': {'size': 1536,\n    'distance': <Distance.DOT: 'Dot'>,\n    'hnsw_config': None,\n    'quantization_config': None,\n    'on_disk': True},\n   'shard_number': 1,\n   'sharding_method': None,\n   'replication_factor': 1,\n   'write_consistency_factor': 1,\n   'read_fan_out_factor': None,\n   'on_disk_payload': True,\n   'sparse_vectors': None},\n  'hnsw_config': {'m': 16,\n   'ef_construct': 100,\n   'full_scan_threshold': 10000,\n   'max_indexing_threads': 0,\n   'on_disk': False,\n   'payload_m': None},\n  'optimizer_config': {'deleted_threshold': 0.2,\n   'vacuum_min_vector_number': 1000,\n   'default_segment_number': 0,\n   'max_segment_size': None,\n   'memmap_threshold': None,\n   'indexing_threshold': 20000,\n   'flush_interval_sec': 5,\n   'max_optimization_threads': None},\n  'wal_config': {'wal_capacity_mb': 32, 'wal_segments_ahead': 0},\n  'quantization_config': {'binary': {'always_ram': True}}},\n 'payload_schema': {}}"
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "collection_info = client.get_collection(collection_name=f\"{collection_name}\")\n",
+    "collection_info.dict()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Oversampling vs Recall\n",
+    "\n",
+    "### Preparing a query dataset\n",
+    "\n",
+    "For the purpose of this illustration, we'll take a few vectors which we know are already in the index and query them. We should get the same vectors back as results from the Qdrant index. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.453626Z",
+     "start_time": "2024-04-01T16:34:13.391567Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "[89391,\n 79659,\n 12006,\n 80978,\n 87219,\n 97885,\n 83155,\n 67504,\n 4645,\n 82711,\n 48395,\n 57375,\n 69208,\n 14136,\n 89515,\n 59880,\n 78730,\n 36952,\n 49620,\n 96486,\n 55473,\n 58179,\n 18926,\n 6489,\n 11931,\n 54146,\n 9850,\n 71259,\n 37825,\n 47331,\n 84964,\n 92399,\n 56669,\n 77042,\n 73744,\n 47993,\n 83780,\n 92429,\n 75114,\n 4463,\n 69030,\n 81185,\n 27950,\n 66217,\n 54652,\n 8260,\n 1151,\n 993,\n 85954,\n 66863,\n 47303,\n 8992,\n 92688,\n 76030,\n 29472,\n 3077,\n 42454,\n 46120,\n 69140,\n 20877,\n 2844,\n 95423,\n 1770,\n 28568,\n 96448,\n 94227,\n 40837,\n 91684,\n 29785,\n 66936,\n 85121,\n 39546,\n 81910,\n 5514,\n 37068,\n 35731,\n 93990,\n 26685,\n 63076,\n 18762,\n 27922,\n 34916,\n 80976,\n 83189,\n 6328,\n 57508,\n 58860,\n 13758,\n 72976,\n 85030,\n 332,\n 34963,\n 85009,\n 31344,\n 11560,\n 58108,\n 85163,\n 17064,\n 44712,\n 45962]"
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "query_indices = random.sample(range(len(dataset)), 100)\n",
+    "query_dataset = dataset[query_indices]\n",
+    "query_indices"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.453928Z",
+     "start_time": "2024-04-01T16:34:13.452405Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "## Add Gaussian noise to any vector\n",
+    "\n",
+    "\n",
+    "def add_noise(vector, noise=0.05):\n",
+    "    return vector + noise * np.random.randn(*vector.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:34:13.457839Z",
+     "start_time": "2024-04-01T16:34:13.455431Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "def correct(results, text):\n",
+    "    return text in [x.payload[\"text\"] for x in results]\n",
+    "\n",
+    "\n",
+    "def count_correct(query_dataset, limit=1, oversampling=1, rescore=False):\n",
+    "    correct_results = 0\n",
+    "    for query_vector, text in zip(query_dataset[\"openai\"], query_dataset[\"text\"]):\n",
+    "        results = client.search(\n",
+    "            collection_name=collection_name,\n",
+    "            query_vector=add_noise(np.array(query_vector)),\n",
+    "            limit=limit,\n",
+    "            search_params=models.SearchParams(\n",
+    "                quantization=models.QuantizationSearchParams(\n",
+    "                    rescore=rescore,\n",
+    "                    oversampling=oversampling,\n",
+    "                )\n",
+    "            ),\n",
+    "        )\n",
+    "        correct_results += correct(results, text)\n",
+    "    return correct_results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:40:48.347002Z",
+     "start_time": "2024-04-01T16:40:42.228551Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "limit_grid = [1, 3, 10, 20, 50]\n",
+    "oversampling_grid = [1.0, 3.0, 5.0]\n",
+    "rescore_grid = [False, True]\n",
+    "results = []\n",
+    "\n",
+    "for limit in limit_grid:\n",
+    "    for oversampling in oversampling_grid:\n",
+    "        for rescore in rescore_grid:\n",
+    "            start = time.perf_counter()\n",
+    "            correct_results = count_correct(\n",
+    "                query_dataset, limit=limit, oversampling=oversampling, rescore=rescore\n",
+    "            )\n",
+    "            end = time.perf_counter()\n",
+    "            results.append(\n",
+    "                {\n",
+    "                    \"limit\": limit,\n",
+    "                    \"oversampling\": oversampling,\n",
+    "                    \"candidates\": int(oversampling * limit),\n",
+    "                    \"rescore\": rescore,\n",
+    "                    \"accuracy\": correct_results / 100,\n",
+    "                    \"total queries\": len(query_dataset[\"text\"]),\n",
+    "                    \"time\": end - start,\n",
+    "                }\n",
+    "            )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-04-01T16:41:55.445405Z",
+     "start_time": "2024-04-01T16:41:55.442687Z"
+    }
+   },
+   "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>candidates</th>\n      <th>rescore</th>\n      <th>accuracy</th>\n      <th>time</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>1</td>\n      <td>False</td>\n      <td>0.90</td>\n      <td>0.221826</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>1</td>\n      <td>True</td>\n      <td>0.91</td>\n      <td>0.134167</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>3</td>\n      <td>False</td>\n      <td>0.88</td>\n      <td>0.115299</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>3</td>\n      <td>True</td>\n      <td>0.97</td>\n      <td>0.209320</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>5</td>\n      <td>False</td>\n      <td>0.84</td>\n      <td>0.154485</td>\n    </tr>\n    <tr>\n      <th>5</th>\n      <td>5</td>\n      <td>True</td>\n      <td>0.91</td>\n      <td>0.124424</td>\n    </tr>\n    <tr>\n      <th>6</th>\n      <td>3</td>\n      <td>False</td>\n      <td>0.99</td>\n      <td>0.121695</td>\n    </tr>\n    <tr>\n      <th>7</th>\n      <td>3</td>\n      <td>True</td>\n      <td>0.96</td>\n      <td>0.123257</td>\n    </tr>\n    <tr>\n      <th>8</th>\n      <td>9</td>\n      <td>False</td>\n      <td>0.94</td>\n      <td>0.119629</td>\n    </tr>\n    <tr>\n      <th>9</th>\n      <td>9</td>\n      <td>True</td>\n      <td>0.98</td>\n      <td>0.119372</td>\n    </tr>\n    <tr>\n      <th>10</th>\n      <td>15</td>\n      <td>False</td>\n      <td>0.90</td>\n      <td>0.121621</td>\n    </tr>\n    <tr>\n      <th>11</th>\n      <td>15</td>\n      <td>True</td>\n      <td>0.97</td>\n      <td>0.125466</td>\n    </tr>\n    <tr>\n      <th>12</th>\n      <td>10</td>\n      <td>False</td>\n      <td>0.93</td>\n      <td>0.135910</td>\n    </tr>\n    <tr>\n      <th>13</th>\n      <td>10</td>\n      <td>True</td>\n      <td>0.95</td>\n      <td>0.138135</td>\n    </tr>\n    <tr>\n      <th>14</th>\n      <td>30</td>\n      <td>False</td>\n      <td>0.94</td>\n      <td>0.177928</td>\n    </tr>\n    <tr>\n      <th>15</th>\n      <td>30</td>\n      <td>True</td>\n      <td>0.98</td>\n      <td>0.254588</td>\n    </tr>\n    <tr>\n      <th>16</th>\n      <td>50</td>\n      <td>False</td>\n      <td>0.94</td>\n      <td>0.268659</td>\n    </tr>\n    <tr>\n      <th>17</th>\n      <td>50</td>\n      <td>True</td>\n      <td>0.96</td>\n      <td>0.269792</td>\n    </tr>\n    <tr>\n      <th>18</th>\n      <td>20</td>\n      <td>False</td>\n      <td>0.96</td>\n      <td>0.249941</td>\n    </tr>\n    <tr>\n      <th>19</th>\n      <td>20</td>\n      <td>True</td>\n      <td>0.96</td>\n      <td>0.247138</td>\n    </tr>\n    <tr>\n      <th>20</th>\n      <td>60</td>\n      <td>False</td>\n      <td>0.97</td>\n      <td>0.251301</td>\n    </tr>\n    <tr>\n      <th>21</th>\n      <td>60</td>\n      <td>True</td>\n      <td>0.98</td>\n      <td>0.256504</td>\n    </tr>\n    <tr>\n      <th>22</th>\n      <td>100</td>\n      <td>False</td>\n      <td>0.98</td>\n      <td>0.270049</td>\n    </tr>\n    <tr>\n      <th>23</th>\n      <td>100</td>\n      <td>True</td>\n      <td>0.97</td>\n      <td>0.248972</td>\n    </tr>\n    <tr>\n      <th>24</th>\n      <td>50</td>\n      <td>False</td>\n      <td>0.97</td>\n      <td>0.306356</td>\n    </tr>\n    <tr>\n      <th>25</th>\n      <td>50</td>\n      <td>True</td>\n      <td>0.98</td>\n      <td>0.257544</td>\n    </tr>\n    <tr>\n      <th>26</th>\n      <td>150</td>\n      <td>False</td>\n      <td>0.98</td>\n      <td>0.238811</td>\n    </tr>\n    <tr>\n      <th>27</th>\n      <td>150</td>\n      <td>True</td>\n      <td>0.99</td>\n      <td>0.263939</td>\n    </tr>\n    <tr>\n      <th>28</th>\n      <td>250</td>\n      <td>False</td>\n      <td>0.99</td>\n      <td>0.256558</td>\n    </tr>\n    <tr>\n      <th>29</th>\n      <td>250</td>\n      <td>True</td>\n      <td>1.00</td>\n      <td>0.335823</td>\n    </tr>\n  </tbody>\n</table>\n</div>",
+      "text/plain": "    candidates  rescore  accuracy      time\n0            1    False      0.90  0.221826\n1            1     True      0.91  0.134167\n2            3    False      0.88  0.115299\n3            3     True      0.97  0.209320\n4            5    False      0.84  0.154485\n5            5     True      0.91  0.124424\n6            3    False      0.99  0.121695\n7            3     True      0.96  0.123257\n8            9    False      0.94  0.119629\n9            9     True      0.98  0.119372\n10          15    False      0.90  0.121621\n11          15     True      0.97  0.125466\n12          10    False      0.93  0.135910\n13          10     True      0.95  0.138135\n14          30    False      0.94  0.177928\n15          30     True      0.98  0.254588\n16          50    False      0.94  0.268659\n17          50     True      0.96  0.269792\n18          20    False      0.96  0.249941\n19          20     True      0.96  0.247138\n20          60    False      0.97  0.251301\n21          60     True      0.98  0.256504\n22         100    False      0.98  0.270049\n23         100     True      0.97  0.248972\n24          50    False      0.97  0.306356\n25          50     True      0.98  0.257544\n26         150    False      0.98  0.238811\n27         150     True      0.99  0.263939\n28         250    False      0.99  0.256558\n29         250     True      1.00  0.335823"
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df = pd.DataFrame(results)\n",
+    "df[[\"candidates\", \"rescore\", \"accuracy\", \"time\"]]\n",
+    "# df.to_csv(\"candidates-rescore-time.csv\", index=False)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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": 4
+}

docs/qdrant/Retrieval_with_FastEmbed.ipynb

@@ -21,7 +21,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -37,13 +37,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
     "from typing import List\n",
     "import numpy as np\n",
-    "from fastembed.embedding import FlagEmbedding as Embedding"
+    "from fastembed import TextEmbedding"
    ]
   },
   {
@@ -58,7 +58,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -84,7 +84,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\")\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 +105,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -124,65 +124,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.06002192,  0.04322132, -0.00545516, -0.04419701, -0.00542277],\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 +175,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.10.13"
   },
   "orig_nbformat": 4
  },

docs/qdrant/Usage_With_Qdrant.ipynb

@@ -28,16 +28,7 @@
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\u001b[33mDEPRECATION: pytorch-lightning 1.6.5 has a non-standard dependency specifier torch>=1.8.*. pip 23.3 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
-      "\u001b[0m"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "!pip install 'qdrant-client[fastembed]' --quiet --upgrade"
    ]
@@ -56,8 +47,6 @@
    "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"
    ]
   },
@@ -117,22 +106,22 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
+      "100%|██████████| 77.7M/77.7M [00:05<00:00, 14.6MiB/s]\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "['77e1e4724dd243b08608f57d5692f6aa',\n",
-       " '74841e5dc3594646bda2c6a6d2795dbd',\n",
-       " '6ef39a9445604d0da84d04f760cd7cf7',\n",
-       " 'e659503d3b3748ef90f23c778274835b',\n",
-       " 'b999675068cd413f93faa0cc890c3819',\n",
-       " '8e452f2935cf4e4b80d8eea68c2aad58',\n",
-       " '28ed4fd4592c48c9a0519618d51bb86e',\n",
-       " '59378c784c5f49109bef65fdc4061334',\n",
-       " 'a78c9b598f7942749156334283a6f24f',\n",
-       " 'f72bb24701c64fabb0182c9e757b581b']"
+       "['4fa8b10c78da4b18ba0830ba8a57367a',\n",
+       " '2eae04b515ee4e9185a9a0e6be812bba',\n",
+       " 'c6039f88486f47f1835ae3b069c5823c',\n",
+       " 'c2c8c51e305144d1917b373125fb4d95',\n",
+       " '79fd23b9ec0648cdab38d1947c6b933e',\n",
+       " '036aa200d8c3492b8a438e4f825f5e7f',\n",
+       " 'c35c77f3ea37460a9a13723fb77b7367',\n",
+       " '6ebccbca571b40d0ab6e83e5e0f2f562',\n",
+       " '38048c2ccc1d4962a4f8f1bd89c8357a',\n",
+       " 'c6b09308360140c7b4f106af3658a31e']"
       ]
      },
      "execution_count": 4,
@@ -186,12 +175,14 @@
     "ids = [42, 2]\n",
     "\n",
     "# Use the new add method\n",
-    "client.add(\n",
-    "    collection_name=\"demo_collection\",\n",
-    "    documents=docs,\n",
-    "    metadata=metadata,\n",
-    "    ids=ids\n",
-    ")"
+    "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."
    ]
   },
   {
@@ -203,15 +194,12 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[QueryResponse(id='42', embedding=None, metadata={'document': 'Qdrant has Langchain integrations', 'source': 'Langchain-docs'}, document='Qdrant has Langchain integrations', score=0.8496814051311954), QueryResponse(id='2', embedding=None, metadata={'document': 'Qdrant also has Llama Index integrations', 'source': 'Linkedin-docs'}, document='Qdrant also has Llama Index integrations', score=0.8478494193031256)]\n"
+      "[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": [
-    "search_result = client.query(\n",
-    "    collection_name=\"demo_collection\",\n",
-    "    query_text=[\"This is a query document\"]\n",
-    ")\n",
+    "search_result = client.query(collection_name=\"demo_collection\", query_text=\"This is a query document\")\n",
     "print(search_result)"
    ]
   },
@@ -245,7 +233,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.11.5"
   },
   "orig_nbformat": 4
  },

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,7 @@
+import importlib.metadata
+
+from fastembed.text import TextEmbedding
+from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
+
+__version__ = importlib.metadata.version("fastembed")
+__all__ = ["TextEmbedding", "SparseTextEmbedding", "SparseEmbedding"]

fastembed/common/model_management.py

@@ -0,0 +1,239 @@
+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
+
+
+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.lower() == model["model"].lower():
+                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,
+        extra_patterns: Optional[List[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.
+            extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
+                includes the required model files.
+        Returns:
+            Path: The path to the model directory.
+        """
+        allow_patterns = [
+            "config.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+        ]
+        if extra_patterns is not None:
+            allow_patterns.extend(extra_patterns)
+
+        return snapshot_download(
+            repo_id=hf_source_repo,
+            allow_patterns=allow_patterns,
+            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"
+
+        if model_tar_gz.exists():
+            model_tar_gz.unlink()
+
+        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:
+            extra_patterns = [model["model_file"]]
+            extra_patterns.extend(model.get("additional_files", []))
+
+            try:
+                return Path(
+                    cls.download_files_from_huggingface(
+                        hf_source, cache_dir=str(cache_dir), extra_patterns=extra_patterns
+                    )
+                )
+            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,54 @@
+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,163 @@
+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,
+    Sequence,
+)
+
+import numpy as np
+import onnxruntime as ort
+
+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,
+        model_file: str,
+        threads: Optional[int],
+        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
+    ) -> None:
+        model_path = model_dir / model_file
+
+        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+
+        onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
+        available_providers = ort.get_available_providers()
+        for provider in onnx_providers:
+            # check providers available
+            provider_name = provider if isinstance(provider, str) else provider[0]
+            if provider_name not in available_providers:
+                raise ValueError(
+                    f"Provider {provider_name} is not available. Available providers: {available_providers}"
+                )
+
+        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)
+        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,404 +1,24 @@
-import json
-import os
-import shutil
-import tarfile
-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Dict, Iterable, List, Union
+from typing import Optional
 
-import onnxruntime as ort
-import numpy as np
-import requests
-from tokenizers import Tokenizer, AddedToken
-from tqdm import tqdm
+from loguru import logger
+
+from fastembed 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, 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 list_supported_models(cls) -> List[Dict[str, Union[str, int]]]:
-        """
-        Lists the supported models.
-        """
-        return [
-            {
-                "model": "BAAI/bge-small-en",
-                "dim": 384,
-                "description": "Fast and Default English model",
-            },
-            {
-                "model": "BAAI/bge-base-en",
-                "dim": 768,
-                "description": "Base English model",
-            },
-            {
-                "model": "sentence-transformers/all-MiniLM-L6-v2",
-                "dim": 384,
-                "description": "Sentence Transformer model, MiniLM-L6-v2",
-            },
-            {
-                "model": "intfloat/multilingual-e5-large",
-                "dim": 1024,
-                "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
-            },
-        ]
-
-    @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_
-    """
-
-    @classmethod
-    def load_tokenizer(cls, model_dir: Path, max_length: int = 512) -> Tokenizer:
-        config_path = model_dir / "config.json"
-        if not config_path.exists():
-            raise ValueError(f"Could not find config.json in {model_dir}")
-
-        tokenizer_path = model_dir / "tokenizer.json"
-        if not tokenizer_path.exists():
-            raise ValueError(f"Could not find tokenizer.json in {model_dir}")
-
-        tokenizer_config_path = model_dir / "tokenizer_config.json"
-        if not tokenizer_config_path.exists():
-            raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
-
-        tokens_map_path = model_dir / "special_tokens_map.json"
-        if not tokens_map_path.exists():
-            raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
-
-        config = json.load(open(str(config_path)))
-        tokenizer_config = json.load(open(str(tokenizer_config_path)))
-        tokens_map = json.load(open(str(tokens_map_path)))
-
-        tokenizer = Tokenizer.from_file(str(tokenizer_path))
-        tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
-        tokenizer.enable_padding(pad_id=config["pad_token_id"], pad_token=tokenizer_config["pad_token"])
-
-        for token in tokens_map.values():
-            if isinstance(token, str):
-                tokenizer.add_special_tokens([token])
-            elif isinstance(token, dict):
-                tokenizer.add_special_tokens([AddedToken(**token)])
-
-        return tokenizer
-
+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.
-            cache_dir (str, optional): The path to the cache directory. Defaults to `local_cache` in the current directory.
-
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
-        """
-        self.model_name = model_name
-
-        if cache_dir is None:
-            cache_dir = Path(".").resolve() / "local_cache"
-            cache_dir.mkdir(parents=True, exist_ok=True)
-
-        model_dir = self.retrieve_model(model_name, cache_dir)
-
-        model_path = model_dir / "model.onnx"
-        optimized_model_path = model_dir / "model_optimized.onnx"
-
-        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
-        onnx_providers = ["CPUExecutionProvider"]
-
-        if not model_path.exists():
-            # Rename file model_optimized.onnx to model.onnx if it exists
-            if optimized_model_path.exists():
-                optimized_model_path.rename(model_path)
-            else:
-                raise ValueError(f"Could not find model.onnx in {model_dir}")
-
-        # Hacky support for multilingual model
-        self.exclude_token_type_ids = False
-        if model_name == "intfloat/multilingual-e5-large":
-            self.exclude_token_type_ids = True
-
-        so = ort.SessionOptions()
-        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-
-        self.tokenizer = self.load_tokenizer(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]) -> np.ndarray:
-        encoded = self.tokenizer.encode_batch(documents)
-        input_ids = np.array([e.ids for e in encoded])
-        attention_mask = np.array([e.attention_mask for e in encoded])
-
-        onnx_input = {
-            "input_ids": np.array(input_ids, dtype=np.int64),
-            "attention_mask": np.array(attention_mask, dtype=np.int64),
-        }
-
-        if not self.exclude_token_type_ids:
-            onnx_input["token_type_ids"] = np.array(
-                [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
-            )
-
-        model_output = self.model.run(None, onnx_input)
-        last_hidden_state = model_output[0][:, 0]
-        embeddings = normalize(last_hidden_state).astype(np.float32)
-        return embeddings
-
-    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
-        """
-        if type(documents) == str:
-            documents = [documents]
-        # 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]
-                yield from self.onnx_embed(batch)
-        else:
-            vectors = self.onnx_embed(documents)
-            yield from vectors
-
-
-class DefaultEmbedding(FlagEmbedding):
-    """
-    Implementation of the default Flag Embedding model.
-
-    Args:
-        FlagEmbedding (_type_): _description_
-    """
-
-    def __init__(
-        self,
-        model_name: str = "BAAI/bge-small-en",
-        max_length: int = 512,
-        cache_dir: str = None,
-    ):
-        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
+        super().__init__(model_name, cache_dir, threads, **kwargs)

fastembed/parallel_processor.py

@@ -0,0 +1,209 @@
+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/__init__.py

@@ -0,0 +1,4 @@
+from fastembed.sparse.sparse_embedding_base import SparseEmbedding
+from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
+
+__all__ = ["SparseEmbedding", "SparseTextEmbedding"]

fastembed/sparse/sparse_embedding_base.py

@@ -0,0 +1,43 @@
+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.lower() == model["model"].lower() 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,134 @@
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type
+
+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": "Misspelled version of the model. Retained for backward compatibility. Independent Implementation of SPLADE++ Model for English",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "Qdrant/SPLADE_PP_en_v1",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "prithivida/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",
+        },
+        "model_file": "model.onnx",
+    },
+]
+
+
+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)
+
+        scores = 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
+        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)
+
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)
+
+        model_dir = self.download_model(model_description, cache_dir)
+
+        self.load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_description["model_file"],
+            threads=threads,
+        )
+
+    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,
+        )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[EmbeddingWorker]:
+        return SpladePPEmbeddingWorker
+
+
+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/__init__.py

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

fastembed/text/e5_onnx_embedding.py

@@ -0,0 +1,62 @@
+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_file": "model.onnx",
+        "additional_files": ["model.onnx_data"],
+    },
+    {
+        "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.00,
+        "sources": {
+            "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+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,67 @@
+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.52,
+        "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-small-en",
+        "dim": 512,
+        "description": "English embedding model supporting 8192 sequence length",
+        "size_in_GB": 0.12,
+        "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+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,294 @@
+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.42,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-base-en-v1.5",
+        "dim": 768,
+        "description": "Base English model, v1.5",
+        "size_in_GB": 0.21,
+        "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_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-large-en-v1.5",
+        "dim": 1024,
+        "description": "Large English model, v1.5",
+        "size_in_GB": 1.20,
+        "sources": {
+            "hf": "qdrant/bge-large-en-v1.5-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "BAAI/bge-small-en",
+        "dim": 384,
+        "description": "Fast English model",
+        "size_in_GB": 0.13,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-small-en-v1.5",
+        "dim": 384,
+        "description": "Fast and Default English model",
+        "size_in_GB": 0.067,
+        "sources": {
+            "hf": "qdrant/bge-small-en-v1.5-onnx-q",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-small-zh-v1.5",
+        "dim": 512,
+        "description": "Fast and recommended Chinese model",
+        "size_in_GB": 0.09,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "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_file": "model.onnx",
+    },
+    {
+        "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+        "dim": 384,
+        "description": "Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2",
+        "size_in_GB": 0.22,
+        "sources": {
+            "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
+        "dim": 768,
+        "description": "Quantized 8192 context length english model",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model_quantized.onnx",
+    },
+    {
+        "model": "thenlper/gte-large",
+        "dim": 1024,
+        "description": "Large general text embeddings model",
+        "size_in_GB": 1.20,
+        "sources": {
+            "hf": "qdrant/gte-large-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "mixedbread-ai/mxbai-embed-large-v1",
+        "dim": 1024,
+        "description": "MixedBread Base sentence embedding model, does well on MTEB",
+        "size_in_GB": 0.64,
+        "sources": {
+            "hf": "mixedbread-ai/mxbai-embed-large-v1",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-xs",
+        "dim": 384,
+        "description": "Based on all-MiniLM-L6-v2 model with only 22m parameters, ideal for latency/TCO budgets.",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-xs",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-s",
+        "dim": 384,
+        "description": "Based on infloat/e5-small-unsupervised, does not trade off retrieval accuracy for its small size.",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-s",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m",
+        "dim": 768,
+        "description": "Based on intfloat/e5-base-unsupervised model, provides the best retrieval without slowing down inference.",
+        "size_in_GB": 0.43,
+        "sources": {
+            "hf": "Snowflake/snowflake-arctic-embed-m",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m-long",
+        "dim": 768,
+        "description": "Based on nomic-ai/nomic-embed-text-v1-unsupervised model, 8192 context-length model",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-m-long",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-l",
+        "dim": 1024,
+        "description": "Based on intfloat/e5-large-unsupervised, large model for most accurate retrieval.",
+        "size_in_GB": 1.02,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-l",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+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)
+
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)
+        model_dir = self.download_model(model_description, cache_dir)
+
+        self.load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_description["model_file"],
+            threads=threads,
+        )
+
+    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.lower() == model["model"].lower() 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,59 @@
+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.5"
+version = "0.2.6"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -11,34 +11,46 @@ repository = "https://github.com/qdrant/fastembed"
 keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-transformers"]
 
 [tool.poetry.dependencies]
-python = ">=3.8.0,<3.12"
-onnx = "^1.11"
-onnxruntime = "^1.15"
-tqdm = "^4.65"
-requests = "^2.31"
-tokenizers = "^0.13"
+python = ">=3.8.0,<3.13"
+onnx = [
+    {version = "^1.15.0", optional = true, markers = "extra != 'gpu'"}
+]
+onnxruntime = [
+    {version = "^1.17.0", optional = true, markers = "extra != 'gpu'"}
+]
+onnxruntime-gpu = [
+    { version = "^1.17.0", optional = true, python = "<3.13", markers = "extra == 'gpu'" }
+]
 
-[tool.poetry.dev-dependencies]
-ruff = "^0.0.277"
-isort = "^5.12.0"
-black = "^23.7.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.group.dev.dependencies]
+pytest = "^7.4.2"
+ruff = "^0.3.1"
 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.extras]
+gpu = ["onnxruntime-gpu"]
 
 [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
+line-length = 99

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,35 +0,0 @@
-import numpy as np
-
-from fastembed.embedding import DefaultEmbedding, Embedding
-
-
-CANONICAL_VECTOR_VALUES = {
-    "BAAI/bge-small-en": np.array([-0.0232, -0.0255,  0.0174, -0.0639, -0.0006]),
-    "BAAI/bge-base-en": np.array([0.0115,  0.0372,  0.0295,  0.0121,  0.0346]),
-    "sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259,  0.0058,  0.0114,  0.0380, -0.0233]),
-    "intfloat/multilingual-e5-large": np.array([0.0098,  0.0045,  0.0066, -0.0354,  0.0070]),
-}
-
-
-def test_default_embedding():
-    for model_desc in Embedding.list_supported_models():
-        dim = model_desc["dim"]
-        model = DefaultEmbedding(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)
-
-
-def test_batch_embedding():
-    model = DefaultEmbedding()
-
-    docs = ["hello world", "flag embedding"] * 100
-    embeddings = list(model.embed(docs, batch_size=10))
-    embeddings = np.stack(embeddings, axis=0)
-
-    assert embeddings.shape == (200, 384)

tests/test_sparse_embeddings.py

@@ -0,0 +1,98 @@
+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]
+
+
+def test_single_embedding():
+    docs_to_embed = docs
+
+    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]
+
+
+def test_parallel_processing():
+    import numpy as np
+
+    model = SparseTextEmbedding(
+        model_name="prithivida/Splade_PP_en_v1",
+    )
+    docs = ["hello world", "flag embedding"] * 30
+    sparse_embeddings_duo = list(model.embed(docs, batch_size=10, parallel=2))
+    sparse_embeddings_all = list(model.embed(docs, batch_size=10, parallel=0))
+    sparse_embeddings = list(model.embed(docs, batch_size=10, parallel=None))
+
+    assert len(sparse_embeddings) == len(sparse_embeddings_duo) == len(sparse_embeddings_all) == len(docs)
+
+    for sparse_embedding, sparse_embedding_duo, sparse_embedding_all in zip(
+        sparse_embeddings, sparse_embeddings_duo, sparse_embeddings_all
+    ):
+        assert (
+            sparse_embedding.indices.tolist()
+            == sparse_embedding_duo.indices.tolist()
+            == sparse_embedding_all.indices.tolist()
+        )
+        assert np.allclose(sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3)
+        assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3)

tests/test_text_onnx_embeddings.py

@@ -0,0 +1,94 @@
+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-en-v1.5-quantized": 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]
+    ),
+    "nomic-ai/nomic-embed-text-v1.5-Q": np.array(
+        [-0.01554983,  0.0129992 , -0.17909265, -0.01062993,  0.00512859]
+    ),
+    "thenlper/gte-large": np.array([-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]),
+    "mixedbread-ai/mxbai-embed-large-v1": np.array([0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]),
+    "snowflake/snowflake-arctic-embed-xs": np.array([0.0092, 0.0619, 0.0196, 0.009, -0.0114]),
+    "snowflake/snowflake-arctic-embed-s": np.array([-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]),
+    "snowflake/snowflake-arctic-embed-m": np.array([-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]),
+    "snowflake/snowflake-arctic-embed-m-long": np.array(
+        [0.0080, -0.0266, -0.0335, 0.0282, 0.0143]
+    ),
+    "snowflake/snowflake-arctic-embed-l": np.array([0.0189, -0.0673, 0.0183, 0.0124, 0.0146]),
+}
+
+
+def test_embedding():
+    is_ci = os.getenv("CI")
+
+    for model_desc in TextEmbedding.list_supported_models():
+        if not is_ci 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)