new: update dev workflow

fix links (#215 )
new: remove slurm environment variables (#213 )
2026-09-24 06:57:51 -05:00 · 2024-05-02 20:34:47 +02:00 · 2024-04-27 22:24:29 +05:30 · 2024-04-26 21:45:02 +02:00 · 2024-04-26 17:27:16 +02:00 · 2024-04-26 10:28:37 +05:30
57 changed files with 7295 additions and 4955 deletions
@@ -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
@@ -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.
@@ -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 }}
@@ -1,8 +1,9 @@
 name: Tests
+run-name: Tests (dev)

 on:
  push:
-    branches: [ master, main ]
+    branches: [ dev ]
  pull_request:

 env:
@@ -10,7 +11,6 @@ env:

 jobs:
  test:
-
    strategy:
      matrix:
        python-version:
@@ -18,6 +18,7 @@ jobs:
          - '3.9.x'
          - '3.10.x'
          - '3.11.x'
+          - '3.12.x'
        os:
          - ubuntu-latest
          - macos-latest
@@ -28,18 +29,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: |
-          export IS_UBUNTU_CI=$(test "${{ matrix.os }}" = "ubuntu-latest" && echo "true" || echo "false")
-          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'
@@ -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/*
@@ -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 ]
@@ -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.
@@ -1,43 +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_with_FastEmbed/) 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/qdrant/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/qdrant/Usage_With_Qdrant/).

-1. Light & Fast
-    - Quantized model weights
-    - ONNX Runtime, no PyTorch dependency
-    - CPU-first design
-    - Data-parallelism for encoding of large datasets
+## 📈 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
-import numpy as np

+# 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)) # Note the list() call - this is a generator 
+
+# 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
@@ -48,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,
@@ -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,43 +33,115 @@
   "metadata": {},
   "outputs": [
    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|██████████| 76.7M/76.7M [00:05<00:00, 15.0MiB/s]\n",
-      "100%|██████████| 3/3 [00:00<00:00, 455.37it/s]"
-     ]
+     "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": [
-      "(384,)\n"
+      "The model BAAI/bge-small-en-v1.5 is ready to use.\n"
     ]
    },
    {
-     "name": "stderr",
+     "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": [
-      "\n"
+      "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\n",
-    "embedding_model = DefaultEmbedding()\n",
-    "embeddings: List[np.ndarray] = list(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"
   ]
  },
  {
@@ -83,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"
-   ]
-  },
-  {
-   "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": 3,
-   "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."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "272c8915",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "embedding_model = DefaultEmbedding()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "5549d501",
-   "metadata": {},
-   "source": [
-    "### Embed your Documents\n",
-    "\n",
-    "Use the embed method of the embedding model to transform the documents into a List of np.array. The method returns a generator, so we cast it to a list to get the embeddings."
+    "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": "8013eee9",
+   "id": "2e9c8766",
   "metadata": {},
   "outputs": [
    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|██████████| 4/4 [00:00<00:00, 361.82it/s]\n"
-     ]
+     "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": [
-    "embeddings: List[np.ndarray] = list(embedding_model.embed(documents))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e5b5a6ad",
-   "metadata": {},
-   "source": [
-    "You can print the shape of the embeddings to understand their dimensions. Typically, the shape will indicate the number of dimensions in the vector."
+    "multilingual_large_model = TextEmbedding(\"intfloat/multilingual-e5-large\")  # This can take a few minutes to download"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
-   "id": "0d8c8e08",
+   "id": "a9e70f0e",
   "metadata": {},
   "outputs": [
    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(384,)\n"
-     ]
+     "data": {
+      "text/plain": [
+       "(4, 1024)"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
    }
   ],
   "source": [
-    "print(embeddings[0].shape)  # (384,) or similar output"
+    "np.array(\n",
+    "    list(multilingual_large_model.embed([\"Hello, world!\", \"你好世界\", \"¡Hola Mundo!\", \"नमस्ते!\"]))\n",
+    ").shape  # Vector of 1024 dimensions"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "64fe20ed",
+   "metadata": {},
+   "source": [
+    "Next: Checkout how to use FastEmbed with Qdrant for similarity search: [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/)"
   ]
  }
 ],
@@ -238,7 +236,7 @@
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.10.13"
  }
 },
 "nbformat": 4,
@@ -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
+}
@@ -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": [
    {
@@ -35,125 +68,311 @@
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
-       "      <td>BAAI/bge-small-en</td>\n",
-       "      <td>384</td>\n",
-       "      <td>Fast English model</td>\n",
-       "      <td>0.20</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
       "      <td>BAAI/bge-small-en-v1.5</td>\n",
       "      <td>384</td>\n",
       "      <td>Fast and Default English model</td>\n",
-       "      <td>0.13</td>\n",
+       "      <td>0.067</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>2</th>\n",
+       "      <th>1</th>\n",
       "      <td>BAAI/bge-small-zh-v1.5</td>\n",
       "      <td>512</td>\n",
       "      <td>Fast and recommended Chinese model</td>\n",
-       "      <td>0.10</td>\n",
+       "      <td>0.090</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>BAAI/bge-base-en</td>\n",
-       "      <td>768</td>\n",
-       "      <td>Base English model</td>\n",
-       "      <td>0.50</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>BAAI/bge-base-en-v1.5</td>\n",
-       "      <td>768</td>\n",
-       "      <td>Base English model, v1.5</td>\n",
-       "      <td>0.44</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>5</th>\n",
+       "      <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.09</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>intfloat/multilingual-e5-large</td>\n",
-       "      <td>1024</td>\n",
-       "      <td>Multilingual model, e5-large. Recommend using this model for non-English languages</td>\n",
-       "      <td>2.24</td>\n",
+       "      <td>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>jinaai/jina-embeddings-v2-base-en</td>\n",
+       "      <td>BAAI/bge-base-en-v1.5</td>\n",
       "      <td>768</td>\n",
-       "      <td>English embedding model supporting 8192 sequence length</td>\n",
-       "      <td>0.55</td>\n",
+       "      <td>Base English model, v1.5</td>\n",
+       "      <td>0.210</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
-       "      <td>jinaai/jina-embeddings-v2-small-en</td>\n",
-       "      <td>512</td>\n",
-       "      <td>English embedding model supporting 8192 sequence length</td>\n",
-       "      <td>0.13</td>\n",
+       "      <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 ...</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-small-en-v1.5   384   \n",
-       "2                  BAAI/bge-small-zh-v1.5   512   \n",
-       "3                        BAAI/bge-base-en   768   \n",
-       "4                   BAAI/bge-base-en-v1.5   768   \n",
-       "5  sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "6          intfloat/multilingual-e5-large  1024   \n",
-       "7       jinaai/jina-embeddings-v2-base-en   768   \n",
-       "8      jinaai/jina-embeddings-v2-small-en   512   \n",
+       "                                                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 English model   \n",
-       "1                                                      Fast and Default English model   \n",
-       "2                                                  Fast and recommended Chinese model   \n",
-       "3                                                                  Base English model   \n",
-       "4                                                            Base English model, v1.5   \n",
-       "5                                            Sentence Transformer model, MiniLM-L6-v2   \n",
-       "6  Multilingual model, e5-large. Recommend using this model for non-English languages   \n",
-       "7                             English embedding model supporting 8192 sequence length   \n",
-       "8                             English embedding model supporting 8192 sequence length   \n",
-       "\n",
-       "   size_in_GB  \n",
-       "0        0.20  \n",
-       "1        0.13  \n",
-       "2        0.10  \n",
-       "3        0.50  \n",
-       "4        0.44  \n",
-       "5        0.09  \n",
-       "6        2.24  \n",
-       "7        0.55  \n",
-       "8        0.13  "
+       "                                          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"
  },
@@ -167,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
@@ -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",
@@ -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)
-```
+```
@@ -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() }}

@@ -24,4 +24,4 @@
        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 %}
@@ -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
+}
@@ -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
 },
@@ -47,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"
   ]
  },
@@ -104,19 +102,26 @@
   "execution_count": 4,
   "metadata": {},
   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 77.7M/77.7M [00:05<00:00, 14.6MiB/s]\n"
+     ]
+    },
    {
     "data": {
      "text/plain": [
-       "['6e8fcf7e0ecc407b9b6bb011d169f629',\n",
-       " 'c9d26e7e0ea741b2b1082d097796b28b',\n",
-       " 'cf05747e7eb34d2490b1df1f8be94049',\n",
-       " '208c197266d547a880dfb65e46738b19',\n",
-       " '27bd985c5d6f49d68fc2cf73dac74199',\n",
-       " 'c5e929c8837f4370818c97f63996f8ef',\n",
-       " 'c12213c6cdac470aa2471f2d30dc4041',\n",
-       " '974e64a7d8624f6e9824fa7b9c94f99d',\n",
-       " '0129fae193c740eba092512d8e53ab4a',\n",
-       " '492cad6e741e4aeebb196bd818a97d17']"
+       "['4fa8b10c78da4b18ba0830ba8a57367a',\n",
+       " '2eae04b515ee4e9185a9a0e6be812bba',\n",
+       " 'c6039f88486f47f1835ae3b069c5823c',\n",
+       " 'c2c8c51e305144d1917b373125fb4d95',\n",
+       " '79fd23b9ec0648cdab38d1947c6b933e',\n",
+       " '036aa200d8c3492b8a438e4f825f5e7f',\n",
+       " 'c35c77f3ea37460a9a13723fb77b7367',\n",
+       " '6ebccbca571b40d0ab6e83e5e0f2f562',\n",
+       " '38048c2ccc1d4962a4f8f1bd89c8357a',\n",
+       " 'c6b09308360140c7b4f106af3658a31e']"
      ]
     },
     "execution_count": 4,
@@ -170,12 +175,7 @@
    "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)"
   ]
  },
  {
@@ -194,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)"
   ]
  },
@@ -236,7 +233,7 @@
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.11.5"
  },
  "orig_nbformat": 4
 },
@@ -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",
@@ -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
+}
@@ -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)
@@ -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])
@@ -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"]
@@ -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.")
@@ -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
@@ -0,0 +1,141 @@
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, TypeVar, Union
+
+import numpy as np
+import onnxruntime as ort
+
+from fastembed.common.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],
+    ) -> None:
+        model_path = model_dir / model_file
+
+        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+        onnx_providers = ["CPUExecutionProvider"]
+
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        if threads is not None:
+            so.intra_op_num_threads = threads
+            so.inter_op_num_threads = threads
+
+        self.tokenizer = load_tokenizer(model_dir=model_dir)
+        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)
@@ -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
@@ -1,689 +1,24 @@
-import json
-import os
-import shutil
-import tarfile
-import tempfile
-from abc import ABC, abstractmethod
-from itertools import islice
-from multiprocessing import get_all_start_methods
-from pathlib import Path
-from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple, Union
+from typing import Optional

-import numpy as np
-import onnxruntime as ort
-import requests
-from tokenizers import AddedToken, Tokenizer
-from tqdm import tqdm
+from loguru import logger

-from fastembed.parallel_processor import ParallelWorkerPool, Worker
-from fastembed.splitter import FastEmbedRecursiveSplitter
+from fastembed import TextEmbedding

-from .models import TextSplitterConfig
+logger.warning(
+    "DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated."
+    "Use from fastembed import TextEmbedding instead."
+)

+DefaultEmbedding = TextEmbedding
+FlagEmbedding = TextEmbedding

-def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
-    """
-    >>> list(iter_batch([1,2,3,4,5], 3))
-    [[1, 2, 3], [4, 5]]
-    """
-    source_iter = iter(iterable)
-    while source_iter:
-        b = list(islice(source_iter, size))
-        if len(b) == 0:
-            break
-        yield b

-
-def normalize(input_array, p=2, dim=1, eps=1e-12):
-    # Calculate the Lp norm along the specified dimension
-    norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
-    norm = np.maximum(norm, eps)  # Avoid division by zero
-    normalized_array = input_array / norm
-    return normalized_array
-
-
-class EmbeddingModel:
-    @classmethod
-    def load_tokenizer(cls, model_dir: Path, max_length: int = 512) -> Tokenizer:
-        config_path = model_dir / "config.json"
-        if not config_path.exists():
-            raise ValueError(f"Could not find config.json in {model_dir}")
-
-        tokenizer_path = model_dir / "tokenizer.json"
-        if not tokenizer_path.exists():
-            raise ValueError(f"Could not find tokenizer.json in {model_dir}")
-
-        tokenizer_config_path = model_dir / "tokenizer_config.json"
-        if not tokenizer_config_path.exists():
-            raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
-
-        tokens_map_path = model_dir / "special_tokens_map.json"
-        if not tokens_map_path.exists():
-            raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
-
-        config = json.load(open(str(config_path)))
-        tokenizer_config = json.load(open(str(tokenizer_config_path)))
-        tokens_map = json.load(open(str(tokens_map_path)))
-
-        tokenizer = Tokenizer.from_file(str(tokenizer_path))
-        tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
-        # TODO: Truncation should be enabled by default, but we need to disable this for when being used with chunking
-        tokenizer.enable_padding(pad_id=config["pad_token_id"], pad_token=tokenizer_config["pad_token"])
-
-        for token in tokens_map.values():
-            if isinstance(token, str):
-                tokenizer.add_special_tokens([token])
-            elif isinstance(token, dict):
-                tokenizer.add_special_tokens([AddedToken(**token)])
-
-        return tokenizer
-
-    def __init__(
-        self,
-        path: Path,
-        model_name: str,
-        max_length: int = 512,
-        max_threads: int = None,
-        splitter_config: Optional[TextSplitterConfig] = None,
-    ):
-        self.path = path
-        self.model_name = model_name
-        model_path = self.path / "model.onnx"
-        optimized_model_path = self.path / "model_optimized.onnx"
-
-        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
-        onnx_providers = ["CPUExecutionProvider"]
-
-        if not model_path.exists():
-            # Rename file model_optimized.onnx to model.onnx if it exists
-            if optimized_model_path.exists():
-                optimized_model_path.rename(model_path)
-            else:
-                raise ValueError(f"Could not find model.onnx in {self.path}")
-
-        # Hacky support for multilingual model
-        self.exclude_token_type_ids = False
-        if model_name == "intfloat/multilingual-e5-large":
-            self.exclude_token_type_ids = True
-
-        so = ort.SessionOptions()
-        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-
-        if max_threads is not None:
-            so.intra_op_num_threads = max_threads
-            so.inter_op_num_threads = max_threads
-
-        self.tokenizer = self.load_tokenizer(self.path, max_length=max_length)
-        self.model = ort.InferenceSession(str(model_path), providers=onnx_providers, sess_options=so)
-        if splitter_config:
-            splitter_config.tokenizer = self.tokenizer
-            self.splitter = FastEmbedRecursiveSplitter(config=splitter_config)
-
-    def onnx_embed(self, documents: List[str]) -> Tuple[np.ndarray, np.ndarray]:
-        encoded = self.tokenizer.encode_batch(documents)
-        input_ids = np.array([e.ids for e in encoded])
-        attention_mask = np.array([e.attention_mask for e in encoded])
-
-        onnx_input = {
-            "input_ids": np.array(input_ids, dtype=np.int64),
-            "attention_mask": np.array(attention_mask, dtype=np.int64),
-        }
-
-        if not self.exclude_token_type_ids:
-            onnx_input["token_type_ids"] = np.array(
-                [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
-            )
-
-        model_output = self.model.run(None, onnx_input)
-        embeddings = model_output[0]
-        return embeddings, attention_mask
-
-    def split_text(self, text: str, chunk_size: Optional[int] = None, chunk_overlap: Optional[int] = None) -> List[str]:
-        return self.splitter.split_text(text=text, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-
-
-class EmbeddingWorker(Worker):
-    def __init__(
-        self,
-        path: Path,
-        model_name: str,
-        max_length: int = 512,
-    ):
-        self.model = EmbeddingModel(path=path, model_name=model_name, max_length=max_length, max_threads=1)
-
-    @classmethod
-    def start(cls, path: Path, model_name: str, max_length: int = 512, **kwargs: Any) -> "EmbeddingWorker":
-        return cls(
-            path=path,
-            model_name=model_name,
-            max_length=max_length,
-        )
-
-    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
-        for idx, batch in items:
-            embeddings, attn_mask = self.model.onnx_embed(batch)
-            yield idx, (embeddings, attn_mask)
-
-
-class Embedding(ABC):
-    """
-    Abstract class for embeddings.
-
-    Inherits:
-        ABC: Abstract base class
-
-    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_
-    """
-
-    model: EmbeddingModel
-
-    @abstractmethod
-    def embed(self, texts: Iterable[str], batch_size: int = 256, parallel: int = None) -> List[np.ndarray]:
-        raise NotImplementedError
-
-    @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Union[str, Union[int, float]]]]:
-        """
-        Lists the supported models.
-        """
-        return [
-            {"model": "BAAI/bge-small-en", "dim": 384, "description": "Fast English model", "size_in_GB": 0.2},
-            {
-                "model": "BAAI/bge-small-en-v1.5",
-                "dim": 384,
-                "description": "Fast and Default English model",
-                "size_in_GB": 0.13,
-            },
-            {
-                "model": "BAAI/bge-small-zh-v1.5",
-                "dim": 512,
-                "description": "Fast and recommended Chinese model",
-                "size_in_GB": 0.1,
-            },
-            {"model": "BAAI/bge-base-en", "dim": 768, "description": "Base English model", "size_in_GB": 0.5},
-            {
-                "model": "BAAI/bge-base-en-v1.5",
-                "dim": 768,
-                "description": "Base English model, v1.5",
-                "size_in_GB": 0.44,
-            },
-            {
-                "model": "sentence-transformers/all-MiniLM-L6-v2",
-                "dim": 384,
-                "description": "Sentence Transformer model, MiniLM-L6-v2",
-                "size_in_GB": 0.09,
-            },
-            {
-                "model": "intfloat/multilingual-e5-large",
-                "dim": 1024,
-                "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
-                "size_in_GB": 2.24,
-            },
-            {
-                "model": "jinaai/jina-embeddings-v2-base-en",
-                "dim": 768,
-                "description": " English embedding model supporting 8192 sequence length",
-                "size_in_GB": 0.55,
-            },
-            {
-                "model": "jinaai/jina-embeddings-v2-small-en",
-                "dim": 512,
-                "description": " English embedding model supporting 8192 sequence length",
-                "size_in_GB": 0.13,
-            },
-        ]
-
-    @classmethod
-    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
-        """
-        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 download_files_from_huggingface(cls, repod_id: str, cache_dir: Optional[str] = None) -> str:
-        """
-        Downloads a model from HuggingFace Hub.
-        Args:
-            repod_id (str): The HF hub id (name) of the model to retrieve.
-            cache_dir (Optional[str]): The path to the cache directory.
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. "jinaai/jina-embeddings-v2-small-en".
-        Returns:
-            Path: The path to the model directory.
-        """
-        from huggingface_hub import snapshot_download
-
-        return snapshot_download(
-            repo_id=repod_id, ignore_patterns=["model.safetensors", "pytorch_model.bin"], cache_dir=cache_dir
-        )
-
-    @classmethod
-    def decompress_to_cache(cls, targz_path: str, cache_dir: str):
-        """
-        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_gcs(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 retrieve_model_hf(self, model_name: str, cache_dir: str) -> Path:
-        """
-        Retrieves a model from HuggingFace Hub.
-        Args:
-            model_name (str): The name of the model to retrieve.
-            cache_dir (str): The path to the cache directory.
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
-        Returns:
-            Path: The path to the model directory.
-        """
-
-        assert (
-            "/" in model_name
-        ), "model_name must be in the format <org>/<model> e.g. jinaai/jina-embeddings-v2-small-en"
-
-        return Path(self.download_files_from_huggingface(repod_id=model_name, cache_dir=cache_dir))
-
-    def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
-        """
-        Embeds a list of text passages into a list of embeddings.
-
-        Args:
-            texts (Iterable[str]): The list of texts to embed.
-            **kwargs: Additional keyword argument to pass to the embed method.
-
-        Yields:
-            Iterable[np.ndarray]: The embeddings.
-        """
-
-        yield from self.embed((f"passage: {t}" for t in texts), **kwargs)
-
-    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])
-        return query_embedding
-
-    def split_text(self, text: str, chunk_size: Optional[int] = None, chunk_overlap: Optional[int] = None) -> List[str]:
-        """Splits text into chunks based on the tokenizer encoding size.
-
-        Args:
-            text (str): The text to split.
-            chunk_size (Optional[int], optional): Maximum size of chunks based on the tokenizer encoding.
-            chunk_overlap (Optional[int], optional): Allowed overlap in characters between chunks.
-
-        Returns:
-            List[str]: The list of strings.
-        """
-        return self.model.split_text(text, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-
-
-class FlagEmbedding(Embedding):
-    """
-    Implementation of the Flag Embedding model.
-
-    Args:
-        Embedding (_type_): _description_
-    """
-
-    def __init__(
-        self,
-        model_name: str = "BAAI/bge-small-en-v1.5",
-        max_length: int = 512,
-        cache_dir: str = None,
-        threads: int = None,
-        splitter_config: Optional[TextSplitterConfig] = None,
-    ):
-        """
-        Args:
-            model_name (str): The name of the model to use.
-            max_length (int, optional): The maximum number of tokens. Defaults to 512. Unknown behavior for values > 512.
-            cache_dir (str, optional): The path to the cache directory.
-                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
-                                       Defaults to `fastembed_cache` in the system's temp directory.
-            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
-            splitter_config (Optional[TextSplitterConfig], optional): The configuration for the text splitter. Defaults to None.
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
-        """
-        self.model_name = model_name
-
-        if cache_dir is None:
-            default_cache_dir = os.path.join(tempfile.gettempdir(), "fastembed_cache")
-            cache_dir = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
-            cache_dir.mkdir(parents=True, exist_ok=True)
-
-        self._cache_dir = cache_dir
-        self._model_dir = self.retrieve_model_gcs(model_name, cache_dir)
-        self._max_length = max_length
-
-        self.model = EmbeddingModel(
-            self._model_dir, self.model_name, max_length=max_length, max_threads=threads, splitter_config=splitter_config
-        )
-
-    def embed(
-        self, documents: Union[str, Iterable[str]], batch_size: int = 256, parallel: int = None
-    ) -> Iterable[np.ndarray]:
-        """
-        Encode a list of documents into list of embeddings.
-        We use mean pooling with attention so that the model can handle variable-length inputs.
-
-        Args:
-            documents: Iterator of documents or single document to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-            parallel:
-                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
-                If 0, use all available cores.
-                If None, don't use data-parallel processing, use default onnxruntime threading instead.
-
-        Returns:
-            List of embeddings, one per document
-        """
-        is_small = False
-
-        if isinstance(documents, str):
-            documents = [documents]
-            is_small = True
-
-        if isinstance(documents, list):
-            if len(documents) < batch_size:
-                is_small = True
-
-        if parallel == 0:
-            parallel = os.cpu_count()
-
-        if parallel is None or is_small:
-            for batch in iter_batch(documents, batch_size):
-                embeddings, _ = self.model.onnx_embed(batch)
-                yield from normalize(embeddings[:, 0]).astype(np.float32)
-        else:
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            params = {
-                "path": self._model_dir,
-                "model_name": self.model_name,
-                "max_length": self._max_length,
-            }
-            pool = ParallelWorkerPool(parallel, EmbeddingWorker, start_method=start_method)
-            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
-                embeddings, _ = batch
-                yield from normalize(embeddings[:, 0]).astype(np.float32)
-
-    @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Union[str, Union[int, float]]]]:
-        """
-        Lists the supported models.
-        """
-        # jina models are not supported by this class
-        return [model for model in super().list_supported_models() if not model["model"].startswith("jinaai")]
-
-
-class DefaultEmbedding(FlagEmbedding):
-    """
-    FastEmbed Default Embedding model is BAAI/bge-small-en-v1.5. This is the recommended model for English.
-
-    Args:
-        FlagEmbedding (Embedding): The Flag Embedding model implementation.
-    """
-
-    def __init__(
-        self,
-        model_name: str = "BAAI/bge-small-en-v1.5",
-        max_length: int = 512,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        splitter_config: Optional[TextSplitterConfig] = None,
-    ):
-        super().__init__(
-            model_name, max_length=max_length, cache_dir=cache_dir, threads=threads, splitter_config=splitter_config
-        )
-
-
-class OpenAIEmbedding(Embedding):
-    def __init__(self):
-        # Initialize your OpenAI model here
-        # self.model = ...
-        ...
-
-    def embed(self, texts, batch_size: int = 256, parallel: int = None):
-        # Use your OpenAI model to embed the texts
-        # return self.model.embed(texts)
-        raise NotImplementedError
-
-
-class JinaEmbedding(Embedding):
+class JinaEmbedding(TextEmbedding):
    def __init__(
        self,
        model_name: str = "jinaai/jina-embeddings-v2-base-en",
-        max_length: int = 512,
-        cache_dir: str = None,
-        threads: int = None,
-        splitter_config: Optional[TextSplitterConfig] = None,
+        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.
-                                       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.
-            splitter_config (Optional[TextSplitterConfig], optional): The configuration for the text splitter. Defaults to None.
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
-        """
-        self.model_name = model_name
-
-        if cache_dir is None:
-            default_cache_dir = os.path.join(tempfile.gettempdir(), "fastembed_cache")
-            cache_dir = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
-            cache_dir.mkdir(parents=True, exist_ok=True)
-
-        self._cache_dir = cache_dir
-        self._model_dir = self.retrieve_model_hf(model_name, cache_dir)
-        self._max_length = max_length
-
-        self.model = EmbeddingModel(
-            self._model_dir, self.model_name, max_length=max_length, max_threads=threads, splitter_config=splitter_config 
-        )
-
-    def embed(
-        self, documents: Union[str, Iterable[str]], batch_size: int = 256, parallel: int = None
-    ) -> Iterable[np.ndarray]:
-        """
-        Encode a list of documents into list of embeddings.
-        We use mean pooling with attention so that the model can handle variable-length inputs.
-        Args:
-            documents: Iterator of documents or single document to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-            parallel:
-                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
-                If 0, use all available cores.
-                If None, don't use data-parallel processing, use default onnxruntime threading instead.
-        Returns:
-            List of embeddings, one per document
-        """
-        is_small = False
-
-        if isinstance(documents, str):
-            documents = [documents]
-            is_small = True
-
-        if isinstance(documents, list):
-            if len(documents) < batch_size:
-                is_small = True
-
-        if parallel == 0:
-            parallel = os.cpu_count()
-
-        if parallel is None or is_small:
-            for batch in iter_batch(documents, batch_size):
-                embeddings, attn_mask = self.model.onnx_embed(batch)
-                yield from normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
-        else:
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            params = {
-                "path": self._model_dir,
-                "model_name": self.model_name,
-                "max_length": self._max_length,
-            }
-            pool = ParallelWorkerPool(parallel, EmbeddingWorker, start_method=start_method)
-            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
-                embeddings, attn_mask = batch
-                yield from normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
-
-    @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Union[str, Union[int, float]]]]:
-        """
-        Lists the supported models.
-        """
-        # only jina models are supported by this class
-        return [model for model in Embedding.list_supported_models() if model["model"].startswith("jinaai")]
-
-    @staticmethod
-    def mean_pooling(model_output, attention_mask):
-        token_embeddings = model_output
-        input_mask_expanded = (np.expand_dims(attention_mask, axis=-1)).astype(float)
-
-        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
-        mask_sum = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
-
-        return sum_embeddings / mask_sum
+        super().__init__(model_name, cache_dir, threads, **kwargs)
@@ -1,70 +0,0 @@
-from typing import Callable, List, Optional
-
-from pydantic import BaseModel, ConfigDict, model_validator
-from tokenizers import Tokenizer
-
-"""
-For how to use validators in Pydantic, see: https://docs.pydantic.dev/latest/concepts/validators/
-"""
-
-
-class TextSplitterConfig(BaseModel):
-    """
-    Configuration for TextSplitter.
-
-    Attributes:
-        chunk_size (int): Size of the chunks, measured in tokens.
-        chunk_overlap (int): Overlap between chunks, measured in tokens.
-        length_function (Callable[[str], int]): Function that measures the length of given chunks. Default is len.
-        keep_separator (bool): Whether to keep the separator or not. Default is False.
-        strip_whitespace (bool): Whether to strip whitespace or not. Default is True.
-        tokenizer (Optional[Tokenizer]): Tokenizer used in FastEmbedRecursiveSplitter. Default is None.
-        is_separator_regex (bool): Whether the separator is a regular expression. Default is False.
-        separators (Optional[List[str]]): List of separators. Default is None.
-    """
-    model_config: ConfigDict = ConfigDict(arbitrary_types_allowed=True)
-
-    chunk_size: Optional[int] = None
-    chunk_overlap: Optional[int] = None
-    length_function: Callable[[str], int] = len
-    keep_separator: bool = False
-    strip_whitespace: bool = True
-    tokenizer: Optional[Tokenizer] = None
-    is_separator_regex: bool = False
-    separators: Optional[List[str]] = None
-
-    @model_validator(mode="after")
-    def _check_chunk_size_overlap(self) -> "TextSplitterConfig":
-        """
-        Validates chunk_size and chunk_overlap.
-
-        Raises:
-            ValueError: If chunk_size is not greater than 0, chunk_overlap is negative, or chunk_overlap is greater than chunk_size.
-        """
-        chunk_size = self.chunk_size
-        chunk_overlap = self.chunk_overlap
-        if chunk_size <= 0:
-            raise ValueError("Invalid value for chunk_size. It must be greater than 0.")
-        if chunk_overlap < 0:
-            raise ValueError("Invalid value for chunk_overlap. It must be greater than or equal to 0.")
-        if chunk_overlap > chunk_size:
-            raise ValueError("Invalid value for chunk_overlap. It must be smaller than chunk_size.")
-        return self
-
-    @model_validator(mode="after")
-    def _check_tokenizer(self) -> "TextSplitterConfig":
-        """
-        Validates tokenizer.
-
-        Raises:
-            ValueError: If tokenizer is not an instance of Tokenizer or its max_length is not greater than or equal to chunk_size.
-        """
-        tokenizer = self.tokenizer
-        chunk_size = self.chunk_size
-        if tokenizer is None:
-            return self
-        if not isinstance(tokenizer, Tokenizer):
-            raise ValueError("Invalid tokenizer. It must be an instance of tokenizers.Tokenizer.")
-        if tokenizer.model_max_length < chunk_size:
-            raise ValueError("Invalid chunk size. It must be smaller than or equal to tokenizer's max_length.")
-        return self
@@ -7,7 +7,7 @@ from multiprocessing.context import BaseContext
 from multiprocessing.process import BaseProcess
 from multiprocessing.sharedctypes import Synchronized as BaseValue
 from queue import Empty
-from typing import Any, Dict, Iterable, List, Optional, Type, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Type

 # Single item should be processed in less than:
 processing_timeout = 10 * 60  # seconds
@@ -128,7 +128,9 @@ class ParallelWorkerPool:
                yield buffer.pop(next_expected)
                next_expected += 1

-    def semi_ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Tuple[int, Any]]:
+    def semi_ordered_map(
+        self, stream: Iterable[Any], *args: Any, **kwargs: Any
+    ) -> Iterable[Tuple[int, Any]]:
        try:
            self.start(**kwargs)

@@ -0,0 +1,4 @@
+from fastembed.sparse.sparse_embedding_base import SparseEmbedding
+from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
+
+__all__ = ["SparseEmbedding", "SparseTextEmbedding"]
@@ -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()
@@ -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)
@@ -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)
@@ -1,190 +0,0 @@
-# Custom implementation based on the Langchain text splitter
-# Provided under standard MIT license by LangChain
-# https://github.com/hwchase17/langchain
-
-import logging
-import re
-from abc import ABC, abstractmethod
-from typing import (
-    Iterable,
-    List,
-    Optional,
-)
-
-from tokenizers import Tokenizer
-
-from .models import TextSplitterConfig
-
-logger = logging.getLogger(__name__)
-
-
-def _split_text_with_regex(text: str, separator: str, keep_separator: bool) -> List[str]:
-    # Now that we have the separator, split the text
-    if separator:
-        if keep_separator:
-            # The parentheses in the pattern keep the delimiters in the result.
-            _splits = re.split(f"({separator})", text)
-            splits = [_splits[i] + _splits[i + 1] for i in range(1, len(_splits), 2)]
-            if len(_splits) % 2 == 0:
-                splits += _splits[-1:]
-            splits = [_splits[0]] + splits
-        else:
-            splits = re.split(separator, text)
-    else:
-        splits = list(text)
-    return [s for s in splits if s != ""]
-
-
-class TextSplitter(ABC):
-    """Interface for splitting text into chunks."""
-
-    def __init__(
-        self,
-        config: TextSplitterConfig,
-    ) -> None:
-        """Create a new TextSplitter.
-
-        Args:
-            chunk_size: Maximum token length per chunk
-            chunk_overlap: Overlap in tokens between chunks
-            length_function: Function that measures the length of given chunks
-            keep_separator: Whether to keep the separator in the chunks
-            strip_whitespace: If `True`, strips whitespace from the start and end of
-                              every document
-        """
-        self._config = config
-        self._chunk_size = config.chunk_size
-        self._chunk_overlap = config.chunk_overlap
-        self._length_function = config.length_function
-        self._keep_separator = config.keep_separator
-        self._strip_whitespace = config.strip_whitespace
-
-    @abstractmethod
-    def split_text(self, text: str) -> List[str]:
-        """Split text into multiple components."""
-
-    def _join_docs(self, docs: List[str], separator: str) -> Optional[str]:
-        text = separator.join(docs)
-        if self._strip_whitespace:
-            text = text.strip()
-        if text == "":
-            return None
-        else:
-            return text
-
-    def _merge_splits(
-        self,
-        splits: Iterable[str],
-        separator: str,
-        chunk_size: Optional[int] = None,
-        chunk_overlap: Optional[int] = None,
-    ) -> List[str]:
-        # We now want to combine these smaller pieces into medium size
-        # chunks to send to the LLM.
-        separator_len = self._length_function(separator)
-
-        chunk_size = chunk_size or self._chunk_size
-        chunk_overlap = chunk_overlap or self._chunk_overlap
-
-        docs = []
-        current_doc: List[str] = []
-        total = 0
-        for d in splits:
-            _len = self._length_function(d)
-            if total + _len + (separator_len if len(current_doc) > 0 else 0) > chunk_size:
-                if total > chunk_size:
-                    logger.warning(
-                        f"Created a chunk of size {total}, " f"which is longer than the specified {chunk_size}"
-                    )
-                if len(current_doc) > 0:
-                    doc = self._join_docs(current_doc, separator)
-                    if doc is not None:
-                        docs.append(doc)
-                    # Keep on popping if:
-                    # - we have a larger chunk than in the chunk overlap
-                    # - or if we still have any chunks and the length is long
-                    while total > chunk_overlap or (
-                        total + _len + (separator_len if len(current_doc) > 0 else 0) > chunk_size and total > 0
-                    ):
-                        total -= self._length_function(current_doc[0]) + (separator_len if len(current_doc) > 1 else 0)
-                        current_doc = current_doc[1:]
-            current_doc.append(d)
-            total += _len + (separator_len if len(current_doc) > 1 else 0)
-        doc = self._join_docs(current_doc, separator)
-        if doc is not None:
-            docs.append(doc)
-        return docs
-
-
-class FastEmbedRecursiveSplitter(TextSplitter):
-    """
-    Splitting text into chunks recursively.
-
-    The splitter splits text into chunks of a maximum size, with a given overlap.
-    """
-
-    def __init__(
-        self,
-        config: TextSplitterConfig,
-    ) -> None:
-        """Create a new TextSplitter."""
-        tokenizer = config.tokenizer
-        if not isinstance(tokenizer, Tokenizer):
-            raise ValueError("Tokenizer received was not an instance of tokenizers.Tokenizer")
-
-        def _tokenizer_length(text: str) -> int:
-            return len(tokenizer.encode(text))
-
-        config.length_function = _tokenizer_length
-        super().__init__(config)
-        self._separators = config.separators or ["\n\n", "\n", " ", ""]
-        self._is_separator_regex = config.is_separator_regex
-
-    def _split_text(
-        self, text: str, separators: List[str], chunk_size: Optional[int] = None, chunk_overlap: Optional[int] = None
-    ) -> List[str]:
-        """Split incoming text and return chunks."""
-
-        chunk_size = chunk_size or self._chunk_size
-        chunk_overlap = chunk_overlap or self._chunk_overlap
-
-        final_chunks = []
-        # Get appropriate separator to use
-        separator = separators[-1]
-        new_separators = []
-        for i, _s in enumerate(separators):
-            _separator = _s if self._is_separator_regex else re.escape(_s)
-            if _s == "":
-                separator = _s
-                break
-            if re.search(_separator, text):
-                separator = _s
-                new_separators = separators[i + 1 :]
-                break
-
-        _separator = separator if self._is_separator_regex else re.escape(separator)
-        splits = _split_text_with_regex(text, _separator, self._keep_separator)
-
-        # Now go merging things, recursively splitting longer texts.
-        _good_splits = []
-        _separator = "" if self._keep_separator else separator
-        for s in splits:
-            if self._length_function(s) < chunk_size:
-                _good_splits.append(s)
-            else:
-                if _good_splits:
-                    merged_text = self._merge_splits(_good_splits, _separator, chunk_size, chunk_overlap)
-                    final_chunks.extend(merged_text)
-                    _good_splits = []
-                if not new_separators:
-                    final_chunks.append(s)
-                else:
-                    other_info = self._split_text(s, new_separators)
-                    final_chunks.extend(other_info)
-        if _good_splits:
-            merged_text = self._merge_splits(_good_splits, _separator)
-            final_chunks.extend(merged_text)
-        return final_chunks
-
-    def split_text(self, text: str, chunk_size: Optional[int] = None, chunk_overlap: Optional[int] = None) -> List[str]:
-        return self._split_text(text, self._separators, chunk_size, chunk_overlap)
@@ -0,0 +1,3 @@
+from fastembed.text.text_embedding import TextEmbedding
+
+__all__ = ["TextEmbedding"]
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.1.3"
+version = "0.2.6"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -11,24 +11,29 @@ 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"
+python = ">=3.8.0,<3.13"
+onnx = "^1.15.0"
+onnxruntime = "^1.17.0"
+tqdm = "^4.66"
 requests = "^2.31"
-tokenizers = "^0.15.0"
-huggingface-hub = "0.19.4"
-pydantic = "^2.5.3"
+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.1.11"
-isort = "^5.12.0"
-black = "^23.7.0"
+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"

@@ -37,11 +42,5 @@ mknotebooks = "^0.8.0"
 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
@@ -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)
-
-
@@ -1,723 +0,0 @@
-Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.  
-
-Last year COVID-19 kept us apart. This year we are finally together again. 
-
-Tonight, we meet as Democrats Republicans and Independents. But most importantly as Americans. 
-
-With a duty to one another to the American people to the Constitution. 
-
-And with an unwavering resolve that freedom will always triumph over tyranny. 
-
-Six days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated. 
-
-He thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined. 
-
-He met the Ukrainian people. 
-
-From President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world. 
-
-Groups of citizens blocking tanks with their bodies. Everyone from students to retirees teachers turned soldiers defending their homeland. 
-
-In this struggle as President Zelenskyy said in his speech to the European Parliament “Light will win over darkness.” The Ukrainian Ambassador to the United States is here tonight. 
-
-Let each of us here tonight in this Chamber send an unmistakable signal to Ukraine and to the world. 
-
-Please rise if you are able and show that, Yes, we the United States of America stand with the Ukrainian people. 
-
-Throughout our history we’ve learned this lesson when dictators do not pay a price for their aggression they cause more chaos.   
-
-They keep moving.   
-
-And the costs and the threats to America and the world keep rising.   
-
-That’s why the NATO Alliance was created to secure peace and stability in Europe after World War 2. 
-
-The United States is a member along with 29 other nations. 
-
-It matters. American diplomacy matters. American resolve matters. 
-
-Putin’s latest attack on Ukraine was premeditated and unprovoked. 
-
-He rejected repeated efforts at diplomacy. 
-
-He thought the West and NATO wouldn’t respond. And he thought he could divide us at home. Putin was wrong. We were ready.  Here is what we did.   
-
-We prepared extensively and carefully. 
-
-We spent months building a coalition of other freedom-loving nations from Europe and the Americas to Asia and Africa to confront Putin. 
-
-I spent countless hours unifying our European allies. We shared with the world in advance what we knew Putin was planning and precisely how he would try to falsely justify his aggression.  
-
-We countered Russia’s lies with truth.   
-
-And now that he has acted the free world is holding him accountable. 
-
-Along with twenty-seven members of the European Union including France, Germany, Italy, as well as countries like the United Kingdom, Canada, Japan, Korea, Australia, New Zealand, and many others, even Switzerland. 
-
-We are inflicting pain on Russia and supporting the people of Ukraine. Putin is now isolated from the world more than ever. 
-
-Together with our allies –we are right now enforcing powerful economic sanctions. 
-
-We are cutting off Russia’s largest banks from the international financial system.  
-
-Preventing Russia’s central bank from defending the Russian Ruble making Putin’s $630 Billion “war fund” worthless.   
-
-We are choking off Russia’s access to technology that will sap its economic strength and weaken its military for years to come.  
-
-Tonight I say to the Russian oligarchs and corrupt leaders who have bilked billions of dollars off this violent regime no more. 
-
-The U.S. Department of Justice is assembling a dedicated task force to go after the crimes of Russian oligarchs.  
-
-We are joining with our European allies to find and seize your yachts your luxury apartments your private jets. We are coming for your ill-begotten gains. 
-
-And tonight I am announcing that we will join our allies in closing off American air space to all Russian flights – further isolating Russia – and adding an additional squeeze –on their economy. The Ruble has lost 30% of its value. 
-
-The Russian stock market has lost 40% of its value and trading remains suspended. Russia’s economy is reeling and Putin alone is to blame. 
-
-Together with our allies we are providing support to the Ukrainians in their fight for freedom. Military assistance. Economic assistance. Humanitarian assistance. 
-
-We are giving more than $1 Billion in direct assistance to Ukraine. 
-
-And we will continue to aid the Ukrainian people as they defend their country and to help ease their suffering.  
-
-Let me be clear, our forces are not engaged and will not engage in conflict with Russian forces in Ukraine.  
-
-Our forces are not going to Europe to fight in Ukraine, but to defend our NATO Allies – in the event that Putin decides to keep moving west.  
-
-For that purpose we’ve mobilized American ground forces, air squadrons, and ship deployments to protect NATO countries including Poland, Romania, Latvia, Lithuania, and Estonia. 
-
-As I have made crystal clear the United States and our Allies will defend every inch of territory of NATO countries with the full force of our collective power.  
-
-And we remain clear-eyed. The Ukrainians are fighting back with pure courage. But the next few days weeks, months, will be hard on them.  
-
-Putin has unleashed violence and chaos.  But while he may make gains on the battlefield – he will pay a continuing high price over the long run. 
-
-And a proud Ukrainian people, who have known 30 years  of independence, have repeatedly shown that they will not tolerate anyone who tries to take their country backwards.  
-
-To all Americans, I will be honest with you, as I’ve always promised. A Russian dictator, invading a foreign country, has costs around the world. 
-
-And I’m taking robust action to make sure the pain of our sanctions  is targeted at Russia’s economy. And I will use every tool at our disposal to protect American businesses and consumers. 
-
-Tonight, I can announce that the United States has worked with 30 other countries to release 60 Million barrels of oil from reserves around the world.  
-
-America will lead that effort, releasing 30 Million barrels from our own Strategic Petroleum Reserve. And we stand ready to do more if necessary, unified with our allies.  
-
-These steps will help blunt gas prices here at home. And I know the news about what’s happening can seem alarming. 
-
-But I want you to know that we are going to be okay. 
-
-When the history of this era is written Putin’s war on Ukraine will have left Russia weaker and the rest of the world stronger. 
-
-While it shouldn’t have taken something so terrible for people around the world to see what’s at stake now everyone sees it clearly. 
-
-We see the unity among leaders of nations and a more unified Europe a more unified West. And we see unity among the people who are gathering in cities in large crowds around the world even in Russia to demonstrate their support for Ukraine.  
-
-In the battle between democracy and autocracy, democracies are rising to the moment, and the world is clearly choosing the side of peace and security. 
-
-This is a real test. It’s going to take time. So let us continue to draw inspiration from the iron will of the Ukrainian people. 
-
-To our fellow Ukrainian Americans who forge a deep bond that connects our two nations we stand with you. 
-
-Putin may circle Kyiv with tanks, but he will never gain the hearts and souls of the Ukrainian people. 
-
-He will never extinguish their love of freedom. He will never weaken the resolve of the free world. 
-
-We meet tonight in an America that has lived through two of the hardest years this nation has ever faced. 
-
-The pandemic has been punishing. 
-
-And so many families are living paycheck to paycheck, struggling to keep up with the rising cost of food, gas, housing, and so much more. 
-
-I understand. 
-
-I remember when my Dad had to leave our home in Scranton, Pennsylvania to find work. I grew up in a family where if the price of food went up, you felt it. 
-
-That’s why one of the first things I did as President was fight to pass the American Rescue Plan.  
-
-Because people were hurting. We needed to act, and we did. 
-
-Few pieces of legislation have done more in a critical moment in our history to lift us out of crisis. 
-
-It fueled our efforts to vaccinate the nation and combat COVID-19. It delivered immediate economic relief for tens of millions of Americans.  
-
-Helped put food on their table, keep a roof over their heads, and cut the cost of health insurance. 
-
-And as my Dad used to say, it gave people a little breathing room. 
-
-And unlike the $2 Trillion tax cut passed in the previous administration that benefitted the top 1% of Americans, the American Rescue Plan helped working people—and left no one behind. 
-
-And it worked. It created jobs. Lots of jobs. 
-
-In fact—our economy created over 6.5 Million new jobs just last year, more jobs created in one year  
-than ever before in the history of America. 
-
-Our economy grew at a rate of 5.7% last year, the strongest growth in nearly 40 years, the first step in bringing fundamental change to an economy that hasn’t worked for the working people of this nation for too long.  
-
-For the past 40 years we were told that if we gave tax breaks to those at the very top, the benefits would trickle down to everyone else. 
-
-But that trickle-down theory led to weaker economic growth, lower wages, bigger deficits, and the widest gap between those at the top and everyone else in nearly a century. 
-
-Vice President Harris and I ran for office with a new economic vision for America. 
-
-Invest in America. Educate Americans. Grow the workforce. Build the economy from the bottom up  
-and the middle out, not from the top down.  
-
-Because we know that when the middle class grows, the poor have a ladder up and the wealthy do very well. 
-
-America used to have the best roads, bridges, and airports on Earth. 
-
-Now our infrastructure is ranked 13th in the world. 
-
-We won’t be able to compete for the jobs of the 21st Century if we don’t fix that. 
-
-That’s why it was so important to pass the Bipartisan Infrastructure Law—the most sweeping investment to rebuild America in history. 
-
-This was a bipartisan effort, and I want to thank the members of both parties who worked to make it happen. 
-
-We’re done talking about infrastructure weeks. 
-
-We’re going to have an infrastructure decade. 
-
-It is going to transform America and put us on a path to win the economic competition of the 21st Century that we face with the rest of the world—particularly with China.  
-
-As I’ve told Xi Jinping, it is never a good bet to bet against the American people. 
-
-We’ll create good jobs for millions of Americans, modernizing roads, airports, ports, and waterways all across America. 
-
-And we’ll do it all to withstand the devastating effects of the climate crisis and promote environmental justice. 
-
-We’ll build a national network of 500,000 electric vehicle charging stations, begin to replace poisonous lead pipes—so every child—and every American—has clean water to drink at home and at school, provide affordable high-speed internet for every American—urban, suburban, rural, and tribal communities. 
-
-4,000 projects have already been announced. 
-
-And tonight, I’m announcing that this year we will start fixing over 65,000 miles of highway and 1,500 bridges in disrepair. 
-
-When we use taxpayer dollars to rebuild America – we are going to Buy American: buy American products to support American jobs. 
-
-The federal government spends about $600 Billion a year to keep the country safe and secure. 
-
-There’s been a law on the books for almost a century 
-to make sure taxpayers’ dollars support American jobs and businesses. 
-
-Every Administration says they’ll do it, but we are actually doing it. 
-
-We will buy American to make sure everything from the deck of an aircraft carrier to the steel on highway guardrails are made in America. 
-
-But to compete for the best jobs of the future, we also need to level the playing field with China and other competitors. 
-
-That’s why it is so important to pass the Bipartisan Innovation Act sitting in Congress that will make record investments in emerging technologies and American manufacturing. 
-
-Let me give you one example of why it’s so important to pass it. 
-
-If you travel 20 miles east of Columbus, Ohio, you’ll find 1,000 empty acres of land. 
-
-It won’t look like much, but if you stop and look closely, you’ll see a “Field of dreams,” the ground on which America’s future will be built. 
-
-This is where Intel, the American company that helped build Silicon Valley, is going to build its $20 billion semiconductor “mega site”. 
-
-Up to eight state-of-the-art factories in one place. 10,000 new good-paying jobs. 
-
-Some of the most sophisticated manufacturing in the world to make computer chips the size of a fingertip that power the world and our everyday lives. 
-
-Smartphones. The Internet. Technology we have yet to invent. 
-
-But that’s just the beginning. 
-
-Intel’s CEO, Pat Gelsinger, who is here tonight, told me they are ready to increase their investment from  
-$20 billion to $100 billion. 
-
-That would be one of the biggest investments in manufacturing in American history. 
-
-And all they’re waiting for is for you to pass this bill. 
-
-So let’s not wait any longer. Send it to my desk. I’ll sign it.  
-
-And we will really take off. 
-
-And Intel is not alone. 
-
-There’s something happening in America. 
-
-Just look around and you’ll see an amazing story. 
-
-The rebirth of the pride that comes from stamping products “Made In America.” The revitalization of American manufacturing.   
-
-Companies are choosing to build new factories here, when just a few years ago, they would have built them overseas. 
-
-That’s what is happening. Ford is investing $11 billion to build electric vehicles, creating 11,000 jobs across the country. 
-
-GM is making the largest investment in its history—$7 billion to build electric vehicles, creating 4,000 jobs in Michigan. 
-
-All told, we created 369,000 new manufacturing jobs in America just last year. 
-
-Powered by people I’ve met like JoJo Burgess, from generations of union steelworkers from Pittsburgh, who’s here with us tonight. 
-
-As Ohio Senator Sherrod Brown says, “It’s time to bury the label “Rust Belt.” 
-
-It’s time. 
-
-But with all the bright spots in our economy, record job growth and higher wages, too many families are struggling to keep up with the bills.  
-
-Inflation is robbing them of the gains they might otherwise feel. 
-
-I get it. That’s why my top priority is getting prices under control. 
-
-Look, our economy roared back faster than most predicted, but the pandemic meant that businesses had a hard time hiring enough workers to keep up production in their factories. 
-
-The pandemic also disrupted global supply chains. 
-
-When factories close, it takes longer to make goods and get them from the warehouse to the store, and prices go up. 
-
-Look at cars. 
-
-Last year, there weren’t enough semiconductors to make all the cars that people wanted to buy. 
-
-And guess what, prices of automobiles went up. 
-
-So—we have a choice. 
-
-One way to fight inflation is to drive down wages and make Americans poorer.  
-
-I have a better plan to fight inflation. 
-
-Lower your costs, not your wages. 
-
-Make more cars and semiconductors in America. 
-
-More infrastructure and innovation in America. 
-
-More goods moving faster and cheaper in America. 
-
-More jobs where you can earn a good living in America. 
-
-And instead of relying on foreign supply chains, let’s make it in America. 
-
-Economists call it “increasing the productive capacity of our economy.” 
-
-I call it building a better America. 
-
-My plan to fight inflation will lower your costs and lower the deficit. 
-
-17 Nobel laureates in economics say my plan will ease long-term inflationary pressures. Top business leaders and most Americans support my plan. And here’s the plan: 
-
-First – cut the cost of prescription drugs. Just look at insulin. One in ten Americans has diabetes. In Virginia, I met a 13-year-old boy named Joshua Davis.  
-
-He and his Dad both have Type 1 diabetes, which means they need insulin every day. Insulin costs about $10 a vial to make.  
-
-But drug companies charge families like Joshua and his Dad up to 30 times more. I spoke with Joshua’s mom. 
-
-Imagine what it’s like to look at your child who needs insulin and have no idea how you’re going to pay for it.  
-
-What it does to your dignity, your ability to look your child in the eye, to be the parent you expect to be. 
-
-Joshua is here with us tonight. Yesterday was his birthday. Happy birthday, buddy.  
-
-For Joshua, and for the 200,000 other young people with Type 1 diabetes, let’s cap the cost of insulin at $35 a month so everyone can afford it.  
-
-Drug companies will still do very well. And while we’re at it let Medicare negotiate lower prices for prescription drugs, like the VA already does. 
-
-Look, the American Rescue Plan is helping millions of families on Affordable Care Act plans save $2,400 a year on their health care premiums. Let’s close the coverage gap and make those savings permanent. 
-
-Second – cut energy costs for families an average of $500 a year by combatting climate change.  
-
-Let’s provide investments and tax credits to weatherize your homes and businesses to be energy efficient and you get a tax credit; double America’s clean energy production in solar, wind, and so much more;  lower the price of electric vehicles, saving you another $80 a month because you’ll never have to pay at the gas pump again. 
-
-Third – cut the cost of child care. Many families pay up to $14,000 a year for child care per child.  
-
-Middle-class and working families shouldn’t have to pay more than 7% of their income for care of young children.  
-
-My plan will cut the cost in half for most families and help parents, including millions of women, who left the workforce during the pandemic because they couldn’t afford child care, to be able to get back to work. 
-
-My plan doesn’t stop there. It also includes home and long-term care. More affordable housing. And Pre-K for every 3- and 4-year-old.  
-
-All of these will lower costs. 
-
-And under my plan, nobody earning less than $400,000 a year will pay an additional penny in new taxes. Nobody.  
-
-The one thing all Americans agree on is that the tax system is not fair. We have to fix it.  
-
-I’m not looking to punish anyone. But let’s make sure corporations and the wealthiest Americans start paying their fair share. 
-
-Just last year, 55 Fortune 500 corporations earned $40 billion in profits and paid zero dollars in federal income tax.  
-
-That’s simply not fair. That’s why I’ve proposed a 15% minimum tax rate for corporations. 
-
-We got more than 130 countries to agree on a global minimum tax rate so companies can’t get out of paying their taxes at home by shipping jobs and factories overseas. 
-
-That’s why I’ve proposed closing loopholes so the very wealthy don’t pay a lower tax rate than a teacher or a firefighter.  
-
-So that’s my plan. It will grow the economy and lower costs for families. 
-
-So what are we waiting for? Let’s get this done. And while you’re at it, confirm my nominees to the Federal Reserve, which plays a critical role in fighting inflation.  
-
-My plan will not only lower costs to give families a fair shot, it will lower the deficit. 
-
-The previous Administration not only ballooned the deficit with tax cuts for the very wealthy and corporations, it undermined the watchdogs whose job was to keep pandemic relief funds from being wasted. 
-
-But in my administration, the watchdogs have been welcomed back. 
-
-We’re going after the criminals who stole billions in relief money meant for small businesses and millions of Americans.  
-
-And tonight, I’m announcing that the Justice Department will name a chief prosecutor for pandemic fraud. 
-
-By the end of this year, the deficit will be down to less than half what it was before I took office.  
-
-The only president ever to cut the deficit by more than one trillion dollars in a single year. 
-
-Lowering your costs also means demanding more competition. 
-
-I’m a capitalist, but capitalism without competition isn’t capitalism. 
-
-It’s exploitation—and it drives up prices. 
-
-When corporations don’t have to compete, their profits go up, your prices go up, and small businesses and family farmers and ranchers go under. 
-
-We see it happening with ocean carriers moving goods in and out of America. 
-
-During the pandemic, these foreign-owned companies raised prices by as much as 1,000% and made record profits. 
-
-Tonight, I’m announcing a crackdown on these companies overcharging American businesses and consumers. 
-
-And as Wall Street firms take over more nursing homes, quality in those homes has gone down and costs have gone up.  
-
-That ends on my watch. 
-
-Medicare is going to set higher standards for nursing homes and make sure your loved ones get the care they deserve and expect. 
-
-We’ll also cut costs and keep the economy going strong by giving workers a fair shot, provide more training and apprenticeships, hire them based on their skills not degrees. 
-
-Let’s pass the Paycheck Fairness Act and paid leave.  
-
-Raise the minimum wage to $15 an hour and extend the Child Tax Credit, so no one has to raise a family in poverty. 
-
-Let’s increase Pell Grants and increase our historic support of HBCUs, and invest in what Jill—our First Lady who teaches full-time—calls America’s best-kept secret: community colleges. 
-
-And let’s pass the PRO Act when a majority of workers want to form a union—they shouldn’t be stopped.  
-
-When we invest in our workers, when we build the economy from the bottom up and the middle out together, we can do something we haven’t done in a long time: build a better America. 
-
-For more than two years, COVID-19 has impacted every decision in our lives and the life of the nation. 
-
-And I know you’re tired, frustrated, and exhausted. 
-
-But I also know this. 
-
-Because of the progress we’ve made, because of your resilience and the tools we have, tonight I can say  
-we are moving forward safely, back to more normal routines.  
-
-We’ve reached a new moment in the fight against COVID-19, with severe cases down to a level not seen since last July.  
-
-Just a few days ago, the Centers for Disease Control and Prevention—the CDC—issued new mask guidelines. 
-
-Under these new guidelines, most Americans in most of the country can now be mask free.   
-
-And based on the projections, more of the country will reach that point across the next couple of weeks. 
-
-Thanks to the progress we have made this past year, COVID-19 need no longer control our lives.  
-
-I know some are talking about “living with COVID-19”. Tonight – I say that we will never just accept living with COVID-19. 
-
-We will continue to combat the virus as we do other diseases. And because this is a virus that mutates and spreads, we will stay on guard. 
-
-Here are four common sense steps as we move forward safely.  
-
-First, stay protected with vaccines and treatments. We know how incredibly effective vaccines are. If you’re vaccinated and boosted you have the highest degree of protection. 
-
-We will never give up on vaccinating more Americans. Now, I know parents with kids under 5 are eager to see a vaccine authorized for their children. 
-
-The scientists are working hard to get that done and we’ll be ready with plenty of vaccines when they do. 
-
-We’re also ready with anti-viral treatments. If you get COVID-19, the Pfizer pill reduces your chances of ending up in the hospital by 90%.  
-
-We’ve ordered more of these pills than anyone in the world. And Pfizer is working overtime to get us 1 Million pills this month and more than double that next month.  
-
-And we’re launching the “Test to Treat” initiative so people can get tested at a pharmacy, and if they’re positive, receive antiviral pills on the spot at no cost.  
-
-If you’re immunocompromised or have some other vulnerability, we have treatments and free high-quality masks. 
-
-We’re leaving no one behind or ignoring anyone’s needs as we move forward. 
-
-And on testing, we have made hundreds of millions of tests available for you to order for free.   
-
-Even if you already ordered free tests tonight, I am announcing that you can order more from covidtests.gov starting next week. 
-
-Second – we must prepare for new variants. Over the past year, we’ve gotten much better at detecting new variants. 
-
-If necessary, we’ll be able to deploy new vaccines within 100 days instead of many more months or years.  
-
-And, if Congress provides the funds we need, we’ll have new stockpiles of tests, masks, and pills ready if needed. 
-
-I cannot promise a new variant won’t come. But I can promise you we’ll do everything within our power to be ready if it does.  
-
-Third – we can end the shutdown of schools and businesses. We have the tools we need. 
-
-It’s time for Americans to get back to work and fill our great downtowns again.  People working from home can feel safe to begin to return to the office.   
-
-We’re doing that here in the federal government. The vast majority of federal workers will once again work in person. 
-
-Our schools are open. Let’s keep it that way. Our kids need to be in school. 
-
-And with 75% of adult Americans fully vaccinated and hospitalizations down by 77%, most Americans can remove their masks, return to work, stay in the classroom, and move forward safely. 
-
-We achieved this because we provided free vaccines, treatments, tests, and masks. 
-
-Of course, continuing this costs money. 
-
-I will soon send Congress a request. 
-
-The vast majority of Americans have used these tools and may want to again, so I expect Congress to pass it quickly.   
-
-Fourth, we will continue vaccinating the world.     
-
-We’ve sent 475 Million vaccine doses to 112 countries, more than any other nation. 
-
-And we won’t stop. 
-
-We have lost so much to COVID-19. Time with one another. And worst of all, so much loss of life. 
-
-Let’s use this moment to reset. Let’s stop looking at COVID-19 as a partisan dividing line and see it for what it is: A God-awful disease.  
-
-Let’s stop seeing each other as enemies, and start seeing each other for who we really are: Fellow Americans.  
-
-We can’t change how divided we’ve been. But we can change how we move forward—on COVID-19 and other issues we must face together. 
-
-I recently visited the New York City Police Department days after the funerals of Officer Wilbert Mora and his partner, Officer Jason Rivera. 
-
-They were responding to a 9-1-1 call when a man shot and killed them with a stolen gun. 
-
-Officer Mora was 27 years old. 
-
-Officer Rivera was 22. 
-
-Both Dominican Americans who’d grown up on the same streets they later chose to patrol as police officers. 
-
-I spoke with their families and told them that we are forever in debt for their sacrifice, and we will carry on their mission to restore the trust and safety every community deserves. 
-
-I’ve worked on these issues a long time. 
-
-I know what works: Investing in crime preventionand community police officers who’ll walk the beat, who’ll know the neighborhood, and who can restore trust and safety. 
-
-So let’s not abandon our streets. Or choose between safety and equal justice. 
-
-Let’s come together to protect our communities, restore trust, and hold law enforcement accountable. 
-
-That’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers. 
-
-That’s why the American Rescue Plan provided $350 Billion that cities, states, and counties can use to hire more police and invest in proven strategies like community violence interruption—trusted messengers breaking the cycle of violence and trauma and giving young people hope.  
-
-We should all agree: The answer is not to Defund the police. The answer is to FUND the police with the resources and training they need to protect our communities. 
-
-I ask Democrats and Republicans alike: Pass my budget and keep our neighborhoods safe.  
-
-And I will keep doing everything in my power to crack down on gun trafficking and ghost guns you can buy online and make at home—they have no serial numbers and can’t be traced. 
-
-And I ask Congress to pass proven measures to reduce gun violence. Pass universal background checks. Why should anyone on a terrorist list be able to purchase a weapon? 
-
-Ban assault weapons and high-capacity magazines. 
-
-Repeal the liability shield that makes gun manufacturers the only industry in America that can’t be sued. 
-
-These laws don’t infringe on the Second Amendment. They save lives. 
-
-The most fundamental right in America is the right to vote – and to have it counted. And it’s under assault. 
-
-In state after state, new laws have been passed, not only to suppress the vote, but to subvert entire elections. 
-
-We cannot let this happen. 
-
-Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. 
-
-Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. 
-
-One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. 
-
-And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence. 
-
-A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. 
-
-And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. 
-
-We can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling.  
-
-We’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers.  
-
-We’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. 
-
-We’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders. 
-
-We can do all this while keeping lit the torch of liberty that has led generations of immigrants to this land—my forefathers and so many of yours. 
-
-Provide a pathway to citizenship for Dreamers, those on temporary status, farm workers, and essential workers. 
-
-Revise our laws so businesses have the workers they need and families don’t wait decades to reunite. 
-
-It’s not only the right thing to do—it’s the economically smart thing to do. 
-
-That’s why immigration reform is supported by everyone from labor unions to religious leaders to the U.S. Chamber of Commerce. 
-
-Let’s get it done once and for all. 
-
-Advancing liberty and justice also requires protecting the rights of women. 
-
-The constitutional right affirmed in Roe v. Wade—standing precedent for half a century—is under attack as never before. 
-
-If we want to go forward—not backward—we must protect access to health care. Preserve a woman’s right to choose. And let’s continue to advance maternal health care in America. 
-
-And for our LGBTQ+ Americans, let’s finally get the bipartisan Equality Act to my desk. The onslaught of state laws targeting transgender Americans and their families is wrong. 
-
-As I said last year, especially to our younger transgender Americans, I will always have your back as your President, so you can be yourself and reach your God-given potential. 
-
-While it often appears that we never agree, that isn’t true. I signed 80 bipartisan bills into law last year. From preventing government shutdowns to protecting Asian-Americans from still-too-common hate crimes to reforming military justice. 
-
-And soon, we’ll strengthen the Violence Against Women Act that I first wrote three decades ago. It is important for us to show the nation that we can come together and do big things. 
-
-So tonight I’m offering a Unity Agenda for the Nation. Four big things we can do together.  
-
-First, beat the opioid epidemic. 
-
-There is so much we can do. Increase funding for prevention, treatment, harm reduction, and recovery.  
-
-Get rid of outdated rules that stop doctors from prescribing treatments. And stop the flow of illicit drugs by working with state and local law enforcement to go after traffickers. 
-
-If you’re suffering from addiction, know you are not alone. I believe in recovery, and I celebrate the 23 million Americans in recovery. 
-
-Second, let’s take on mental health. Especially among our children, whose lives and education have been turned upside down.  
-
-The American Rescue Plan gave schools money to hire teachers and help students make up for lost learning.  
-
-I urge every parent to make sure your school does just that. And we can all play a part—sign up to be a tutor or a mentor. 
-
-Children were also struggling before the pandemic. Bullying, violence, trauma, and the harms of social media. 
-
-As Frances Haugen, who is here with us tonight, has shown, we must hold social media platforms accountable for the national experiment they’re conducting on our children for profit. 
-
-It’s time to strengthen privacy protections, ban targeted advertising to children, demand tech companies stop collecting personal data on our children. 
-
-And let’s get all Americans the mental health services they need. More people they can turn to for help, and full parity between physical and mental health care. 
-
-Third, support our veterans. 
-
-Veterans are the best of us. 
-
-I’ve always believed that we have a sacred obligation to equip all those we send to war and care for them and their families when they come home. 
-
-My administration is providing assistance with job training and housing, and now helping lower-income veterans get VA care debt-free.  
-
-Our troops in Iraq and Afghanistan faced many dangers. 
-
-One was stationed at bases and breathing in toxic smoke from “burn pits” that incinerated wastes of war—medical and hazard material, jet fuel, and more. 
-
-When they came home, many of the world’s fittest and best trained warriors were never the same. 
-
-Headaches. Numbness. Dizziness. 
-
-A cancer that would put them in a flag-draped coffin. 
-
-I know. 
-
-One of those soldiers was my son Major Beau Biden. 
-
-We don’t know for sure if a burn pit was the cause of his brain cancer, or the diseases of so many of our troops. 
-
-But I’m committed to finding out everything we can. 
-
-Committed to military families like Danielle Robinson from Ohio. 
-
-The widow of Sergeant First Class Heath Robinson.  
-
-He was born a soldier. Army National Guard. Combat medic in Kosovo and Iraq. 
-
-Stationed near Baghdad, just yards from burn pits the size of football fields. 
-
-Heath’s widow Danielle is here with us tonight. They loved going to Ohio State football games. He loved building Legos with their daughter. 
-
-But cancer from prolonged exposure to burn pits ravaged Heath’s lungs and body. 
-
-Danielle says Heath was a fighter to the very end. 
-
-He didn’t know how to stop fighting, and neither did she. 
-
-Through her pain she found purpose to demand we do better. 
-
-Tonight, Danielle—we are. 
-
-The VA is pioneering new ways of linking toxic exposures to diseases, already helping more veterans get benefits. 
-
-And tonight, I’m announcing we’re expanding eligibility to veterans suffering from nine respiratory cancers. 
-
-I’m also calling on Congress: pass a law to make sure veterans devastated by toxic exposures in Iraq and Afghanistan finally get the benefits and comprehensive health care they deserve. 
-
-And fourth, let’s end cancer as we know it. 
-
-This is personal to me and Jill, to Kamala, and to so many of you. 
-
-Cancer is the #2 cause of death in America–second only to heart disease. 
-
-Last month, I announced our plan to supercharge  
-the Cancer Moonshot that President Obama asked me to lead six years ago. 
-
-Our goal is to cut the cancer death rate by at least 50% over the next 25 years, turn more cancers from death sentences into treatable diseases.  
-
-More support for patients and families. 
-
-To get there, I call on Congress to fund ARPA-H, the Advanced Research Projects Agency for Health. 
-
-It’s based on DARPA—the Defense Department project that led to the Internet, GPS, and so much more.  
-
-ARPA-H will have a singular purpose—to drive breakthroughs in cancer, Alzheimer’s, diabetes, and more. 
-
-A unity agenda for the nation. 
-
-We can do this. 
-
-My fellow Americans—tonight , we have gathered in a sacred space—the citadel of our democracy. 
-
-In this Capitol, generation after generation, Americans have debated great questions amid great strife, and have done great things. 
-
-We have fought for freedom, expanded liberty, defeated totalitarianism and terror. 
-
-And built the strongest, freest, and most prosperous nation the world has ever known. 
-
-Now is the hour. 
-
-Our moment of responsibility. 
-
-Our test of resolve and conscience, of history itself. 
-
-It is in this moment that our character is formed. Our purpose is found. Our future is forged. 
-
-Well I know this nation.  
-
-We will meet the test. 
-
-To protect freedom and liberty, to expand fairness and opportunity. 
-
-We will save democracy. 
-
-As hard as these times have been, I am more optimistic about America today than I have been my whole life. 
-
-Because I see the future that is within our grasp. 
-
-Because I know there is simply nothing beyond our capacity. 
-
-We are the only nation on Earth that has always turned every crisis we have faced into an opportunity. 
-
-The only nation that can be defined by a single word: possibilities. 
-
-So on this night, in our 245th year as a nation, I have come to report on the State of the Union. 
-
-And my report is this: the State of the Union is strong—because you, the American people, are strong. 
-
-We are stronger today than we were a year ago. 
-
-And we will be stronger a year from now than we are today. 
-
-Now is our moment to meet and overcome the challenges of our time. 
-
-And we will, as one people. 
-
-One America. 
-
-The United States of America. 
-
-May God bless you all. May God protect our troops.
@@ -1,67 +0,0 @@
-import os
-import pytest
-import numpy as np
-
-from fastembed.embedding import DefaultEmbedding, JinaEmbedding
-
-CANONICAL_VECTOR_VALUES = {
-    "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
-    "BAAI/bge-small-en-v1.5": np.array([0.01522374, -0.02271799, 0.00860278, -0.07424029, 0.00386434]),
-    "BAAI/bge-small-zh-v1.5": np.array([-0.01023294, 0.07634465, 0.0691722, -0.04458365, -0.03160762]),
-    "BAAI/bge-base-en": np.array([0.0115, 0.0372, 0.0295, 0.0121, 0.0346]),
-    "BAAI/bge-base-en-v1.5": np.array([0.01129394, 0.05493144, 0.02615099, 0.00328772, 0.02996045]),
-    "sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259, 0.0058, 0.0114, 0.0380, -0.0233]),
-    "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
-    "jinaai/jina-embeddings-v2-small-en": np.array([-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]),
-    "jinaai/jina-embeddings-v2-base-en": np.array([-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]),
-}
-
-
-@pytest.mark.parametrize('embedding_class', [DefaultEmbedding, JinaEmbedding])
-def test_embedding(embedding_class):
-    is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
-
-    for model_desc in embedding_class.list_supported_models():
-        if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
-            continue
-
-        dim = model_desc["dim"]
-        model = embedding_class(model_name=model_desc["model"])
-
-        docs = ["hello world", "flag embedding"]
-        embeddings = list(model.embed(docs))
-        embeddings = np.stack(embeddings, axis=0)
-        assert embeddings.shape == (2, dim)
-
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
-        assert np.allclose(embeddings[0, :canonical_vector.shape[0]], canonical_vector, atol=1e-3), model_desc["model"]
-
-
-@pytest.mark.parametrize('n_dims,embedding_class', [(384, DefaultEmbedding), (768, JinaEmbedding)])
-def test_batch_embedding(n_dims, embedding_class):
-    model = embedding_class()
-
-    docs = ["hello world", "flag embedding"] * 100
-    embeddings = list(model.embed(docs, batch_size=10))
-    embeddings = np.stack(embeddings, axis=0)
-
-    assert embeddings.shape == (200, n_dims)
-
-
-@pytest.mark.parametrize('n_dims,embedding_class', [(384, DefaultEmbedding), (768, JinaEmbedding)])
-def test_parallel_processing(n_dims, embedding_class):
-    model = embedding_class()
-
-    docs = ["hello world", "flag embedding"] * 100
-    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
-    embeddings = np.stack(embeddings, axis=0)
-
-    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
-    embeddings_2 = np.stack(embeddings_2, axis=0)
-
-    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
-    embeddings_3 = np.stack(embeddings_3, axis=0)
-
-    assert embeddings.shape == (200, n_dims)
-    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
-    assert np.allclose(embeddings, embeddings_3, atol=1e-3)
@@ -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)
@@ -1,27 +0,0 @@
-import os
-from pathlib import Path
-
-import pytest
-
-from fastembed.embedding import DefaultEmbedding
-from fastembed.models import TextSplitterConfig
-
-
-@pytest.mark.parametrize(["chunk_size", "chunk_overlap"], [[500, 50], [1000, 100]])
-def test_embedding(chunk_size: int, chunk_overlap: int):
-    is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
-
-    for model_desc in DefaultEmbedding.list_supported_models():
-        if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
-            continue
-
-        p = Path(__file__).with_name("state_of_the_union.txt")
-        with open(p, encoding='utf-8') as f:
-            text = f.read()
-            embedding = DefaultEmbedding(
-                model_name=model_desc["model"],
-                splitter_config=TextSplitterConfig(chunk_size=chunk_size, chunk_overlap=chunk_overlap),
-            )
-            texts = embedding.split_text(text)
-            for text in texts:
-                assert len(embedding.model.tokenizer.encode(text)) <= chunk_size
@@ -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)
Author	SHA1	Message	Date
George Panchuk	5afa07b921	new: update dev workflow	2024-05-02 20:34:47 +02:00
Arun	432da42c11	fix links (#215 )	2024-04-27 22:24:29 +05:30
George	5cde2898bc	new: remove slurm environment variables (#213 )	2024-04-26 21:45:02 +02:00
AnushandGeorge	ab7a99a748	feat: Quantized models (#201 ) * 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>	2024-04-26 17:27:16 +02:00
Anush	466886a317	ci: Schedule python-tests.yml (#211 ) * ci: Schedule python-tests.yml * ci: use emojis * ci: Bump action versions python-tests.yml * ci: python-tests.yml	2024-04-26 10:28:37 +05:30
Anush	cc4112d859	feat: Snowflake models (#207 ) * feat: Snowflake models * Added snowflake/snowflake-arctic-embed-m * docs: snowflake/snowflake-arctic-embed-m	2024-04-19 19:00:10 +05:30
Nirant	864217a7d9	Re-run of identical hardware and generate graphs (#205 )	2024-04-18 11:13:18 +05:30
Andrew Green	9bad44368e	Workaround for running on SLURM (#198 ) * Workaround for running on SLURM onnxruntime would usually get the number of threads from OMP_NUM_THREADS, but that isn't set on SLURM which handles the number of threads differently. This addition tries to figure out if we're running under SLURM, and if so sets the session options accordingly using SLURM environment variables instead. Tested working with latest versions of onnxruntime and fastembed on slurm 23.02.07 * onnxruntime requires number of threads be an integer Caused by me having mis-matched version from my machine and the slurm cluster :( * `os.getenv` returns None for unset environment variables, fix logic Instead of empty string, which I thought it did	2024-04-15 19:14:11 +05:30
dependabot[bot]andNirant	a5cab7a31a	build(deps): bump idna from 3.6 to 3.7 (#195 ) Bumps [idna](https://github.com/kjd/idna) from 3.6 to 3.7. - [Release notes](https://github.com/kjd/idna/releases) - [Changelog](https://github.com/kjd/idna/blob/master/HISTORY.rst) - [Commits](https://github.com/kjd/idna/compare/v3.6...v3.7) --- updated-dependencies: - dependency-name: idna dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-04-15 06:04:54 +05:30
Anush	e55c145924	chore: Exclude unused model repo files (#196 ) * chore: Exclude unused model files * fix: blob pattern	2024-04-12 20:15:02 +05:30
dependabot[bot]andNirant	ad02f60eea	build(deps-dev): bump pillow from 10.2.0 to 10.3.0 (#186 ) Bumps [pillow](https://github.com/python-pillow/Pillow) from 10.2.0 to 10.3.0. - [Release notes](https://github.com/python-pillow/Pillow/releases) - [Changelog](https://github.com/python-pillow/Pillow/blob/main/CHANGES.rst) - [Commits](https://github.com/python-pillow/Pillow/compare/10.2.0...10.3.0) --- updated-dependencies: - dependency-name: pillow dependency-type: direct:development ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-04-08 16:21:32 +05:30
NirantandGeorge	2fcec07f1b	Add Python version to the ISSUE_TEMPLATE (#188 ) * Add Python version * Update .github/ISSUE_TEMPLATE/bug-report.yml Co-authored-by: George <george.panchuk@qdrant.tech> --------- Co-authored-by: George <george.panchuk@qdrant.tech>	2024-04-08 16:05:29 +05:30
George	f340a73a3e	refactoring: update binary quantization notebook (#180 )	2024-04-02 14:45:27 +05:30
Nirant	335f673f3a	Update bug-report.yml	2024-04-02 13:51:51 +05:30
Nirant	ee7ba0a536	Update bug-report.yml	2024-04-02 13:51:27 +05:30
George	912e95e5c8	fix: remove model archive if model extraction was not finished correctly (#179 )	2024-04-01 20:11:32 +02:00
Nirant	c09909773e	Add Execution Counts (#177 ) * Update notebooks * Update FastEmbed usage across docs * Refactor code for better readability and maintainability * Clean outputs * Change dataset * Add numbers inline in output * Remove inline outputs since I used :memory: * Fix syntax error in Hindi_Tamil_RAG_with_Navarasa7B.ipynb	2024-04-01 20:48:26 +05:30
Nirant	3d2254d215	Bump version to 0.2.6 in pyproject.toml (#175 )	2024-04-01 17:16:07 +05:30
GeorgeandNirant Kasliwal	e11f66bb55	refactoring: update imports in notebooks (#173 ) * new: simplify imports * refactoring: update import * refactoring: update imports in notebooks * fix: fix notebook output * Re-run notebook with revised imports --------- Co-authored-by: Nirant Kasliwal <nirant.bits@gmail.com>	2024-04-01 17:02:13 +05:30
Nirant	d6f9ace425	Update size_in_GB for BAAI/bge-small-en-v1.5 model (#176 )	2024-04-01 16:40:45 +05:30
George Nirant Nirant Kasliwal	8b7b8476a5	fix: fix model sizes in supported models lists (#167 ) * fix: fix model sizes in supported models lists * fix: remove redundant comment * fix: fix test * fix: update supported models notebook * Consistentcy around quantization in supported_onnx_models --------- Co-authored-by: Nirant <NirantK@users.noreply.github.com> Co-authored-by: Nirant Kasliwal <nirant.bits@gmail.com>	2024-04-01 16:18:49 +05:30
George	ae35a96bcb	new: simplify imports (#171 ) * new: simplify imports * refactoring: update import	2024-04-01 12:58:26 +05:30
GeorgeandNirant	25671ec349	Update ruff (#172 ) * refactoring: reduce max line-length * new: update ruff --------- Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-04-01 12:55:15 +05:30
George	ce98631b9a	Fix spladepp parallelism (#169 ) * fix: add get_worker_class implementation to spladepp * fix: add tests for parallel embed for spladepp	2024-04-01 10:22:24 +05:30
George	0a4ed42b58	fix: unify existing patterns, remove redundant (#168 )	2024-03-30 14:40:19 +05:30
NirantandAnush	e3d2e1dc44	Hybrid Search Tutorial (#165 ) * Re-organize docs * Rename notebooks * Move nbs * Working Sparse and Dense Search * Add RRF * Refactor code to improve performance and readability * Add ESCI label for the RRF results * Update docs/examples/Hybrid_Search.ipynb Co-authored-by: Anush <anushshetty90@gmail.com> * Update docs/examples/Hybrid_Search.ipynb Co-authored-by: Anush <anushshetty90@gmail.com> * Remove unnecessary code and update vector format --------- Co-authored-by: Anush <anushshetty90@gmail.com>	2024-03-29 21:10:14 +05:30
Nirant	62c21b0237	Add misspelled version of SPLADE++ model for English (#161 )	2024-03-22 21:06:03 +05:30
Anush	d791f38704	chore: case-insensitive model_management.py (#160 )	2024-03-22 21:02:31 +05:30
Nirant Kasliwal	c651b2b539	Update SPLADE notebook with new sections	2024-03-22 15:18:34 +05:30
Yuvraj WaleandNirant	bc1e23849b	feat: support mixedbread-ai/mxbai-embed-large-v1 (#158 ) * add: support mixedbread-ai/mxbai-embed-large-v1 * refactor: canonical vector * Update fastembed/text/onnx_embedding.py --------- Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-03-22 06:23:54 +05:30
Nirant	4db4839995	[PyPi Publish] Bump version to 0.2.5 in pyproject.toml (#156 ) * Bump version to 0.2.5 in pyproject.toml * chore: case insensitive check (#157)	2024-03-20 19:11:44 +05:30
NirantandAnush	96a2a9097a	Fix model name typo + Add SPLADE notebook (#155 ) * Rename model + Add SPLADE notebook * Update docs/examples/SPLADE_with_FastEmbed.ipynb Co-authored-by: Anush <anushshetty90@gmail.com> * Update docs/examples/SPLADE_with_FastEmbed.ipynb Co-authored-by: Anush <anushshetty90@gmail.com> * Update docs/examples/SPLADE_with_FastEmbed.ipynb Co-authored-by: Anush <anushshetty90@gmail.com> * Update CANONICAL_COLUMN_VALUES in test_sparse_embeddings.py --------- Co-authored-by: Anush <anushshetty90@gmail.com>	2024-03-20 18:15:53 +05:30
NirantandKumar Shivendu	256b2265d5	Move CONTRIBUTING.md + Add Test for Adding New Models (#154 ) * Move CONTRIBUTING.md + Ad Test for Adding New Models * Update CONTRIBUTING.md Co-authored-by: Kumar Shivendu <kshivendu1@gmail.com> --------- Co-authored-by: Kumar Shivendu <kshivendu1@gmail.com>	2024-03-18 19:58:02 +05:30
NirantandKumar Shivendu	1e91c8d165	Fix Issue Template forms (#152 ) * Add CONTRIBUTING.md file with guidelines for contributing to FastEmbed * Add code linting and pre-commit info to CONTRIBUTING * Update CONTRIBUTING.md Co-authored-by: Kumar Shivendu <kshivendu1@gmail.com> * Add bug/new model issue template and move CONTRIBUTING.md * Re-organize issue templates --------- Co-authored-by: Kumar Shivendu <kshivendu1@gmail.com>	2024-03-14 15:11:23 +05:30
Nirant	a761f456b2	Add import statement for version debugging (#151 )	2024-03-14 15:02:15 +05:30
NirantandKumar Shivendu	287e19c494	Add CONTRIBUTING.md file with guidelines for contributing to FastEmbed (#150 ) * Add CONTRIBUTING.md file with guidelines for contributing to FastEmbed * Add code linting and pre-commit info to CONTRIBUTING * Update CONTRIBUTING.md Co-authored-by: Kumar Shivendu <kshivendu1@gmail.com> * Add bug/new model issue template and move CONTRIBUTING.md --------- Co-authored-by: Kumar Shivendu <kshivendu1@gmail.com>	2024-03-14 14:50:48 +05:30
Anush	6a94994038	Release v0.2.4 (#149 )	2024-03-13 23:56:01 +05:30
Andrey Vasnetsov	361f674e47	avoid changing output dimentionality for a single input (#148 )	2024-03-13 23:38:51 +05:30
Nirant	041a606285	Merge pull request #146 from qdrant:v0.2.3 Publish to PyPi with SPLADE models	2024-03-13 18:17:41 +05:30
Nirant Kasliwal	5b937c29f6	Update poetry install command to exclude docs	2024-03-13 18:15:51 +05:30
Nirant Kasliwal	8d368889c0	Update version and add pre-commit dependency	2024-03-13 18:07:31 +05:30
Nirant generall Anush008	d817da2e01	Add Splade v1 (#144 ) * Add SPLADE v1 * WIP SPLADE Export errors * add ONNX model to HF hub and use that * Update sentences in Converting_SPLADE_to_ONNX.ipynb * Remove unnecessary files and directories * Rename var in TextEmbedding class to use EMBEDDING_MODEL_TYPE * Add SPLADE to list of text embeddings * Add SPLADE model support for text embedding * Fix deprecation warning in embedding.py * Add test for batch embedding with sparse embeddings * Refactor import statement in test_sparse_embeddings.py * Rename nbs * Update vocab size in SPLADE model * Fix canonical vector lookup in test_text_onnx_embeddings.py * review refactoring * restore list_supported_models in OnnxTextEmbedding * Remove unused method _preprocess_onnx_input() in SpladePP class * Update SPLADE_PP_en_v1 source in splade_pp.py * Refactor onnx_model.py to change base model behavior * extend tests to sparse values as well as indicies * chore: pre-commit hooks --------- Co-authored-by: generall <andrey@vasnetsov.com> Co-authored-by: Anush008 <anushshetty90@gmail.com>	2024-03-13 18:04:44 +05:30
Artem Daineko	68635efa3f	Fix link to Optimus in docs (#143 )	2024-03-10 11:45:44 +05:30
Nirant	9ffde58df6	Getting Started Improvements (#138 ) * Update FastEmbed README.md * Rewrite GettingStarted to use TextEmbedding instead of DefaultEmbedding * Improve grammar * Update Getting Started.ipynb with model information and document format	2024-03-07 17:01:11 +05:30
Nirant Kasliwal	5603fbe1fb	Fix naming typo	2024-03-06 14:30:54 +05:30
Nirant Kasliwal	9ed4486d9c	Rename typo in notebook	2024-03-06 14:24:48 +05:30
Nirant Kasliwal	e36a39c388	Update author information in notebook	2024-03-06 14:18:48 +05:30
Nirant Kasliwal	97c359c1b9	Add Navrasa LLM download and explanation sections	2024-03-06 14:17:41 +05:30
Nirant Kasliwal	9fd51425fe	Add author information and Colab link to notebook	2024-03-06 14:12:04 +05:30
Nirant Kasliwal	0dec22d02d	Remove inline outputs	2024-03-06 14:06:46 +05:30
Nirant Kasliwal	8a3b746b71	Rename notebook	2024-03-06 14:05:13 +05:30
Nirant	337ad9c93f	Hindi RAG with Qdrant and FastEmbed (#135 ) * Add workingnb * Add A100 Colab * Remove old checkpoint * Refactor code to separate HF Token	2024-03-06 14:02:32 +05:30
NirantandAnush	74062e8607	Add attention export functionality to experiments (#134 ) * Add attention export functionality * Update experiments/attention_export.py Co-authored-by: Anush <anushshetty90@gmail.com> --------- Co-authored-by: Anush <anushshetty90@gmail.com>	2024-03-04 16:30:17 +05:30
Anush	1e298a00b3	feat: Added gte-large, nomic-text 1.5, cleanup (#130 )	2024-02-21 17:54:32 +05:30
Nathan LeRoyandNirant	38c4eb1cc5	Check for existing files in cache dir before instantiating a model (#128 ) * check for existing files * (chore: model_management.py): Add comment --------- Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-02-21 13:56:14 +05:30
Armaghan	406f432edc	feat: Support sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 (#129 ) * feat: Support sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 * test: Include sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 * docs: supported models update	2024-02-21 13:35:49 +05:30
Nirant	defb6183c1	* feat(pyproject): updated version to '0.2.2' (#124 ) * chore(pyproject): updated dev dependencies versions	2024-02-19 14:46:01 +05:30
AnushandNirant	98141cc8d3	feat: Added nomic-embed-text-v1 support + formatting changes + import fixes (#118 ) * feat: Added nomic-embed-text-v1 support * chore: xenova/nomic-embed-text-v1 -> nomic-ai/nomic-embed-text-v1 --------- Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-02-19 14:17:33 +05:30
Nirant	b1f5e7a989	Add import statement to the warning message (#119 )	2024-02-15 12:40:11 +05:30
Kumar Shivendu	558a837531	Merge pull request #112 from qdrant/KShivendu-patch-1 docs: Describe how to change the model and how to just create embeddings	2024-02-13 10:33:41 +05:30
Nirant	b81e40c95d	Merge branch 'main' into KShivendu-patch-1	2024-02-13 08:25:51 +05:30
Nirant	81bab0cd1d	Make 0.2.1 Release + Update docs (#116 ) * Update version from 0.2.0 (yanked) to 0.2.1 * Update text embedding to include prefix for passages and queries * Update supported models to use the latest API * * fix(text_embedding_base.py): remove unnecessary prefix from texts in embed method * feat(text_embedding_base.py): update query_embed method to updated instruction for the v1.5 model * Remove comparison, since the ranking is identical even with varying embedding * Refactor text embedding query handling	2024-02-08 09:06:00 +05:30
Nirant	973da354ae	Fix query to align with Qdrant mixin usage (#115 ) * fix: query in text_embedding_base to work with both Iterable and str as users might supply both * Fix Qdrant query to align with future usage * * refactor(text_embedding_base.py): change query parameter type from str to Union[str, Iterable[str]] in query_embed method * Update return type of query_embed method * Update return type in TextEmbeddingBase	2024-02-07 22:01:31 +05:30
Zihong Lin Kumar Shivendu Nirant	46968181ad	Simplify imports: #110 (#113 ) * Simplify imports: #110 * Update fastembed/__init__.py Co-authored-by: Kumar Shivendu <kshivendu1@gmail.com> * Remove outdated import * Remove outdated import --------- Co-authored-by: Kumar Shivendu <kshivendu1@gmail.com> Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-02-07 21:12:46 +05:30
Nirant Kasliwal	3948f0db2e	Update fastembed v0.2.0	2024-02-07 20:36:46 +05:30
Nirant	cf66d0e5e1	Update Python and dependency versions (#111 )	2024-02-05 12:37:52 +01:00
Kumar Shivendu	c11ba70fbc	Update README.md	2024-02-05 07:17:10 +01:00
Kumar Shivendu	ea3ef26fa2	Improve README	2024-02-05 07:07:41 +01:00
Kumar Shivendu	4b1ffb47f0	docs: Describe how to change the model and how to just create embeddings	2024-02-05 11:24:47 +05:30
Kumar Shivendu	d3f5f29ee0	docs: Improve README (#109 )	2024-02-05 10:49:45 +05:30
Kumar Shivendu	05885a36dd	refactor: Introduce experiments dir (#108 )	2024-02-05 10:49:12 +05:30
Andrey Vasnetsov	a3bc73c556	Merge pull request #105 from qdrant/refactoring-off-everything Refactoring of internal structure	2024-02-02 16:48:04 +01:00
generall	fcdc5690b9	rename models	2024-02-02 15:51:19 +01:00
generall	7883fa3c41	new multilingual models	2024-02-02 15:32:42 +01:00
generall	5dbd0073e2	rename flag -> onnx	2024-02-02 15:12:57 +01:00
generall	8b800da7bc	ruff	2024-02-02 15:12:57 +01:00
generall	3e6f69e2eb	review fixes	2024-02-02 15:12:57 +01:00
generallandGeorge Panchuk	4813b18854	refactoring Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>	2024-02-02 15:12:57 +01:00
Anush	96f7d83d33	feat: Support xenova/multilingual-e5-large, xenova/paraphrase-multili… (#103 ) * feat: Support xenova/multilingual-e5-large, xenova/paraphrase-multilingual-mpnet-base-v2 * chore: updated exclude_token_type_ids check * docs: supported models update	2024-02-02 15:16:35 +05:30
Nirant	2e3e5508c0	Update poetry lock to latest versions (#98 ) * Update poetry lock to latest versions * Update poetry lock to latest versions	2024-01-30 21:39:52 +05:30
Anush	87decb0d53	chore: port to Xenova Jina source (#102 ) * chore: xenova jina * chore: try recusive model location * chore: updated doc string, blob pattern	2024-01-30 21:27:16 +05:30
Anush	ede507e2cf	feat: HuggingFace download support for FlagEmbedding (#94 ) * feat: HF support for FlagEmbedding * chore: update docstring embedding.py * refactor: GCS URLs models.json * chore: toLower() models.json * chore: update tqdm declarative * chore: exclude keys list_supported_models * chore: review changes	2024-01-23 12:47:55 +05:30
David Janes	f87330fcd1	use "with" to open JSON files (#96 )	2024-01-22 12:20:06 +05:30
Nirant	3b32619a4c	Update Python version and add pre-commit dependency (#93 ) * Update Python version and add pre-commit dependency * Remove Python 3.8.x from matrix * Update Python version and pre-commit configuration	2024-01-16 19:30:47 +05:30
AnushandNirant Kasliwal	9b63427118	chore: pre-commit formatting (#91 ) * chore: formatting * chore: formatting * chore: remove other hooks * Update poetry lock --------- Co-authored-by: Nirant Kasliwal <nirant.bits@gmail.com>	2024-01-16 15:06:54 +05:30