new: gpu package (#224)

* fix: stage GCS downloads instead of deleting the caller's cache dir

* fix: verify archive integrity and give each download its own staging dir

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -0,0 +1,59 @@
+name: Bug
+description: File a bug report
+title: "[Bug]: "
+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: Describe the error you encountered.
+      placeholder: <Description>
+    validations:
+      required: true
+  - type: textarea
+    id: expected
+    attributes:
+      label: What is the expected behaviour?
+      description: Describe the way you expected the code to behave.
+      placeholder: <Description>
+  - type: textarea
+    id: code-snippet
+    attributes:
+      label: A minimal reproducible example
+      description: It would really help us to fix the problem if you could provide a code snippet that reproduces the issue.
+      placeholder: <Code snippet>
+  - 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: textarea
+    id: version
+    attributes:
+      label: FastEmbed 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.
+      placeholder: v0.7.4
+    validations:
+      required: true
+  - type: dropdown
+    id: os
+    attributes:
+      label: What os are you seeing the problem on?
+      multiple: true
+      options:
+        - Linux
+        - MacOS
+        - Windows
+  - type: textarea
+    id: logs
+    attributes:
+      label: Relevant stack traces and/or logs
+      description: Please copy and paste any relevant raised exceptions. This will be automatically formatted into code, so no need for backticks.
+      render: shell

.github/ISSUE_TEMPLATE/config.yaml

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

.github/ISSUE_TEMPLATE/feature-request.yml

@@ -0,0 +1,22 @@
+name: Feature
+description: New functionality request
+title: "[Feature]: "
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Thanks for taking the time to fill out this report!
+  - type: textarea
+    id: feature-description
+    attributes:
+      label: What feature would you like to request?
+      description: Please provide the description of the feature you would like to request.
+      placeholder: <Description>
+    validations:
+      required: true
+  - type: textarea
+    id: additional-info
+    attributes:
+      label: Is there any additional information you would like to provide?
+      description: Please provide any additional information that you think might be useful.
+      placeholder: <Info>

.github/ISSUE_TEMPLATE/model-request.yml

@@ -0,0 +1,22 @@
+name: Model
+description: Request a new model
+title: "[Model]: "
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Thanks for taking the time to fill out this report!
+  - type: textarea
+    id: model-name
+    attributes:
+      label: Which model would you like to support?
+      description: Please provide the name of the model you would like to see supported.
+      placeholder: Link to the model (e.g. on HuggingFace)
+    validations:
+      required: true
+  - type: textarea
+    id: motivation
+    attributes:
+      label: What are the main advantages of this model?
+      description: Please describe the main advantages of this model comparing to the existing ones and provide links to benchmarks if there are any.
+      placeholder: <Description>

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,19 @@
+### All Submissions:
+
+* [ ] Have you followed the guidelines in our Contributing document?
+* [ ] Have you checked to ensure there aren't other open [Pull Requests](../../../pulls) for the same update/change?
+
+<!-- You can erase any parts of this template not applicable to your Pull Request. -->
+
+### New Feature Submissions:
+
+* [ ] Does your submission pass the existing tests?
+* [ ] Have you added tests for your feature?
+* [ ] Have you installed `pre-commit` with `pip3 install pre-commit` and set up hooks with `pre-commit install`?
+
+### New models submission:
+
+* [ ] Have you added an explanation of why it's important to include this model?
+* [ ] Have you added tests for the new model? Were canonical values for tests computed via the original model?
+* [ ] Have you added the code snippet for how canonical values were computed?
+* [ ] Have you successfully ran tests with your changes locally?

.github/dependabot.yml

@@ -0,0 +1,40 @@
+version: 2
+updates:
+  - package-ecosystem: "pip"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    open-pull-requests-limit: 5
+    groups:
+      security-updates:
+        applies-to: security-updates
+        patterns:
+          - "*"
+      version-updates:
+        applies-to: version-updates
+        update-types:
+          - "minor"
+          - "patch"
+        patterns:
+          - "*"
+    cooldown:
+      default-days: 7
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    open-pull-requests-limit: 5
+    groups:
+      security-updates:
+        applies-to: security-updates
+        patterns:
+          - "*"
+      version-updates:
+        applies-to: version-updates
+        update-types:
+          - "minor"
+          - "patch"
+        patterns:
+          - "*"
+    cooldown:
+      default-days: 7

.github/workflows/ci.yml

@@ -1,8 +1,8 @@
-name: ci 
+name: ci
 on:
   push:
     branches:
-      - master 
+      - master
       - main
 permissions:
   contents: write
@@ -10,16 +10,16 @@ jobs:
   deploy:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
         with:
           python-version: 3.x
-      - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV 
-      - uses: actions/cache@v3
+      - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
+      - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
         with:
           key: mkdocs-material-${{ env.cache_id }}
           path: .cache
           restore-keys: |
             mkdocs-material-
-      - run: pip install mkdocs-material mkdocstrings pillow cairosvg mknotebooks
+      - run: pip install mkdocs-material mkdocstrings==0.27.0 pillow cairosvg mknotebooks
       - run: mkdocs gh-deploy --force

.github/workflows/python-publish.yml

@@ -15,18 +15,17 @@ on:
     tags:
       - 'v*'           # Push events to every version tag
 
-
 jobs:
   deploy:
 
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
       with:
-        python-version: '3.9.x'
+        python-version: '3.10.x'
     - name: Install dependencies
       run: |
         python -m pip install poetry
@@ -34,7 +33,7 @@ jobs:
     - name: Build package
       run: poetry build
     - name: Publish package
-      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2
       with:
         user: __token__
         password: ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/python-tests.yml

@@ -1,9 +1,11 @@
 name: Tests
+run-name: Tests (gpu)
 
 on:
-  push:
-    branches: [ master, main ]
   pull_request:
+    branches: [ master, main, gpu ]
+  workflow_dispatch:
+
 
 env:
   CARGO_TERM_COLOR: always
@@ -14,32 +16,31 @@ jobs:
     strategy:
       matrix:
         python-version:
-          - '3.8.x'
-          - '3.9.x'
           - '3.10.x'
           - '3.11.x'
+          - '3.12.x'
+          - '3.13.x'
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
 
     runs-on: ${{ matrix.os }}
 
     name: Python ${{ matrix.python-version }} on ${{ matrix.os }} test
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
         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
+          poetry install --no-interaction --no-ansi --without dev,docs
+
+      - name: Run pytest
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          export IS_UBUNTU_CI=$(test "${{ matrix.os }}" = "ubuntu-latest" && echo "true" || echo "false")
-          pytest
-        shell: bash
+          poetry run pytest

.github/workflows/type-checkers.yml

@@ -0,0 +1,38 @@
+name: type-checkers
+
+on: [push]
+
+jobs:
+  build:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: true
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        os: [ubuntu-latest]
+
+    name: Python ${{ matrix.python-version }} test
+
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip poetry
+          poetry install --no-interaction --no-ansi --without dev,docs,test
+
+      - name: mypy
+        run: |
+            poetry run mypy fastembed \
+              --disallow-incomplete-defs \
+              --disallow-untyped-defs \
+              --disable-error-code=import-untyped
+
+      - name: pyright
+        run: |
+          poetry run pyright tests/type_stub.py

.gitignore

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

.pre-commit-config.yaml

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

CONTRIBUTING.md

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

LICENSE

@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright [yyyy] [name of copyright owner]
+   Copyright 2026 Qdrant Solutions GmbH
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

NOTICE

@@ -0,0 +1,24 @@
+Copyright 2024 Qdrant
+
+This product includes software developed by Qdrant
+
+This distribution includes the following Jina AI models, each with its respective license:
+- jinaai/jina-colbert-v2
+    - License: cc-by-nc-4.0
+- jinaai/jina-reranker-v2-base-multilingual
+    - License: cc-by-nc-4.0
+- jinaai/jina-embeddings-v3
+    - License: cc-by-nc-4.0
+
+These models are developed by Jina (https://jina.ai/) and are subject to Jina AI's licensing terms.
+
+This distribution includes the following Google models, each with its respective license:
+- vidore/colpali-v1.3
+    - License: gemma
+- google/embeddinggemma-300m
+    - License: gemma
+
+Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms
+
+Additional Notes:
+This project also includes third-party libraries with their respective licenses. Please refer to the documentation of each library for details regarding its usage and licensing terms.

README.md

@@ -1,41 +1,232 @@
 # ⚡️ What is FastEmbed?
 
-FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a Github issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
+FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
-The default embedding supports "query" and "passage" prefixes for the input text. The default model is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval%20with%20FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/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 [support](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 best. You can install it with or without GPU support:
 
 ```bash
 pip install fastembed
+
+# or with GPU support
+
+pip install fastembed-gpu
 ```
 
-## 📖 Usage
+## 📖 Quickstart
 
 ```python
-from fastembed.embedding import FlagEmbedding as Embedding
+from fastembed import TextEmbedding
 
-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
+
+# Example list of documents
+documents: list[str] = [
+    "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
+```
+
+Fastembed supports a variety of models for different tasks and modalities.
+The list of all the available models can be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/)
+### 🎒 Dense text embeddings
+
+```python
+from fastembed import TextEmbedding
+
+model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
+embeddings = list(model.embed(documents))
+
+# [
+#   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
+#   array([-0.1019,  0.0635, -0.0332,  0.0522, ...], dtype=float32)
+# ]
+
+```
+
+Dense text embedding can also be extended with models which are not in the list of supported models.
+
+```python
+from fastembed import TextEmbedding
+from fastembed.common.model_description import PoolingType, ModelSource
+
+TextEmbedding.add_custom_model(
+    model="intfloat/multilingual-e5-small",
+    pooling=PoolingType.MEAN,
+    normalization=True,
+    sources=ModelSource(hf="intfloat/multilingual-e5-small"),  # can be used with an `url` to load files from a private storage
+    dim=384,
+    model_file="onnx/model.onnx",  # can be used to load an already supported model with another optimization or quantization, e.g. onnx/model_O4.onnx
+)
+model = TextEmbedding(model_name="intfloat/multilingual-e5-small")
+embeddings = list(model.embed(documents))
+```
+
+
+### 🔱 Sparse text embeddings
+
+* SPLADE++
+
+```python
+from fastembed import SparseTextEmbedding
+
+model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
+embeddings = list(model.embed(documents))
+
+# [
+#   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
+#   SparseEmbedding(indices=[ 38,  12,  91, ... ], values=[0.11, 0.22, 0.39, ...])
+# ]
+```
+
+<!--
+* BM42 - ([link](ToDo))
+
+```
+from fastembed import SparseTextEmbedding
+
+model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
+embeddings = list(model.embed(documents))
+
+# [
+#   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
+#   SparseEmbedding(indices=[ 38,  12,  91, ... ], values=[0.11, 0.22, 0.39, ...])
+# ]
+```
+-->
+
+### 🦥 Late interaction models (aka ColBERT)
+
+
+```python
+from fastembed import LateInteractionTextEmbedding
+
+model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
+embeddings = list(model.embed(documents))
+
+# [
+#   array([
+#       [-0.1115,  0.0097,  0.0052,  0.0195, ...],
+#       [-0.1019,  0.0635, -0.0332,  0.0522, ...],
+#   ]),
+#   array([
+#       [-0.9019,  0.0335, -0.0032,  0.0991, ...],
+#       [-0.2115,  0.8097,  0.1052,  0.0195, ...],
+#   ]),  
+# ]
+```
+
+### 🖼️ Image embeddings
+
+```python
+from fastembed import ImageEmbedding
+
+images = [
+    "./path/to/image1.jpg",
+    "./path/to/image2.jpg",
+]
+
+model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
+embeddings = list(model.embed(images))
+
+# [
+#   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
+#   array([-0.1019,  0.0635, -0.0332,  0.0522, ...], dtype=float32)
+# ]
+```
+
+### Late interaction multimodal models (ColPali)
+
+```python
+from fastembed import LateInteractionMultimodalEmbedding
+
+doc_images = [
+    "./path/to/qdrant_pdf_doc_1_screenshot.jpg",
+    "./path/to/colpali_pdf_doc_2_screenshot.jpg",
+]
+
+query = "What is Qdrant?"
+
+model = LateInteractionMultimodalEmbedding(model_name="Qdrant/colpali-v1.3-fp16")
+doc_images_embeddings = list(model.embed_image(doc_images))
+# shape (2, 1030, 128)
+# [array([[-0.03353882, -0.02090454, ..., -0.15576172, -0.07678223]], dtype=float32)]
+query_embedding = model.embed_text(query)
+# shape (1, 20, 128)
+# [array([[-0.00218201,  0.14758301, ...,  -0.02207947,  0.16833496]], dtype=float32)]
+```
+
+### 🔄 Rerankers
+```python
+from fastembed.rerank.cross_encoder import TextCrossEncoder
+
+query = "Who is maintaining Qdrant?"
+documents: list[str] = [
+    "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.",
+]
+encoder = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-6-v2")
+scores = list(encoder.rerank(query, documents))
+
+# [-11.48061752319336, 5.472434997558594]
+```
+
+Text cross encoders can also be extended with models which are not in the list of supported models.
+
+```python
+from fastembed.rerank.cross_encoder import TextCrossEncoder 
+from fastembed.common.model_description import ModelSource
+
+TextCrossEncoder.add_custom_model(
+    model="Xenova/ms-marco-MiniLM-L-4-v2",
+    model_file="onnx/model.onnx",
+    sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-4-v2"),
+)
+model = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-4-v2")
+scores = list(model.rerank_pairs(
+    [("What is AI?", "Artificial intelligence is ..."), ("What is ML?", "Machine learning is ..."),]
+))
+```
+
+## ⚡️ FastEmbed on a GPU
+
+FastEmbed supports running on GPU devices.
+It requires installation of the `fastembed-gpu` package.
+
+```bash
+pip install fastembed-gpu
+```
+
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for detailed instructions, CUDA 12.x support and troubleshooting of the common issues.
+
+```python
+from fastembed import TextEmbedding
+
+embedding_model = TextEmbedding(
+    model_name="BAAI/bge-small-en-v1.5", 
+    providers=["CUDAExecutionProvider"]
+)
+print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
+
 ```
 
 ## Usage with Qdrant
@@ -46,37 +237,45 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```
 
-Might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh. 
+or 
 
-```python
-from qdrant_client import QdrantClient
-
-# Initialize the client
-client = QdrantClient(":memory:")  # or QdrantClient(path="path/to/db")
-
-# Prepare your documents, metadata, and IDs
-docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
-metadata = [
-    {"source": "Langchain-docs"},
-    {"source": "Linkedin-docs"},
-]
-ids = [42, 2]
-
-# Use the new add method
-client.add(
-    collection_name="demo_collection",
-    documents=docs,
-    metadata=metadata,
-    ids=ids
-)
-
-search_result = client.query(
-    collection_name="demo_collection",
-    query_text="This is a query document"
-)
-print(search_result)
+```bash
+pip install qdrant-client[fastembed-gpu]
 ```
 
-#### Similar Work
+You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
-Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.
+```python
+from qdrant_client import QdrantClient, models
+
+# Initialize the client
+client = QdrantClient("localhost", port=6333) # For production
+# client = QdrantClient(":memory:") # For experimentation
+
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
+payload = [
+    {"document": "Qdrant has Langchain integrations", "source": "Langchain-docs", },
+    {"document": "Qdrant also has Llama Index integrations", "source": "LlamaIndex-docs"},
+]
+docs = [models.Document(text=data["document"], model=model_name) for data in payload]
+ids = [42, 2]
+
+client.create_collection(
+    "demo_collection",
+    vectors_config=models.VectorParams(
+        size=client.get_embedding_size(model_name), distance=models.Distance.COSINE)
+)
+
+client.upload_collection(
+    collection_name="demo_collection",
+    vectors=docs,
+    ids=ids,
+    payload=payload,
+)
+
+search_result = client.query_points(
+    collection_name="demo_collection",
+    query=models.Document(text="This is a query document", model=model_name)
+).points
+print(search_result)
+```

RELEASE.md

@@ -0,0 +1,41 @@
+# Releasing FastEmbed
+
+This is a guide how to release `fastembed` and `fastembed-gpu` packages.
+
+## How to
+
+1. Accumulate changes in the `main` branch.
+2. Bump the version in `pyproject.toml`
+
+3. Rebase the `gpu` branch on `main` and resolve conflicts if occurred:
+
+```bash
+git checkout gpu
+git rebase main
+git push -f origin gpu
+```
+
+4. Draft release notes
+5. Checkout to `main` and create a tag, e.g.:
+
+```bash
+git checkout main
+git tag -a v0.1.0 -m "Release v0.1.0"
+```
+
+6. Checkout `gpu` and create a tag, e.g.:
+
+```bash
+git checkout gpu
+git tag -a v0.1.0-gpu -m "Release v0.1.0"
+```
+
+7. Push tags:
+
+```bash
+git push --tags
+```
+
+8. Verify that both packages have been published successfully on PyPI. Try installing them and verify imports.
+9. Create a release on GitHub with the written release notes.
+

docs/Getting Started.ipynb

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

docs/examples/ColBERT_with_FastEmbed.ipynb

@@ -0,0 +1,421 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d14d29ebd3592ecb",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "# Late Interaction Text Embedding Models\n",
+    "\n",
+    "As of version 0.3.0 FastEmbed supports Late Interaction Text Embedding Models and currently available with one of the most popular embedding model of the family - ColBERT.\n",
+    "\n",
+    "## What is a Late Interaction Text Embedding Model?\n",
+    "\n",
+    "Late Interaction Text Embedding Model is a kind of information retrieval model which performs query and documents interactions at the scoring stage.\n",
+    "In order to better understand it, we can compare it to the models without interaction.  \n",
+    "For instance, if you take a sentence-transformer model, compute embeddings for your documents, compute embeddings for your queries, and just compare them by cosine similarity, then you're retrieving points without interaction.\n",
+    "\n",
+    "It is a pretty much easy and straightforward approach, however we might be sacrificing some precision due to its simplicity. It is caused by several facts: \n",
+    "- there is no interaction between queries and documents at the early stage (embedding generation) nor at the late stage (during scoring). \n",
+    "- we are trying to encapsulate all the document information in only one pooled embedding, and obviously, some information might be lost.\n",
+    "\n",
+    "Late Interaction Text Embedding models are trying to address it by computing embeddings for each token in queries and documents, and then finding the most similar ones via model specific operation, e.g. ColBERT (Contextual Late Interaction over BERT) uses MaxSim operation.\n",
+    "With this approach we can have not only a better representation of the documents, but also make queries and documents more aware one of another.\n",
+    "\n",
+    "For more information on ColBERT and MaxSim operation, you can check out [this blogpost](https://jina.ai/news/what-is-colbert-and-late-interaction-and-why-they-matter-in-search/) by Jina AI.\n",
+    "\n",
+    "## ColBERT in FastEmbed\n",
+    "\n",
+    "FastEmbed provides a simple way to use ColBERT model, similar to the ones it has with `TextEmbedding`.\n",
+    " "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "7f1053b17c810be5",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-03T17:20:26.927643Z",
+     "start_time": "2024-06-03T17:20:25.128994Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/joein/work/qdrant/fastembed/venv/lib/python3.10/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"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[{'model': 'colbert-ir/colbertv2.0',\n",
+       "  'dim': 128,\n",
+       "  'description': 'Late interaction model',\n",
+       "  'size_in_GB': 0.44,\n",
+       "  'sources': {'hf': 'colbert-ir/colbertv2.0'},\n",
+       "  'model_file': 'model.onnx'}]"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from fastembed import LateInteractionTextEmbedding\n",
+    "\n",
+    "LateInteractionTextEmbedding.list_supported_models()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "c2c15893df422631",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-03T17:23:35.764183Z",
+     "start_time": "2024-06-03T17:23:21.630277Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Fetching 5 files:   0%|          | 0/5 [00:00<?, ?it/s]\n",
+      "config.json: 100%|██████████| 743/743 [00:00<00:00, 4.56MB/s]\n",
+      "\n",
+      "tokenizer_config.json: 100%|██████████| 405/405 [00:00<00:00, 3.34MB/s]\n",
+      "Fetching 5 files:  20%|██        | 1/5 [00:00<00:01,  3.64it/s]\n",
+      "tokenizer.json:   0%|          | 0.00/466k [00:00<?, ?B/s]\u001b[A\n",
+      "\n",
+      "special_tokens_map.json: 100%|██████████| 112/112 [00:00<00:00, 727kB/s]\n",
+      "\n",
+      "tokenizer.json: 100%|██████████| 466k/466k [00:00<00:00, 1.48MB/s]\u001b[A\n",
+      "\n",
+      "model.onnx:   0%|          | 0.00/436M [00:00<?, ?B/s]\u001b[A\n",
+      "model.onnx:   2%|▏         | 10.5M/436M [00:00<00:34, 12.2MB/s]\u001b[A\n",
+      "model.onnx:   5%|▍         | 21.0M/436M [00:01<00:20, 20.3MB/s]\u001b[A\n",
+      "model.onnx:   7%|▋         | 31.5M/436M [00:01<00:15, 25.7MB/s]\u001b[A\n",
+      "model.onnx:  10%|▉         | 41.9M/436M [00:01<00:13, 29.4MB/s]\u001b[A\n",
+      "model.onnx:  12%|█▏        | 52.4M/436M [00:01<00:12, 31.9MB/s]\u001b[A\n",
+      "model.onnx:  14%|█▍        | 62.9M/436M [00:02<00:11, 33.7MB/s]\u001b[A\n",
+      "model.onnx:  17%|█▋        | 73.4M/436M [00:02<00:10, 34.9MB/s]\u001b[A\n",
+      "model.onnx:  19%|█▉        | 83.9M/436M [00:02<00:09, 35.5MB/s]\u001b[A\n",
+      "model.onnx:  22%|██▏       | 94.4M/436M [00:03<00:09, 36.1MB/s]\u001b[A\n",
+      "model.onnx:  24%|██▍       | 105M/436M [00:03<00:09, 36.6MB/s] \u001b[A\n",
+      "model.onnx:  26%|██▋       | 115M/436M [00:03<00:08, 36.9MB/s]\u001b[A\n",
+      "model.onnx:  29%|██▉       | 126M/436M [00:03<00:08, 37.1MB/s]\u001b[A\n",
+      "model.onnx:  31%|███▏      | 136M/436M [00:04<00:08, 37.3MB/s]\u001b[A\n",
+      "model.onnx:  34%|███▎      | 147M/436M [00:04<00:07, 37.4MB/s]\u001b[A\n",
+      "model.onnx:  36%|███▌      | 157M/436M [00:04<00:07, 37.4MB/s]\u001b[A\n",
+      "model.onnx:  38%|███▊      | 168M/436M [00:05<00:07, 37.5MB/s]\u001b[A\n",
+      "model.onnx:  41%|████      | 178M/436M [00:05<00:06, 37.6MB/s]\u001b[A\n",
+      "model.onnx:  43%|████▎     | 189M/436M [00:05<00:06, 37.6MB/s]\u001b[A\n",
+      "model.onnx:  46%|████▌     | 199M/436M [00:05<00:06, 37.6MB/s]\u001b[A\n",
+      "model.onnx:  48%|████▊     | 210M/436M [00:06<00:06, 37.5MB/s]\u001b[A\n",
+      "model.onnx:  50%|█████     | 220M/436M [00:06<00:05, 37.5MB/s]\u001b[A\n",
+      "model.onnx:  53%|█████▎    | 231M/436M [00:06<00:05, 37.6MB/s]\u001b[A\n",
+      "model.onnx:  55%|█████▌    | 241M/436M [00:06<00:05, 37.6MB/s]\u001b[A\n",
+      "model.onnx:  58%|█████▊    | 252M/436M [00:07<00:04, 37.6MB/s]\u001b[A\n",
+      "model.onnx:  60%|██████    | 262M/436M [00:07<00:04, 37.7MB/s]\u001b[A\n",
+      "model.onnx:  63%|██████▎   | 273M/436M [00:07<00:04, 37.7MB/s]\u001b[A\n",
+      "model.onnx:  65%|██████▍   | 283M/436M [00:08<00:04, 36.0MB/s]\u001b[A\n",
+      "model.onnx:  67%|██████▋   | 294M/436M [00:08<00:03, 36.4MB/s]\u001b[A\n",
+      "model.onnx:  70%|██████▉   | 304M/436M [00:08<00:03, 36.8MB/s]\u001b[A\n",
+      "model.onnx:  72%|███████▏  | 315M/436M [00:08<00:03, 37.0MB/s]\u001b[A\n",
+      "model.onnx:  75%|███████▍  | 325M/436M [00:09<00:02, 37.3MB/s]\u001b[A\n",
+      "model.onnx:  77%|███████▋  | 336M/436M [00:09<00:03, 30.8MB/s]\u001b[A\n",
+      "model.onnx:  79%|███████▉  | 346M/436M [00:10<00:02, 32.6MB/s]\u001b[A\n",
+      "model.onnx:  82%|████████▏ | 357M/436M [00:10<00:02, 33.9MB/s]\u001b[A\n",
+      "model.onnx:  84%|████████▍ | 367M/436M [00:10<00:01, 34.8MB/s]\u001b[A\n",
+      "model.onnx:  87%|████████▋ | 377M/436M [00:10<00:01, 35.7MB/s]\u001b[A\n",
+      "model.onnx:  89%|████████▉ | 388M/436M [00:11<00:01, 36.2MB/s]\u001b[A\n",
+      "model.onnx:  91%|█████████▏| 398M/436M [00:11<00:01, 36.6MB/s]\u001b[A\n",
+      "model.onnx:  94%|█████████▍| 409M/436M [00:11<00:00, 36.9MB/s]\u001b[A\n",
+      "model.onnx:  96%|█████████▌| 419M/436M [00:11<00:00, 37.1MB/s]\u001b[A\n",
+      "model.onnx:  99%|█████████▊| 430M/436M [00:12<00:00, 37.3MB/s]\u001b[A\n",
+      "model.onnx: 100%|██████████| 436M/436M [00:12<00:00, 35.1MB/s]\u001b[A\n",
+      "Fetching 5 files: 100%|██████████| 5/5 [00:13<00:00,  2.68s/it]\n"
+     ]
+    }
+   ],
+   "source": [
+    "embedding_model = LateInteractionTextEmbedding(\"colbert-ir/colbertv2.0\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "e560b5fa7d63bea3",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-03T17:39:33.400876Z",
+     "start_time": "2024-06-03T17:39:33.397431Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "documents = [\n",
+    "    \"ColBERT is a late interaction text embedding model, however, there are also other models such as TwinBERT.\",\n",
+    "    \"On the contrary to the late interaction models, the early interaction models contains interaction steps at embedding generation process\",\n",
+    "]\n",
+    "queries = [\n",
+    "    \"Are there any other late interaction text embedding models except ColBERT?\",\n",
+    "    \"What is the difference between late interaction and early interaction text embedding models?\",\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "347ad924a3449743",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "*NOTE*: ColBERT computes query and documents embeddings differently, make sure to use the corresponding methods."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "496fbf51e4eaaae",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-03T17:39:34.379885Z",
+     "start_time": "2024-06-03T17:39:34.316257Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "document_embeddings = list(\n",
+    "    embedding_model.embed(documents)\n",
+    ")  # embed and qury_embed return generators,\n",
+    "# which we need to evaluate by writing them to a list\n",
+    "query_embeddings = list(embedding_model.query_embed(queries))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "50595bb0498f0c7c",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-03T17:39:34.793528Z",
+     "start_time": "2024-06-03T17:39:34.788545Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "((26, 128), (32, 128))"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "document_embeddings[0].shape, query_embeddings[0].shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13e43f2c24a7d5fc",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "Don't worry about query embeddings having the bigger shape in this case. \n",
+    "ColBERT authors recommend to pad queries with [MASK] tokens to 32 tokens.\n",
+    "They also recommends to truncate queries to 32 tokens, however we don't do that in FastEmbed, so you can put some straight into the queries."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bb1a4011effd3699",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "## MaxSim operator"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e9ea4cf82521f2de",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "Qdrant will support ColBERT as of the next version (v1.10), however, at the moment, you can compute embedding similarities manually.   "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "f84392f63d2c6076",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-03T17:39:36.431622Z",
+     "start_time": "2024-06-03T17:39:36.427363Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "\n",
+    "def compute_relevance_scores(\n",
+    "    query_embedding: np.array, document_embeddings: np.array, k: int\n",
+    ") -> list[int]:\n",
+    "    \"\"\"\n",
+    "    Compute relevance scores for top-k documents given a query.\n",
+    "\n",
+    "    :param query_embedding: Numpy array representing the query embedding, shape: [num_query_terms, embedding_dim]\n",
+    "    :param document_embeddings: Numpy array representing embeddings for documents, shape: [num_documents, max_doc_length, embedding_dim]\n",
+    "    :param k: Number of top documents to return\n",
+    "    :return: Indices of the top-k documents based on their relevance scores\n",
+    "    \"\"\"\n",
+    "    # Compute batch dot-product of query_embedding and document_embeddings\n",
+    "    # Resulting shape: [num_documents, num_query_terms, max_doc_length]\n",
+    "    scores = np.matmul(query_embedding, document_embeddings.transpose(0, 2, 1))\n",
+    "\n",
+    "    # Apply max-pooling across document terms (axis=2) to find the max similarity per query term\n",
+    "    # Shape after max-pool: [num_documents, num_query_terms]\n",
+    "    max_scores_per_query_term = np.max(scores, axis=2)\n",
+    "\n",
+    "    # Sum the scores across query terms to get the total score for each document\n",
+    "    # Shape after sum: [num_documents]\n",
+    "    total_scores = np.sum(max_scores_per_query_term, axis=1)\n",
+    "\n",
+    "    # Sort the documents based on their total scores and get the indices of the top-k documents\n",
+    "    sorted_indices = np.argsort(total_scores)[::-1][:k]\n",
+    "\n",
+    "    return sorted_indices"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "c61d07bed7b60e35",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-03T17:39:37.053383Z",
+     "start_time": "2024-06-03T17:39:37.050926Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Sorted document indices: [0 1]\n"
+     ]
+    }
+   ],
+   "source": [
+    "sorted_indices = compute_relevance_scores(\n",
+    "    np.array(query_embeddings[0]), np.array(document_embeddings), k=3\n",
+    ")\n",
+    "print(\"Sorted document indices:\", sorted_indices)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "b24df2569970d9e8",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-03T17:40:52.276846Z",
+     "start_time": "2024-06-03T17:40:52.273789Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Query: Are there any other late interaction text embedding models except ColBERT?\n",
+      "Document: ColBERT is a late interaction text embedding model, however, there are also other models such as TwinBERT.\n",
+      "Document: On the contrary to the late interaction models, the early interaction models contains interaction steps at embedding generation process\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(f\"Query: {queries[0]}\")\n",
+    "for index in sorted_indices:\n",
+    "    print(f\"Document: {documents[index]}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6de537c37aff3927",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "## Use-case recommendation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "37e3525d3259cd2b",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "Despite ColBERT allows to compute embeddings independently and spare some workload offline, it still computes more resources than no interaction models. Due to this, it might be more reasonable to use ColBERT not as a first-stage retriever, but as a re-ranker.\n",
+    "\n",
+    "The first-stage retriever would then be a no-interaction model, which e.g. retrieves first 100 or 500 examples, and leave the final ranking to the ColBERT model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cfa922793454b4ad",
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/examples/FastEmbed_GPU.ipynb

@@ -0,0 +1,540 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "ntGNDuSCeAR2"
+   },
+   "source": [
+    "# FastEmbed on GPU\n",
+    "\n",
+    "As of version 0.2.7 FastEmbed supports GPU acceleration.\n",
+    "\n",
+    "This notebook covers the installation process and usage of fastembed on GPU.\n",
+    "\n",
+    "## Installation\n",
+    "\n",
+    "Fastembed depends on `onnxruntime` and inherits its scheme of GPU support.\n",
+    "\n",
+    "In order to use GPU with onnx models, you would need to have `onnxruntime-gpu` package, which substitutes all the `onnxruntime` functionality.\n",
+    "Fastembed mimics this behavior and requires `fastembed-gpu` package to be installed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "GK2XADwUeEK7"
+   },
+   "outputs": [],
+   "source": [
+    "!pip install fastembed-gpu"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "3aiGPqjCeGzo"
+   },
+   "source": [
+    "**NOTE**: `onnxruntime-gpu` and `onnxruntime` can't be installed in the same environment. If you have `onnxruntime` installed, you would need to uninstall it before installing `onnxruntime-gpu`. Same is true for `fastembed` and `fastembed-gpu`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "3xx3r-9jgAMi"
+   },
+   "source": [
+    "### CUDA 12.x support\n",
+    "You can check your CUDA version using such commands as `nvidia-smi` or `nvcc --version`\n",
+    "\n",
+    "Starting from version 1.19.0, onnxruntime-gpu ships with support for CUDA 12.x by default.\n",
+    "\n",
+    "Google Colab notebooks have by default CUDA 12.x and CuDNN 8.x.\n",
+    "\n",
+    "Latest version of `onnxruntime-gpu` requires CuDNN 9.x, in order to install it you can run the following command: "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!sudo apt install cudnn9\n",
+    "!pip install fastembed-gpu -qqq"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If it necessary to work with CuDNN 8, you can consider locking `onnxruntime-gpu` to 1.18.0 with CUDA 12.x by this command:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install onnxruntime-gpu==1.18.0 -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ -qq\n",
+    "!pip install fastembed-gpu -qqq"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### CUDA 11.x support\n",
+    "To use latest version of `onnxruntime-gpu` with CUDA 11.x, you can run the following command:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install onnxruntime-gpu -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-11/pypi/simple/ -qq"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**NOTE**: Ensure that CuDNN 9.x is installed when working with the latest `onnxruntime-gpu`, whether using CUDA 11.x or 12.x."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Igv5RXhSeO68"
+   },
+   "source": [
+    "### CUDA drivers\n",
+    "\n",
+    "FastEmbed does not include CUDA drivers and CuDNN libraries.\n",
+    "You would need to take care of the environment setup on your own.\n",
+    "The dependencies required for the chosen onnxruntime version are listed in the [CUDA Execution Provider requirements](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Setting up fastembed-gpu on GCP\n",
+    "\n",
+    "#### CUDA drivers\n",
+    "[CUDA 11.8 toolkit](https://developer.nvidia.com/cuda-11-8-0-download-archive) or [CUDA 12.x toolkit](https://developer.nvidia.com/cuda-downloads) has to be installed if they haven't yet been set up.\n",
+    "\n",
+    "#### Example of setting up CUDA 12.x on Ubuntu 22.04\n",
+    "Make sure to download an archive which has been created for your particular platform, CPU architecture and OS distribution.\n",
+    "\n",
+    "For Ubuntu 22.04 with x86_64 CPU architecture the following [archive](https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=22.04&target_type=deb_network) has to be downloaded.\n",
+    "\n",
+    "```bash\n",
+    "wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb\n",
+    "sudo dpkg -i cuda-keyring_1.1-1_all.deb\n",
+    "sudo apt-get update\n",
+    "sudo apt-get -y install cuda\n",
+    "```\n",
+    "**NOTE**: Specific CUDA libraries can be found in the [meta packages section](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/#meta-packages) in the CUDA installation guide.\n",
+    "\n",
+    "**NOTE**: When installing CUDA, the environment variable might not be set by default. Make sure to add the following line to your environment variables:\n",
+    "```bash\n",
+    "LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH\n",
+    "```\n",
+    "This will ensure that the CUDA libraries are properly linked.\n",
+    "\n",
+    "#### CuDNN 9.x\n",
+    "CuDNN 9.x library can be installed via the following [archive](https://developer.nvidia.com/rdp/cudnn-archive).\n",
+    "\n",
+    "#### Example of setting up CuDNN 9.x on Ubuntu 22.04\n",
+    "CuDNN 9.x for Ubuntu 22.04 x86_64 [archive](https://developer.nvidia.com/cudnn-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=22.04&target_type=deb_network) can be downloaded and installed in the following way:\n",
+    "```bash\n",
+    "wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb\n",
+    "sudo dpkg -i cuda-keyring_1.1-1_all.deb\n",
+    "sudo apt-get update\n",
+    "sudo apt-get -y install cudnn\n",
+    "```\n",
+    "**NOTE**: When installing CuDNN, you can choose specific version, cudnn-cuda-11 or cudnn-cuda-12"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Common issues\n",
+    "\n",
+    "The following are some common issues that may arise while using `fastembed-gpu` if not installed properly:\n",
+    "\n",
+    "CUDA library is not installed:\n",
+    "```bash\n",
+    "FAIL : Failed to load library libonnxruntime_providers_cuda.so with error: libcublasLt.so.x: cannot open shared object file: No such file or directory\n",
+    "```\n",
+    "\n",
+    "\n",
+    "CuDNN library is not installed:\n",
+    "```bash\n",
+    "FAIL : Failed to load library libonnxruntime_providers_cuda.so with error: libcudnn.so.x: cannot open shared object file: No such file or directory\n",
+    "```\n",
+    "\n",
+    "\n",
+    "CUDA library path is not set:\n",
+    "```bash\n",
+    "FAIL : Failed to load library libonnxruntime_providers_cuda.so with error: libcufft.so.x: failed to map segment from shared object\n",
+    "```\n",
+    "\n",
+    "Make sure to add the following line to your environment variables:\n",
+    "```bash\n",
+    "LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Usage"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 334,
+     "referenced_widgets": [
+      "aacf08a7aa444b64a2efad1967d28a53",
+      "5606aa785de74d65a9928b31c0be8a53",
+      "d4ec9d3b74ec4412894da2161ed2bddf",
+      "8edd544c3e074ec1813e5b9d1aef43d9",
+      "9898890f8a75468ea20e3ce319d0b6e2",
+      "da3b18abb16241a0a7191ee9afcb0510",
+      "258a619168824253a6a329efdc51ebe6",
+      "53c7cdc967d24faba0b5c659c94c50b8",
+      "e9348d8be28d408e8e760c71b21ab294",
+      "0ba06e0816714f2fbdec8260f160abc0",
+      "0b96563334964d449dd34f35b6b3e715",
+      "11c2eec490e8479b944eec7f30cb1ca2",
+      "91463da0d1c5466795e06ab586002259",
+      "30f4f7833406474f89ef0700b00a33aa",
+      "4302c304ec6a4b5985797e300bd7e353",
+      "2605640c7b824ed7aa137d404e14b774",
+      "b02efe3a33d04f06aa8938719ab35671",
+      "50408e5d052343b1a1b44a0fae0f801d",
+      "1be01c95d9e84f8ea88367c987a72fdc",
+      "a109c13bc93a449186424542dc330be8",
+      "4adce304ce1947b5a01dde10bbb3bb8c",
+      "a761366a37e44837a25e0f25b18efed2",
+      "94512b9055e546389471197b76ad5449",
+      "072dca00bd7b4918a178f90ccabf698a",
+      "48a856c59ef74cc3834521b1bf616541",
+      "c020c503aeaa464cad643ade5ee3ae24",
+      "a4e7e40c0bbd4f878c20a9f65fe3a048",
+      "e8c0a1c339fd47668d944a9defad79d4",
+      "cd782d35c6bd40c0a60d57b1828a7251",
+      "04f638ab08da4d20928644c4ba03f8ef",
+      "17f20477fc79475f97adf1c1f64a4192",
+      "96f7b5a2e224462e9fcffd03f906a593",
+      "755cd32d9fc9407c80a160f45c802d1e",
+      "a886258e7cd14c048b58391d7b772901",
+      "bc3e48f826a74840867a6209e622b75e",
+      "125b2ac0f78043bba7eca53474ca44c4",
+      "82f186d1ffb4435d94a6c7e9025242ef",
+      "77000333e5ca4094be291ad82d4a627a",
+      "7fe64fb53055431488d002c76c8e331e",
+      "7de59ae9919f4a5bb2b6e601a3c02412",
+      "97a69423a6644eab87fc636e182f23a4",
+      "4df936d1065b41f4bf02ed394fdf7b7e",
+      "3918bd1affa3454e8e9044a418a056ea",
+      "163b27ae0bce41e5b48efcb4b3fd780d",
+      "94631fd6e0744085bc79c3121de4a9f7",
+      "31cd98d66bc54418b35e70fbbc0fa3c0",
+      "6d21627a638b4ddca6fe7bfb80a621b5",
+      "b37bed9dc4fe45c08b8397288fe5b1a9",
+      "164fef95d1414177a40d563f5682f6a3",
+      "1a9a0ea53448413a8e4b360b7bb69e26",
+      "dd1a4483b4b045c6929e3d2cf1338f63",
+      "496ddd8e05f949cd8cbba8e677f476ac",
+      "2813be951d7f48b2aad1dd4a444ce3eb",
+      "8e9a2c2dd21942edbdfecb3b7dffc70b",
+      "08a10fe247f1425db044cfc13f2fb384",
+      "b8786aded92d421592bc7623c5c7899e",
+      "c91a20a9433d4016ba2db69fa50e0b4d",
+      "e997820738594c6dadb061908d7afdc1",
+      "a5fc751f81ae498f9aa55ece0e6853b2",
+      "2aee4fc8cda64c5eb8722be81e48e0ca",
+      "3a53e8624dff48b3959875ef58ee99ce",
+      "50a70044f77542108fe188598e70797e",
+      "13cf998b35ae4507a63e797f6fa3eada",
+      "6209eb6a68cf4a378767ef34d0d9216d",
+      "7395db766b944af9b41d6b56c9ada0b1",
+      "42122c317ec648688f0164a1adb5df28"
+     ]
+    },
+    "id": "Ttf4YggPeQQK",
+    "outputId": "aa75129d-9e2d-4c88-cf03-251dd43a11b1"
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:88: UserWarning: \n",
+      "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
+      "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
+      "You will be able to reuse this secret in all of your notebooks.\n",
+      "Please note that authentication is recommended but still optional to access public models or datasets.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "aacf08a7aa444b64a2efad1967d28a53",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Fetching 5 files:   0%|          | 0/5 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "11c2eec490e8479b944eec7f30cb1ca2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "tokenizer_config.json:   0%|          | 0.00/1.24k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "94512b9055e546389471197b76ad5449",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "config.json:   0%|          | 0.00/706 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a886258e7cd14c048b58391d7b772901",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "special_tokens_map.json:   0%|          | 0.00/695 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "94631fd6e0744085bc79c3121de4a9f7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "tokenizer.json:   0%|          | 0.00/711k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b8786aded92d421592bc7623c5c7899e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model_optimized.onnx:   0%|          | 0.00/66.5M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "['CUDAExecutionProvider', 'CPUExecutionProvider']"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "from fastembed import TextEmbedding\n",
+    "\n",
+    "embedding_model_gpu = TextEmbedding(\n",
+    "    model_name=\"BAAI/bge-small-en-v1.5\", providers=[\"CUDAExecutionProvider\"]\n",
+    ")\n",
+    "embedding_model_gpu.model.model.get_providers()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "id": "iPtoHf7GeV-i"
+   },
+   "outputs": [],
+   "source": "documents: list[str] = list(np.repeat(\"Demonstrating GPU acceleration in fastembed\", 500))"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "islhyLf4ed-H",
+    "outputId": "8c8ed09b-9eac-438f-97bc-578751975148"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "43.4 ms ± 2.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%timeit\n",
+    "list(embedding_model_gpu.embed(documents))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 67,
+     "referenced_widgets": [
+      "9c306ce5188c45feb8dfb9089592591c",
+      "296ff54c6e61441f978084df59626598",
+      "d6d42b4f245a49b7ba7769e23a3202fc",
+      "39ce7754480147759c16a3089d8105af",
+      "8253960a069d4106863a75faae54b90d",
+      "7ccf959452af4c0b873c7567747f0816",
+      "ac9d0b5a5b1f401e90a1cc9ffe6d4b4c",
+      "0aada067dec3472f9aba1772d6b775a5",
+      "07597b1287e04653b80c47a771549376",
+      "054be1dd9f084cae911745b692ccd929",
+      "ab19e8e831694e308a4b79f05aff728e"
+     ]
+    },
+    "id": "bOKVUvWJegYJ",
+    "outputId": "dde74917-08b0-4ce2-9a2b-cc31e02cafb2"
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9c306ce5188c45feb8dfb9089592591c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Fetching 5 files:   0%|          | 0/5 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "['CPUExecutionProvider']"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "embedding_model_cpu = TextEmbedding(model_name=\"BAAI/bge-small-en-v1.5\")\n",
+    "embedding_model_cpu.model.model.get_providers()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "0NJj9RvSfASP",
+    "outputId": "526f5280-99bd-454e-8af8-6a860ad96e54"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4.33 s ± 591 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%timeit\n",
+    "list(embedding_model_cpu.embed(documents))"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "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.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 1
+}

docs/examples/FastEmbed_Multi_GPU.ipynb

@@ -0,0 +1,88 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Fastembed Multi-GPU Tutorial\n",
+    "This tutorial demonstrates how to leverage multi-GPU support in Fastembed. Fastembed supports embedding text and images utilizing modern GPUs for acceleration. Let's explore how to use Fastembed with multiple GPUs step by step."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Prerequisites\n",
+    "To get started, ensure you have the following installed:\n",
+    "- Python 3.9 or later\n",
+    "- Fastembed (`pip install fastembed-gpu`)\n",
+    "- Refer to [this](https://github.com/qdrant/fastembed/blob/main/docs/examples/FastEmbed_GPU.ipynb) tutorial if you have issues with GPU dependencies\n",
+    "- Access to a multi-GPU server"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Multi-GPU using cuda argument with TextEmbedding Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from fastembed import TextEmbedding\n",
+    "\n",
+    "# define the documents to embed\n",
+    "docs = [\"hello world\", \"flag embedding\"] * 100\n",
+    "\n",
+    "# define gpu ids\n",
+    "device_ids = [0, 1]\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    # initialize a TextEmbedding model using CUDA\n",
+    "    text_model = TextEmbedding(\n",
+    "        model_name=\"sentence-transformers/all-MiniLM-L6-v2\",\n",
+    "        cuda=True,\n",
+    "        device_ids=device_ids,\n",
+    "        lazy_load=True,\n",
+    "    )\n",
+    "\n",
+    "    # generate embeddings\n",
+    "    text_embeddings = list(text_model.embed(docs, batch_size=2, parallel=len(device_ids)))\n",
+    "    print(text_embeddings)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In this snippet:\n",
+    "- `cuda=True` enables GPU acceleration.\n",
+    "- `device_ids=[0, 1]` specifies GPUs to use. Replace `[0, 1]` with available GPU IDs.\n",
+    "- `lazy_load=True`\n",
+    "\n",
+    "**NOTE**: When using multi-GPU settings, it is important to configure `parallel` and `lazy_load` properly to avoid inefficiencies:\n",
+    "\n",
+    "`parallel`: This parameter enables multi-GPU support by spawning child processes for each GPU specified in device_ids. To ensure proper utilization, the value of `parallel` must match the number of GPUs in device_ids. If using a single GPU, this parameter is not necessary.\n",
+    "\n",
+    "`lazy_load`: Enabling `lazy_load` prevents redundant memory usage. Without `lazy_load`, the model is initially loaded into the memory of the first GPU by the main process. When child processes are spawned for each GPU, the model is reloaded on the first GPU, causing redundant memory consumption and inefficiencies."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/examples/Image_Embedding.ipynb

@@ -0,0 +1,128 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "aa0a86859809102",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "# Image Embedding\n",
+    "As of version 0.3.0 fastembed supports computation of image embeddings.\n",
+    "\n",
+    "The process is as easy and straightforward as with text embeddings. Let's see how it works."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "cea8fd5c019571fe",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-02T11:35:40.126023Z",
+     "start_time": "2024-06-02T11:35:39.864701Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Fetching 3 files: 100%|██████████| 3/3 [00:00<00:00, 47482.69it/s]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": "[array([0.        , 0.        , 0.        , ..., 0.        , 0.01139933,\n        0.        ], dtype=float32),\n array([0.02169187, 0.        , 0.        , ..., 0.        , 0.00848291,\n        0.        ], dtype=float32)]"
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from fastembed import ImageEmbedding\n",
+    "\n",
+    "model = ImageEmbedding(\"Qdrant/resnet50-onnx\")\n",
+    "\n",
+    "embeddings_generator = model.embed(\n",
+    "    [\"../../tests/misc/image.jpeg\", \"../../tests/misc/small_image.jpeg\"]\n",
+    ")\n",
+    "embeddings_list = list(embeddings_generator)\n",
+    "embeddings_list"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3f838f18523ad1e0",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "## Preprocessing\n",
+    "\n",
+    "Preprocessing is encapsulated in the ImageEmbedding class, applied operations are identical to the ones provided by [Hugging Face Transformers](https://huggingface.co/docs/transformers/en/index).\n",
+    "You don't need to think about batching, opening/closing files, resizing images, etc., Fastembed will take care of it."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "894b33ff9b385d72",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "## Supported models\n",
+    "\n",
+    "List of supported image embedding models can either be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-image-embedding-models) or by calling the `ImageEmbedding.list_supported_models()` method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "6d6a4cbbd2200d14",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-02T11:40:19.313226Z",
+     "start_time": "2024-06-02T11:40:19.309845Z"
+    },
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "[{'model': 'Qdrant/clip-ViT-B-32-vision',\n  'dim': 512,\n  'description': 'CLIP vision encoder based on ViT-B/32',\n  'size_in_GB': 0.34,\n  'sources': {'hf': 'Qdrant/clip-ViT-B-32-vision'},\n  'model_file': 'model.onnx'},\n {'model': 'Qdrant/resnet50-onnx',\n  'dim': 2048,\n  'description': 'ResNet-50 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`__.',\n  'size_in_GB': 0.1,\n  'sources': {'hf': 'Qdrant/resnet50-onnx'},\n  'model_file': 'model.onnx'}]"
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ImageEmbedding.list_supported_models()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/examples/SPLADE_with_FastEmbed.ipynb

@@ -0,0 +1,389 @@
+{
+ "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": null,
+   "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"
+   ]
+  },
+  {
+   "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": null,
+   "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(\n",
+    "        f\"Token at index {sparse_embeddings_list[0].indices[i]} has weight {sparse_embeddings_list[0].values[i]}\"\n",
+    "    )"
+   ]
+  },
+  {
+   "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(\n",
+    "    SparseTextEmbedding.list_supported_models()[0][\"sources\"][\"hf\"]\n",
+    ")"
+   ]
+  },
+  {
+   "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(\n",
+    "        sorted(token_weight_dict.items(), key=lambda item: item[1], reverse=True)\n",
+    "    )\n",
+    "    return token_weight_dict\n",
+    "\n",
+    "\n",
+    "# Test the function with the first SparseEmbedding\n",
+    "print(json.dumps(get_tokens_and_weights(sparse_embeddings_list[index], tokenizer), indent=4))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Observations and Model Design Choices\n",
+    "\n",
+    "1. The relative order of importance is quite useful. The most important tokens in the sentence have the highest weights.\n",
+    "1. **Term Expansion**: The model can expand the terms in the document. This means that the model can generate weights for tokens that are not present in the document but are related to the tokens in the document. This is a powerful feature that allows the model to capture the context of the document. Here, you'll see that the model has added the tokens '3' from 'third' and 'moon' from 'lunar' to the sparse vector.\n",
+    "\n",
+    "### Design Choices\n",
+    "\n",
+    "1. The weights are not normalized. This means that the sum of the weights is not 1 or 100. This is a common practice in sparse embeddings, as it allows the model to capture the importance of each token in the document.\n",
+    "1. Tokens are included in the sparse vector only if they are present in the model's vocabulary. This means that the model will not generate a weight for tokens that it has not seen during training.\n",
+    "1. Tokens do not map to words directly -- allowing you to gracefully handle typo errors and out-of-vocabulary tokens."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "fst",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/examples/Supported_Models.ipynb

@@ -2,11 +2,132 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-13T09:01:03.324551Z",
+     "start_time": "2024-11-13T09:01:03.234711Z"
+    }
+   },
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The autoreload extension is already loaded. To reload it, use:\n",
+      "  %reload_ext autoreload\n"
+     ]
+    }
+   ],
+   "execution_count": 10
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-13T09:01:04.505772Z",
+     "start_time": "2024-11-13T09:01:04.493296Z"
+    }
+   },
+   "source": [
+    "import pandas as pd\n",
+    "\n",
+    "from fastembed import (\n",
+    "    SparseTextEmbedding,\n",
+    "    TextEmbedding,\n",
+    "    LateInteractionTextEmbedding,\n",
+    "    ImageEmbedding,\n",
+    ")\n",
+    "from fastembed.rerank.cross_encoder import TextCrossEncoder"
+   ],
+   "outputs": [],
+   "execution_count": 11
+  },
+  {
+   "cell_type": "markdown",
    "metadata": {},
+   "source": [
+    "## Supported Text Embedding Models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-13T09:01:05.812271Z",
+     "start_time": "2024-11-13T09:01:05.795846Z"
+    }
+   },
+   "source": [
+    "supported_models = (\n",
+    "    pd.DataFrame(TextEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")\n",
+    "supported_models"
+   ],
    "outputs": [
     {
      "data": {
+      "text/plain": [
+       "                                                model   dim  \\\n",
+       "0                              BAAI/bge-small-en-v1.5   384   \n",
+       "1                              BAAI/bge-small-zh-v1.5   512   \n",
+       "2                 snowflake/snowflake-arctic-embed-xs   384   \n",
+       "3              sentence-transformers/all-MiniLM-L6-v2   384   \n",
+       "4                  jinaai/jina-embeddings-v2-small-en   512   \n",
+       "5                                   BAAI/bge-small-en   384   \n",
+       "6                  snowflake/snowflake-arctic-embed-s   384   \n",
+       "7                    nomic-ai/nomic-embed-text-v1.5-Q   768   \n",
+       "8                               BAAI/bge-base-en-v1.5   768   \n",
+       "9   sentence-transformers/paraphrase-multilingual-...   384   \n",
+       "10                          Qdrant/clip-ViT-B-32-text   512   \n",
+       "11                  jinaai/jina-embeddings-v2-base-de   768   \n",
+       "12                                   BAAI/bge-base-en   768   \n",
+       "13                 snowflake/snowflake-arctic-embed-m   768   \n",
+       "14                     nomic-ai/nomic-embed-text-v1.5   768   \n",
+       "15                  jinaai/jina-embeddings-v2-base-en   768   \n",
+       "16                       nomic-ai/nomic-embed-text-v1   768   \n",
+       "17            snowflake/snowflake-arctic-embed-m-long   768   \n",
+       "18                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
+       "19                jinaai/jina-embeddings-v2-base-code   768   \n",
+       "20  sentence-transformers/paraphrase-multilingual-...   768   \n",
+       "21                 snowflake/snowflake-arctic-embed-l  1024   \n",
+       "22                                 thenlper/gte-large  1024   \n",
+       "23                             BAAI/bge-large-en-v1.5  1024   \n",
+       "24                     intfloat/multilingual-e5-large  1024   \n",
+       "\n",
+       "                                          description     license  size_in_GB  \n",
+       "0   Text embeddings, Unimodal (text), English, 512...         mit       0.067  \n",
+       "1   Text embeddings, Unimodal (text), Chinese, 512...         mit       0.090  \n",
+       "2   Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.090  \n",
+       "3   Text embeddings, Unimodal (text), English, 256...  apache-2.0       0.090  \n",
+       "4   Text embeddings, Unimodal (text), English, 819...  apache-2.0       0.120  \n",
+       "5   Text embeddings, Unimodal (text), English, 512...         mit       0.130  \n",
+       "6   Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.130  \n",
+       "7   Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.130  \n",
+       "8   Text embeddings, Unimodal (text), English, 512...         mit       0.210  \n",
+       "9   Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.220  \n",
+       "10  Text embeddings, Multimodal (text&image), Engl...         mit       0.250  \n",
+       "11  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.320  \n",
+       "12  Text embeddings, Unimodal (text), English, 512...         mit       0.420  \n",
+       "13  Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.430  \n",
+       "14  Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.520  \n",
+       "15  Text embeddings, Unimodal (text), English, 819...  apache-2.0       0.520  \n",
+       "16  Text embeddings, Multimodal (text, image), Eng...  apache-2.0       0.520  \n",
+       "17  Text embeddings, Unimodal (text), English, 204...  apache-2.0       0.540  \n",
+       "18  Text embeddings, Unimodal (text), English, 512...  apache-2.0       0.640  \n",
+       "19  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       0.640  \n",
+       "20  Text embeddings, Unimodal (text), Multilingual...  apache-2.0       1.000  \n",
+       "21  Text embeddings, Unimodal (text), English, 512...  apache-2.0       1.020  \n",
+       "22  Text embeddings, Unimodal (text), English, 512...         mit       1.200  \n",
+       "23  Text embeddings, Unimodal (text), English, 512...         mit       1.200  \n",
+       "24  Text embeddings, Unimodal (text), Multilingual...         mit       2.240  "
+      ],
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
@@ -29,70 +150,703 @@
        "      <th>model</th>\n",
        "      <th>dim</th>\n",
        "      <th>description</th>\n",
+       "      <th>license</th>\n",
+       "      <th>size_in_GB</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
-       "      <td>BAAI/bge-small-en</td>\n",
+       "      <td>BAAI/bge-small-en-v1.5</td>\n",
        "      <td>384</td>\n",
-       "      <td>Fast and Default English model</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>0.067</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
-       "      <td>BAAI/bge-base-en</td>\n",
-       "      <td>768</td>\n",
-       "      <td>Base English model</td>\n",
+       "      <td>BAAI/bge-small-zh-v1.5</td>\n",
+       "      <td>512</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Chinese, 512...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>0.090</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
-       "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
+       "      <td>snowflake/snowflake-arctic-embed-xs</td>\n",
        "      <td>384</td>\n",
-       "      <td>Sentence Transformer model, MiniLM-L6-v2</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.090</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
+       "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 256...</td>\n",
+       "      <td>apache-2.0</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>Text embeddings, Unimodal (text), English, 819...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.120</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>BAAI/bge-small-en</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>0.130</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-s</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.130</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1.5-Q</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Multimodal (text, image), Eng...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.130</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>8</th>\n",
+       "      <td>BAAI/bge-base-en-v1.5</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>0.210</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>9</th>\n",
+       "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.220</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>10</th>\n",
+       "      <td>Qdrant/clip-ViT-B-32-text</td>\n",
+       "      <td>512</td>\n",
+       "      <td>Text embeddings, Multimodal (text&amp;image), Engl...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>0.250</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>11</th>\n",
+       "      <td>jinaai/jina-embeddings-v2-base-de</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.320</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>12</th>\n",
+       "      <td>BAAI/bge-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>0.420</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>13</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-m</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.430</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>14</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Multimodal (text, image), Eng...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>15</th>\n",
+       "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 819...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>16</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Multimodal (text, image), Eng...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>17</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-m-long</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 204...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.540</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>18</th>\n",
+       "      <td>mixedbread-ai/mxbai-embed-large-v1</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.640</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>19</th>\n",
+       "      <td>jinaai/jina-embeddings-v2-base-code</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.640</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>20</th>\n",
+       "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>1.000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>21</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-l</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>1.020</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>22</th>\n",
+       "      <td>thenlper/gte-large</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>1.200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>23</th>\n",
+       "      <td>BAAI/bge-large-en-v1.5</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Text embeddings, Unimodal (text), English, 512...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>1.200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>24</th>\n",
        "      <td>intfloat/multilingual-e5-large</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Multilingual model, e5-large. Recommend using this model for non-English languages. Recommend using this via Torch implementation of FastEmbed</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>2.240</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
-      ],
-      "text/plain": [
-       "                                    model   dim  \\\n",
-       "0                       BAAI/bge-small-en   384   \n",
-       "1                        BAAI/bge-base-en   768   \n",
-       "2  sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "3          intfloat/multilingual-e5-large  1024   \n",
-       "\n",
-       "                                                                                                                                      description  \n",
-       "0                                                                                                                  Fast and Default English model  \n",
-       "1                                                                                                                              Base English model  \n",
-       "2                                                                                                        Sentence Transformer model, MiniLM-L6-v2  \n",
-       "3  Multilingual model, e5-large. Recommend using this model for non-English languages. Recommend using this via Torch implementation of FastEmbed  "
       ]
      },
-     "execution_count": 1,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
+   "execution_count": 12
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
    "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 Sparse Text Embedding Models"
    ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-13T09:01:07.038954Z",
+     "start_time": "2024-11-13T09:01:07.019656Z"
+    }
+   },
+   "source": [
+    "(\n",
+    "    pd.DataFrame(SparseTextEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\", \"additional_files\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")"
+   ],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "                                     model  vocab_size  \\\n",
+       "0                              Qdrant/bm25         NaN   \n",
+       "1  Qdrant/bm42-all-minilm-l6-v2-attentions     30522.0   \n",
+       "2               prithivida/Splade_PP_en_v1     30522.0   \n",
+       "3                prithvida/Splade_PP_en_v1     30522.0   \n",
+       "\n",
+       "                                         description     license  size_in_GB  \\\n",
+       "0  BM25 as sparse embeddings meant to be used wit...  apache-2.0       0.010   \n",
+       "1  Light sparse embedding model, which assigns an...  apache-2.0       0.090   \n",
+       "2  Independent Implementation of SPLADE++ Model f...  apache-2.0       0.532   \n",
+       "3  Independent Implementation of SPLADE++ Model f...  apache-2.0       0.532   \n",
+       "\n",
+       "  requires_idf  \n",
+       "0         True  \n",
+       "1         True  \n",
+       "2          NaN  \n",
+       "3          NaN  "
+      ],
+      "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>license</th>\n",
+       "      <th>size_in_GB</th>\n",
+       "      <th>requires_idf</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Qdrant/bm25</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>BM25 as sparse embeddings meant to be used wit...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.010</td>\n",
+       "      <td>True</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Qdrant/bm42-all-minilm-l6-v2-attentions</td>\n",
+       "      <td>30522.0</td>\n",
+       "      <td>Light sparse embedding model, which assigns an...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.090</td>\n",
+       "      <td>True</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>prithivida/Splade_PP_en_v1</td>\n",
+       "      <td>30522.0</td>\n",
+       "      <td>Independent Implementation of SPLADE++ Model f...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.532</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>prithvida/Splade_PP_en_v1</td>\n",
+       "      <td>30522.0</td>\n",
+       "      <td>Independent Implementation of SPLADE++ Model f...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.532</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 13
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "## Supported Late Interaction Text Embedding Models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-11-13T09:01:08.074442Z",
+     "start_time": "2024-11-13T09:01:08.056138Z"
+    }
+   },
+   "source": [
+    "(\n",
+    "    pd.DataFrame(LateInteractionTextEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")"
+   ],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "                                   model  dim  \\\n",
+       "0  answerdotai/answerai-colbert-small-v1   96   \n",
+       "1                 colbert-ir/colbertv2.0  128   \n",
+       "2                 jinaai/jina-colbert-v2  128   \n",
+       "\n",
+       "                                         description       license  \\\n",
+       "0  Text embeddings, Unimodal (text), Multilingual...    apache-2.0   \n",
+       "1                             Late interaction model           mit   \n",
+       "2  New model that expands capabilities of colbert...  cc-by-nc-4.0   \n",
+       "\n",
+       "   size_in_GB        additional_files  \n",
+       "0        0.13                     NaN  \n",
+       "1        0.44                     NaN  \n",
+       "2        2.24  [onnx/model.onnx_data]  "
+      ],
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>model</th>\n",
+       "      <th>dim</th>\n",
+       "      <th>description</th>\n",
+       "      <th>license</th>\n",
+       "      <th>size_in_GB</th>\n",
+       "      <th>additional_files</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>answerdotai/answerai-colbert-small-v1</td>\n",
+       "      <td>96</td>\n",
+       "      <td>Text embeddings, Unimodal (text), Multilingual...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.13</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>colbert-ir/colbertv2.0</td>\n",
+       "      <td>128</td>\n",
+       "      <td>Late interaction model</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>0.44</td>\n",
+       "      <td>NaN</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>jinaai/jina-colbert-v2</td>\n",
+       "      <td>128</td>\n",
+       "      <td>New model that expands capabilities of colbert...</td>\n",
+       "      <td>cc-by-nc-4.0</td>\n",
+       "      <td>2.24</td>\n",
+       "      <td>[onnx/model.onnx_data]</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 14
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "## Supported Image Embedding Models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-11-13T09:01:09.171647Z",
+     "start_time": "2024-11-13T09:01:09.150940Z"
+    }
+   },
+   "source": [
+    "(\n",
+    "    pd.DataFrame(ImageEmbedding.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")"
+   ],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "                         model   dim  \\\n",
+       "0         Qdrant/resnet50-onnx  2048   \n",
+       "1  Qdrant/clip-ViT-B-32-vision   512   \n",
+       "2       Qdrant/Unicom-ViT-B-32   512   \n",
+       "3       Qdrant/Unicom-ViT-B-16   768   \n",
+       "\n",
+       "                                         description     license  size_in_GB  \n",
+       "0      Image embeddings, Unimodal (image), 2016 year  apache-2.0        0.10  \n",
+       "1  Image embeddings, Multimodal (text&image), 202...         mit        0.34  \n",
+       "2  Image embeddings, Multimodal (text&image), 202...  apache-2.0        0.48  \n",
+       "3  Image embeddings (more detailed than Unicom-Vi...  apache-2.0        0.82  "
+      ],
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>model</th>\n",
+       "      <th>dim</th>\n",
+       "      <th>description</th>\n",
+       "      <th>license</th>\n",
+       "      <th>size_in_GB</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Qdrant/resnet50-onnx</td>\n",
+       "      <td>2048</td>\n",
+       "      <td>Image embeddings, Unimodal (image), 2016 year</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.10</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Qdrant/clip-ViT-B-32-vision</td>\n",
+       "      <td>512</td>\n",
+       "      <td>Image embeddings, Multimodal (text&amp;image), 202...</td>\n",
+       "      <td>mit</td>\n",
+       "      <td>0.34</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>Qdrant/Unicom-ViT-B-32</td>\n",
+       "      <td>512</td>\n",
+       "      <td>Image embeddings, Multimodal (text&amp;image), 202...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.48</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>Qdrant/Unicom-ViT-B-16</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Image embeddings (more detailed than Unicom-Vi...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "      <td>0.82</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 15
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Supported Rerank Cross Encoder Models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-13T09:01:10.313943Z",
+     "start_time": "2024-11-13T09:01:10.298428Z"
+    }
+   },
+   "source": [
+    "(\n",
+    "    pd.DataFrame(TextCrossEncoder.list_supported_models())\n",
+    "    .sort_values(\"size_in_GB\")\n",
+    "    .drop(columns=[\"sources\", \"model_file\"])\n",
+    "    .reset_index(drop=True)\n",
+    ")"
+   ],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "                                       model  size_in_GB  \\\n",
+       "0              Xenova/ms-marco-MiniLM-L-6-v2        0.08   \n",
+       "1             Xenova/ms-marco-MiniLM-L-12-v2        0.12   \n",
+       "2            jinaai/jina-reranker-v1-tiny-en        0.13   \n",
+       "3           jinaai/jina-reranker-v1-turbo-en        0.15   \n",
+       "4                     BAAI/bge-reranker-base        1.04   \n",
+       "5  jinaai/jina-reranker-v2-base-multilingual        1.11   \n",
+       "\n",
+       "                                         description       license  \n",
+       "0  MiniLM-L-6-v2 model optimized for re-ranking t...    apache-2.0  \n",
+       "1  MiniLM-L-12-v2 model optimized for re-ranking ...    apache-2.0  \n",
+       "2  Designed for blazing-fast re-ranking with 8K c...    apache-2.0  \n",
+       "3  Designed for blazing-fast re-ranking with 8K c...    apache-2.0  \n",
+       "4  BGE reranker base model for cross-encoder re-r...           mit  \n",
+       "5  A multi-lingual reranker model for cross-encod...  cc-by-nc-4.0  "
+      ],
+      "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>size_in_GB</th>\n",
+       "      <th>description</th>\n",
+       "      <th>license</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Xenova/ms-marco-MiniLM-L-6-v2</td>\n",
+       "      <td>0.08</td>\n",
+       "      <td>MiniLM-L-6-v2 model optimized for re-ranking t...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Xenova/ms-marco-MiniLM-L-12-v2</td>\n",
+       "      <td>0.12</td>\n",
+       "      <td>MiniLM-L-12-v2 model optimized for re-ranking ...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>jinaai/jina-reranker-v1-tiny-en</td>\n",
+       "      <td>0.13</td>\n",
+       "      <td>Designed for blazing-fast re-ranking with 8K c...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>jinaai/jina-reranker-v1-turbo-en</td>\n",
+       "      <td>0.15</td>\n",
+       "      <td>Designed for blazing-fast re-ranking with 8K c...</td>\n",
+       "      <td>apache-2.0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>BAAI/bge-reranker-base</td>\n",
+       "      <td>1.04</td>\n",
+       "      <td>BGE reranker base model for cross-encoder re-r...</td>\n",
+       "      <td>mit</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>jinaai/jina-reranker-v2-base-multilingual</td>\n",
+       "      <td>1.11</td>\n",
+       "      <td>A multi-lingual reranker model for cross-encod...</td>\n",
+       "      <td>cc-by-nc-4.0</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 16
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": ""
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "fst",
+   "display_name": "Python 3.8.18 ('base')",
    "language": "python",
    "name": "python3"
   },
@@ -106,9 +860,14 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.11.8"
   },
-  "orig_nbformat": 4
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "c4a27af61e455bc18dcf16f5867a2ff0402fa12b01dd0f6ce3a79ae73ad15e91"
+   }
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2

docs/experimental/Binary Quantization from Scratch.ipynb

@@ -3,22 +3,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": [
-    "# Binary Quantization of OpenAI Embedding\n",
-    "---\n",
-    "\n",
-    "In the world of large-scale data retrieval and processing, efficiency is crucial. With the exponential growth of data, the ability to retrieve information quickly and accurately can significantly affect system performance. This blog post explores a technique known as binary quantization applied to OpenAI embeddings, demonstrating how it can enhance **retrieval latency by 20x** or more.\n",
-    "\n",
-    "## What Are OpenAI Embeddings?\n",
-    "OpenAI embeddings are numerical representations of textual information. They transform text into a vector space where semantically similar texts are mapped close together. This mathematical representation enables computers to understand and process human language more effectively.\n",
-    "\n",
-    "## Binary Quantization\n",
-    "Binary quantization is a method which converts continuous numerical values into binary values (0 or 1). It simplifies the data structure, allowing faster computations. Here's a brief overview of the binary quantization process applied to OpenAI embeddings:\n",
-    "\n",
-    "1. **Load Embeddings**: OpenAI embeddings are loaded from parquet files.\n",
-    "2. **Binary Transformation**: The continuous valued vectors are converted into binary form. Here, values greater than 0 are set to 1, and others remain 0.\n",
-    "3. **Comparison & Retrieval**: Binary vectors are used for comparison using logical XOR operations and other efficient algorithms."
-   ]
+   "source": []
   },
   {
    "cell_type": "markdown",
@@ -29,24 +14,33 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
+   "execution_count": 1,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:00:06.460001Z",
+     "start_time": "2024-06-06T17:00:04.214098Z"
+    }
+   },
    "outputs": [],
    "source": [
-    "!pip install matplotlib tqdm pandas numpy --quiet"
+    "!pip install matplotlib tqdm pandas numpy datasets --quiet --upgrade"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 2,
    "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:00:07.041784Z",
+     "start_time": "2024-06-06T17:00:06.461658Z"
+    },
     "id": "WBVTItUX4yyr"
    },
    "outputs": [],
    "source": [
     "import numpy as np\n",
     "import pandas as pd\n",
-    "import matplotlib.pyplot as plt\n",
+    "from datasets import load_dataset\n",
     "from tqdm import tqdm"
    ]
   },
@@ -68,8 +62,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 3,
    "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:09.343230Z",
+     "start_time": "2024-06-06T17:00:07.042526Z"
+    },
     "colab": {
      "base_uri": "https://localhost:8080/",
      "height": 250
@@ -77,58 +75,24 @@
     "id": "REJpFqkG7EG2",
     "outputId": "7a43c0ae-fbcc-45fe-fd58-bfe691297b22"
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|██████████| 26/26 [00:10<00:00,  2.45it/s]\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "(1000000, 1536)"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "def get_openai_vectors(force_download: bool = False):\n",
-    "    res = []\n",
-    "    for i in tqdm(range(26)):\n",
-    "        if force_download:\n",
-    "            !wget https://huggingface.co/api/datasets/KShivendu/dbpedia-entities-openai-1M/parquet/KShivendu--dbpedia-entities-openai-1M/train/{i}.parquet\n",
-    "        df = pd.read_parquet(f\"{i}.parquet\", engine=\"pyarrow\")\n",
-    "        res.append(np.stack(df.openai))\n",
-    "        del df\n",
-    "\n",
-    "    openai_vectors = np.concatenate(res)\n",
-    "    del res\n",
-    "    return openai_vectors\n",
-    "\n",
-    "\n",
-    "openai_vectors = get_openai_vectors(force_download=False)\n",
-    "openai_vectors.shape"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## ㆓ Binary Conversion\n",
-    "\n",
-    "Here, we will use 0 as the threshold for the binary conversion. All values greater than 0 will be set to 1, and others will remain 0. This is a simple and effective way to convert continuous values into binary values for OpenAI embeddings."
+    "# Download from Huggingface Hub\n",
+    "ds = load_dataset(\n",
+    "    \"Qdrant/dbpedia-entities-openai3-text-embedding-3-large-3072-100K\", split=\"train\"\n",
+    ")\n",
+    "openai_vectors = np.array(ds[\"text-embedding-3-large-3072-embedding\"])\n",
+    "del ds"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 4,
    "metadata": {
-    "id": "0JM2-Bj2Jkab"
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:10.900963Z",
+     "start_time": "2024-06-06T17:01:09.344842Z"
+    }
    },
    "outputs": [],
    "source": [
@@ -136,6 +100,30 @@
     "openai_bin[openai_vectors > 0] = 1"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:10.906827Z",
+     "start_time": "2024-06-06T17:01:10.901820Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "3072"
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "n_dim = openai_vectors.shape[1]\n",
+    "n_dim"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -147,8 +135,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 6,
    "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:10.909730Z",
+     "start_time": "2024-06-06T17:01:10.908166Z"
+    },
     "id": "FqshI-GlIERd"
    },
    "outputs": [],
@@ -157,7 +149,7 @@
     "    scores = np.dot(openai_vectors, openai_vectors[idx])\n",
     "    dot_results = np.argsort(scores)[-limit:][::-1]\n",
     "\n",
-    "    bin_scores = 1536 - np.logical_xor(openai_bin, openai_bin[idx]).sum(axis=1)\n",
+    "    bin_scores = n_dim - np.logical_xor(openai_bin, openai_bin[idx]).sum(axis=1)\n",
     "    bin_results = np.argsort(bin_scores)[-(limit * oversampling) :][::-1]\n",
     "\n",
     "    return len(set(dot_results).intersection(set(bin_results))) / limit"
@@ -172,8 +164,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 7,
    "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:25.206592Z",
+     "start_time": "2024-06-06T17:01:10.911971Z"
+    },
     "colab": {
      "base_uri": "https://localhost:8080/"
     },
@@ -185,110 +181,128 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "  0%|          | 0/4 [00:00<?, ?it/s]"
+      "  0%|          | 0/4 [00:00<?, ?it/s]\n",
+      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      " 50%|█████     | 1/2 [00:02<00:02,  2.05s/it]\u001b[A"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 1, 'limit': 10, 'recall': 0.8}\n"
+      "{'sampling_rate': 1, 'limit': 3, 'mean_acc': 0.9}\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 2/2 [00:33<00:00, 16.98s/it]\n",
-      " 25%|██▌       | 1/4 [00:33<01:41, 33.96s/it]"
+      "\n",
+      "100%|██████████| 2/2 [00:04<00:00,  2.02s/it]\u001b[A\n",
+      " 25%|██▌       | 1/4 [00:04<00:12,  4.05s/it]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 1, 'limit': 100, 'recall': 0.708}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": []
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'sampling_rate': 2, 'limit': 10, 'recall': 0.95}\n"
+      "{'sampling_rate': 1, 'limit': 10, 'mean_acc': 0.8300000000000001}\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 2/2 [00:32<00:00, 16.38s/it]\n",
-      " 50%|█████     | 2/4 [01:06<01:06, 33.26s/it]"
+      "\n",
+      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      " 50%|█████     | 1/2 [00:01<00:01,  1.72s/it]\u001b[A"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 2, 'limit': 100, 'recall': 0.877}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": []
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'sampling_rate': 3, 'limit': 10, 'recall': 0.96}\n"
+      "{'sampling_rate': 2, 'limit': 3, 'mean_acc': 1.0}\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 2/2 [00:32<00:00, 16.49s/it]\n",
-      " 75%|███████▌  | 3/4 [01:39<00:33, 33.13s/it]"
+      "\n",
+      "100%|██████████| 2/2 [00:03<00:00,  1.76s/it]\u001b[A\n",
+      " 50%|█████     | 2/4 [00:07<00:07,  3.75s/it]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 3, 'limit': 100, 'recall': 0.937}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": []
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'sampling_rate': 5, 'limit': 10, 'recall': 0.9800000000000001}\n"
+      "{'sampling_rate': 2, 'limit': 10, 'mean_acc': 0.9700000000000001}\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 2/2 [00:32<00:00, 16.47s/it]\n",
-      "100%|██████████| 4/4 [02:12<00:00, 33.17s/it]"
+      "\n",
+      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      " 50%|█████     | 1/2 [00:01<00:01,  1.72s/it]\u001b[A"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 5, 'limit': 100, 'recall': 0.977}\n"
+      "{'sampling_rate': 3, 'limit': 3, 'mean_acc': 1.0}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 2/2 [00:03<00:00,  1.69s/it]\u001b[A\n",
+      " 75%|███████▌  | 3/4 [00:10<00:03,  3.58s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 3, 'limit': 10, 'mean_acc': 0.9800000000000001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      " 50%|█████     | 1/2 [00:01<00:01,  1.68s/it]\u001b[A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 5, 'limit': 3, 'mean_acc': 1.0}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 2/2 [00:03<00:00,  1.65s/it]\u001b[A\n",
+      "100%|██████████| 4/4 [00:14<00:00,  3.57s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 5, 'limit': 10, 'mean_acc': 0.99}\n"
      ]
     },
     {
@@ -301,117 +315,53 @@
    ],
    "source": [
     "number_of_samples = 10\n",
-    "limits = [10, 100]\n",
+    "limits = [3, 10]\n",
     "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",
+    "    return np.mean(\n",
+    "        [accuracy(i, limit=limit, oversampling=sampling_rate) for i in range(number_of_samples)]\n",
+    "    )\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",
+    "        result = {\n",
+    "            \"sampling_rate\": i,\n",
+    "            \"limit\": j,\n",
+    "            \"mean_acc\": mean_accuracy(number_of_samples, j, i),\n",
+    "        }\n",
     "        print(result)\n",
     "        results.append(result)"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 19,
+   "cell_type": "markdown",
    "metadata": {},
+   "source": [
+    "## ㆓ Binary Conversion\n",
+    "\n",
+    "Here, we will use 0 as the threshold for the binary conversion. All values greater than 0 will be set to 1, and others will remain 0. This is a simple and effective way to convert continuous values into binary values for OpenAI embeddings."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:25.247495Z",
+     "start_time": "2024-06-06T17:01:25.213508Z"
+    }
+   },
    "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>sampling_rate</th>\n",
-       "      <th>limit</th>\n",
-       "      <th>recall</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>1</td>\n",
-       "      <td>10</td>\n",
-       "      <td>0.800</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>1</td>\n",
-       "      <td>100</td>\n",
-       "      <td>0.708</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>2</td>\n",
-       "      <td>10</td>\n",
-       "      <td>0.950</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>2</td>\n",
-       "      <td>100</td>\n",
-       "      <td>0.877</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>3</td>\n",
-       "      <td>10</td>\n",
-       "      <td>0.960</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>5</th>\n",
-       "      <td>3</td>\n",
-       "      <td>100</td>\n",
-       "      <td>0.937</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6</th>\n",
-       "      <td>5</td>\n",
-       "      <td>10</td>\n",
-       "      <td>0.980</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>7</th>\n",
-       "      <td>5</td>\n",
-       "      <td>100</td>\n",
-       "      <td>0.977</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "   sampling_rate  limit  recall\n",
-       "0              1     10   0.800\n",
-       "1              1    100   0.708\n",
-       "2              2     10   0.950\n",
-       "3              2    100   0.877\n",
-       "4              3     10   0.960\n",
-       "5              3    100   0.937\n",
-       "6              5     10   0.980\n",
-       "7              5    100   0.977"
-      ]
+      "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>sampling_rate</th>\n      <th>limit</th>\n      <th>mean_acc</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>1</td>\n      <td>3</td>\n      <td>0.90</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>1</td>\n      <td>10</td>\n      <td>0.83</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>2</td>\n      <td>3</td>\n      <td>1.00</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>2</td>\n      <td>10</td>\n      <td>0.97</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>3</td>\n      <td>3</td>\n      <td>1.00</td>\n    </tr>\n    <tr>\n      <th>5</th>\n      <td>3</td>\n      <td>10</td>\n      <td>0.98</td>\n    </tr>\n    <tr>\n      <th>6</th>\n      <td>5</td>\n      <td>3</td>\n      <td>1.00</td>\n    </tr>\n    <tr>\n      <th>7</th>\n      <td>5</td>\n      <td>10</td>\n      <td>0.99</td>\n    </tr>\n  </tbody>\n</table>\n</div>",
+      "text/plain": "   sampling_rate  limit  mean_acc\n0              1      3      0.90\n1              1     10      0.83\n2              2      3      1.00\n3              2     10      0.97\n4              3      3      1.00\n5              3     10      0.98\n6              5      3      1.00\n7              5     10      0.99"
      },
-     "execution_count": 19,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -422,22 +372,13 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "| sampling_rate | limit | accuracy |\n",
-    "|---------------|-------|----------|\n",
-    "| 1             | 10    | 0.800    |\n",
-    "| 1             | 100   | 0.708    |\n",
-    "| 2             | 10    | 0.950    |\n",
-    "| 2             | 100   | 0.877    |\n",
-    "| 4             | 10    | 0.970    |\n",
-    "| 4             | 100   | 0.956    |\n",
-    "| 8             | 10    | 0.990    |\n",
-    "| 8             | 100   | 0.990    |\n",
-    "| 16            | 10    | 1.000    |\n",
-    "| 16            | 100   | 0.998    |"
-   ]
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -446,7 +387,8 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
    "name": "python3"
   },
   "language_info": {
@@ -459,7 +401,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

docs/index.md

@@ -2,20 +2,20 @@
 
 FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a Github issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
-The default embedding supports "query" and "passage" prefixes for the input text. The default model is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval%20with%20FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
-
 1. Light & Fast
     - Quantized model weights
-    - ONNX Runtime for inference via [Optimum](github.com/huggingface/optimum)
+    - ONNX Runtime for inference
 
 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
+    - Default is Flag Embedding, which has shown good results on the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard
     - List of [supported models](https://qdrant.github.io/fastembed/examples/Supported_Models/) - including multilingual models
 
+Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval%20with%20FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+
 ## 🚀 Installation
 
-To install the FastEmbed library, pip works: 
+To install the FastEmbed library, pip works:
 
 ```bash
 pip install fastembed
@@ -24,16 +24,16 @@ pip install fastembed
 ## 📖 Usage
 
 ```python
-from fastembed.embedding import FlagEmbedding as Embedding
+from fastembed import TextEmbedding
 
-documents: List[str] = [
+documents: list[str] = [
     "passage: Hello, World!",
-    "query: Hello, World!", # these are two different embedding
+    "query: Hello, World!",
     "passage: This is an example passage.",
-    "fastembed is supported by and maintained by Qdrant." # You can leave out the prefix but it's recommended
+    "fastembed is supported by and maintained by Qdrant."
 ]
-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 = TextEmbedding()
+embeddings: list[np.ndarray] = embedding_model.embed(documents)
 ```
 
 ## Usage with Qdrant
@@ -44,23 +44,22 @@ 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
 
 # Initialize the client
-client = QdrantClient(":memory:")  # or QdrantClient(path="path/to/db")
+client = QdrantClient(":memory:")  # Using an in-process Qdrant
 
 # 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
 client.add(
     collection_name="demo_collection",
     documents=docs,
@@ -73,4 +72,4 @@ search_result = client.query(
     query_text="This is a query document"
 )
 print(search_result)
-```
+```

docs/overrides/main.html

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

docs/qdrant/Retrieval_with_FastEmbed.ipynb

@@ -21,7 +21,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -37,13 +37,12 @@
   },
   {
    "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 +57,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -71,7 +70,7 @@
    ],
    "source": [
     "# Example list of documents\n",
-    "documents: List[str] = [\n",
+    "documents: list[str] = [\n",
     "    \"Maharana Pratap was a Rajput warrior king from Mewar\",\n",
     "    \"He fought against the Mughal Empire led by Akbar\",\n",
     "    \"The Battle of Haldighati in 1576 was his most famous battle\",\n",
@@ -84,10 +83,10 @@
     "    \"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",
+    "embeddings: list[np.ndarray] = list(\n",
     "    embedding_model.passage_embed(documents)\n",
     ")  # notice that we are casting the generator to a list\n",
     "\n",
@@ -105,7 +104,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -124,65 +123,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 +174,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.10.13"
   },
   "orig_nbformat": 4
  },

docs/qdrant/Usage_With_Qdrant.ipynb

@@ -28,16 +28,7 @@
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\u001b[33mDEPRECATION: pytorch-lightning 1.6.5 has a non-standard dependency specifier torch>=1.8.*. pip 23.3 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
-      "\u001b[0m"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "!pip install 'qdrant-client[fastembed]' --quiet --upgrade"
    ]
@@ -55,9 +46,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from typing import List\n",
-    "import numpy as np\n",
-    "from fastembed.embedding import FlagEmbedding as Embedding\n",
     "from qdrant_client import QdrantClient"
    ]
   },
@@ -78,7 +66,7 @@
    "outputs": [],
    "source": [
     "# Example list of documents\n",
-    "documents: List[str] = [\n",
+    "documents: list[str] = [\n",
     "    \"Maharana Pratap was a Rajput warrior king from Mewar\",\n",
     "    \"He fought against the Mughal Empire led by Akbar\",\n",
     "    \"The Battle of Haldighati in 1576 was his most famous battle\",\n",
@@ -117,22 +105,22 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
+      "100%|██████████| 77.7M/77.7M [00:05<00:00, 14.6MiB/s]\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "['77e1e4724dd243b08608f57d5692f6aa',\n",
-       " '74841e5dc3594646bda2c6a6d2795dbd',\n",
-       " '6ef39a9445604d0da84d04f760cd7cf7',\n",
-       " 'e659503d3b3748ef90f23c778274835b',\n",
-       " 'b999675068cd413f93faa0cc890c3819',\n",
-       " '8e452f2935cf4e4b80d8eea68c2aad58',\n",
-       " '28ed4fd4592c48c9a0519618d51bb86e',\n",
-       " '59378c784c5f49109bef65fdc4061334',\n",
-       " 'a78c9b598f7942749156334283a6f24f',\n",
-       " 'f72bb24701c64fabb0182c9e757b581b']"
+       "['4fa8b10c78da4b18ba0830ba8a57367a',\n",
+       " '2eae04b515ee4e9185a9a0e6be812bba',\n",
+       " 'c6039f88486f47f1835ae3b069c5823c',\n",
+       " 'c2c8c51e305144d1917b373125fb4d95',\n",
+       " '79fd23b9ec0648cdab38d1947c6b933e',\n",
+       " '036aa200d8c3492b8a438e4f825f5e7f',\n",
+       " 'c35c77f3ea37460a9a13723fb77b7367',\n",
+       " '6ebccbca571b40d0ab6e83e5e0f2f562',\n",
+       " '38048c2ccc1d4962a4f8f1bd89c8357a',\n",
+       " 'c6b09308360140c7b4f106af3658a31e']"
       ]
      },
      "execution_count": 4,
@@ -186,12 +174,14 @@
     "ids = [42, 2]\n",
     "\n",
     "# Use the new add method\n",
-    "client.add(\n",
-    "    collection_name=\"demo_collection\",\n",
-    "    documents=docs,\n",
-    "    metadata=metadata,\n",
-    "    ids=ids\n",
-    ")"
+    "client.add(collection_name=\"demo_collection\", documents=docs, metadata=metadata, ids=ids)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Behind the scenes, Qdrant Client uses the FastEmbed library to make a passage embedding and then uses the Qdrant API to upsert the documents with metadata, put together as a Points into the collection."
    ]
   },
   {
@@ -203,14 +193,13 @@
      "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",
+    "    collection_name=\"demo_collection\", query_text=\"This is a query document\"\n",
     ")\n",
     "print(search_result)"
    ]
@@ -245,7 +234,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.11.5"
   },
   "orig_nbformat": 4
  },

experiments/01_ONNX_Port.ipynb

@@ -19,7 +19,7 @@
    "outputs": [],
    "source": [
     "from pathlib import Path\n",
-    "from typing import List, Tuple, Any\n",
+    "from typing import Any\n",
     "\n",
     "import numpy as np\n",
     "import time\n",
@@ -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"
    ]
@@ -92,9 +91,11 @@
     "    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]\n",
     "\n",
     "\n",
-    "def hf_embed(model_id: str, inputs: List[str]):\n",
+    "def hf_embed(model_id: str, inputs: list[str]):\n",
     "    # Tokenize the input texts\n",
-    "    batch_dict = hf_tokenizer(inputs, max_length=512, padding=True, truncation=True, return_tensors=\"pt\")\n",
+    "    batch_dict = hf_tokenizer(\n",
+    "        inputs, max_length=512, padding=True, truncation=True, return_tensors=\"pt\"\n",
+    "    )\n",
     "\n",
     "    outputs = hf_model(**batch_dict)\n",
     "    embeddings = average_pool(outputs.last_hidden_state, batch_dict[\"attention_mask\"])\n",
@@ -134,7 +135,9 @@
     "optimization_config = AutoOptimizationConfig.O4()\n",
     "optimizer = ORTOptimizer.from_pretrained(model)\n",
     "\n",
-    "optimizer.optimize(save_dir=save_dir, optimization_config=optimization_config, use_external_data_format=True)\n",
+    "optimizer.optimize(\n",
+    "    save_dir=save_dir, optimization_config=optimization_config, use_external_data_format=True\n",
+    ")\n",
     "model = ORTModelForFeatureExtraction.from_pretrained(save_dir)\n",
     "\n",
     "tokenizer.save_pretrained(save_dir)\n",
@@ -149,7 +152,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 +164,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,8 +175,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
-    "def measure_pipeline_time(pipeline, input_texts: List[str], num_runs=10, **kwargs: Any) -> Tuple[float, float]:\n",
+    "def measure_pipeline_time(\n",
+    "    pipeline, input_texts: list[str], num_runs=10, **kwargs: Any\n",
+    ") -> tuple[float, float]:\n",
     "    \"\"\"Measures the time it takes to run the pipeline on the input texts.\"\"\"\n",
     "    times = []\n",
     "    total_chars = sum(len(text) for text in input_texts)\n",
@@ -256,6 +261,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 +310,9 @@
     }
    ],
    "source": [
-    "import os\n",
     "from google.cloud import storage\n",
     "\n",
+    "\n",
     "def upload(bucket_name, source_file_path):\n",
     "    storage_client = storage.Client(project=\"main\")\n",
     "    bucket = storage_client.bucket(bucket_name)\n",

experiments/02_SPLADE_to_ONNX.ipynb

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

experiments/Example. Convert Resnet50 to ONNX.ipynb

@@ -0,0 +1,122 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "4bdb2a91-fa2a-4cee-ad5a-176cc957394d",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-05-23T12:15:28.171586Z",
+     "start_time": "2024-05-23T12:15:28.076314Z"
+    }
+   },
+   "outputs": [
+    {
+     "ename": "ModuleNotFoundError",
+     "evalue": "No module named 'torch'",
+     "output_type": "error",
+     "traceback": [
+      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
+      "\u001B[0;31mModuleNotFoundError\u001B[0m                       Traceback (most recent call last)",
+      "Cell \u001B[0;32mIn[1], line 1\u001B[0m\n\u001B[0;32m----> 1\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorch\u001B[39;00m\n\u001B[1;32m      2\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorch\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01monnx\u001B[39;00m\n\u001B[1;32m      3\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorchvision\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mmodels\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m \u001B[38;5;21;01mmodels\u001B[39;00m\n",
+      "\u001B[0;31mModuleNotFoundError\u001B[0m: No module named 'torch'"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "import torch.onnx\n",
+    "import torchvision.models as models\n",
+    "import torchvision.transforms as transforms\n",
+    "from PIL import Image\n",
+    "import numpy as np\n",
+    "from tests.config import TEST_MISC_DIR\n",
+    "\n",
+    "# Load pre-trained ResNet-50 model\n",
+    "resnet = models.resnet50(pretrained=True)\n",
+    "resnet = torch.nn.Sequential(*(list(resnet.children())[:-1]))  # Remove the last fully connected layer\n",
+    "resnet.eval()\n",
+    "\n",
+    "# Define preprocessing transform\n",
+    "preprocess = transforms.Compose([\n",
+    "    transforms.Resize(256),\n",
+    "    transforms.CenterCrop(224),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
+    "])\n",
+    "\n",
+    "# Load and preprocess the image\n",
+    "def preprocess_image(image_path):\n",
+    "    input_image = Image.open(image_path)\n",
+    "    input_tensor = preprocess(input_image)\n",
+    "    input_batch = input_tensor.unsqueeze(0)  # Add batch dimension\n",
+    "    return input_batch\n",
+    "\n",
+    "# Example input for exporting\n",
+    "input_image = preprocess_image('example.jpg')\n",
+    "\n",
+    "# Export the model to ONNX with dynamic axes\n",
+    "torch.onnx.export(\n",
+    "    resnet, \n",
+    "    input_image, \n",
+    "    \"model.onnx\", \n",
+    "    export_params=True, \n",
+    "    opset_version=9, \n",
+    "    input_names=['input'], \n",
+    "    output_names=['output'],\n",
+    "    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}\n",
+    ")\n",
+    "\n",
+    "# Load ONNX model\n",
+    "import onnx\n",
+    "import onnxruntime as ort\n",
+    "\n",
+    "onnx_model = onnx.load(\"model.onnx\")\n",
+    "ort_session = ort.InferenceSession(\"model.onnx\")\n",
+    "\n",
+    "# Run inference and extract feature vectors\n",
+    "def extract_feature_vectors(image_paths):\n",
+    "    input_images = [preprocess_image(image_path) for image_path in image_paths]\n",
+    "    input_batch = torch.cat(input_images, dim=0)  # Combine images into a single batch\n",
+    "    ort_inputs = {ort_session.get_inputs()[0].name: input_batch.numpy()}\n",
+    "    ort_outs = ort_session.run(None, ort_inputs)\n",
+    "    return ort_outs[0]\n",
+    "\n",
+    "# Example usage\n",
+    "images = [TEST_MISC_DIR / \"image.jpeg\", str(TEST_MISC_DIR / \"small_image.jpeg\")]  # Replace with your image paths\n",
+    "feature_vectors = extract_feature_vectors(images)\n",
+    "print(\"Feature vector shape:\", feature_vectors.shape)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "baa650c4cb3e0e6d"
+  }
+ ],
+ "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.12.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

experiments/attention_export.py

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

experiments/if_splade_to_onnx.py

@@ -0,0 +1,134 @@
+"""Export an inference-free SPLADE document encoder to ONNX.
+
+Converts `opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte` (an MLM head
+over a GTE backbone) into an onnx model producing token logits, and assembles a model dir
+with everything fastembed's `IfSplade` needs: model.onnx, tokenizer files and idf.json.
+
+Usage:
+    python experiments/if_splade_to_onnx.py --output-dir models/opensearch-neural-sparse-encoding-doc-v3-gte
+"""
+
+import argparse
+import shutil
+from pathlib import Path
+
+import torch
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+MODEL_ID = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte"
+# revision of the remote modeling code (Alibaba-NLP/new-impl), pinned in the model card
+CODE_REVISION = "40ced75c3017eb27626c9d4ea981bde21a2662f4"
+
+TOKENIZER_FILES = [
+    "config.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+    "vocab.txt",
+    "idf.json",
+]
+
+
+class LogitsOnly(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module):
+        super().__init__()
+        self.model = model
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        return self.model(input_ids=input_ids, attention_mask=attention_mask).logits
+
+
+def export(model_id: str, output_dir: Path, opset: int = 14) -> Path:
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    model = AutoModelForMaskedLM.from_pretrained(
+        model_id, trust_remote_code=True, code_revision=CODE_REVISION
+    )
+    model.eval()
+    wrapped = LogitsOnly(model)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    dummy = tokenizer(
+        ["fastembed is a library", "onnx export"],
+        padding=True,
+        truncation=True,
+        return_tensors="pt",
+        return_token_type_ids=False,
+    )
+
+    onnx_path = output_dir / "model.onnx"
+    with torch.inference_mode():
+        torch.onnx.export(
+            wrapped,
+            (dummy["input_ids"], dummy["attention_mask"]),
+            f=onnx_path.as_posix(),
+            input_names=["input_ids", "attention_mask"],
+            output_names=["logits"],
+            dynamic_axes={
+                "input_ids": {0: "batch_size", 1: "sequence_length"},
+                "attention_mask": {0: "batch_size", 1: "sequence_length"},
+                "logits": {0: "batch_size", 1: "sequence_length"},
+            },
+            do_constant_folding=True,
+            opset_version=opset,
+            dynamo=False,
+        )
+
+    for file_name in TOKENIZER_FILES:
+        local_path = hf_hub_download(repo_id=model_id, filename=file_name)
+        shutil.copy(local_path, output_dir / file_name)
+
+    return onnx_path
+
+
+def parity_check(model_id: str, output_dir: Path) -> None:
+    import numpy as np
+    import onnxruntime as ort
+
+    model = AutoModelForMaskedLM.from_pretrained(
+        model_id, trust_remote_code=True, code_revision=CODE_REVISION
+    )
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+    documents = [
+        "Currently New York is rainy.",
+        "fastembed is a lightweight library for generating embeddings",
+        "hello world",
+    ]
+    features = tokenizer(
+        documents, padding=True, truncation=True, return_tensors="pt", return_token_type_ids=False
+    )
+
+    with torch.inference_mode():
+        torch_logits = model(**features).logits.numpy()
+
+    session = ort.InferenceSession(output_dir / "model.onnx")
+    onnx_logits = session.run(
+        ["logits"],
+        {
+            "input_ids": features["input_ids"].numpy(),
+            "attention_mask": features["attention_mask"].numpy(),
+        },
+    )[0]
+
+    max_diff = np.abs(torch_logits - onnx_logits).max()
+    print(f"max |torch - onnx| logits diff: {max_diff}")
+    assert max_diff < 1e-3, "onnx export does not match the torch model"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-id", default=MODEL_ID)
+    parser.add_argument("--output-dir", default=f"models/{MODEL_ID.replace('/', '_')}", type=Path)
+    parser.add_argument("--opset", default=14, type=int)
+    args = parser.parse_args()
+
+    onnx_path = export(args.model_id, args.output_dir, args.opset)
+    print(f"Exported to {onnx_path}")
+    parity_check(args.model_id, args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

experiments/try_attention_export.py

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

fastembed/__init__.py

@@ -0,0 +1,22 @@
+import importlib.metadata
+
+from fastembed.image import ImageEmbedding
+from fastembed.late_interaction import LateInteractionTextEmbedding
+from fastembed.late_interaction_multimodal import LateInteractionMultimodalEmbedding
+from fastembed.sparse import SparseEmbedding, SparseTextEmbedding
+from fastembed.text import TextEmbedding
+
+try:
+    version = importlib.metadata.version("fastembed")
+except importlib.metadata.PackageNotFoundError as _:
+    version = importlib.metadata.version("fastembed-gpu")
+
+__version__ = version
+__all__ = [
+    "TextEmbedding",
+    "SparseTextEmbedding",
+    "SparseEmbedding",
+    "ImageEmbedding",
+    "LateInteractionTextEmbedding",
+    "LateInteractionMultimodalEmbedding",
+]

fastembed/common/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.common.types import ImageInput, OnnxProvider, PathInput
+
+__all__ = ["OnnxProvider", "ImageInput", "PathInput"]

fastembed/common/model_description.py

@@ -0,0 +1,53 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ModelSource:
+    hf: str | None = None
+    url: str | None = None
+    _deprecated_tar_struct: bool = False
+
+    @property
+    def deprecated_tar_struct(self) -> bool:
+        return self._deprecated_tar_struct
+
+    def __post_init__(self) -> None:
+        if self.hf is None and self.url is None:
+            raise ValueError(
+                f"At least one source should be set, current sources: hf={self.hf}, url={self.url}"
+            )
+
+
+@dataclass(frozen=True)
+class BaseModelDescription:
+    model: str
+    sources: ModelSource
+    model_file: str
+    description: str
+    license: str
+    size_in_GB: float
+    additional_files: list[str] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class DenseModelDescription(BaseModelDescription):
+    dim: int | None = None
+    tasks: dict[str, Any] | None = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        assert self.dim is not None, "dim is required for dense model description"
+
+
+@dataclass(frozen=True)
+class SparseModelDescription(BaseModelDescription):
+    requires_idf: bool | None = None
+    vocab_size: int | None = None
+
+
+class PoolingType(str, Enum):
+    CLS = "CLS"
+    MEAN = "MEAN"
+    LAST_TOKEN = "LAST_TOKEN"
+    DISABLED = "DISABLED"

fastembed/common/model_management.py

@@ -0,0 +1,536 @@
+import os
+import time
+import gzip
+import json
+import shutil
+import tarfile
+import tempfile
+import contextlib
+from copy import deepcopy
+from pathlib import Path, PureWindowsPath
+from typing import Any, TypeVar, Generic
+
+import requests
+from huggingface_hub import snapshot_download, model_info, list_repo_tree
+from huggingface_hub.hf_api import RepoFile
+from huggingface_hub.utils import (
+    RepositoryNotFoundError,
+    disable_progress_bars,
+    enable_progress_bars,
+)
+from loguru import logger
+from tqdm import tqdm
+from fastembed.common.model_description import BaseModelDescription
+
+T = TypeVar("T", bound=BaseModelDescription)
+
+_DOWNLOAD_CHUNK_SIZE = 256 * 1024
+
+
+class ModelManagement(Generic[T]):
+    METADATA_FILE = "files_metadata.json"
+
+    @classmethod
+    def list_supported_models(cls) -> list[dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            list[T]: A list of dictionaries containing the model information.
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def add_custom_model(
+        cls,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Add a custom model to the existing embedding classes based on the passed model descriptions
+
+        Model description dict should contain the fields same as in one of the model descriptions presented
+         in fastembed.common.model_description
+
+         E.g. for BaseModelDescription:
+              model: str
+              sources: ModelSource
+              model_file: str
+              description: str
+              license: str
+              size_in_GB: float
+              additional_files: list[str]
+
+        Returns:
+            None
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def _list_supported_models(cls) -> list[T]:
+        raise NotImplementedError()
+
+    @classmethod
+    def _get_model_description(cls, model_name: str) -> T:
+        """
+        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:
+            T: 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.
+        """
+
+        response = requests.get(url, stream=True, timeout=(10, 120))
+
+        # 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."
+            )
+        # Otherwise an error page gets written out as though it were the archive.
+        response.raise_for_status()
+
+        # 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 = bool(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=_DOWNLOAD_CHUNK_SIZE):
+                    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: str,
+        extra_patterns: list[str],
+        local_files_only: bool = False,
+        **kwargs: Any,
+    ) -> 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 (list[str]): extra patterns to allow in the snapshot download, typically
+                includes the required model files.
+            local_files_only (bool, optional): Whether to only use local files. Defaults to False.
+        Returns:
+            Path: The path to the model directory.
+        """
+
+        def _verify_files_from_metadata(
+            model_dir: Path, stored_metadata: dict[str, Any], repo_files: list[RepoFile]
+        ) -> bool:
+            try:
+                for rel_path, meta in stored_metadata.items():
+                    file_path = model_dir / rel_path
+
+                    if not file_path.exists():
+                        return False
+
+                    if repo_files:  # online verification
+                        file_info = next((f for f in repo_files if f.path == file_path.name), None)
+                        if (
+                            not file_info
+                            or file_info.size != meta["size"]
+                            or file_info.blob_id != meta["blob_id"]
+                        ):
+                            return False
+
+                    else:  # offline verification
+                        if file_path.stat().st_size != meta["size"]:
+                            return False
+                return True
+            except (OSError, KeyError) as e:
+                logger.error(f"Error verifying files: {str(e)}")
+                return False
+
+        def _collect_file_metadata(
+            model_dir: Path, repo_files: list[RepoFile]
+        ) -> dict[str, dict[str, int | str]]:
+            meta: dict[str, dict[str, int | str]] = {}
+            file_info_map = {f.path: f for f in repo_files}
+            for file_path in model_dir.rglob("*"):
+                if file_path.is_file() and file_path.name != cls.METADATA_FILE:
+                    repo_file = file_info_map.get(file_path.name)
+                    if repo_file:
+                        meta[str(file_path.relative_to(model_dir))] = {
+                            "size": repo_file.size,
+                            "blob_id": repo_file.blob_id,
+                        }
+            return meta
+
+        def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) -> None:
+            try:
+                if not model_dir.exists():
+                    model_dir.mkdir(parents=True, exist_ok=True)
+                (model_dir / cls.METADATA_FILE).write_text(json.dumps(meta))
+            except (OSError, ValueError) as e:
+                logger.warning(f"Error saving metadata: {str(e)}")
+
+        allow_patterns = [
+            "config.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+            "preprocessor_config.json",
+        ]
+
+        allow_patterns.extend(extra_patterns)
+
+        snapshot_dir = Path(cache_dir) / f"models--{hf_source_repo.replace('/', '--')}"
+        metadata_file = snapshot_dir / cls.METADATA_FILE
+
+        if local_files_only:
+            disable_progress_bars()
+            if metadata_file.exists():
+                metadata = json.loads(metadata_file.read_text())
+                verified = _verify_files_from_metadata(snapshot_dir, metadata, repo_files=[])
+                if not verified:
+                    logger.warning(
+                        "Local file sizes do not match the metadata."
+                    )  # do not raise, still make an attempt to load the model
+            result = snapshot_download(
+                repo_id=hf_source_repo,
+                allow_patterns=allow_patterns,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+                **kwargs,
+            )
+            return result
+
+        repo_revision = model_info(hf_source_repo).sha
+        repo_tree = list(list_repo_tree(hf_source_repo, revision=repo_revision, repo_type="model"))
+
+        allowed_extensions = {".json", ".onnx", ".txt"}
+        repo_files = (
+            [
+                f
+                for f in repo_tree
+                if isinstance(f, RepoFile) and Path(f.path).suffix in allowed_extensions
+            ]
+            if repo_tree
+            else []
+        )
+
+        verified_metadata = False
+
+        if snapshot_dir.exists() and metadata_file.exists():
+            metadata = json.loads(metadata_file.read_text())
+            verified_metadata = _verify_files_from_metadata(snapshot_dir, metadata, repo_files)
+
+        if verified_metadata:
+            disable_progress_bars()
+
+        result = snapshot_download(
+            repo_id=hf_source_repo,
+            allow_patterns=allow_patterns,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+            **kwargs,
+        )
+
+        if (
+            not verified_metadata
+        ):  # metadata is not up-to-date, update it and check whether the files have been
+            # downloaded correctly
+            metadata = _collect_file_metadata(snapshot_dir, repo_files)
+
+            download_successful = _verify_files_from_metadata(
+                snapshot_dir, metadata, repo_files=[]
+            )  # offline verification
+            if not download_successful:
+                raise ValueError(
+                    "Files have been corrupted during downloading process. "
+                    "Please check your internet connection and try again."
+                )
+            _save_file_metadata(snapshot_dir, metadata)
+
+        return result
+
+    @classmethod
+    def decompress_to_cache(cls, targz_path: str, cache_dir: str) -> str:
+        """
+        Decompresses a .tar.gz file to a cache directory.
+
+        Nothing is deleted on failure, since `cache_dir` may hold more than this archive.
+        Cleaning up a partial extraction is the caller's job.
+
+        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.
+
+        Raises:
+            ValueError: If the archive is missing, corrupt, or holds an unsafe member.
+        """
+        # 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:
+                if hasattr(tarfile, "data_filter"):
+                    tar.extractall(path=cache_dir, filter="data")
+                else:
+                    # No PEP 706 filter before 3.10.12, so vet the members by hand.
+                    members = tar.getmembers()
+                    for member in members:
+                        cls._validate_tar_member(member)
+                    tar.extractall(path=cache_dir, members=members)
+                # tarfile stops at the end-of-archive marker, short of the gzip trailer, so
+                # the CRC is only checked if the rest of the stream is read.
+                while tar.fileobj.read(1 << 20):
+                    pass
+        except (tarfile.TarError, ValueError, EOFError, gzip.BadGzipFile) as e:
+            # gzip raises EOFError for a truncated stream and BadGzipFile for a corrupted one.
+            raise ValueError(f"An error occurred while decompressing {targz_path}: {e}") from e
+
+        return cache_dir
+
+    @staticmethod
+    def _is_unsafe_tar_path(path: str) -> bool:
+        """Checks whether a tar member name or link target may escape the extraction dir.
+
+        Lexical on purpose: resolving against the extraction directory is unsound before
+        extraction, since `link/../escape` only escapes once an earlier member has been
+        written as a symlink. Any `..` component is therefore rejected outright.
+        """
+        # PureWindowsPath splits on both separators, so `root` covers POSIX "/evil" as
+        # well as "\\evil", which escapes on Windows without being absolute.
+        windows_path = PureWindowsPath(path)
+        return bool(windows_path.drive or windows_path.root) or ".." in windows_path.parts
+
+    @classmethod
+    def _validate_tar_member(cls, member: tarfile.TarInfo) -> None:
+        """Raises ValueError if a member could write outside the extraction directory."""
+        if cls._is_unsafe_tar_path(member.name):
+            raise ValueError(f"Unsafe tar member path: {member.name}")
+
+        if member.issym() or member.islnk():
+            if cls._is_unsafe_tar_path(member.linkname):
+                raise ValueError(f"Unsafe tar link target: {member.name} -> {member.linkname}")
+        elif not (member.isfile() or member.isdir()):
+            # Devices, fifos and the like have no place in a model archive.
+            raise ValueError(f"Unsupported tar member type: {member.name}")
+
+    @classmethod
+    def retrieve_model_gcs(
+        cls,
+        model_name: str,
+        source_url: str,
+        cache_dir: str,
+        deprecated_tar_struct: bool = False,
+        local_files_only: bool = False,
+    ) -> Path:
+        fast_model_name = f"{'fast-' if deprecated_tar_struct else ''}{model_name.split('/')[-1]}"
+        cache_tmp_dir = Path(cache_dir) / "tmp"
+        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 local_files_only:
+            logger.error(
+                f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
+            )
+            raise ValueError(
+                f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
+            )
+
+        if cache_tmp_dir.is_symlink():
+            raise ValueError(
+                f"{cache_tmp_dir} is a symlink, refusing to stage downloads through it"
+            )
+        cache_tmp_dir.mkdir(parents=True, exist_ok=True)
+
+        # The archive and everything extracted from it go in a directory of this attempt's own,
+        # so removing it undoes the attempt without touching any other download of the model.
+        staging_dir = Path(tempfile.mkdtemp(dir=cache_tmp_dir, prefix=f"{fast_model_name}-"))
+        try:
+            model_tar_gz = staging_dir / f"{fast_model_name}.tar.gz"
+            cls.download_file_from_gcs(
+                source_url,
+                output_path=str(model_tar_gz),
+            )
+
+            cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(staging_dir))
+
+            model_tmp_dir = staging_dir / fast_model_name
+            if not model_tmp_dir.is_dir() or model_tmp_dir.is_symlink():
+                raise ValueError(
+                    f"The archive from {source_url} has no {fast_model_name} directory"
+                )
+
+            # Replace a stale empty model_dir, which Windows will not rename onto. rmdir leaves
+            # anything else alone, including one another download has just filled.
+            with contextlib.suppress(OSError):
+                model_dir.rmdir()
+
+            try:
+                # Rename from the staging dir to the final name is atomic
+                model_tmp_dir.rename(model_dir)
+            except OSError:
+                # Another download of the same model finished first, so keep its copy.
+                if not (model_dir.is_dir() and any(model_dir.iterdir())):
+                    raise
+        finally:
+            shutil.rmtree(staging_dir, ignore_errors=True)
+
+        return model_dir
+
+    @classmethod
+    def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: Any) -> Path:
+        """
+        Downloads a model from HuggingFace Hub or Google Cloud Storage.
+
+        Args:
+            model (T): 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.
+            retries: (int): The number of times to retry (including the first attempt)
+
+        Returns:
+            Path: The path to the downloaded model directory.
+        """
+        local_files_only = kwargs.get("local_files_only", False)
+        hf_offline = os.environ.get("HF_HUB_OFFLINE", "").strip().upper()
+        if not local_files_only and hf_offline in {"1", "TRUE", "YES", "ON"}:
+            local_files_only = True
+            kwargs["local_files_only"] = True
+        specific_model_path: str | None = kwargs.pop("specific_model_path", None)
+        if specific_model_path:
+            return Path(specific_model_path)
+        retries = 1 if local_files_only else retries
+        hf_source = model.sources.hf
+        url_source = model.sources.url
+
+        extra_patterns = [model.model_file]
+        extra_patterns.extend(model.additional_files)
+
+        if hf_source:
+            try:
+                cache_kwargs = deepcopy(kwargs)
+                cache_kwargs["local_files_only"] = True
+                resolved_path = Path(
+                    cls.download_files_from_huggingface(
+                        hf_source,
+                        cache_dir=cache_dir,
+                        extra_patterns=extra_patterns,
+                        **cache_kwargs,
+                    )
+                )
+                if (resolved_path / model.model_file).exists() and all(
+                    (resolved_path / file).exists() for file in extra_patterns
+                ):
+                    return resolved_path
+            except Exception:
+                pass
+            finally:
+                enable_progress_bars()
+
+        sleep = 3.0
+        while retries > 0:
+            retries -= 1
+
+            if hf_source and not local_files_only:
+                # we have already tried loading with `local_files_only=True` via hf and we failed
+                try:
+                    return Path(
+                        cls.download_files_from_huggingface(
+                            hf_source,
+                            cache_dir=cache_dir,
+                            extra_patterns=extra_patterns,
+                            **kwargs,
+                        )
+                    )
+                except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
+                    if not local_files_only:
+                        logger.error(
+                            f"Could not download model from HuggingFace: {e} "
+                            "Falling back to other sources."
+                        )
+                finally:
+                    enable_progress_bars()
+            if url_source or local_files_only:
+                try:
+                    return cls.retrieve_model_gcs(
+                        model.model,
+                        str(url_source),
+                        str(cache_dir),
+                        deprecated_tar_struct=model.sources.deprecated_tar_struct,
+                        local_files_only=local_files_only,
+                    )
+                except Exception:
+                    if not local_files_only:
+                        logger.error(f"Could not download model from url: {url_source}")
+
+            if local_files_only:
+                logger.error("Could not find model in cache_dir")
+                break
+            else:
+                logger.error(
+                    f"Could not download model from either source, sleeping for {sleep} seconds, {retries} retries left."
+                )
+                time.sleep(sleep)
+                sleep *= 3
+
+        raise ValueError(f"Could not load model {model.model} from any source.")

fastembed/common/onnx_model.py

@@ -0,0 +1,200 @@
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Generic, Iterable, Sequence, Type, TypeVar
+
+import numpy as np
+import onnxruntime as ort
+
+from numpy.typing import NDArray
+from tokenizers import Tokenizer
+
+from fastembed.common.types import OnnxProvider, NumpyArray, Device
+from fastembed.parallel_processor import Worker
+
+# Holds type of the embedding result
+T = TypeVar("T")
+
+
+@dataclass
+class OnnxOutputContext:
+    model_output: NumpyArray
+    attention_mask: NDArray[np.int64] | None = None
+    input_ids: NDArray[np.int64] | None = None
+    metadata: dict[str, Any] | None = None
+
+
+class OnnxModel(Generic[T]):
+    EXPOSED_SESSION_OPTIONS = ("enable_cpu_mem_arena",)
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _get_worker_init_kwargs(self) -> dict[str, Any]:
+        """Additional kwargs a worker process needs to reconstruct this model.
+
+        Workers are started with `spawn`/`forkserver`, hence they don't inherit class-level state
+        which has been set up in runtime, e.g. models registered via `add_custom_model`.
+        Such state has to be shipped to the workers explicitly.
+
+        Returns:
+            dict[str, Any]: kwargs to pass to `_get_worker_class().init_embedding`.
+        """
+        return {}
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[T]: Post-processed output as an iterable of type T.
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def __init__(self) -> None:
+        self.model: ort.InferenceSession | None = None
+        self.tokenizer: Tokenizer | None = None
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def _load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
+    ) -> None:
+        model_path = model_dir / model_file
+        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+        available_providers = ort.get_available_providers()
+        cuda_available = "CUDAExecutionProvider" in available_providers
+        explicit_cuda = cuda is True or cuda == Device.CUDA
+
+        if explicit_cuda and providers is not None:
+            warnings.warn(
+                f"`cuda` and `providers` are mutually exclusive parameters, "
+                f"cuda: {cuda}, providers: {providers}. If you'd like to use providers, cuda should be one of "
+                f"[False, Device.CPU, Device.AUTO].",
+                category=UserWarning,
+                stacklevel=6,
+            )
+
+        if providers is not None:
+            onnx_providers = list(providers)
+        elif explicit_cuda or (cuda == Device.AUTO and cuda_available):
+            if device_id is None:
+                onnx_providers = ["CUDAExecutionProvider"]
+            else:
+                onnx_providers = [("CUDAExecutionProvider", {"device_id": device_id})]
+        else:
+            onnx_providers = ["CPUExecutionProvider"]
+
+        requested_provider_names: list[str] = []
+        for provider in onnx_providers:
+            # check providers available
+            provider_name = provider if isinstance(provider, str) else provider[0]
+            requested_provider_names.append(provider_name)
+            if provider_name not in available_providers:
+                raise ValueError(
+                    f"Provider {provider_name} is not available. Available providers: {available_providers}"
+                )
+
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        if threads is not None:
+            so.intra_op_num_threads = threads
+            so.inter_op_num_threads = threads
+
+        if extra_session_options is not None:
+            self.add_extra_session_options(so, extra_session_options)
+
+        self.model = ort.InferenceSession(
+            str(model_path), providers=onnx_providers, sess_options=so
+        )
+        if "CUDAExecutionProvider" in requested_provider_names:
+            assert self.model is not None
+            current_providers = self.model.get_providers()
+            if "CUDAExecutionProvider" not in current_providers:
+                warnings.warn(
+                    f"Attempt to set CUDAExecutionProvider failed. Current providers: {current_providers}."
+                    "If you are using CUDA 12.x, install onnxruntime-gpu via "
+                    "`pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/`",
+                    RuntimeWarning,
+                )
+
+    @classmethod
+    def _select_exposed_session_options(cls, model_kwargs: dict[str, Any]) -> dict[str, Any]:
+        """A convenience method to select the exposed session options in models
+
+        Args:
+            model_kwargs (dict[str, Any]): The model kwargs.
+
+        Returns:
+            dict[str, Any]: a dict with filtered exposed session options.
+        """
+        return {k: v for k, v in model_kwargs.items() if k in cls.EXPOSED_SESSION_OPTIONS}
+
+    @classmethod
+    def add_extra_session_options(
+        cls, session_options: ort.SessionOptions, extra_options: dict[str, Any]
+    ) -> None:
+        """Add extra session options to the existing options object in-place
+
+        Args:
+            session_options (ort.SessionOptions): The existing session options object.
+            extra_options (dict[str, Any]): The extra session options available in cls.EXPOSED_SESSION_OPTIONS.
+
+        Returns:
+            None
+        """
+        for option in extra_options:
+            assert (
+                option in cls.EXPOSED_SESSION_OPTIONS
+            ), f"{option} is unknown or not exposed (exposed options: {cls.EXPOSED_SESSION_OPTIONS})"
+        if "enable_cpu_mem_arena" in extra_options:
+            session_options.enable_cpu_mem_arena = extra_options["enable_cpu_mem_arena"]
+
+    def load_onnx_model(self) -> None:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def onnx_embed(self, *args: Any, **kwargs: Any) -> OnnxOutputContext:
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class EmbeddingWorker(Worker, Generic[T]):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxModel[T]:
+        raise NotImplementedError()
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ):
+        self.model = self.init_embedding(model_name, cache_dir, **kwargs)
+
+    @classmethod
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker[T]":
+        return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
+
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+        raise NotImplementedError("Subclasses must implement this method")

fastembed/common/preprocessor_utils.py

@@ -0,0 +1,151 @@
+import json
+import sys
+from typing import Any, Iterator
+from pathlib import Path
+
+from tokenizers import AddedToken, Tokenizer
+
+from fastembed.image.transform.operators import Compose
+
+
+def load_special_tokens(model_dir: Path) -> dict[str, Any]:
+    """Read special_tokens_map.json, treating an absent file as an empty map."""
+    tokens_map_path = model_dir / "special_tokens_map.json"
+    if not tokens_map_path.exists():
+        return {}
+
+    with open(str(tokens_map_path)) as tokens_map_file:
+        tokens_map = json.load(tokens_map_file)
+
+    return tokens_map
+
+
+def iter_special_tokens(tokens_map: dict[str, Any]) -> Iterator[str | dict[str, Any]]:
+    """Yield the individual tokens declared in a special tokens map.
+
+    Most keys hold one token, but `additional_special_tokens` holds a list of them.
+    """
+    for value in tokens_map.values():
+        if isinstance(value, list):
+            yield from value
+        else:
+            yield value
+
+
+def _valid_context(value: Any) -> int | None:
+    """Return `value` if it can be used as a truncation limit, `None` otherwise.
+
+    Config files do not always carry a real limit: transformers writes `model_max_length` as
+    1e30 when the value is unknown, and some repos ship a 0 or a null. `enable_truncation`
+    raises an `OverflowError` on the former and silently produces empty encodings on the
+    latter, so both are rejected here rather than passed through.
+    """
+    if isinstance(value, bool) or not isinstance(value, int):
+        return None
+    if not 0 < value <= sys.maxsize:
+        return None
+    return value
+
+
+def _resolve_max_context(tokenizer_config: dict[str, Any], model_dir: Path) -> int:
+    """Pick the truncation limit, preferring the stricter of the two tokenizer config keys.
+
+    `config.json:max_position_embeddings` deliberately is not used as a fallback: it is the size
+    of the position table, not the usable context, and the two differ per architecture, e.g.
+    roberta reports 514 for a usable 512.
+    """
+    candidates = [
+        context
+        for context in (
+            _valid_context(tokenizer_config.get("model_max_length")),
+            _valid_context(tokenizer_config.get("max_length")),
+        )
+        if context is not None
+    ]
+    if not candidates:
+        raise ValueError(
+            f"Could not determine the maximum context length for {model_dir}. Set a positive "
+            "`model_max_length` or `max_length` in tokenizer_config.json."
+        )
+
+    return min(candidates)
+
+
+def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
+    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}")
+
+    # config.json is optional: transformers v5 no longer writes it for every model.
+    config_path = model_dir / "config.json"
+    config: dict[str, Any] = {}
+    if config_path.exists():
+        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)
+
+    max_context = _resolve_max_context(tokenizer_config, model_dir)
+
+    tokens_map = load_special_tokens(model_dir)
+
+    tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    tokenizer.enable_truncation(max_length=max_context)
+
+    # Registered before the padding is resolved: the map may name a pad token that
+    # tokenizer.json does not carry, and it only gets an id once it is added.
+    for token in iter_special_tokens(tokens_map):
+        if isinstance(token, str):
+            tokenizer.add_special_tokens([token])
+        elif isinstance(token, dict):
+            tokenizer.add_special_tokens([AddedToken(**token)])
+
+    # Padding is always normalized to batch-longest. A serialized fixed length shorter than the
+    # truncation limit leaves longer encodings untouched, which produces ragged batches, and a
+    # fixed length equal to it pads every batch to the maximum. Direction and pad token metadata
+    # are taken from the serialized settings, since some models pad on the left.
+    padding = tokenizer.padding or {}
+    pad_token = padding.get("pad_token") or tokenizer_config.get("pad_token")
+    if pad_token is None:
+        raise ValueError(f"Could not find a pad token for {model_dir}")
+
+    # The vocabulary is the last resort, not a hardcoded 0: that silently disagrees with
+    # `pad_token` for every model whose pad token is not the first entry.
+    pad_id = padding.get("pad_id", config.get("pad_token_id"))
+    if pad_id is None:
+        pad_id = tokenizer.token_to_id(pad_token)
+    if pad_id is None:
+        raise ValueError(f"Could not resolve an id for the pad token {pad_token!r} in {model_dir}")
+
+    tokenizer.enable_padding(
+        direction=padding.get("direction", "right"),
+        pad_id=pad_id,
+        pad_type_id=padding.get("pad_type_id", 0),
+        pad_token=pad_token,
+        pad_to_multiple_of=padding.get("pad_to_multiple_of"),
+        length=None,
+    )
+
+    special_token_to_id = {
+        token.content: token_id
+        for token_id, token in tokenizer.get_added_tokens_decoder().items()
+        if token.special
+    }
+
+    return tokenizer, special_token_to_id
+
+
+def load_preprocessor(model_dir: Path) -> Compose:
+    preprocessor_config_path = model_dir / "preprocessor_config.json"
+    if not preprocessor_config_path.exists():
+        raise ValueError(f"Could not find preprocessor_config.json in {model_dir}")
+
+    with open(str(preprocessor_config_path)) as preprocessor_config_file:
+        preprocessor_config = json.load(preprocessor_config_file)
+        transforms = Compose.from_config(preprocessor_config)
+    return transforms

fastembed/common/types.py

@@ -0,0 +1,27 @@
+from enum import Enum
+from pathlib import Path
+from typing import Any, TypeAlias
+
+import numpy as np
+from numpy.typing import NDArray
+from PIL import Image
+
+
+class Device(str, Enum):
+    CPU = "cpu"
+    CUDA = "cuda"
+    AUTO = "auto"
+
+
+PathInput: TypeAlias = str | Path
+ImageInput: TypeAlias = PathInput | Image.Image
+
+OnnxProvider: TypeAlias = str | tuple[str, dict[Any, Any]]
+NumpyArray: TypeAlias = (
+    NDArray[np.float64]
+    | NDArray[np.float32]
+    | NDArray[np.float16]
+    | NDArray[np.int8]
+    | NDArray[np.int64]
+    | NDArray[np.int32]
+)

fastembed/common/utils.py

@@ -0,0 +1,79 @@
+import os
+import sys
+import re
+import tempfile
+import unicodedata
+from pathlib import Path
+from itertools import islice
+from typing import Iterable, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+from fastembed.common.types import NumpyArray
+
+T = TypeVar("T")
+
+
+def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> NumpyArray:
+    # 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
+
+
+def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
+    input_mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(np.int64)
+    input_mask_expanded = np.tile(input_mask_expanded, (1, 1, input_array.shape[-1]))
+    sum_embeddings = np.sum(input_array * input_mask_expanded, axis=1)
+    sum_mask = np.sum(input_mask_expanded, axis=1)
+    pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+    return pooled_embeddings
+
+
+def last_token_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
+    """Take the embedding of the last non-padding token of each sequence.
+
+    Locates the last position the attention mask marks as real, so it holds whichever
+    side the tokenizer pads on.
+    """
+    last_token_indices = attention_mask.shape[1] - 1 - np.argmax(attention_mask[:, ::-1], axis=1)
+    return input_array[np.arange(input_array.shape[0]), last_token_indices]
+
+
+def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
+    """
+    >>> 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: str | None = 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
+
+
+def get_all_punctuation() -> set[str]:
+    return set(
+        chr(i) for i in range(sys.maxunicode) if unicodedata.category(chr(i)).startswith("P")
+    )
+
+
+def remove_non_alphanumeric(text: str) -> str:
+    return re.sub(r"[^\w\s]", " ", text, flags=re.UNICODE)

fastembed/embedding.py

@@ -1,504 +1,24 @@
-import json
-import os
-import shutil
-import tarfile
-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 Any
 
-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 import TextEmbedding
 
+logger.warning(
+    "DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated."
+    "Use from fastembed import TextEmbedding instead."
+)
 
-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
+DefaultEmbedding = TextEmbedding
+FlagEmbedding = TextEmbedding
 
 
-def normalize(input_array, p=2, dim=1, eps=1e-12):
-    # Calculate the Lp norm along the specified dimension
-    norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
-    norm = np.maximum(norm, eps)  # Avoid division by zero
-    normalized_array = input_array / norm
-    return normalized_array
-
-
-class EmbeddingModel(ABC):
-    @classmethod
-    def load_tokenizer(cls, model_dir: Path, max_length: int = 512) -> Tokenizer:
-        config_path = model_dir / "config.json"
-        if not config_path.exists():
-            raise ValueError(f"Could not find config.json in {model_dir}")
-
-        tokenizer_path = model_dir / "tokenizer.json"
-        if not tokenizer_path.exists():
-            raise ValueError(f"Could not find tokenizer.json in {model_dir}")
-
-        tokenizer_config_path = model_dir / "tokenizer_config.json"
-        if not tokenizer_config_path.exists():
-            raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
-
-        tokens_map_path = model_dir / "special_tokens_map.json"
-        if not tokens_map_path.exists():
-            raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
-
-        config = json.load(open(str(config_path)))
-        tokenizer_config = json.load(open(str(tokenizer_config_path)))
-        tokens_map = json.load(open(str(tokens_map_path)))
-
-        tokenizer = Tokenizer.from_file(str(tokenizer_path))
-        tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
-        tokenizer.enable_padding(pad_id=config["pad_token_id"], pad_token=tokenizer_config["pad_token"])
-
-        for token in tokens_map.values():
-            if isinstance(token, str):
-                tokenizer.add_special_tokens([token])
-            elif isinstance(token, dict):
-                tokenizer.add_special_tokens([AddedToken(**token)])
-
-        return tokenizer
-
-    def __init__(
-            self,
-            path: Path,
-            model_name: str,
-            max_length: int = 512,
-            max_threads: int = None,
-    ):
-        self.path = path
-        self.model_name = model_name
-        model_path = self.path / "model.onnx"
-        optimized_model_path = self.path / "model_optimized.onnx"
-
-        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
-        onnx_providers = ["CPUExecutionProvider"]
-
-        if not model_path.exists():
-            # Rename file model_optimized.onnx to model.onnx if it exists
-            if optimized_model_path.exists():
-                optimized_model_path.rename(model_path)
-            else:
-                raise ValueError(f"Could not find model.onnx in {self.path}")
-
-        # Hacky support for multilingual model
-        self.exclude_token_type_ids = False
-        if model_name == "intfloat/multilingual-e5-large":
-            self.exclude_token_type_ids = True
-
-        so = ort.SessionOptions()
-        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-
-        if max_threads is not None:
-            so.intra_op_num_threads = max_threads
-            so.inter_op_num_threads = max_threads
-
-        self.tokenizer = self.load_tokenizer(self.path, max_length=max_length)
-        self.model = ort.InferenceSession(str(model_path), providers=onnx_providers, sess_options=so)
-
-    def onnx_embed(self, documents: List[str]) -> np.ndarray:
-        encoded = self.tokenizer.encode_batch(documents)
-        input_ids = np.array([e.ids for e in encoded])
-        attention_mask = np.array([e.attention_mask for e in encoded])
-
-        onnx_input = {
-            "input_ids": np.array(input_ids, dtype=np.int64),
-            "attention_mask": np.array(attention_mask, dtype=np.int64),
-        }
-
-        if not self.exclude_token_type_ids:
-            onnx_input["token_type_ids"] = np.array(
-                [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
-            )
-
-        model_output = self.model.run(None, onnx_input)
-        last_hidden_state = model_output[0][:, 0]
-        embeddings = normalize(last_hidden_state).astype(np.float32)
-        return embeddings
-
-
-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 = self.model.onnx_embed(batch)
-            yield idx, embeddings
-
-
-class Embedding(ABC):
-    """
-    Abstract class for embeddings.
-
-    Args:
-        ABC ():
-
-    Raises:
-        NotImplementedError: Raised when you call an abstract method that has not been implemented.
-        PermissionError: _description_
-        ValueError: Several possible reasons: 1) targz_path does not exist or is not a file, 2) targz_path is not a .tar.gz file, 3) An error occurred while decompressing targz_path, 4) Could not find model_dir in cache_dir, 5) Could not find tokenizer.json in model_dir, 6) Could not find model.onnx in model_dir.
-        NotImplementedError: _description_
-
-    Returns:
-        _type_: _description_
-
-    Yields:
-        _type_: _description_
-    """
-
-    @abstractmethod
-    def embed(self, texts: List[str]) -> List[np.ndarray]:
-        raise NotImplementedError
-
-    @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Union[str, 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-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
-            },
-        ]
-
-    @classmethod
-    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
-        """
-        Downloads a file from Google Cloud Storage.
-
-        Args:
-            url (str): The URL to download the file from.
-            output_path (str): The path to save the downloaded file to.
-            show_progress (bool, optional): Whether to show a progress bar. Defaults to True.
-
-        Returns:
-            str: The path to the downloaded file.
-        """
-
-        if os.path.exists(output_path):
-            return output_path
-        response = requests.get(url, stream=True)
-
-        # Handle HTTP errors
-        if response.status_code == 403:
-            raise PermissionError(
-                "Authentication Error: You do not have permission to access this resource. Please check your credentials."
-            )
-
-        # Get the total size of the file
-        total_size_in_bytes = int(response.headers.get("content-length", 0))
-
-        # Warn if the total size is zero
-        if total_size_in_bytes == 0:
-            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
-
-        # Initialize the progress bar
-        progress_bar = (
-            tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
-            if total_size_in_bytes and show_progress
-            else None
-        )
-
-        # Attempt to download the file
-        try:
-            with open(output_path, "wb") as file:
-                for chunk in response.iter_content(chunk_size=1024):  # Adjust chunk size to your preference
-                    if chunk:  # Filter out keep-alive new chunks
-                        if progress_bar is not None:
-                            progress_bar.update(len(chunk))
-                        file.write(chunk)
-        except Exception as e:
-            print(f"An error occurred while trying to download the file: {str(e)}")
-            return
-        finally:
-            if progress_bar is not None:
-                progress_bar.close()
-        return output_path
-
-    @classmethod
-    def decompress_to_cache(cls, targz_path: str, cache_dir: str):
-        """
-        Decompresses a .tar.gz file to a cache directory.
-
-        Args:
-            targz_path (str): Path to the .tar.gz file.
-            cache_dir (str): Path to the cache directory.
-
-        Returns:
-            cache_dir (str): Path to the cache directory.
-        """
-        # Check if targz_path exists and is a file
-        if not os.path.isfile(targz_path):
-            raise ValueError(f"{targz_path} does not exist or is not a file.")
-
-        # Check if targz_path is a .tar.gz file
-        if not targz_path.endswith(".tar.gz"):
-            raise ValueError(f"{targz_path} is not a .tar.gz file.")
-
-        try:
-            # Open the tar.gz file
-            with tarfile.open(targz_path, "r:gz") as tar:
-                # Extract all files into the cache directory
-                tar.extractall(path=cache_dir)
-        except tarfile.TarError as e:
-            # If any error occurs while opening or extracting the tar.gz file,
-            # delete the cache directory (if it was created in this function)
-            # and raise the error again
-            if "tmp" in cache_dir:
-                shutil.rmtree(cache_dir)
-            raise ValueError(f"An error occurred while decompressing {targz_path}: {e}")
-
-        return cache_dir
-
-    def retrieve_model(self, model_name: str, cache_dir: str) -> Path:
-        """
-        Retrieves a model from Google Cloud Storage.
-
-        Args:
-            model_name (str): The name of the model to retrieve.
-            cache_dir (str): The path to the cache directory.
-
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
-
-        Returns:
-            Path: The path to the model directory.
-        """
-
-        assert "/" in model_name, "model_name must be in the format <org>/<model> e.g. BAAI/bge-base-en"
-
-        fast_model_name = f"fast-{model_name.split('/')[-1]}"
-
-        model_dir = Path(cache_dir) / fast_model_name
-        if model_dir.exists():
-            return model_dir
-
-        model_tar_gz = Path(cache_dir) / f"{fast_model_name}.tar.gz"
-        try:
-            self.download_file_from_gcs(
-                f"https://storage.googleapis.com/qdrant-fastembed/{fast_model_name}.tar.gz",
-                output_path=str(model_tar_gz),
-            )
-        except PermissionError:
-            simple_model_name = model_name.replace("/", "-")
-            print(f"Was not able to download {fast_model_name}.tar.gz, trying {simple_model_name}.tar.gz")
-            self.download_file_from_gcs(
-                f"https://storage.googleapis.com/qdrant-fastembed/{simple_model_name}.tar.gz",
-                output_path=str(model_tar_gz),
-            )
-
-        self.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=cache_dir)
-        assert model_dir.exists(), f"Could not find {model_dir} in {cache_dir}"
-
-        model_tar_gz.unlink()
-
-        return model_dir
-
-    def passage_embed(self, texts: List[str], batch_size: int = 256) -> Iterable[np.ndarray]:
-        """
-        Embeds a list of text passages into a list of embeddings.
-
-        Args:
-            texts (List[str]): The list of texts to embed.
-            batch_size (int, optional): The batch size. Defaults to 256.
-
-        Yields:
-            Iterable[np.ndarray]: The embeddings.
-        """
-
-        for i in range(0, len(texts), batch_size):
-            # Prepend "passage: " to each text
-            yield from self.embed([f"passage: {t}" for t in texts[i: i + batch_size]])
-
-    def query_embed(self, query: str) -> Iterable[np.ndarray]:
-        """
-        Embeds a query
-
-        Args:
-            query (str): The query to search for.
-
-        Returns:
-            Iterable[np.ndarray]: The embeddings.
-        """
-
-        # Prepend "query: " to the query
-        query = f"query: {query}"
-        # Embed the query
-        query_embedding = self.embed([query])
-        # Compute the cosine similarity between the query embedding and the document embeddings
-        return query_embedding
-
-
-class FlagEmbedding(Embedding):
-    """
-    Implementation of the Flag Embedding model.
-
-    Args:
-        Embedding (_type_): _description_
-    """
-
-    def __init__(
-            self,
-            model_name: str = "BAAI/bge-small-en",
-            max_length: int = 512,
-            cache_dir: str = None,
-            threads: int = None,
-    ):
-        """
-        Args:
-            model_name (str): The name of the model to use.
-            max_length (int, optional): The maximum number of tokens. Defaults to 512. Unknown behavior for values > 512.
-            cache_dir (str, optional): The path to the cache directory. Defaults to `local_cache` in the current directory.
-            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
-
-        Raises:
-            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
-        """
-        self.model_name = model_name
-
-        if cache_dir is None:
-            cache_dir = Path(".").resolve() / "local_cache"
-            cache_dir.mkdir(parents=True, exist_ok=True)
-
-        self._cache_dir = cache_dir
-        self._model_dir = self.retrieve_model(model_name, cache_dir)
-        self._max_length = max_length
-
-        self.model = EmbeddingModel(self._model_dir, self.model_name, max_length=max_length,
-                                    max_threads=threads)
-
-    def embed(
-            self, documents: Union[str, Iterable[str]], batch_size: int = 256, parallel: int = None
-    ) -> Iterable[np.ndarray]:
-        """
-        Encode a list of documents into list of embeddings.
-        We use mean pooling with attention so that the model can handle variable-length inputs.
-
-        Args:
-            documents: Iterator of documents or single document to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
-            parallel:
-                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
-                If 0, use all available cores.
-                If None, don't use data-parallel processing, use default onnxruntime threading instead.
-
-        Returns:
-            List of embeddings, one per document
-        """
-        is_small = False
-
-        if isinstance(documents, str):
-            documents = [documents]
-            is_small = True
-
-        if isinstance(documents, list):
-            if len(documents) < batch_size:
-                is_small = True
-
-        if parallel == 0:
-            parallel = os.cpu_count()
-
-        if parallel is None or is_small:
-            for batch in iter_batch(documents, batch_size):
-                yield from self.model.onnx_embed(batch)
-        else:
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            params = {
-                "path": self._model_dir,
-                "model_name": self.model_name,
-                "max_length": self._max_length,
-            }
-            pool = ParallelWorkerPool(parallel, EmbeddingWorker, start_method=start_method)
-            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
-                yield from batch
-
-
-class DefaultEmbedding(FlagEmbedding):
-    """
-    Implementation of the default Flag Embedding model.
-
-    Args:
-        FlagEmbedding (_type_): _description_
-    """
-
+class JinaEmbedding(TextEmbedding):
     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,
+        model_name: str = "jinaai/jina-embeddings-v2-base-en",
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        **kwargs: Any,
     ):
-        super().__init__(model_name, max_length=max_length, cache_dir=cache_dir, threads=threads)
-
-
-class OpenAIEmbedding(Embedding):
-    def __init__(self):
-        # Initialize your OpenAI model here
-        # self.model = ...
-        ...
-
-    def embed(self, texts):
-        # Use your OpenAI model to embed the texts
-        # return self.model.embed(texts)
-        raise NotImplementedError
+        super().__init__(model_name, cache_dir, threads, **kwargs)

fastembed/image/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.image.image_embedding import ImageEmbedding
+
+__all__ = ["ImageEmbedding"]

fastembed/image/image_embedding.py

@@ -0,0 +1,141 @@
+from typing import Any, Iterable, Sequence, Type
+from dataclasses import asdict
+
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common import ImageInput, OnnxProvider
+from fastembed.image.image_embedding_base import ImageEmbeddingBase
+from fastembed.image.onnx_embedding import OnnxImageEmbedding
+from fastembed.image.normalized_embedding import NormalizedEmbedding
+from fastembed.image.siglip_embedding import SiglipOnnxImageEmbedding
+from fastembed.common.model_description import DenseModelDescription
+
+
+class ImageEmbedding(ImageEmbeddingBase):
+    EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [
+        OnnxImageEmbedding,
+        NormalizedEmbedding,
+        SiglipOnnxImageEmbedding,
+    ]
+
+    @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": "Qdrant/clip-ViT-B-32-vision",
+                        "dim": 512,
+                        "description": "CLIP vision encoder based on ViT-B/32",
+                        "license": "mit",
+                        "size_in_GB": 0.33,
+                        "sources": {
+                            "hf": "Qdrant/clip-ViT-B-32-vision",
+                        },
+                        "model_file": "model.onnx",
+                    }
+                ]
+                ```
+        """
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        result: list[DenseModelDescription] = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding._list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        **kwargs: Any,
+    ):
+        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=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **kwargs,
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in ImageEmbedding."
+            "Please check the supported models using `ImageEmbedding.list_supported_models()`"
+        )
+
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: int | None = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
+    def embed(
+        self,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 16,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of images into list of embeddings.
+
+        Args:
+            images: Iterator of image paths or single image path 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(images, batch_size, parallel, **kwargs)

fastembed/image/image_embedding_base.py

@@ -0,0 +1,55 @@
+from typing import Iterable, Any
+
+from fastembed.common.model_description import DenseModelDescription
+from fastembed.common.types import NumpyArray
+from fastembed.common.model_management import ModelManagement
+from fastembed.common.types import ImageInput
+
+
+class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        **kwargs: Any,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: int | None = None
+
+    def embed(
+        self,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 16,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Embeds a list of images into a list of embeddings.
+
+        Args:
+            images: The list of image paths to preprocess and embed.
+            batch_size: Batch size for encoding
+            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.
+            **kwargs: Additional keyword argument to pass to the embed method.
+
+        Yields:
+            Iterable[NdArray]: The embeddings.
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the chosen model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")

fastembed/image/normalized_embedding.py

@@ -0,0 +1,69 @@
+from typing import Any, Iterable, Type
+
+
+from fastembed.common.types import NumpyArray
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import normalize
+from fastembed.image.onnx_embedding import OnnxImageEmbedding
+from fastembed.image.onnx_image_model import ImageEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_normalized_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="nomic-ai/nomic-embed-vision-v1.5",
+        dim=768,
+        description="Image embeddings, Multimodal (text&image), 2024 year",
+        license="apache-2.0",
+        size_in_GB=0.37,
+        sources=ModelSource(hf="nomic-ai/nomic-embed-vision-v1.5"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="nomic-ai/nomic-embed-vision-v1.5-Q",
+        dim=768,
+        description="Image embeddings, Multimodal (text&image), 2024 year",
+        license="apache-2.0",
+        size_in_GB=0.1,
+        sources=ModelSource(hf="nomic-ai/nomic-embed-vision-v1.5"),
+        model_file="onnx/model_quantized.onnx",
+    ),
+]
+
+
+class NormalizedEmbedding(OnnxImageEmbedding):
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """
+        Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_normalized_models
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[NumpyArray]"]:
+        return NormalizedEmbeddingWorker
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        # The model emits last_hidden_state, which onnx_embed flattens to (batch, tokens * dim).
+        # Recover the token axis, take the CLS token (index 0) and normalize, matching the reference
+        # F.normalize(last_hidden_state[:, 0], p=2, dim=1).
+        dim = self.model_description.dim
+        assert dim is not None, "Model description is missing the embedding dim"
+        hidden_states = output.model_output.reshape(output.model_output.shape[0], -1, dim)
+        return normalize(hidden_states[:, 0])
+
+
+class NormalizedEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs: Any
+    ) -> NormalizedEmbedding:
+        return NormalizedEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/image/onnx_embedding.py

@@ -0,0 +1,216 @@
+from typing import Any, Iterable, Sequence, Type
+
+
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common import ImageInput, OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import define_cache_dir, normalize
+from fastembed.image.image_embedding_base import ImageEmbeddingBase
+from fastembed.image.onnx_image_model import ImageEmbeddingWorker, OnnxImageModel
+
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_onnx_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="Qdrant/clip-ViT-B-32-vision",
+        dim=512,
+        description="Image embeddings, Multimodal (text&image), 2021 year",
+        license="mit",
+        size_in_GB=0.34,
+        sources=ModelSource(hf="Qdrant/clip-ViT-B-32-vision"),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="Qdrant/resnet50-onnx",
+        dim=2048,
+        description="Image embeddings, Unimodal (image), 2016 year",
+        license="apache-2.0",
+        size_in_GB=0.1,
+        sources=ModelSource(hf="Qdrant/resnet50-onnx"),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="Qdrant/Unicom-ViT-B-16",
+        dim=768,
+        description="Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
+        license="apache-2.0",
+        size_in_GB=0.82,
+        sources=ModelSource(hf="Qdrant/Unicom-ViT-B-16"),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="Qdrant/Unicom-ViT-B-32",
+        dim=512,
+        description="Image embeddings, Multimodal (text&image), 2023 year",
+        license="apache-2.0",
+        size_in_GB=0.48,
+        sources=ModelSource(hf="Qdrant/Unicom-ViT-B-32"),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="jinaai/jina-clip-v1",
+        dim=768,
+        description="Image embeddings, Multimodal (text&image), 2024 year",
+        license="apache-2.0",
+        size_in_GB=0.34,
+        sources=ModelSource(hf="jinaai/jina-clip-v1"),
+        model_file="onnx/vision_model.onnx",
+    ),
+]
+
+
+class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: int | None = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        """
+        Load the onnx model.
+        """
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """
+        Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_onnx_models
+
+    def embed(
+        self,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 16,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of images into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+
+        Args:
+            images: Iterator of image paths or single image path 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_images(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            images=images,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[NumpyArray]"]:
+        return OnnxImageEmbeddingWorker
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Preprocess the onnx input.
+        """
+
+        return onnx_input
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        return normalize(output.model_output)
+
+
+class OnnxImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> OnnxImageEmbedding:
+        return OnnxImageEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/image/onnx_image_model.py

@@ -0,0 +1,158 @@
+import contextlib
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+from PIL import Image
+
+from fastembed.image.transform.operators import Compose
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common import ImageInput, OnnxProvider
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
+from fastembed.common.preprocessor_utils import load_preprocessor
+from fastembed.common.utils import iter_batch
+from fastembed.parallel_processor import ParallelWorkerPool
+
+# Holds type of the embedding result
+
+
+class OnnxImageModel(OnnxModel[T]):
+    ONNX_OUTPUT_NAMES: list[str] | None = None
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[T]: Post-processed output as an iterable of type T.
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.processor: Compose | None = None
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def _load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
+    ) -> None:
+        super()._load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_id=device_id,
+            extra_session_options=extra_session_options,
+        )
+        self.processor = load_preprocessor(model_dir=model_dir)
+
+    def load_onnx_model(self) -> None:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _build_onnx_input(self, encoded: NumpyArray) -> dict[str, NumpyArray]:
+        input_name = self.model.get_inputs()[0].name  # type: ignore[union-attr]
+        return {input_name: encoded}
+
+    def onnx_embed(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
+        with contextlib.ExitStack() as stack:
+            image_files = [
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
+                for image in images
+            ]
+            assert self.processor is not None, "Processor is not initialized"
+            encoded = np.array(self.processor(image_files))
+        onnx_input = self._build_onnx_input(encoded)
+        onnx_input = self._preprocess_onnx_input(onnx_input)
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        embeddings = model_output[0].reshape(len(images), -1)
+        return OnnxOutputContext(model_output=embeddings)
+
+    def _embed_images(
+        self,
+        model_name: str,
+        cache_dir: str,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        local_files_only: bool = False,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Iterable[T]:
+        is_small = False
+
+        if isinstance(images, (str, Path, Image.Image)):
+            images = [images]
+            is_small = True
+
+        if isinstance(images, list) and len(images) < batch_size:
+            is_small = True
+
+        if parallel is None or is_small:
+            if not hasattr(self, "model") or self.model is None:
+                self.load_onnx_model()
+
+            for batch in iter_batch(images, batch_size):
+                yield from self._post_process_onnx_output(self.onnx_embed(batch), **kwargs)
+        else:
+            if parallel == 0:
+                parallel = os.cpu_count()
+
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
+                **kwargs,
+            }
+
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
+            pool = ParallelWorkerPool(
+                num_workers=parallel or 1,
+                worker=self._get_worker_class(),
+                cuda=cuda,
+                device_ids=device_ids,
+                start_method=start_method,
+            )
+            for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
+                yield from self._post_process_onnx_output(batch, **kwargs)  # type: ignore
+
+
+class ImageEmbeddingWorker(EmbeddingWorker[T]):
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+        for idx, batch in items:
+            embeddings = self.model.onnx_embed(batch)
+            yield idx, embeddings

fastembed/image/siglip_embedding.py

@@ -0,0 +1,44 @@
+from typing import Any, Type
+
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.image.onnx_embedding import OnnxImageEmbedding, OnnxImageEmbeddingWorker
+
+supported_siglip_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="google/siglip2-base-patch16-224",
+        dim=768,
+        description="Image embeddings, Multimodal (text&image), 2025 year",
+        license="apache-2.0",
+        size_in_GB=0.37,
+        sources=ModelSource(hf="onnx-community/siglip2-base-patch16-224-ONNX"),
+        model_file="onnx/vision_model.onnx",
+    ),
+]
+
+
+class SiglipOnnxImageEmbedding(OnnxImageEmbedding):
+    """SigLIP vision tower.
+
+    The exported graph returns both the per-patch `last_hidden_state` and the pooled
+    `pooler_output`; only the latter is the image embedding, so it must be selected explicitly.
+    """
+
+    ONNX_OUTPUT_NAMES = ["pooler_output"]
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["OnnxImageEmbeddingWorker"]:
+        return SiglipImageEmbeddingWorker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        return supported_siglip_models
+
+
+class SiglipImageEmbeddingWorker(OnnxImageEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> OnnxImageEmbedding:
+        return SiglipOnnxImageEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/image/transform/functional.py

@@ -0,0 +1,234 @@
+import numpy as np
+from PIL import Image
+
+from fastembed.common.types import NumpyArray
+
+
+def convert_to_rgb(image: Image.Image) -> Image.Image:
+    if image.mode == "RGB":
+        return image
+
+    image = image.convert("RGB")
+    return image
+
+
+def center_crop(
+    image: Image.Image | NumpyArray,
+    size: tuple[int, int],
+) -> NumpyArray:
+    if isinstance(image, np.ndarray):
+        _, orig_height, orig_width = image.shape
+    else:
+        orig_height, orig_width = image.height, image.width
+        # (H, W, C) -> (C, H, W)
+        image = np.array(image).transpose((2, 0, 1))
+
+    crop_height, crop_width = size
+
+    # left upper corner (0, 0)
+    top = (orig_height - crop_height) // 2
+    bottom = top + crop_height
+    left = (orig_width - crop_width) // 2
+    right = left + crop_width
+
+    # Check if cropped area is within image boundaries
+    if top >= 0 and bottom <= orig_height and left >= 0 and right <= orig_width:
+        image = image[..., top:bottom, left:right]
+        return image
+
+    # Padding with zeros
+    new_height = max(crop_height, orig_height)
+    new_width = max(crop_width, orig_width)
+    new_shape = image.shape[:-2] + (new_height, new_width)
+    new_image = np.zeros_like(image, shape=new_shape, dtype=np.float32)
+
+    top_pad = (new_height - orig_height) // 2
+    bottom_pad = top_pad + orig_height
+    left_pad = (new_width - orig_width) // 2
+    right_pad = left_pad + orig_width
+    new_image[..., top_pad:bottom_pad, left_pad:right_pad] = image
+
+    top += top_pad
+    bottom += top_pad
+    left += left_pad
+    right += left_pad
+
+    new_image = new_image[
+        ..., max(0, top) : min(new_height, bottom), max(0, left) : min(new_width, right)
+    ]
+
+    return new_image
+
+
+def normalize(
+    image: NumpyArray,
+    mean: float | list[float],
+    std: float | list[float],
+) -> NumpyArray:
+    if image.ndim < 3:
+        raise ValueError(f"image must be (C, H, W) or (N, C, H, W), got shape {image.shape}")
+
+    # Channels sit on the third axis from the end, which covers (C, H, W) and
+    # (N, C, H, W) alike. Transposing instead reversed every axis, which put the
+    # batch dimension where the channels were meant to be.
+    num_channels = image.shape[-3]
+
+    if not np.issubdtype(image.dtype, np.floating):
+        image = image.astype(np.float32)
+
+    mean_list = mean if isinstance(mean, list) else [mean] * num_channels
+
+    if len(mean_list) != num_channels:
+        raise ValueError(
+            f"mean must have the same number of channels as the image, image has {num_channels} channels, got "
+            f"{len(mean_list)}"
+        )
+
+    # (C, 1, 1) lines the channels up with the trailing (C, H, W) axes under numpy
+    # broadcasting, whatever batch dimensions lead them.
+    mean_arr = np.array(mean_list, dtype=np.float32).reshape(-1, 1, 1)
+
+    std_list = std if isinstance(std, list) else [std] * num_channels
+    if len(std_list) != num_channels:
+        raise ValueError(
+            f"std must have the same number of channels as the image, image has {num_channels} channels, got {len(std_list)}"
+        )
+
+    std_arr = np.array(std_list, dtype=np.float32).reshape(-1, 1, 1)
+
+    image_upd = (image - mean_arr) / std_arr
+    return image_upd
+
+
+def resize(
+    image: Image.Image,
+    size: int | tuple[int, int],
+    resample: int | Image.Resampling = Image.Resampling.BILINEAR,
+) -> Image.Image:
+    if isinstance(size, tuple):
+        # fastembed keeps sizes as (height, width) — `Compose.from_config` builds the
+        # tuple as (size["height"], size["width"]) — while Pillow's resize takes
+        # (width, height). The two agree for square sizes, so this only shows up on a
+        # non-square image processor configuration.
+        height, width = size
+        return image.resize((width, height), resample)
+
+    height, width = image.height, image.width
+    short, long = (width, height) if width <= height else (height, width)
+
+    new_short, new_long = size, int(size * long / short)
+    if width <= height:
+        new_size = (new_short, new_long)
+    else:
+        new_size = (new_long, new_short)
+    return image.resize(new_size, resample)
+
+
+def rescale(image: NumpyArray, scale: float, dtype: type = np.float32) -> NumpyArray:
+    return (image * scale).astype(dtype)
+
+
+def pil2ndarray(image: Image.Image | NumpyArray) -> NumpyArray:
+    if isinstance(image, Image.Image):
+        return np.asarray(image).transpose((2, 0, 1))
+    return image
+
+
+def pad2square(
+    image: Image.Image,
+    size: int,
+    fill_color: str | int | tuple[int, ...] = 0,
+) -> Image.Image:
+    height, width = image.height, image.width
+
+    left, right = 0, width
+    top, bottom = 0, height
+
+    crop_required = False
+    if width > size:
+        left = (width - size) // 2
+        right = left + size
+        crop_required = True
+
+    if height > size:
+        top = (height - size) // 2
+        bottom = top + size
+        crop_required = True
+
+    new_image = Image.new(mode="RGB", size=(size, size), color=fill_color)
+    new_image.paste(image.crop((left, top, right, bottom)) if crop_required else image)
+    return new_image
+
+
+def resize_longest_edge(
+    image: Image.Image,
+    max_size: int,
+    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
+) -> Image.Image:
+    height, width = image.height, image.width
+    aspect_ratio = width / height
+
+    if width >= height:
+        # Width is longer
+        new_width = max_size
+        new_height = int(new_width / aspect_ratio)
+    else:
+        # Height is longer
+        new_height = max_size
+        new_width = int(new_height * aspect_ratio)
+
+    # Ensure even dimensions
+    if new_height % 2 != 0:
+        new_height += 1
+    if new_width % 2 != 0:
+        new_width += 1
+
+    return image.resize((new_width, new_height), resample)
+
+
+def crop_ndarray(
+    image: NumpyArray,
+    x1: int,
+    y1: int,
+    x2: int,
+    y2: int,
+    channel_first: bool = True,
+) -> NumpyArray:
+    if channel_first:
+        # (C, H, W) format
+        return image[:, y1:y2, x1:x2]
+    else:
+        # (H, W, C) format
+        return image[y1:y2, x1:x2, :]
+
+
+def resize_ndarray(
+    image: NumpyArray,
+    size: tuple[int, int],
+    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
+    channel_first: bool = True,
+) -> NumpyArray:
+    # Convert to PIL-friendly format (H, W, C)
+    if channel_first:
+        img_hwc = image.transpose((1, 2, 0))
+    else:
+        img_hwc = image
+
+    # Handle different dtypes
+    if img_hwc.dtype == np.float32 or img_hwc.dtype == np.float64:
+        # Assume normalized, scale to 0-255 for PIL
+        img_hwc_scaled = (img_hwc * 255).astype(np.uint8)
+        pil_img = Image.fromarray(img_hwc_scaled, mode="RGB")
+        resized = pil_img.resize(size, resample)
+        result = np.array(resized).astype(np.float32) / 255.0
+    else:
+        # uint8 or similar
+        pil_img = Image.fromarray(img_hwc.astype(np.uint8), mode="RGB")
+        resized = pil_img.resize(size, resample)
+        result = np.array(resized)
+
+    # Convert back to original format
+    if channel_first:
+        result = result.transpose((2, 0, 1))
+
+    return result

fastembed/image/transform/operators.py

@@ -0,0 +1,499 @@
+from typing import Any
+import math
+
+from PIL import Image
+
+from fastembed.common.types import NumpyArray
+from fastembed.image.transform.functional import (
+    center_crop,
+    convert_to_rgb,
+    crop_ndarray,
+    normalize,
+    pil2ndarray,
+    rescale,
+    resize,
+    resize_longest_edge,
+    resize_ndarray,
+    pad2square,
+)
+
+
+class Transform:
+    def __call__(self, images: list[Any]) -> list[Image.Image] | list[NumpyArray]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class ConvertToRGB(Transform):
+    def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
+        return [convert_to_rgb(image=image) for image in images]
+
+
+class CenterCrop(Transform):
+    def __init__(self, size: tuple[int, int]):
+        self.size = size
+
+    def __call__(self, images: list[Image.Image]) -> list[NumpyArray]:
+        return [center_crop(image=image, size=self.size) for image in images]
+
+
+class Normalize(Transform):
+    def __init__(self, mean: float | list[float], std: float | list[float]):
+        self.mean = mean
+        self.std = std
+
+    def __call__(  # type: ignore[override]
+        self, images: list[NumpyArray] | list[list[NumpyArray]]
+    ) -> list[NumpyArray] | list[list[NumpyArray]]:
+        if images and isinstance(images[0], list):
+            # Nested structure from ImageSplitter
+            return [
+                [normalize(image, mean=self.mean, std=self.std) for image in img_patches]  # type: ignore[arg-type]
+                for img_patches in images
+            ]
+        else:
+            # Flat structure (backward compatibility)
+            return [normalize(image, mean=self.mean, std=self.std) for image in images]  # type: ignore[arg-type]
+
+
+class Resize(Transform):
+    def __init__(
+        self,
+        size: int | tuple[int, int],
+        resample: Image.Resampling = Image.Resampling.BICUBIC,
+    ):
+        self.size = size
+        self.resample = resample
+
+    def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
+        return [resize(image, size=self.size, resample=self.resample) for image in images]
+
+
+class Rescale(Transform):
+    def __init__(self, scale: float = 1 / 255):
+        self.scale = scale
+
+    def __call__(  # type: ignore[override]
+        self, images: list[NumpyArray] | list[list[NumpyArray]]
+    ) -> list[NumpyArray] | list[list[NumpyArray]]:
+        if images and isinstance(images[0], list):
+            # Nested structure from ImageSplitter
+            return [
+                [rescale(image, scale=self.scale) for image in img_patches]  # type: ignore[arg-type]
+                for img_patches in images
+            ]
+        else:
+            # Flat structure (backward compatibility)
+            return [rescale(image, scale=self.scale) for image in images]  # type: ignore[arg-type]
+
+
+class PILtoNDarray(Transform):
+    def __call__(self, images: list[Image.Image | NumpyArray]) -> list[NumpyArray]:
+        return [pil2ndarray(image) for image in images]
+
+
+class PadtoSquare(Transform):
+    def __init__(
+        self,
+        size: int,
+        fill_color: str | int | tuple[int, ...],
+    ):
+        self.size = size
+        self.fill_color = fill_color
+
+    def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
+        return [
+            pad2square(image=image, size=self.size, fill_color=self.fill_color) for image in images
+        ]
+
+
+class ResizeLongestEdge(Transform):
+    """Resize images so the longest edge equals target size, preserving aspect ratio."""
+
+    def __init__(
+        self,
+        size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.size = size
+        self.resample = resample
+
+    def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
+        return [resize_longest_edge(image, self.size, self.resample) for image in images]
+
+
+class ResizeForVisionEncoder(Transform):
+    """
+    Resize both dimensions to be multiples of vision_encoder_max_size.
+    Preserves aspect ratio approximately.
+    Works on numpy arrays in (C, H, W) format.
+    """
+
+    def __init__(
+        self,
+        max_size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.max_size = max_size
+        self.resample = resample
+
+    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
+        result = []
+        for image in images:
+            # Assume (C, H, W) format
+            _, height, width = image.shape
+
+            aspect_ratio = width / height
+
+            if width >= height:
+                # Calculate new width as multiple of max_size
+                new_width = math.ceil(width / self.max_size) * self.max_size
+                new_height = int(new_width / aspect_ratio)
+                new_height = math.ceil(new_height / self.max_size) * self.max_size
+            else:
+                # Calculate new height as multiple of max_size
+                new_height = math.ceil(height / self.max_size) * self.max_size
+                new_width = int(new_height * aspect_ratio)
+                new_width = math.ceil(new_width / self.max_size) * self.max_size
+
+            # Resize using the ndarray resize function
+            resized = resize_ndarray(
+                image,
+                size=(new_width, new_height),  # PIL expects (width, height)
+                resample=self.resample,
+                channel_first=True,
+            )
+            result.append(resized)
+
+        return result
+
+
+class ImageSplitter(Transform):
+    """
+    Split images into grid of patches plus a global view.
+
+    If image dimensions exceed max_size:
+    - Divide into ceil(H/max_size) x ceil(W/max_size) patches
+    - Each patch is cropped from the image
+    - Add a global view (original resized to max_size x max_size)
+
+    If image is smaller than max_size:
+    - Return single image unchanged
+
+    Works on numpy arrays in (C, H, W) format.
+    """
+
+    def __init__(
+        self,
+        max_size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.max_size = max_size
+        self.resample = resample
+
+    def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]:  # type: ignore[override]
+        result = []
+
+        for image in images:
+            # Assume (C, H, W) format
+            _, height, width = image.shape
+            max_height = max_width = self.max_size
+
+            frames = []
+
+            if height > max_height or width > max_width:
+                # Calculate the number of splits needed
+                num_splits_h = math.ceil(height / max_height)
+                num_splits_w = math.ceil(width / max_width)
+
+                # Calculate optimal patch dimensions
+                optimal_height = math.ceil(height / num_splits_h)
+                optimal_width = math.ceil(width / num_splits_w)
+
+                # Generate patches in grid order (row by row)
+                for r in range(num_splits_h):
+                    for c in range(num_splits_w):
+                        # Calculate crop coordinates
+                        start_x = c * optimal_width
+                        start_y = r * optimal_height
+                        end_x = min(start_x + optimal_width, width)
+                        end_y = min(start_y + optimal_height, height)
+
+                        # Crop the patch
+                        cropped = crop_ndarray(
+                            image, x1=start_x, y1=start_y, x2=end_x, y2=end_y, channel_first=True
+                        )
+                        frames.append(cropped)
+
+                # Add global view (resized to max_size x max_size)
+                global_view = resize_ndarray(
+                    image,
+                    size=(max_width, max_height),  # PIL expects (width, height)
+                    resample=self.resample,
+                    channel_first=True,
+                )
+                frames.append(global_view)
+            else:
+                # Image is small enough, no splitting needed
+                frames.append(image)
+
+            # Append (not extend) to preserve per-image grouping
+            result.append(frames)
+
+        return result
+
+
+class SquareResize(Transform):
+    """
+    Resize images to square dimensions (max_size x max_size).
+    Works on numpy arrays in (C, H, W) format.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.size = size
+        self.resample = resample
+
+    def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]:  # type: ignore[override]
+        return [
+            [
+                resize_ndarray(
+                    image, size=(self.size, self.size), resample=self.resample, channel_first=True
+                )
+            ]
+            for image in images
+        ]
+
+
+class Compose:
+    def __init__(self, transforms: list[Transform]):
+        self.transforms = transforms
+
+    def __call__(
+        self, images: list[Image.Image] | list[NumpyArray]
+    ) -> list[NumpyArray] | list[Image.Image]:
+        for transform in self.transforms:
+            images = transform(images)
+        return images
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "Compose":
+        """Creates processor from a config dict.
+        Args:
+            config (dict[str, Any]): Configuration dictionary.
+
+                Valid keys:
+                    - do_resize
+                    - resize_mode
+                    - size
+                    - fill_color
+                    - do_center_crop
+                    - crop_size
+                    - do_rescale
+                    - rescale_factor
+                    - do_normalize
+                    - image_mean
+                    - mean
+                    - image_std
+                    - std
+                    - resample
+                    - interpolation
+                Valid size keys (nested):
+                    - {"height", "width"}
+                    - {"shortest_edge"}
+                    - {"longest_edge"}
+
+        Returns:
+            Compose: Image processor.
+        """
+        transforms: list[Transform] = []
+        cls._get_convert_to_rgb(transforms, config)
+        cls._get_resize(transforms, config)
+        cls._get_pad2square(transforms, config)
+        cls._get_center_crop(transforms, config)
+        cls._get_pil2ndarray(transforms, config)
+        cls._get_image_splitting(transforms, config)
+        cls._get_rescale(transforms, config)
+        cls._get_normalize(transforms, config)
+        return cls(transforms=transforms)
+
+    @staticmethod
+    def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]) -> None:
+        transforms.append(ConvertToRGB())
+
+    @classmethod
+    def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]) -> None:
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+        if mode in ("CLIPImageProcessor", "SiglipImageProcessor"):
+            if config.get("do_resize", False):
+                size = config["size"]
+                if "shortest_edge" in size:
+                    size = size["shortest_edge"]
+                elif "height" in size and "width" in size:
+                    size = (size["height"], size["width"])
+                else:
+                    raise ValueError(
+                        "Size must contain either 'shortest_edge' or 'height' and 'width'."
+                    )
+                transforms.append(
+                    Resize(
+                        size=size,
+                        resample=config.get("resample", Image.Resampling.BICUBIC),
+                    )
+                )
+        elif mode == "ConvNextFeatureExtractor":
+            if "size" in config and "shortest_edge" not in config["size"]:
+                raise ValueError(
+                    f"Size dictionary must contain 'shortest_edge' key. Got {config['size'].keys()}"
+                )
+            shortest_edge = config["size"]["shortest_edge"]
+            crop_pct = config.get("crop_pct", 0.875)
+            if shortest_edge < 384:
+                # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
+                resize_shortest_edge = int(shortest_edge / crop_pct)
+                transforms.append(
+                    Resize(
+                        size=resize_shortest_edge,
+                        resample=config.get("resample", Image.Resampling.BICUBIC),
+                    )
+                )
+                transforms.append(CenterCrop(size=(shortest_edge, shortest_edge)))
+            else:
+                transforms.append(
+                    Resize(
+                        size=(shortest_edge, shortest_edge),
+                        resample=config.get("resample", Image.Resampling.BICUBIC),
+                    )
+                )
+        elif mode == "JinaCLIPImageProcessor":
+            interpolation = config.get("interpolation")
+            if isinstance(interpolation, str):
+                resample = cls._interpolation_resolver(interpolation)
+            else:
+                resample = interpolation or Image.Resampling.BICUBIC
+
+            if "size" in config:
+                resize_mode = config.get("resize_mode", "shortest")
+                if resize_mode == "shortest":
+                    transforms.append(
+                        Resize(
+                            size=config["size"],
+                            resample=resample,
+                        )
+                    )
+        elif mode == "Idefics3ImageProcessor":
+            if config.get("do_resize", False):
+                size = config.get("size", {})
+                if "longest_edge" not in size:
+                    raise ValueError(
+                        "Size dictionary must contain 'longest_edge' key for Idefics3ImageProcessor"
+                    )
+
+                # Handle resample parameter - can be int enum or PIL.Image.Resampling
+                resample = config.get("resample", Image.Resampling.LANCZOS)
+                if isinstance(resample, int):
+                    resample = Image.Resampling(resample)
+
+                transforms.append(
+                    ResizeLongestEdge(
+                        size=size["longest_edge"],
+                        resample=resample,
+                    )
+                )
+        else:
+            raise ValueError(f"Preprocessor {mode} is not supported")
+
+    @staticmethod
+    def _get_center_crop(transforms: list[Transform], config: dict[str, Any]) -> None:
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+        if mode in ("CLIPImageProcessor", "SiglipImageProcessor"):
+            if config.get("do_center_crop", False):
+                crop_size_raw = config["crop_size"]
+                crop_size: tuple[int, int]
+                if isinstance(crop_size_raw, int):
+                    crop_size = (crop_size_raw, crop_size_raw)
+                elif isinstance(crop_size_raw, dict):
+                    crop_size = (crop_size_raw["height"], crop_size_raw["width"])
+                else:
+                    raise ValueError(f"Invalid crop size: {crop_size_raw}")
+                transforms.append(CenterCrop(size=crop_size))
+        elif mode == "ConvNextFeatureExtractor":
+            pass
+        elif mode == "JinaCLIPImageProcessor":
+            pass
+        elif mode == "Idefics3ImageProcessor":
+            pass
+        else:
+            raise ValueError(f"Preprocessor {mode} is not supported")
+
+    @staticmethod
+    def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]) -> None:
+        transforms.append(PILtoNDarray())
+
+    @classmethod
+    def _get_image_splitting(cls, transforms: list[Transform], config: dict[str, Any]) -> None:
+        """
+        Add image splitting transforms for Idefics3.
+        Handles conditional logic: splitting vs square resize.
+        Must be called AFTER PILtoNDarray.
+        """
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+
+        if mode == "Idefics3ImageProcessor":
+            do_splitting = config.get("do_image_splitting", False)
+            max_size = config.get("max_image_size", {}).get("longest_edge", 512)
+            resample = config.get("resample", Image.Resampling.LANCZOS)
+            if isinstance(resample, int):
+                resample = Image.Resampling(resample)
+
+            if do_splitting:
+                transforms.append(ResizeForVisionEncoder(max_size, resample))
+                transforms.append(ImageSplitter(max_size, resample))
+            else:
+                transforms.append(SquareResize(max_size, resample))
+
+    @staticmethod
+    def _get_rescale(transforms: list[Transform], config: dict[str, Any]) -> None:
+        if config.get("do_rescale", True):
+            rescale_factor = config.get("rescale_factor", 1 / 255)
+            transforms.append(Rescale(scale=rescale_factor))
+
+    @staticmethod
+    def _get_normalize(transforms: list[Transform], config: dict[str, Any]) -> None:
+        if config.get("do_normalize", False):
+            transforms.append(Normalize(mean=config["image_mean"], std=config["image_std"]))
+        elif "mean" in config and "std" in config:
+            transforms.append(Normalize(mean=config["mean"], std=config["std"]))
+
+    @staticmethod
+    def _get_pad2square(transforms: list[Transform], config: dict[str, Any]) -> None:
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+        if mode == "CLIPImageProcessor":
+            pass
+        elif mode == "ConvNextFeatureExtractor":
+            pass
+        elif mode == "JinaCLIPImageProcessor":
+            transforms.append(
+                PadtoSquare(
+                    size=config["size"],
+                    fill_color=config.get("fill_color", 0),
+                )
+            )
+
+    @staticmethod
+    def _interpolation_resolver(resample: str | None = None) -> Image.Resampling:
+        interpolation_map = {
+            "nearest": Image.Resampling.NEAREST,
+            "lanczos": Image.Resampling.LANCZOS,
+            "bilinear": Image.Resampling.BILINEAR,
+            "bicubic": Image.Resampling.BICUBIC,
+            "box": Image.Resampling.BOX,
+            "hamming": Image.Resampling.HAMMING,
+        }
+
+        if resample and (method := interpolation_map.get(resample.lower())):
+            return method
+
+        raise ValueError(f"Unknown interpolation method: {resample}")

fastembed/late_interaction/__init__.py

@@ -0,0 +1,5 @@
+from fastembed.late_interaction.late_interaction_text_embedding import (
+    LateInteractionTextEmbedding,
+)
+
+__all__ = ["LateInteractionTextEmbedding"]

fastembed/late_interaction/colbert.py

@@ -0,0 +1,301 @@
+import string
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+from tokenizers import Encoding, Tokenizer
+
+from fastembed.common.preprocessor_utils import load_tokenizer
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common import OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.late_interaction.late_interaction_embedding_base import (
+    LateInteractionTextEmbeddingBase,
+)
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_colbert_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="colbert-ir/colbertv2.0",
+        dim=128,
+        description="Text embeddings, Unimodal (text), English, 512 input tokens truncation, 2023 year",
+        license="mit",
+        size_in_GB=0.44,
+        sources=ModelSource(hf="colbert-ir/colbertv2.0"),
+        model_file="model.onnx",
+    ),
+    DenseModelDescription(
+        model="answerdotai/answerai-colbert-small-v1",
+        dim=96,
+        description="Text embeddings, Unimodal (text), English, 512 input tokens truncation, 2024 year",
+        license="apache-2.0",
+        size_in_GB=0.13,
+        sources=ModelSource(hf="answerdotai/answerai-colbert-small-v1"),
+        model_file="vespa_colbert.onnx",
+    ),
+]
+
+
+class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
+    QUERY_MARKER_TOKEN_ID = 1
+    DOCUMENT_MARKER_TOKEN_ID = 2
+    MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning
+    MASK_TOKEN = "[MASK]"
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, is_doc: bool = True, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        if not is_doc:
+            for embedding in output.model_output:
+                yield embedding
+        else:
+            if output.input_ids is None or output.attention_mask is None:
+                raise ValueError(
+                    "input_ids and attention_mask must be provided for document post-processing"
+                )
+
+            for i, token_sequence in enumerate(output.input_ids):
+                for j, token_id in enumerate(token_sequence):  # type: ignore
+                    if token_id in self.skip_list or token_id == self.pad_token_id:
+                        output.attention_mask[i, j] = 0
+
+            output.model_output *= np.expand_dims(output.attention_mask, 2)
+            norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
+            norm_clamped = np.maximum(norm, 1e-12)
+            output.model_output /= norm_clamped
+
+            for embedding, attention_mask in zip(output.model_output, output.attention_mask):
+                yield embedding[attention_mask == 1]
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        marker_token = self.DOCUMENT_MARKER_TOKEN_ID if is_doc else self.QUERY_MARKER_TOKEN_ID
+        onnx_input["input_ids"] = np.insert(
+            onnx_input["input_ids"].astype(np.int64), 1, marker_token, axis=1
+        )
+        onnx_input["attention_mask"] = np.insert(
+            onnx_input["attention_mask"].astype(np.int64), 1, 1, axis=1
+        )
+        return onnx_input
+
+    def tokenize(self, documents: list[str], is_doc: bool = True, **kwargs: Any) -> list[Encoding]:
+        return (
+            self._tokenize_documents(documents=documents)
+            if is_doc
+            else self._tokenize_query(query=next(iter(documents)))
+        )
+
+    def _tokenize_query(self, query: str) -> list[Encoding]:
+        assert self.query_tokenizer is not None
+        encoded = self.query_tokenizer.encode_batch([query])
+        return encoded
+
+    def _tokenize_documents(self, documents: list[str]) -> list[Encoding]:
+        encoded = self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
+        return encoded
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        is_doc: bool = True,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        tokenizer = self.tokenizer if is_doc else self.query_tokenizer
+        assert tokenizer is not None
+        for batch in iter_batch(texts, batch_size):
+            for tokens in tokenizer.encode_batch(batch):
+                if is_doc:
+                    token_num += sum(tokens.attention_mask)
+                else:
+                    attend_count = sum(tokens.attention_mask)
+                    if include_extension:
+                        token_num += max(attend_count, self.MIN_QUERY_LENGTH)
+
+                    else:
+                        token_num += attend_count
+            if include_extension:
+                token_num += len(
+                    batch
+                )  # add 1 for each cls.DOC_MARKER_TOKEN_ID or cls.QUERY_MARKER_TOKEN_ID
+
+        return token_num
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_colbert_models
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: int | None = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+        self.mask_token_id: int | None = None
+        self.pad_token_id: int | None = None
+        self.skip_list: set[int] = set()
+
+        self.query_tokenizer: Tokenizer | None = None
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+        self.query_tokenizer, _ = load_tokenizer(model_dir=self._model_dir)
+
+        assert self.tokenizer is not None
+        self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN]
+        self.pad_token_id = self.tokenizer.padding["pad_id"]
+        self.skip_list = {
+            self.tokenizer.encode(symbol, add_special_tokens=False).ids[0]
+            for symbol in string.punctuation
+        }
+        current_max_length = self.tokenizer.truncation["max_length"]
+        # ensure not to overflow after adding document-marker
+        self.tokenizer.enable_truncation(max_length=current_max_length - 1)
+        self.query_tokenizer.enable_truncation(max_length=current_max_length - 1)
+        self.query_tokenizer.enable_padding(
+            pad_token=self.MASK_TOKEN,
+            pad_id=self.mask_token_id,
+            length=self.MIN_QUERY_LENGTH,
+        )
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        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,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        if isinstance(query, str):
+            query = [query]
+
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()
+
+        for text in query:
+            yield from self._post_process_onnx_output(
+                self.onnx_embed([text], is_doc=False), is_doc=False
+            )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+        return ColbertEmbeddingWorker
+
+
+class ColbertEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Colbert:
+        return Colbert(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/late_interaction/jina_colbert.py

@@ -0,0 +1,58 @@
+from typing import Any, Type
+
+from fastembed.common.types import NumpyArray
+from fastembed.late_interaction.colbert import Colbert, ColbertEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_jina_colbert_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="jinaai/jina-colbert-v2",
+        dim=128,
+        description="New model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 year",
+        license="cc-by-nc-4.0",
+        size_in_GB=2.24,
+        sources=ModelSource(hf="jinaai/jina-colbert-v2"),
+        model_file="onnx/model.onnx",
+        additional_files=["onnx/model.onnx_data"],
+    )
+]
+
+
+class JinaColbert(Colbert):
+    QUERY_MARKER_TOKEN_ID = 250002
+    DOCUMENT_MARKER_TOKEN_ID = 250003
+    MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning
+    MASK_TOKEN = "<mask>"
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]:
+        return JinaColbertEmbeddingWorker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_jina_colbert_models
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        onnx_input = super()._preprocess_onnx_input(onnx_input, is_doc)
+
+        # the attention mask for jina-colbert-v2 is always 1 in queries
+        if not is_doc:
+            onnx_input["attention_mask"][:] = 1
+        return onnx_input
+
+
+class JinaColbertEmbeddingWorker(ColbertEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> JinaColbert:
+        return JinaColbert(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/late_interaction/late_interaction_embedding_base.py

@@ -0,0 +1,80 @@
+from typing import Iterable, Any
+
+from fastembed.common.model_description import DenseModelDescription
+from fastembed.common.types import NumpyArray
+from fastembed.common.model_management import ModelManagement
+
+
+class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        **kwargs: Any,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: int | None = None
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        raise NotImplementedError()
+
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        """
+        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[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: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        """
+        Embeds queries
+
+        Args:
+            query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
+
+        Returns:
+            Iterable[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)
+        else:
+            yield from self.embed(query, **kwargs)
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the chosen model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        **kwargs: Any,
+    ) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -0,0 +1,180 @@
+from typing import Any, Iterable, Sequence, Type
+from dataclasses import asdict
+
+from fastembed.common.model_description import DenseModelDescription
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common import OnnxProvider
+from fastembed.late_interaction.colbert import Colbert
+from fastembed.late_interaction.jina_colbert import JinaColbert
+from fastembed.late_interaction.late_interaction_embedding_base import (
+    LateInteractionTextEmbeddingBase,
+)
+
+
+class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
+    EMBEDDINGS_REGISTRY: list[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert]
+
+    @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": "colbert-ir/colbertv2.0",
+                        "dim": 128,
+                        "description": "Late interaction model",
+                        "license": "mit",
+                        "size_in_GB": 0.44,
+                        "sources": {
+                            "hf": "colbert-ir/colbertv2.0",
+                        },
+                        "model_file": "model.onnx",
+                    },
+                ]
+                ```
+        """
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        result: list[DenseModelDescription] = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding._list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        **kwargs: Any,
+    ):
+        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=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **kwargs,
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in LateInteractionTextEmbedding."
+            "Please check the supported models using `LateInteractionTextEmbedding.list_supported_models()`"
+        )
+
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: int | None = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        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)
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
+        """
+        Embeds queries
+
+        Args:
+            query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
+
+        Returns:
+            Iterable[NdArray]: The embeddings.
+        """
+
+        # This is model-specific, so that different models can have specialized implementations
+        yield from self.model.query_embed(query, **kwargs)
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        is_doc: bool = True,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        """Returns the number of tokens in the texts.
+
+        Args:
+            texts (str | Iterable[str]): The list of texts to embed.
+            batch_size (int): Batch size for encoding
+            is_doc (bool): Whether the texts are documents (disable embedding a query with include_mask=True).
+            include_extension (bool): Turn on to count DOC / QUERY marker tokens, and [MASK] token in query mode.
+
+        Returns:
+            int: Sum of number of tokens in the texts.
+        """
+        return self.model.token_count(
+            texts,
+            batch_size=batch_size,
+            is_doc=is_doc,
+            include_extension=include_extension,
+            **kwargs,
+        )

fastembed/late_interaction/token_embeddings.py

@@ -0,0 +1,83 @@
+from dataclasses import asdict
+from typing import Iterable, Any, Type
+
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray
+from fastembed.late_interaction.late_interaction_embedding_base import (
+    LateInteractionTextEmbeddingBase,
+)
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+
+supported_token_embeddings_models = [
+    DenseModelDescription(
+        model="jinaai/jina-embeddings-v2-small-en-tokens",
+        dim=512,
+        description="Text embeddings, Unimodal (text), English, 8192 input tokens truncation,"
+        " Prefixes for queries/documents: not necessary, 2023 year.",
+        license="apache-2.0",
+        size_in_GB=0.12,
+        sources=ModelSource(hf="xenova/jina-embeddings-v2-small-en"),
+        model_file="onnx/model.onnx",
+    ),
+]
+
+
+class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_token_embeddings_models
+
+    @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 [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+        return TokensEmbeddingWorker
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        # Size: (batch_size, sequence_length, hidden_size)
+        embeddings = output.model_output
+        # Size: (batch_size, sequence_length)
+        assert output.attention_mask is not None
+        masks = output.attention_mask
+
+        # For each document we only select those embeddings that are not masked out
+        for i in range(embeddings.shape[0]):
+            yield embeddings[i, masks[i] == 1]
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        yield from super().embed(documents, batch_size=batch_size, parallel=parallel, **kwargs)
+
+
+class TokensEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs: Any
+    ) -> TokenEmbeddingsModel:
+        return TokenEmbeddingsModel(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/late_interaction_multimodal/__init__.py

@@ -0,0 +1,5 @@
+from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding import (
+    LateInteractionMultimodalEmbedding,
+)
+
+__all__ = ["LateInteractionMultimodalEmbedding"]

fastembed/late_interaction_multimodal/colmodernvbert.py

@@ -0,0 +1,532 @@
+import contextlib
+from typing import Any, Iterable, Type, Optional, Sequence
+import json
+
+import numpy as np
+from tokenizers import Encoding
+from PIL import Image
+
+from fastembed.common import ImageInput
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
+    LateInteractionMultimodalEmbeddingBase,
+)
+from fastembed.late_interaction_multimodal.onnx_multimodal_model import (
+    OnnxMultimodalModel,
+    TextEmbeddingWorker,
+    ImageEmbeddingWorker,
+)
+
+supported_colmodernvbert_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="Qdrant/colmodernvbert",
+        dim=128,
+        description="The late-interaction version of ModernVBERT, CPU friendly, English, 2025.",
+        license="mit",
+        size_in_GB=1.0,
+        sources=ModelSource(hf="Qdrant/colmodernvbert"),
+        additional_files=["processor_config.json"],
+        model_file="model.onnx",
+    ),
+]
+
+
+class ColModernVBERT(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyArray]):
+    """
+    The ModernVBERT/colmodernvbert model implementation. This model uses
+    bidirectional attention, which proves to work better for retrieval.
+
+    See: https://huggingface.co/ModernVBERT/colmodernvbert
+    """
+
+    VISUAL_PROMPT_PREFIX = (
+        "<|begin_of_text|>User:<image>Describe the image.<end_of_utterance>\nAssistant:"
+    )
+    QUERY_AUGMENTATION_TOKEN = "<end_of_utterance>"
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to False.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: Optional[int] = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+        self.mask_token_id = None
+        self.pad_token_id = None
+        self.image_seq_len: Optional[int] = None
+        self.max_image_size: Optional[int] = None
+        self.image_size: Optional[int] = None
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_colmodernvbert_models
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+        # Load image processing configuration
+        processor_config_path = self._model_dir / "processor_config.json"
+        with open(processor_config_path) as f:
+            processor_config = json.load(f)
+            self.image_seq_len = processor_config.get("image_seq_len", 64)
+
+        preprocessor_config_path = self._model_dir / "preprocessor_config.json"
+        with open(preprocessor_config_path) as f:
+            preprocessor_config = json.load(f)
+            self.max_image_size = preprocessor_config.get("max_image_size", {}).get(
+                "longest_edge", 512
+            )
+
+        # Load model configuration
+        config_path = self._model_dir / "config.json"
+        with open(config_path) as f:
+            model_config = json.load(f)
+            vision_config = model_config.get("vision_config", {})
+            self.image_size = vision_config.get("image_size", 512)
+
+    def _preprocess_onnx_text_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        batch_size, seq_length = onnx_input["input_ids"].shape
+        empty_image_placeholder: NumpyArray = np.zeros(
+            (batch_size, seq_length, 3, self.image_size, self.image_size),
+            dtype=np.float32,  # type: ignore[type-var,arg-type,assignment]
+        )
+        onnx_input["pixel_values"] = empty_image_placeholder
+        return onnx_input
+
+    def _post_process_onnx_text_output(
+        self,
+        output: OnnxOutputContext,
+    ) -> Iterable[NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        return output.model_output
+
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        # Add query augmentation tokens (matching process_queries logic from colpali-engine)
+        augmented_queries = [doc + self.QUERY_AUGMENTATION_TOKEN * 10 for doc in documents]
+        encoded = self.tokenizer.encode_batch(augmented_queries)  # type: ignore[union-attr]
+        return encoded
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        assert self.tokenizer is not None
+        tokenize_func = self.tokenize if include_extension else self.tokenizer.encode_batch
+        for batch in iter_batch(texts, batch_size):
+            token_num += sum([sum(encoding.attention_mask) for encoding in tokenize_func(batch)])
+        return token_num
+
+    def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
+        with contextlib.ExitStack() as stack:
+            image_files = [
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
+                for image in images
+            ]
+            assert self.processor is not None, "Processor is not initialized"
+            processed = self.processor(image_files)
+            encoded, attention_mask, metadata = self._process_nested_patches(processed)  # type: ignore[arg-type]
+
+        onnx_input = {"pixel_values": encoded, "attention_mask": attention_mask}
+        onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
+        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
+
+        return OnnxOutputContext(
+            model_output=model_output[0],
+            attention_mask=attention_mask,  # type: ignore[arg-type]
+            metadata=metadata,
+        )
+
+    @staticmethod
+    def _process_nested_patches(
+        processed: list[list[NumpyArray]],
+    ) -> tuple[NumpyArray, NumpyArray, dict[str, Any]]:
+        """
+        Process nested image patches (from ImageSplitter).
+
+        Args:
+            processed: List of patch lists, one per image [[img1_patches], [img2_patches], ...]
+
+        Returns:
+            tuple: (encoded array, attention_mask, metadata)
+                - encoded: (batch_size, max_patches, C, H, W)
+                - attention_mask: (batch_size, max_patches) with 1 for real patches, 0 for padding
+                - metadata: Dict with 'patch_counts' key
+        """
+        patch_counts = [len(patches) for patches in processed]
+        max_patches = max(patch_counts)
+
+        # Get dimensions from first patch
+        channels, height, width = processed[0][0].shape
+        batch_size = len(processed)
+
+        # Create padded array
+        encoded = np.zeros(
+            (batch_size, max_patches, channels, height, width), dtype=processed[0][0].dtype
+        )
+
+        # Create attention mask (1 for real patches, 0 for padding)
+        attention_mask = np.zeros((batch_size, max_patches), dtype=np.int64)
+
+        # Fill in patches and attention mask
+        for i, patches in enumerate(processed):
+            for j, patch in enumerate(patches):
+                encoded[i, j] = patch
+                attention_mask[i, j] = 1
+
+        metadata = {"patch_counts": patch_counts}
+        return encoded, attention_mask, metadata  # type: ignore[return-value]
+
+    def _preprocess_onnx_image_input(
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Add text input placeholders for image data, following Idefics3 processing logic.
+
+        Constructs input_ids dynamically based on the actual number of image patches,
+        using the same token expansion logic as Idefics3Processor.
+
+        Args:
+            onnx_input: Dict with 'pixel_values' (batch, num_patches, C, H, W)
+                        and 'attention_mask' (batch, num_patches) indicating real patches
+            **kwargs: Additional arguments
+
+        Returns:
+            Updated onnx_input with 'input_ids' and updated 'attention_mask' for token sequence
+        """
+        # The attention_mask in onnx_input has a shape of (batch_size, num_patches),
+        # and should be used to create an attention mask matching the input_ids shape.
+        patch_attention_mask = onnx_input["attention_mask"]
+        pixel_values = onnx_input["pixel_values"]
+
+        batch_size = pixel_values.shape[0]
+        batch_input_ids = []
+
+        # Build input_ids for each image based on its actual patch count
+        for i in range(batch_size):
+            # Count real patches (non-padded) from attention mask
+            patch_count = int(np.sum(patch_attention_mask[i]))
+
+            # Compute rows/cols from patch count
+            rows, cols = self._compute_rows_cols_from_patches(patch_count)
+
+            # Build input_ids for this image
+            input_ids = self._build_input_ids_for_image(rows, cols)
+            batch_input_ids.append(input_ids)
+
+        # Pad sequences to max length in batch
+        max_len = max(len(ids) for ids in batch_input_ids)
+
+        # Get padding config from tokenizer
+        padding_direction = self.tokenizer.padding["direction"]  # type: ignore[index,union-attr]
+        pad_token_id = self.tokenizer.padding["pad_id"]  # type: ignore[index,union-attr]
+
+        # Initialize with pad token
+        padded_input_ids = np.full((batch_size, max_len), pad_token_id, dtype=np.int64)
+        attention_mask = np.zeros((batch_size, max_len), dtype=np.int64)
+
+        for i, input_ids in enumerate(batch_input_ids):
+            seq_len = len(input_ids)
+            if padding_direction == "left":
+                # Left padding: place tokens at the END of the array
+                start_idx = max_len - seq_len
+                padded_input_ids[i, start_idx:] = input_ids
+                attention_mask[i, start_idx:] = 1
+            else:
+                # Right padding: place tokens at the START of the array
+                padded_input_ids[i, :seq_len] = input_ids
+                attention_mask[i, :seq_len] = 1
+
+        onnx_input["input_ids"] = padded_input_ids
+        # Update attention_mask with token-level data
+        onnx_input["attention_mask"] = attention_mask
+        return onnx_input
+
+    @staticmethod
+    def _compute_rows_cols_from_patches(patch_count: int) -> tuple[int, int]:
+        if patch_count <= 1:
+            return 0, 0
+
+        # Subtract 1 for the global image
+        grid_patches = patch_count - 1
+
+        # Find rows and cols (assume square or near-square grid)
+        rows = int(grid_patches**0.5)
+        cols = grid_patches // rows
+
+        # Verify the calculation
+        if rows * cols + 1 != patch_count:
+            # Handle non-square grids
+            for r in range(1, grid_patches + 1):
+                if grid_patches % r == 0:
+                    c = grid_patches // r
+                    if r * c + 1 == patch_count:
+                        return r, c
+            # Fallback: treat as unsplit
+            return 0, 0
+
+        return rows, cols
+
+    def _create_single_image_prompt_string(self) -> str:
+        return (
+            "<fake_token_around_image>"
+            + "<global-img>"
+            + "<image>" * self.image_seq_len  # type: ignore[operator]
+            + "<fake_token_around_image>"
+        )
+
+    def _create_split_image_prompt_string(self, rows: int, cols: int) -> str:
+        text_split_images = ""
+
+        # Add tokens for each patch in the grid
+        for n_h in range(rows):
+            for n_w in range(cols):
+                text_split_images += (
+                    "<fake_token_around_image>"
+                    + f"<row_{n_h + 1}_col_{n_w + 1}>"
+                    + "<image>" * self.image_seq_len  # type: ignore[operator]
+                )
+            text_split_images += "\n"
+
+        # Add global image at the end
+        text_split_images += (
+            "\n<fake_token_around_image>"
+            + "<global-img>"
+            + "<image>" * self.image_seq_len  # type: ignore[operator]
+            + "<fake_token_around_image>"
+        )
+
+        return text_split_images
+
+    def _build_input_ids_for_image(self, rows: int, cols: int) -> np.ndarray:
+        # Create the appropriate image prompt string
+        if rows == 0 and cols == 0:
+            image_prompt_tokens = self._create_single_image_prompt_string()
+        else:
+            image_prompt_tokens = self._create_split_image_prompt_string(rows, cols)
+
+        # Replace <image> in visual prompt with expanded tokens
+        # The visual prompt is: "<|begin_of_text|>User:<image>Describe the image.<end_of_utterance>\nAssistant:"
+        expanded_prompt = self.VISUAL_PROMPT_PREFIX.replace("<image>", image_prompt_tokens)
+
+        # Tokenize the complete prompt
+        encoded = self.tokenizer.encode(expanded_prompt)  # type: ignore[union-attr]
+
+        # Convert to numpy array
+        return np.array(encoded.ids, dtype=np.int64)
+
+    def _post_process_onnx_image_output(
+        self,
+        output: OnnxOutputContext,
+    ) -> Iterable[NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        assert self.model_description.dim is not None, "Model dim is not defined"
+        return output.model_output.reshape(
+            output.model_output.shape[0], -1, self.model_description.dim
+        )
+
+    def embed_text(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of documents into list of embeddings.
+
+        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,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    def embed_image(
+        self,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of images into list of embeddings.
+
+        Args:
+            images: Iterator of image paths or single image path 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_images(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            images=images,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    @classmethod
+    def _get_text_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+        return ColModernVBERTTextEmbeddingWorker
+
+    @classmethod
+    def _get_image_worker_class(cls) -> Type[ImageEmbeddingWorker[NumpyArray]]:
+        return ColModernVBERTImageEmbeddingWorker
+
+
+class ColModernVBERTTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColModernVBERT:
+        return ColModernVBERT(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )
+
+
+class ColModernVBERTImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColModernVBERT:
+        return ColModernVBERT(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/late_interaction_multimodal/colpali.py

@@ -0,0 +1,327 @@
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+from tokenizers import Encoding
+
+from fastembed.common import OnnxProvider, ImageInput
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
+    LateInteractionMultimodalEmbeddingBase,
+)
+from fastembed.late_interaction_multimodal.onnx_multimodal_model import (
+    OnnxMultimodalModel,
+    TextEmbeddingWorker,
+    ImageEmbeddingWorker,
+)
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_colpali_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="Qdrant/colpali-v1.3-fp16",
+        dim=128,
+        description="Text embeddings, Multimodal (text&image), English, 50 tokens query length truncation, 2024.",
+        license="mit",
+        size_in_GB=6.5,
+        sources=ModelSource(hf="Qdrant/colpali-v1.3-fp16"),
+        additional_files=["model.onnx_data"],
+        model_file="model.onnx",
+    ),
+]
+
+
+class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyArray]):
+    QUERY_PREFIX = "Query: "
+    BOS_TOKEN = "<s>"
+    PAD_TOKEN = "<pad>"
+    QUERY_MARKER_TOKEN_ID = [2, 5098]
+    IMAGE_PLACEHOLDER_SIZE = (3, 448, 448)
+    EMPTY_TEXT_PLACEHOLDER = np.array(
+        [257152] * 1024 + [2, 50721, 573, 2416, 235265, 108]
+    )  # This is a tokenization of '<image>' * 1024 + '<bos>Describe the image.\n' line which is used as placeholder
+    # while processing an image
+    EVEN_ATTENTION_MASK = np.array([1] * 1030)
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: int | None = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+        self.mask_token_id = None
+        self.pad_token_id = None
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_colpali_models
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+    def _post_process_onnx_image_output(
+        self,
+        output: OnnxOutputContext,
+    ) -> Iterable[NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        assert self.model_description.dim is not None, "Model dim is not defined"
+        return output.model_output.reshape(
+            output.model_output.shape[0], -1, self.model_description.dim
+        )
+
+    def _post_process_onnx_text_output(
+        self,
+        output: OnnxOutputContext,
+    ) -> Iterable[NumpyArray]:
+        """
+        Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+
+        Returns:
+            Iterable[NumpyArray]: Post-processed output as NumPy arrays.
+        """
+        return output.model_output
+
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        texts_query: list[str] = []
+        for query in documents:
+            query = self.BOS_TOKEN + self.QUERY_PREFIX + query + self.PAD_TOKEN * 10
+            query += "\n"
+
+            texts_query.append(query)
+        encoded = self.tokenizer.encode_batch(texts_query)  # type: ignore[union-attr]
+        return encoded
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        assert self.tokenizer is not None
+        tokenize_func = self.tokenize if include_extension else self.tokenizer.encode_batch
+        for batch in iter_batch(texts, batch_size):
+            token_num += sum([sum(encoding.attention_mask) for encoding in tokenize_func(batch)])
+        return token_num
+
+    def _preprocess_onnx_text_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        onnx_input["input_ids"] = np.array(
+            [
+                self.QUERY_MARKER_TOKEN_ID + input_ids[2:].tolist()  # type: ignore[index]
+                for input_ids in onnx_input["input_ids"]
+            ]
+        )
+        empty_image_placeholder: NumpyArray = np.zeros(
+            self.IMAGE_PLACEHOLDER_SIZE, dtype=np.float32
+        )
+        onnx_input["pixel_values"] = np.array(
+            [empty_image_placeholder for _ in onnx_input["input_ids"]],
+        )
+        return onnx_input
+
+    def _preprocess_onnx_image_input(
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Add placeholders for text input when processing image data for ONNX.
+        Args:
+            onnx_input (Dict[str, NumpyArray]): Preprocessed image inputs.
+            **kwargs: Additional arguments.
+        Returns:
+            Dict[str, NumpyArray]: ONNX input with text placeholders.
+        """
+        onnx_input["input_ids"] = np.array(
+            [self.EMPTY_TEXT_PLACEHOLDER for _ in onnx_input["pixel_values"]]
+        )
+        onnx_input["attention_mask"] = np.array(
+            [self.EVEN_ATTENTION_MASK for _ in onnx_input["pixel_values"]]
+        )
+        return onnx_input
+
+    def embed_text(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of documents into list of embeddings.
+
+        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,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    def embed_image(
+        self,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 16,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of images into list of embeddings.
+
+        Args:
+            images: Iterator of image paths or single image path 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_images(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            images=images,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    @classmethod
+    def _get_text_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
+        return ColPaliTextEmbeddingWorker
+
+    @classmethod
+    def _get_image_worker_class(cls) -> Type[ImageEmbeddingWorker[NumpyArray]]:
+        return ColPaliImageEmbeddingWorker
+
+
+class ColPaliTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColPali:
+        return ColPali(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )
+
+
+class ColPaliImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> ColPali:
+        return ColPali(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

@@ -0,0 +1,189 @@
+from typing import Any, Iterable, Sequence, Type
+from dataclasses import asdict
+
+from fastembed.common import OnnxProvider, ImageInput
+from fastembed.common.types import NumpyArray, Device
+from fastembed.late_interaction_multimodal.colpali import ColPali
+from fastembed.late_interaction_multimodal.colmodernvbert import ColModernVBERT
+
+from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
+    LateInteractionMultimodalEmbeddingBase,
+)
+from fastembed.common.model_description import DenseModelDescription
+
+
+class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase):
+    EMBEDDINGS_REGISTRY: list[Type[LateInteractionMultimodalEmbeddingBase]] = [
+        ColPali,
+        ColModernVBERT,
+    ]
+
+    @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": "Qdrant/colpali-v1.3-fp16",
+                         "dim": 128,
+                         "description": "Text embeddings, Unimodal (text), Aligned to image latent space, ColBERT-compatible, 512 tokens max, 2024.",
+                         "license": "mit",
+                         "size_in_GB": 6.06,
+                         "sources": {
+                            "hf": "Qdrant/colpali-v1.3-fp16",
+                         },
+                         "additional_files": [
+                            "model.onnx_data",
+                         ],
+                        "model_file": "model.onnx",
+                    },
+                ]
+                ```
+        """
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        result: list[DenseModelDescription] = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding._list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        **kwargs: Any,
+    ):
+        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=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **kwargs,
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in LateInteractionMultimodalEmbedding."
+            "Please check the supported models using `LateInteractionMultimodalEmbedding.list_supported_models()`"
+        )
+
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: int | None = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
+    def embed_text(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of documents into list of embeddings.
+
+        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_text(documents, batch_size, parallel, **kwargs)
+
+    def embed_image(
+        self,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 16,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of images into list of embeddings.
+
+        Args:
+            images: Iterator of image paths or single image path 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 image
+        """
+        yield from self.model.embed_image(images, batch_size, parallel, **kwargs)
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 1024,
+        include_extension: bool = False,
+        **kwargs: Any,
+    ) -> int:
+        """Returns the number of tokens in the texts.
+
+        Args:
+            texts (str | Iterable[str]): The list of texts to embed.
+            batch_size (int): Batch size for encoding
+            include_extension (bool): Whether to include tokens added by preprocessing
+
+        Returns:
+            int: Sum of number of tokens in the texts.
+        """
+        return self.model.token_count(
+            texts, batch_size=batch_size, include_extension=include_extension, **kwargs
+        )

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -0,0 +1,86 @@
+from typing import Iterable, Any
+
+
+from fastembed.common import ImageInput
+from fastembed.common.model_description import DenseModelDescription
+from fastembed.common.model_management import ModelManagement
+from fastembed.common.types import NumpyArray
+
+
+class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescription]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        **kwargs: Any,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: int | None = None
+
+    def embed_text(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Embeds a list of documents into a list of embeddings.
+
+        Args:
+            documents (Iterable[str]): The list of texts 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.
+            **kwargs: Additional keyword argument to pass to the embed method.
+
+        Yields:
+            Iterable[NumpyArray]: The embeddings.
+        """
+        raise NotImplementedError()
+
+    def embed_image(
+        self,
+        images: ImageInput | Iterable[ImageInput],
+        batch_size: int = 16,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[NumpyArray]:
+        """
+        Encode a list of images into list of embeddings.
+        Args:
+            images: Iterator of image paths or single image path 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 image
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the chosen model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def token_count(
+        self,
+        texts: str | Iterable[str],
+        **kwargs: Any,
+    ) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -0,0 +1,291 @@
+import contextlib
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+from PIL import Image
+from tokenizers import Encoding, Tokenizer
+
+from fastembed.common import OnnxProvider, ImageInput
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
+from fastembed.common.preprocessor_utils import load_tokenizer, load_preprocessor
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common.utils import iter_batch
+from fastembed.image.transform.operators import Compose
+from fastembed.parallel_processor import ParallelWorkerPool
+
+
+class OnnxMultimodalModel(OnnxModel[T]):
+    ONNX_OUTPUT_NAMES: list[str] | None = None
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.tokenizer: Tokenizer | None = None
+        self.processor: Compose | None = None
+        self.special_token_to_id: dict[str, int] = {}
+
+    def _preprocess_onnx_text_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def _preprocess_onnx_image_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    @classmethod
+    def _get_text_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @classmethod
+    def _get_image_worker_class(cls) -> Type["ImageEmbeddingWorker[T]"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _post_process_onnx_image_output(self, output: OnnxOutputContext) -> Iterable[T]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _post_process_onnx_text_output(self, output: OnnxOutputContext) -> Iterable[T]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
+    ) -> None:
+        super()._load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_id=device_id,
+            extra_session_options=extra_session_options,
+        )
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+        assert self.tokenizer is not None
+        self.processor = load_preprocessor(model_dir=model_dir)
+
+    def load_onnx_model(self) -> None:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
+
+    def onnx_embed_text(
+        self,
+        documents: list[str],
+        **kwargs: Any,
+    ) -> OnnxOutputContext:
+        encoded = self.tokenize(documents, **kwargs)
+        input_ids = np.array([e.ids for e in encoded])
+        attention_mask = np.array([e.attention_mask for e in encoded])  # type: ignore[union-attr]
+        input_names = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
+        onnx_input: dict[str, NumpyArray] = {
+            "input_ids": np.array(input_ids, dtype=np.int64),
+        }
+        if "attention_mask" in input_names:
+            onnx_input["attention_mask"] = np.array(attention_mask, dtype=np.int64)
+        if "token_type_ids" in input_names:
+            onnx_input["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_text_input(onnx_input, **kwargs)
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        return OnnxOutputContext(
+            model_output=model_output[0],
+            attention_mask=onnx_input.get("attention_mask", attention_mask),
+            input_ids=onnx_input.get("input_ids", input_ids),
+        )
+
+    def _embed_documents(
+        self,
+        model_name: str,
+        cache_dir: str,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        local_files_only: bool = False,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> 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 is None or is_small:
+            if not hasattr(self, "model") or self.model is None:
+                self.load_onnx_model()
+            for batch in iter_batch(documents, batch_size):
+                yield from self._post_process_onnx_text_output(self.onnx_embed_text(batch))
+        else:
+            if parallel == 0:
+                parallel = os.cpu_count()
+
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
+                **kwargs,
+            }
+
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
+            pool = ParallelWorkerPool(
+                num_workers=parallel or 1,
+                worker=self._get_text_worker_class(),
+                cuda=cuda,
+                device_ids=device_ids,
+                start_method=start_method,
+            )
+            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
+                yield from self._post_process_onnx_text_output(batch)  # type: ignore
+
+    def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
+        with contextlib.ExitStack() as stack:
+            image_files = [
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
+                for image in images
+            ]
+            assert self.processor is not None, "Processor is not initialized"
+            encoded = np.array(self.processor(image_files))
+        onnx_input = {"pixel_values": encoded}
+        onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
+        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
+        embeddings = model_output[0].reshape(len(images), -1)
+        return OnnxOutputContext(model_output=embeddings)
+
+    def _embed_images(
+        self,
+        model_name: str,
+        cache_dir: str,
+        images: Iterable[ImageInput] | ImageInput,
+        batch_size: int = 256,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        local_files_only: bool = False,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Iterable[T]:
+        is_small = False
+
+        if isinstance(images, (str, Path, Image.Image)):
+            images = [images]
+            is_small = True
+
+        if isinstance(images, list) and len(images) < batch_size:
+            is_small = True
+
+        if parallel is None or is_small:
+            if not hasattr(self, "model") or self.model is None:
+                self.load_onnx_model()
+
+            for batch in iter_batch(images, batch_size):
+                yield from self._post_process_onnx_image_output(self.onnx_embed_image(batch))
+        else:
+            if parallel == 0:
+                parallel = os.cpu_count()
+
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
+                **kwargs,
+            }
+
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
+            pool = ParallelWorkerPool(
+                num_workers=parallel or 1,
+                worker=self._get_image_worker_class(),
+                cuda=cuda,
+                device_ids=device_ids,
+                start_method=start_method,
+            )
+            for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
+                yield from self._post_process_onnx_image_output(batch)  # type: ignore
+
+
+class TextEmbeddingWorker(EmbeddingWorker[T]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ):
+        self.model: OnnxMultimodalModel
+        super().__init__(model_name, cache_dir, **kwargs)
+
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxMultimodalModel:
+        raise NotImplementedError()
+
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+        for idx, batch in items:
+            onnx_output = self.model.onnx_embed_text(batch)
+            yield idx, onnx_output
+
+
+class ImageEmbeddingWorker(EmbeddingWorker[T]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ):
+        self.model: OnnxMultimodalModel
+        super().__init__(model_name, cache_dir, **kwargs)
+
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxMultimodalModel:
+        raise NotImplementedError()
+
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+        for idx, batch in items:
+            embeddings = self.model.onnx_embed_image(batch)
+            yield idx, embeddings

fastembed/parallel_processor.py

@@ -1,13 +1,16 @@
 import logging
 import os
 from collections import defaultdict
+from copy import deepcopy
 from enum import Enum
 from multiprocessing import Queue, get_context
 from multiprocessing.context import BaseContext
 from multiprocessing.process import BaseProcess
 from multiprocessing.sharedctypes import Synchronized as BaseValue
 from queue import Empty
-from typing import Any, Dict, Iterable, List, Optional, Type, Tuple
+from typing import Any, Iterable, Type
+
+from fastembed.common.types import Device
 
 # Single item should be processed in less than:
 processing_timeout = 10 * 60  # seconds
@@ -23,10 +26,10 @@ class QueueSignals(str, Enum):
 
 class Worker:
     @classmethod
-    def start(cls, **kwargs: Any) -> "Worker":
+    def start(cls, *args: Any, **kwargs: Any) -> "Worker":
         raise NotImplementedError()
 
-    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         raise NotImplementedError()
 
 
@@ -36,7 +39,7 @@ def _worker(
     output_queue: Queue,
     num_active_workers: BaseValue,
     worker_id: int,
-    kwargs: Optional[Dict[str, Any]] = None,
+    kwargs: dict[str, Any] | None = None,
 ) -> None:
     """
     A worker that pulls data pints off the input queue, and places the execution result on the output queue.
@@ -47,7 +50,9 @@ def _worker(
     if kwargs is None:
         kwargs = {}
 
-    logging.info(f"Reader worker: {worker_id} PID: {os.getpid()}")
+    logging.info(
+        f"Reader worker: {worker_id} PID: {os.getpid()} Device: {kwargs.get('device_id', 'CPU')}"
+    )
     try:
         worker = worker_class.start(**kwargs)
 
@@ -73,7 +78,9 @@ def _worker(
         # See:
         # https://docs.python.org/3.6/library/multiprocessing.html?highlight=process#pipes-and-queues
         # https://docs.python.org/3.6/library/multiprocessing.html?highlight=process#programming-guidelines
+        input_queue.close()
         output_queue.close()
+        input_queue.join_thread()
         output_queue.join_thread()
 
         with num_active_workers.get_lock():
@@ -83,16 +90,25 @@ def _worker(
 
 
 class ParallelWorkerPool:
-    def __init__(self, num_workers: int, worker: Type[Worker], start_method: Optional[str] = None):
+    def __init__(
+        self,
+        num_workers: int,
+        worker: Type[Worker],
+        start_method: str | None = None,
+        device_ids: list[int] | None = None,
+        cuda: bool | Device = Device.AUTO,
+    ):
         self.worker_class = worker
         self.num_workers = num_workers
-        self.input_queue: Optional[Queue] = None
-        self.output_queue: Optional[Queue] = None
+        self.input_queue: Queue | None = None
+        self.output_queue: Queue | None = None
         self.ctx: BaseContext = get_context(start_method)
-        self.processes: List[BaseProcess] = []
+        self.processes: list[BaseProcess] = []
         self.queue_size = self.num_workers * max_internal_batch_size
-
-        self.num_active_workers: Optional[BaseValue] = None
+        self.emergency_shutdown = False
+        self.device_ids = device_ids
+        self.cuda = cuda
+        self.num_active_workers: BaseValue | None = None
 
     def start(self, **kwargs: Any) -> None:
         self.input_queue = self.ctx.Queue(self.queue_size)
@@ -103,6 +119,12 @@ class ParallelWorkerPool:
         self.num_active_workers = ctx_value
 
         for worker_id in range(0, self.num_workers):
+            worker_kwargs = deepcopy(kwargs)
+            if self.device_ids:
+                device_id = self.device_ids[worker_id % len(self.device_ids)]
+                worker_kwargs["device_id"] = device_id
+                worker_kwargs["cuda"] = self.cuda
+
             assert hasattr(self.ctx, "Process")
             process = self.ctx.Process(
                 target=_worker,
@@ -112,14 +134,14 @@ class ParallelWorkerPool:
                     self.output_queue,
                     self.num_active_workers,
                     worker_id,
-                    kwargs.copy(),
+                    worker_kwargs,
                 ),
             )
             process.start()
             self.processes.append(process)
 
     def ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Any]:
-        buffer = defaultdict(Any)
+        buffer: defaultdict[int, Any] = defaultdict(Any)  # type: ignore
         next_expected = 0
 
         for idx, item in self.semi_ordered_map(stream, *args, **kwargs):
@@ -128,7 +150,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)
 
@@ -138,6 +162,7 @@ class ParallelWorkerPool:
             pushed = 0
             read = 0
             for idx, item in enumerate(stream):
+                self.check_worker_health()
                 if pushed - read < self.queue_size:
                     try:
                         out_item = self.output_queue.get_nowait()
@@ -164,6 +189,7 @@ class ParallelWorkerPool:
                 self.input_queue.put(QueueSignals.stop)
 
             while read < pushed:
+                self.check_worker_health()
                 out_item = self.output_queue.get(timeout=processing_timeout)
                 if out_item == QueueSignals.error:
                     self.join_or_terminate()
@@ -173,10 +199,29 @@ class ParallelWorkerPool:
         finally:
             assert self.input_queue is not None, "Input queue is None"
             assert self.output_queue is not None, "Output queue is None"
+            self.join()
             self.input_queue.close()
             self.output_queue.close()
+            if self.emergency_shutdown:
+                self.input_queue.cancel_join_thread()
+                self.output_queue.cancel_join_thread()
+            else:
+                self.input_queue.join_thread()
+                self.output_queue.join_thread()
 
-    def join_or_terminate(self, timeout: Optional[int] = 1) -> None:
+    def check_worker_health(self) -> None:
+        """
+        Checks if any worker process has terminated unexpectedly
+        """
+        for process in self.processes:
+            if not process.is_alive() and process.exitcode != 0:
+                self.emergency_shutdown = True
+                self.join_or_terminate()
+                raise RuntimeError(
+                    f"Worker PID: {process.pid} terminated unexpectedly with code {process.exitcode}"
+                )
+
+    def join_or_terminate(self, timeout: int = 1) -> None:
         """
         Emergency shutdown
         @param timeout:
@@ -204,4 +249,5 @@ class ParallelWorkerPool:
         https://eli.thegreenplace.net/2009/06/12/safely-using-destructors-in-python/.
         """
         for process in self.processes:
-            process.terminate()
+            if process.is_alive():
+                process.terminate()

fastembed/postprocess/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.postprocess.muvera import Muvera
+
+__all__ = ["Muvera"]

fastembed/postprocess/muvera.py

@@ -0,0 +1,362 @@
+import numpy as np
+
+from fastembed.common.types import NumpyArray
+from fastembed.late_interaction.late_interaction_embedding_base import (
+    LateInteractionTextEmbeddingBase,
+)
+from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
+    LateInteractionMultimodalEmbeddingBase,
+)
+
+
+MultiVectorModel = LateInteractionTextEmbeddingBase | LateInteractionMultimodalEmbeddingBase
+MAX_HAMMING_DISTANCE = 65  # 64 bits + 1
+POPCOUNT_LUT = np.array([bin(x).count("1") for x in range(256)], dtype=np.uint8)
+
+
+def hamming_distance_matrix(ids: np.ndarray) -> np.ndarray:
+    """Compute full Hamming distance matrix
+
+    Args:
+    ids: shape (n,) - array of ids, only size of the array matters
+
+    Return:
+        np.ndarray (n, n) - hamming distance matrix
+    """
+    n = len(ids)
+    xor_vals = np.bitwise_xor(ids[:, None], ids[None, :])  # (n, n) uint64
+    bytes_view = xor_vals.view(np.uint8).reshape(n, n, 8)  # (n, n, 8)
+    return POPCOUNT_LUT[bytes_view].sum(axis=2)
+
+
+class SimHashProjection:
+    """
+    SimHash projection component for MUVERA clustering.
+
+    This class implements locality-sensitive hashing using random hyperplanes
+    to partition the vector space into 2^k_sim clusters. Each vector is assigned
+    to a cluster based on which side of k_sim random hyperplanes it falls on.
+
+    Attributes:
+        k_sim (int): Number of SimHash functions (hyperplanes)
+        dim (int): Dimensionality of input vectors
+        simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (dim, k_sim)
+    """
+
+    def __init__(self, k_sim: int, dim: int, random_generator: np.random.Generator):
+        """
+        Initialize SimHash projection with random hyperplanes.
+
+        Args:
+            k_sim (int): Number of SimHash functions, determines 2^k_sim clusters
+            dim (int): Dimensionality of input vectors
+            random_generator (np.random.Generator): Random number generator for reproducibility
+        """
+        self.k_sim = k_sim
+        self.dim = dim
+        # Generate k_sim random hyperplanes (normal vectors) from standard normal distribution
+        self.simhash_vectors = random_generator.normal(size=(dim, k_sim))
+
+    def get_cluster_ids(self, vectors: np.ndarray) -> np.ndarray:
+        """
+        Compute the cluster IDs for a given vector using SimHash.
+
+        The cluster ID is determined by computing the dot product of the vector
+        with each hyperplane normal vector, taking the sign, and interpreting
+        the resulting binary string as an integer.
+
+        Args:
+            vectors (np.ndarray): Input vectors of shape (n, dim,)
+
+        Returns:
+            np.ndarray: Cluster IDs in range [0, 2^k_sim - 1]
+
+        Raises:
+            AssertionError: If a vector shape doesn't match expected dimensionality
+        """
+        dot_product = (
+            vectors @ self.simhash_vectors
+        )  # (token_num, dim) x (dim, k_sim) -> (token_num, k_sim)
+        cluster_ids = (dot_product > 0) @ (1 << np.arange(self.k_sim))
+        return cluster_ids
+
+
+class Muvera:
+    """
+    MUVERA (Multi-Vector Retrieval Architecture) algorithm implementation.
+
+    This class creates Fixed Dimensional Encodings (FDEs) from variable-length
+    sequences of vectors by using SimHash clustering and random projections.
+    The process involves:
+    1. Clustering vectors using multiple SimHash projections
+    2. Computing cluster centers (with different strategies for docs vs queries)
+    3. Applying random projections for dimensionality reduction
+    4. Concatenating results from all projections
+
+    Attributes:
+        k_sim (int): Number of SimHash functions per projection
+        dim (int): Input vector dimensionality
+        dim_proj (int): Output dimensionality after random projection
+        r_reps (int): Number of random projection repetitions
+        random_seed (int): Random seed for consistent random matrix generation
+        simhash_projections (List[SimHashProjection]): SimHash instances for clustering
+        dim_reduction_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj)
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        k_sim: int = 5,
+        dim_proj: int = 16,
+        r_reps: int = 20,
+        random_seed: int = 42,
+    ):
+        """
+        Initialize MUVERA algorithm with specified parameters.
+
+        Args:
+            dim (int): Dimensionality of individual input vectors
+            k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters).
+                                   Defaults to 5.
+            dim_proj (int, optional): Dimensionality after random projection (must be <= dim).
+                                    Defaults to 16.
+            r_reps (int, optional): Number of random projection repetitions for robustness.
+                                    Defaults to 20.
+            random_seed (int, optional): Seed for random number generator to ensure
+                                         reproducible results. Defaults to 42.
+
+        Raises:
+            ValueError: If dim_proj > dim (cannot project to higher dimensionality)
+        """
+        if dim_proj > dim:
+            raise ValueError(
+                f"Cannot project to a higher dimensionality (dim_proj={dim_proj} > dim={dim})"
+            )
+
+        self.k_sim = k_sim
+        self.dim = dim
+        self.dim_proj = dim_proj
+        self.r_reps = r_reps
+        # Create r_reps independent SimHash projections for robustness
+        generator = np.random.default_rng(random_seed)
+        self.simhash_projections = [
+            SimHashProjection(k_sim=self.k_sim, dim=self.dim, random_generator=generator)
+            for _ in range(r_reps)
+        ]
+        # Random projection matrices with entries from {-1, +1} for each repetition
+        self.dim_reduction_projections = generator.choice([-1, 1], size=(r_reps, dim, dim_proj))
+
+    @classmethod
+    def from_multivector_model(
+        cls,
+        model: MultiVectorModel,
+        k_sim: int = 5,
+        dim_proj: int = 16,
+        r_reps: int = 20,  # noqa[naming]
+        random_seed: int = 42,
+    ) -> "Muvera":
+        """
+        Create a Muvera instance from a multi-vector embedding model.
+
+        This class method provides a convenient way to initialize a MUVERA
+        that is compatible with a given multi-vector model by automatically extracting
+        the embedding dimensionality from the model.
+
+        Args:
+            model (MultiVectorModel): A late interaction text or multimodal embedding model
+                                    that provides multi-vector embeddings. Must have an
+                                    `embedding_size` attribute specifying the dimensionality
+                                    of individual vectors.
+            k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters).
+                                   Defaults to 5.
+            dim_proj (int, optional): Dimensionality after random projection (must be <= model's
+                                    embedding_size). Defaults to 16.
+            r_reps (int, optional): Number of random projection repetitions for robustness.
+                                    Defaults to 20.
+            random_seed (int, optional): Seed for random number generator to ensure
+                                         reproducible results. Defaults to 42.
+
+        Returns:
+            Muvera: A configured MUVERA instance ready to process embeddings from the given model.
+
+        Raises:
+            ValueError: If dim_proj > model.embedding_size (cannot project to higher dimensionality)
+
+        Example:
+            >>> from fastembed import LateInteractionTextEmbedding
+            >>> model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
+            >>> muvera = Muvera.from_multivector_model(
+            ...     model=model,
+            ...     k_sim=6,
+            ...     dim_proj=32
+            ... )
+            >>> # Now use postprocessor with embeddings from the model
+            >>> embeddings = np.array(list(model.embed(["sample text"])))
+            >>> fde = muvera.process_document(embeddings[0])
+        """
+        return cls(
+            dim=model.embedding_size,
+            k_sim=k_sim,
+            dim_proj=dim_proj,
+            r_reps=r_reps,
+            random_seed=random_seed,
+        )
+
+    def _get_output_dimension(self) -> int:
+        """
+        Get the output dimension of the MUVERA algorithm.
+
+        Returns:
+            int: Output dimension (r_reps * num_partitions * dim_proj) where b = 2^k_sim
+        """
+        num_partitions = 2**self.k_sim
+        return self.r_reps * num_partitions * self.dim_proj
+
+    @property
+    def embedding_size(self) -> int:
+        return self._get_output_dimension()
+
+    def process_document(self, vectors: NumpyArray) -> NumpyArray:
+        """
+        Encode a document's vectors into a Fixed Dimensional Encoding (FDE).
+
+        Uses document-specific settings: normalizes cluster centers by vector count
+        and fills empty clusters using Hamming distance-based selection.
+
+        Args:
+            vectors (NumpyArray): Document vectors of shape (n_tokens, dim)
+
+        Returns:
+            NumpyArray: Fixed dimensional encodings of shape (r_reps * b * dim_proj,)
+        """
+        return self.process(vectors, fill_empty_clusters=True, normalize_by_count=True)
+
+    def process_query(self, vectors: NumpyArray) -> NumpyArray:
+        """
+        Encode a query's vectors into a Fixed Dimensional Encoding (FDE).
+
+        Uses query-specific settings: no normalization by count and no empty
+        cluster filling to preserve query vector magnitudes.
+
+        Args:
+            vectors (NumpyArray]): Query vectors of shape (n_tokens, dim)
+
+        Returns:
+            NumpyArray: Fixed dimensional encoding of shape (r_reps * b * dim_proj,)
+        """
+        return self.process(vectors, fill_empty_clusters=False, normalize_by_count=False)
+
+    def process(
+        self,
+        vectors: NumpyArray,
+        fill_empty_clusters: bool = True,
+        normalize_by_count: bool = True,
+    ) -> NumpyArray:
+        """
+        Core encoding method that transforms variable-length vector sequences into FDEs.
+
+        The encoding process:
+        1. For each of r_reps random projections:
+           a. Assign vectors to clusters using SimHash
+           b. Compute cluster centers (sum of vectors in each cluster)
+           c. Optionally normalize by cluster size
+           d. Fill empty clusters using Hamming distance if requested
+           e. Apply random projection for dimensionality reduction
+           f. Flatten cluster centers into a vector
+        2. Concatenate all projection results
+
+        Args:
+            vectors (np.ndarray): Input vectors of shape (n_vectors, dim)
+            fill_empty_clusters (bool): Whether to fill empty clusters using nearest
+                                      vectors based on Hamming distance of cluster IDs
+            normalize_by_count (bool): Whether to normalize cluster centers by the
+                                     number of vectors assigned to each cluster
+
+        Returns:
+            np.ndarray: Fixed dimensional encoding of shape (r_reps * b * dim_proj)
+                        where B = 2^k_sim is the number of clusters
+
+        Raises:
+            AssertionError: If input vectors don't have expected dimensionality
+        """
+        assert (
+            vectors.shape[1] == self.dim
+        ), f"Expected vectors of shape (n, {self.dim}), got {vectors.shape}"
+
+        # Store results from each random projection
+        output_vectors = []
+
+        # num of space partitions in SimHash
+        num_partitions = 2**self.k_sim
+        cluster_center_ids = np.arange(num_partitions)
+        precomputed_hamming_matrix = (
+            hamming_distance_matrix(cluster_center_ids) if fill_empty_clusters else None
+        )
+
+        for projection_index, simhash in enumerate(self.simhash_projections):
+            # Initialize cluster centers and count vectors assigned to each cluster
+            cluster_centers = np.zeros((num_partitions, self.dim))
+            cluster_center_id_to_vectors: dict[int, list[int]] = {
+                cluster_center_id: [] for cluster_center_id in cluster_center_ids
+            }
+            cluster_vector_counts = None
+            empty_mask = None
+
+            # Assign each vector to its cluster and accumulate cluster centers
+            vector_cluster_ids = simhash.get_cluster_ids(vectors)
+            for cluster_id, (vec_idx, vec) in zip(vector_cluster_ids, enumerate(vectors)):
+                cluster_centers[cluster_id] += vec
+                cluster_center_id_to_vectors[cluster_id].append(vec_idx)
+
+            if normalize_by_count or fill_empty_clusters:
+                cluster_vector_counts = np.bincount(vector_cluster_ids, minlength=num_partitions)
+                empty_mask = cluster_vector_counts == 0
+
+            if normalize_by_count:
+                assert empty_mask is not None
+                assert cluster_vector_counts is not None
+                non_empty_mask = ~empty_mask
+                cluster_centers[non_empty_mask] /= cluster_vector_counts[non_empty_mask][:, None]
+
+            # Fill empty clusters using vectors with minimum Hamming distance
+            if fill_empty_clusters:
+                assert empty_mask is not None
+                assert precomputed_hamming_matrix is not None
+                masked_hamming = np.where(
+                    empty_mask[None, :], MAX_HAMMING_DISTANCE, precomputed_hamming_matrix
+                )
+                nearest_non_empty = np.argmin(masked_hamming, axis=1)
+                fill_vectors = np.array(
+                    [
+                        vectors[cluster_center_id_to_vectors[cluster_id][0]]
+                        for cluster_id in nearest_non_empty[empty_mask]
+                    ]
+                ).reshape(-1, self.dim)
+                cluster_centers[empty_mask] = fill_vectors
+
+            # Apply random projection for dimensionality reduction if needed
+            if self.dim_proj < self.dim:
+                dim_reduction_projection = self.dim_reduction_projections[
+                    projection_index
+                ]  # Get projection matrix for this repetition
+                projected_centers = (1 / np.sqrt(self.dim_proj)) * (
+                    cluster_centers @ dim_reduction_projection
+                )
+
+                # Flatten cluster centers into a single vector and add to output
+                output_vectors.append(projected_centers.flatten())
+                continue
+
+            # If no projection needed (dim_proj == dim), use original cluster centers
+            output_vectors.append(cluster_centers.flatten())
+
+        # Concatenate results from all R_reps projections into final FDE
+        return np.concatenate(output_vectors)
+
+
+if __name__ == "__main__":
+    v_arrs = np.random.randn(10, 100, 128)
+    muvera = Muvera(128, 4, 8, 20, 42)
+
+    for v_arr in v_arrs:
+        muvera.process(v_arr)  # type: ignore

fastembed/py.typed

@@ -0,0 +1 @@
+partial

fastembed/rerank/cross_encoder/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.rerank.cross_encoder.text_cross_encoder import TextCrossEncoder
+
+__all__ = ["TextCrossEncoder"]

fastembed/rerank/cross_encoder/custom_text_cross_encoder.py

@@ -0,0 +1,78 @@
+from typing import Sequence, Any, Type
+
+from fastembed.common import OnnxProvider
+from fastembed.common.model_description import BaseModelDescription
+from fastembed.common.types import Device
+from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+from fastembed.rerank.cross_encoder.onnx_text_model import TextRerankerWorker
+
+
+class CustomTextCrossEncoder(OnnxTextCrossEncoder):
+    SUPPORTED_MODELS: list[BaseModelDescription] = []
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_ids=device_ids,
+            lazy_load=lazy_load,
+            device_id=device_id,
+            specific_model_path=specific_model_path,
+            **kwargs,
+        )
+
+    @classmethod
+    def _list_supported_models(cls) -> list[BaseModelDescription]:
+        return cls.SUPPORTED_MODELS
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextRerankerWorker]:
+        return CustomTextCrossEncoderWorker
+
+    def _get_worker_init_kwargs(self) -> dict[str, Any]:
+        return {"model_description": self.model_description}
+
+    @classmethod
+    def add_model(
+        cls,
+        model_description: BaseModelDescription,
+    ) -> None:
+        cls.SUPPORTED_MODELS.append(model_description)
+
+
+class CustomTextCrossEncoderWorker(TextRerankerWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        model_description: BaseModelDescription | None = None,
+        **kwargs: Any,
+    ) -> CustomTextCrossEncoder:
+        if model_description is None:
+            raise ValueError(
+                "`model_description` is required to initialize a custom model in a worker "
+                "process, it is provided by `CustomTextCrossEncoder._get_worker_init_kwargs`"
+            )
+        # custom models live in a class-level registry, which spawned workers don't inherit
+        CustomTextCrossEncoder.add_model(model_description)
+        return CustomTextCrossEncoder(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -0,0 +1,239 @@
+from typing import Any, Iterable, Sequence, Type
+
+from loguru import logger
+
+from fastembed.common import OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
+from fastembed.common.utils import define_cache_dir
+from fastembed.rerank.cross_encoder.onnx_text_model import (
+    OnnxCrossEncoderModel,
+    TextRerankerWorker,
+)
+from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
+from fastembed.common.model_description import BaseModelDescription, ModelSource
+
+supported_onnx_models: list[BaseModelDescription] = [
+    BaseModelDescription(
+        model="Xenova/ms-marco-MiniLM-L-6-v2",
+        description="MiniLM-L-6-v2 model optimized for re-ranking tasks.",
+        license="apache-2.0",
+        size_in_GB=0.08,
+        sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-6-v2"),
+        model_file="onnx/model.onnx",
+    ),
+    BaseModelDescription(
+        model="Xenova/ms-marco-MiniLM-L-12-v2",
+        description="MiniLM-L-12-v2 model optimized for re-ranking tasks.",
+        license="apache-2.0",
+        size_in_GB=0.12,
+        sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-12-v2"),
+        model_file="onnx/model.onnx",
+    ),
+    BaseModelDescription(
+        model="BAAI/bge-reranker-base",
+        description="BGE reranker base model for cross-encoder re-ranking.",
+        license="mit",
+        size_in_GB=1.04,
+        sources=ModelSource(hf="BAAI/bge-reranker-base"),
+        model_file="onnx/model.onnx",
+    ),
+    BaseModelDescription(
+        model="jinaai/jina-reranker-v1-tiny-en",
+        description="Designed for blazing-fast re-ranking with 8K context length and fewer parameters than jina-reranker-v1-turbo-en.",
+        license="apache-2.0",
+        size_in_GB=0.13,
+        sources=ModelSource(hf="jinaai/jina-reranker-v1-tiny-en"),
+        model_file="onnx/model.onnx",
+    ),
+    BaseModelDescription(
+        model="jinaai/jina-reranker-v1-turbo-en",
+        description="Designed for blazing-fast re-ranking with 8K context length.",
+        license="apache-2.0",
+        size_in_GB=0.15,
+        sources=ModelSource(hf="jinaai/jina-reranker-v1-turbo-en"),
+        model_file="onnx/model.onnx",
+    ),
+    BaseModelDescription(
+        model="jinaai/jina-reranker-v2-base-multilingual",
+        description="A multi-lingual reranker model for cross-encoder re-ranking with 1K context length and sliding window",
+        license="cc-by-nc-4.0",
+        size_in_GB=1.11,
+        sources=ModelSource(hf="jinaai/jina-reranker-v2-base-multilingual"),
+        model_file="onnx/model.onnx",
+    ),
+]
+
+
+class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
+    @classmethod
+    def _list_supported_models(cls) -> list[BaseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[BaseModelDescription]: A list of BaseModelDescription objects containing the model information.
+        """
+        return supported_onnx_models
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. Xenova/ms-marco-MiniLM-L-6-v2.
+        """
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        if self.device_ids is not None and len(self.device_ids) > 1:
+            logger.warning(
+                "Parallel execution is currently not supported for cross encoders, "
+                f"only the first device will be used for inference: {self.device_ids[0]}."
+            )
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: int | None = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+    def rerank(
+        self,
+        query: str,
+        documents: Iterable[str],
+        batch_size: int = 64,
+        **kwargs: Any,
+    ) -> Iterable[float]:
+        """Reranks documents based on their relevance to a given query.
+
+        Args:
+            query (str): The query string to which document relevance is calculated.
+            documents (Iterable[str]): Iterable of documents to be reranked.
+            batch_size (int, optional): The number of documents processed in each batch. Higher batch sizes improve speed
+                                        but require more memory. Default is 64.
+        Returns:
+            Iterable[float]: An iterable of relevance scores for each document.
+        """
+
+        yield from self._rerank_documents(
+            query=query, documents=documents, batch_size=batch_size, **kwargs
+        )
+
+    def rerank_pairs(
+        self,
+        pairs: Iterable[tuple[str, str]],
+        batch_size: int = 64,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[float]:
+        yield from self._rerank_pairs(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            pairs=pairs,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextRerankerWorker]:
+        return TextCrossEncoderWorker
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[float]:
+        return (float(elem) for elem in output.model_output)
+
+    def token_count(
+        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        """Returns the number of tokens in the pairs.
+
+        Args:
+            pairs: Iterable of tuples, where each tuple contains a query and a document to be tokenized
+            batch_size: Batch size for tokenizing
+
+        Returns:
+            token count: overall number of tokens in the pairs
+        """
+        return self._token_count(pairs, batch_size=batch_size, **kwargs)
+
+
+class TextCrossEncoderWorker(TextRerankerWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxTextCrossEncoder:
+        return OnnxTextCrossEncoder(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -0,0 +1,205 @@
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+from tokenizers import Encoding
+
+from fastembed.common.onnx_model import (
+    EmbeddingWorker,
+    OnnxModel,
+    OnnxOutputContext,
+    OnnxProvider,
+)
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common.preprocessor_utils import load_tokenizer
+from fastembed.common.utils import iter_batch
+from fastembed.parallel_processor import ParallelWorkerPool
+
+
+class OnnxCrossEncoderModel(OnnxModel[float]):
+    ONNX_OUTPUT_NAMES: list[str] | None = None
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["TextRerankerWorker"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
+    ) -> None:
+        super()._load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_id=device_id,
+            extra_session_options=extra_session_options,
+        )
+        self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
+        assert self.tokenizer is not None
+
+    def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
+        return self.tokenizer.encode_batch(pairs)  # type: ignore[union-attr]
+
+    def _build_onnx_input(self, tokenized_input: list[Encoding]) -> dict[str, NumpyArray]:
+        input_names: set[str] = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
+        inputs: dict[str, NumpyArray] = {
+            "input_ids": np.array([enc.ids for enc in tokenized_input], dtype=np.int64),
+        }
+        if "token_type_ids" in input_names:
+            inputs["token_type_ids"] = np.array(
+                [enc.type_ids for enc in tokenized_input], dtype=np.int64
+            )
+        if "attention_mask" in input_names:
+            inputs["attention_mask"] = np.array(
+                [enc.attention_mask for enc in tokenized_input], dtype=np.int64
+            )
+        return inputs
+
+    def onnx_embed(self, query: str, documents: list[str], **kwargs: Any) -> OnnxOutputContext:
+        pairs = [(query, doc) for doc in documents]
+        return self.onnx_embed_pairs(pairs, **kwargs)
+
+    def onnx_embed_pairs(self, pairs: list[tuple[str, str]], **kwargs: Any) -> OnnxOutputContext:
+        tokenized_input = self.tokenize(pairs, **kwargs)
+        inputs = self._build_onnx_input(tokenized_input)
+        onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
+        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        relevant_output = outputs[0]
+        scores: NumpyArray = relevant_output[:, 0]
+        return OnnxOutputContext(model_output=scores)
+
+    def _rerank_documents(
+        self, query: str, documents: Iterable[str], batch_size: int, **kwargs: Any
+    ) -> Iterable[float]:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()
+        for batch in iter_batch(documents, batch_size):
+            yield from self._post_process_onnx_output(self.onnx_embed(query, batch, **kwargs))
+
+    def _rerank_pairs(
+        self,
+        model_name: str,
+        cache_dir: str,
+        pairs: Iterable[tuple[str, str]],
+        batch_size: int,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        local_files_only: bool = False,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Iterable[float]:
+        is_small = False
+
+        if isinstance(pairs, tuple):
+            pairs = [pairs]
+            is_small = True
+
+        if isinstance(pairs, list):
+            if len(pairs) < batch_size:
+                is_small = True
+
+        if parallel is None or is_small:
+            if not hasattr(self, "model") or self.model is None:
+                self.load_onnx_model()
+            for batch in iter_batch(pairs, batch_size):
+                yield from self._post_process_onnx_output(self.onnx_embed_pairs(batch, **kwargs))
+        else:
+            if parallel == 0:
+                parallel = os.cpu_count()
+
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
+                **kwargs,
+                **self._get_worker_init_kwargs(),
+            }
+
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
+            pool = ParallelWorkerPool(
+                num_workers=parallel or 1,
+                worker=self._get_worker_class(),
+                cuda=cuda,
+                device_ids=device_ids,
+                start_method=start_method,
+            )
+            for batch in pool.ordered_map(iter_batch(pairs, batch_size), **params):
+                yield from self._post_process_onnx_output(batch)  # type: ignore
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[float]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[float]: Post-processed output as an iterable of float values.
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def _token_count(
+        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **_: Any
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+
+        token_num = 0
+        assert self.tokenizer is not None
+        for batch in iter_batch(pairs, batch_size):
+            for tokens in self.tokenizer.encode_batch(batch):
+                token_num += sum(tokens.attention_mask)
+
+        return token_num
+
+
+class TextRerankerWorker(EmbeddingWorker[float]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ):
+        self.model: OnnxCrossEncoderModel
+        super().__init__(model_name, cache_dir, **kwargs)
+
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxCrossEncoderModel:
+        raise NotImplementedError()
+
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+        for idx, batch in items:
+            onnx_output = self.model.onnx_embed_pairs(batch)
+            yield idx, onnx_output

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -0,0 +1,178 @@
+from typing import Any, Iterable, Sequence, Type
+from dataclasses import asdict
+
+from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
+from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
+
+from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
+from fastembed.common.model_description import (
+    ModelSource,
+    BaseModelDescription,
+)
+
+
+class TextCrossEncoder(TextCrossEncoderBase):
+    CROSS_ENCODER_REGISTRY: list[Type[TextCrossEncoderBase]] = [
+        OnnxTextCrossEncoder,
+        CustomTextCrossEncoder,
+    ]
+
+    @classmethod
+    def list_supported_models(cls) -> list[dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            list[BaseModelDescription]: A list of dictionaries containing the model information.
+
+            Example:
+                ```
+                [
+                    {
+                        "model": "Xenova/ms-marco-MiniLM-L-6-v2",
+                        "size_in_GB": 0.08,
+                        "sources": {
+                            "hf": "Xenova/ms-marco-MiniLM-L-6-v2",
+                        },
+                        "model_file": "onnx/model.onnx",
+                        "description": "MiniLM-L-6-v2 model optimized for re-ranking tasks.",
+                        "license": "apache-2.0",
+                    }
+                ]
+                ```
+        """
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _list_supported_models(cls) -> list[BaseModelDescription]:
+        result: list[BaseModelDescription] = []
+        for encoder in cls.CROSS_ENCODER_REGISTRY:
+            result.extend(encoder._list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        for CROSS_ENCODER_TYPE in self.CROSS_ENCODER_REGISTRY:
+            supported_models = CROSS_ENCODER_TYPE._list_supported_models()
+            if any(model_name.lower() == model.model.lower() for model in supported_models):
+                self.model = CROSS_ENCODER_TYPE(
+                    model_name=model_name,
+                    cache_dir=cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **kwargs,
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in TextCrossEncoder."
+            "Please check the supported models using `TextCrossEncoder.list_supported_models()`"
+        )
+
+    def rerank(
+        self, query: str, documents: Iterable[str], batch_size: int = 64, **kwargs: Any
+    ) -> Iterable[float]:
+        """Rerank a list of documents based on a query.
+
+        Args:
+            query: Query to rerank the documents against
+            documents: Iterator of documents to rerank
+            batch_size: Batch size for reranking
+
+        Returns:
+            Iterable of scores for each document
+        """
+        yield from self.model.rerank(query, documents, batch_size=batch_size, **kwargs)
+
+    def rerank_pairs(
+        self,
+        pairs: Iterable[tuple[str, str]],
+        batch_size: int = 64,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[float]:
+        """
+        Rerank a list of query-document pairs.
+
+        Args:
+            pairs (Iterable[tuple[str, str]]): An iterable of tuples, where each tuple contains a query and a document
+                to be scored together.
+            batch_size (int, optional): The number of query-document pairs to process in a single batch. Defaults to 64.
+            parallel (Optional[int], optional): The number of parallel processes to use for reranking.
+                If None, parallelization is disabled. Defaults to None.
+            **kwargs (Any): Additional arguments to pass to the underlying reranking model.
+
+        Returns:
+            Iterable[float]: An iterable of scores corresponding to each query-document pair in the input.
+            Higher scores indicate a stronger match between the query and the document.
+
+        Example:
+            >>> encoder = TextCrossEncoder("Xenova/ms-marco-MiniLM-L-6-v2")
+            >>> pairs = [("What is AI?", "Artificial intelligence is ..."), ("What is ML?", "Machine learning is ...")]
+            >>> scores = list(encoder.rerank_pairs(pairs))
+            >>> print(list(map(lambda x: round(x, 2), scores)))
+            [-1.24, -10.6]
+        """
+        yield from self.model.rerank_pairs(
+            pairs, batch_size=batch_size, parallel=parallel, **kwargs
+        )
+
+    @classmethod
+    def add_custom_model(
+        cls,
+        model: str,
+        sources: ModelSource,
+        model_file: str = "onnx/model.onnx",
+        description: str = "",
+        license: str = "",
+        size_in_gb: float = 0.0,
+        additional_files: list[str] | None = None,
+    ) -> None:
+        registered_models = cls._list_supported_models()
+        for registered_model in registered_models:
+            if model == registered_model.model:
+                raise ValueError(
+                    f"Model {model} is already registered in CrossEncoderModel, if you still want to add this model, "
+                    f"please use another model name"
+                )
+
+        CustomTextCrossEncoder.add_model(
+            BaseModelDescription(
+                model=model,
+                sources=sources,
+                model_file=model_file,
+                description=description,
+                license=license,
+                size_in_GB=size_in_gb,
+                additional_files=additional_files or [],
+            )
+        )
+
+    def token_count(
+        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        """Returns the number of tokens in the pairs.
+
+        Args:
+            pairs: Iterable of tuples, where each tuple contains a query and a document to be tokenized
+            batch_size: Batch size for tokenizing
+
+        Returns:
+            token count: overall number of tokens in the pairs
+        """
+        return self.model.token_count(pairs, batch_size=batch_size, **kwargs)

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -0,0 +1,63 @@
+from typing import Any, Iterable
+
+from fastembed.common.model_description import BaseModelDescription
+from fastembed.common.model_management import ModelManagement
+
+
+class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        **kwargs: Any,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+
+    def rerank(
+        self,
+        query: str,
+        documents: Iterable[str],
+        batch_size: int = 64,
+        **kwargs: Any,
+    ) -> Iterable[float]:
+        """Rerank a list of documents given a query.
+
+        Args:
+            query (str): The query to rerank the documents.
+            documents (Iterable[str]): The list of texts to rerank.
+            batch_size (int): The batch size to use for reranking.
+            **kwargs: Additional keyword argument to pass to the rerank method.
+
+        Yields:
+            Iterable[float]: The scores of the reranked the documents.
+        """
+        raise NotImplementedError("This method should be overridden by subclasses")
+
+    def rerank_pairs(
+        self,
+        pairs: Iterable[tuple[str, str]],
+        batch_size: int = 64,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[float]:
+        """Rerank query-document pairs.
+        Args:
+            pairs (Iterable[tuple[str, str]]): Query-document pairs to rerank
+            batch_size (int): The batch size to use for reranking.
+            parallel: 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.
+            **kwargs: Additional keyword argument to pass to the rerank method.
+        Yields:
+            Iterable[float]: Scores for each individual pair
+        """
+        raise NotImplementedError("This method should be overridden by subclasses")
+
+    def token_count(self, pairs: Iterable[tuple[str, str]], **kwargs: Any) -> int:
+        """Returns the number of tokens in the pairs."""
+        raise NotImplementedError("This method should be overridden by subclasses")

fastembed/sparse/__init__.py

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

fastembed/sparse/bm25.py

@@ -0,0 +1,359 @@
+import os
+from collections import defaultdict
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Type
+
+import mmh3
+import numpy as np
+from py_rust_stemmers import SnowballStemmer
+from fastembed.common.utils import (
+    define_cache_dir,
+    iter_batch,
+    get_all_punctuation,
+    remove_non_alphanumeric,
+)
+from fastembed.parallel_processor import ParallelWorkerPool, Worker
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.sparse.utils.tokenizer import SimpleTokenizer
+from fastembed.common.model_description import SparseModelDescription, ModelSource
+
+
+supported_languages = [
+    "arabic",
+    "danish",
+    "dutch",
+    "english",
+    "finnish",
+    "french",
+    "german",
+    "greek",
+    "hungarian",
+    "italian",
+    "norwegian",
+    "portuguese",
+    "romanian",
+    "russian",
+    "spanish",
+    "swedish",
+    "tamil",
+    "turkish",
+]
+
+supported_bm25_models: list[SparseModelDescription] = [
+    SparseModelDescription(
+        model="Qdrant/bm25",
+        vocab_size=0,
+        description="BM25 as sparse embeddings meant to be used with Qdrant",
+        license="apache-2.0",
+        size_in_GB=0.01,
+        sources=ModelSource(hf="Qdrant/bm25"),
+        additional_files=[f"{lang}.txt" for lang in supported_languages],
+        requires_idf=True,
+        model_file="mock.file",
+    ),
+]
+
+
+class Bm25(SparseTextEmbeddingBase):
+    """Implements traditional BM25 in a form of sparse embeddings.
+    Uses a count of tokens in the document to evaluate the importance of the token.
+
+    WARNING: This model is expected to be used with `modifier="idf"` in the sparse vector index of Qdrant.
+
+    BM25 formula:
+
+    score(q, d) = SUM[ IDF(q_i) * (f(q_i, d) * (k + 1)) / (f(q_i, d) + k * (1 - b + b * (|d| / avg_len))) ],
+
+    where IDF is the inverse document frequency, computed on Qdrant's side
+    f(q_i, d) is the term frequency of the token q_i in the document d
+    k, b, avg_len are hyperparameters, described below.
+
+    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.
+        k (float, optional): The k parameter in the BM25 formula. Defines the saturation of the term frequency.
+            I.e. defines how fast the moment when additional terms stop to increase the score. Defaults to 1.2.
+        b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.
+            Defaults to 0.75.
+        avg_len (float, optional): The average length of the documents in the corpus. Defaults to 256.0.
+        language (str): Specifies the language for the stemmer.
+        disable_stemmer (bool): Disable the stemmer.
+    Raises:
+        ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        k: float = 1.2,
+        b: float = 0.75,
+        avg_len: float = 256.0,
+        language: str = "english",
+        token_max_length: int = 40,
+        disable_stemmer: bool = False,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name, cache_dir, **kwargs)
+
+        if language not in supported_languages:
+            raise ValueError(f"{language} language is not supported")
+        else:
+            self.language = language
+
+        self.k = k
+        self.b = b
+        self.avg_len = avg_len
+
+        model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+        self.token_max_length = token_max_length
+        self.punctuation = set(get_all_punctuation())
+        self.disable_stemmer = disable_stemmer
+
+        if disable_stemmer:
+            self.stopwords: set[str] = set()
+            self.stemmer = None
+        else:
+            self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
+            self.stemmer = SnowballStemmer(language)
+
+        self.tokenizer = SimpleTokenizer
+
+    @classmethod
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+        """
+        return supported_bm25_models
+
+    @classmethod
+    def _load_stopwords(cls, model_dir: Path, language: str) -> list[str]:
+        stopwords_path = model_dir / f"{language}.txt"
+        if not stopwords_path.exists():
+            return []
+
+        with open(stopwords_path, "r") as f:
+            return f.read().splitlines()
+
+    def _embed_documents(
+        self,
+        model_name: str,
+        cache_dir: str,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        local_files_only: bool = False,
+        specific_model_path: str | None = None,
+    ) -> Iterable[SparseEmbedding]:
+        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 is None or is_small:
+            for batch in iter_batch(documents, batch_size):
+                yield from self.raw_embed(batch)
+        else:
+            if parallel == 0:
+                parallel = os.cpu_count()
+
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "k": self.k,
+                "b": self.b,
+                "avg_len": self.avg_len,
+                "language": self.language,
+                "token_max_length": self.token_max_length,
+                "disable_stemmer": self.disable_stemmer,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
+            }
+            pool = ParallelWorkerPool(
+                num_workers=parallel or 1,
+                worker=self._get_worker_class(),
+                start_method=start_method,
+            )
+            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
+                for record in batch:
+                    yield record  # type: ignore
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> 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,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+    def _stem(self, tokens: list[str]) -> list[str]:
+        stemmed_tokens: list[str] = []
+        for token in tokens:
+            lower_token = token.lower()
+
+            if token in self.punctuation:
+                continue
+
+            if lower_token in self.stopwords:
+                continue
+
+            if len(token) > self.token_max_length:
+                continue
+
+            stemmed_token = self.stemmer.stem_word(lower_token) if self.stemmer else lower_token
+
+            if stemmed_token:
+                stemmed_tokens.append(stemmed_token)
+        return stemmed_tokens
+
+    def raw_embed(
+        self,
+        documents: list[str],
+    ) -> list[SparseEmbedding]:
+        embeddings: list[SparseEmbedding] = []
+        for document in documents:
+            document = remove_non_alphanumeric(document)
+            tokens = self.tokenizer.tokenize(document)
+            stemmed_tokens = self._stem(tokens)
+            token_id2value = self._term_frequency(stemmed_tokens)
+            embeddings.append(SparseEmbedding.from_dict(token_id2value))
+        return embeddings
+
+    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        for text in texts:
+            document = remove_non_alphanumeric(text)
+            tokens = self.tokenizer.tokenize(document)
+            token_num += len(tokens)
+        return token_num
+
+    def _term_frequency(self, tokens: list[str]) -> dict[int, float]:
+        """Calculate the term frequency part of the BM25 formula.
+
+        (
+            f(q_i, d) * (k + 1)
+        ) / (
+            f(q_i, d) + k * (1 - b + b * (|d| / avg_len))
+        )
+
+        Args:
+            tokens (list[str]): The list of tokens in the document.
+
+        Returns:
+            dict[int, float]: The token_id to term frequency mapping.
+        """
+        tf_map: dict[int, float] = {}
+        counter: defaultdict[str, int] = defaultdict(int)
+        for stemmed_token in tokens:
+            counter[stemmed_token] += 1
+
+        doc_len = len(tokens)
+        for stemmed_token in counter:
+            token_id = self.compute_token_id(stemmed_token)
+            num_occurrences = counter[stemmed_token]
+            tf_map[token_id] = num_occurrences * (self.k + 1)
+            tf_map[token_id] /= num_occurrences + self.k * (
+                1 - self.b + self.b * doc_len / self.avg_len
+            )
+        return tf_map
+
+    @classmethod
+    def compute_token_id(cls, token: str) -> int:
+        return abs(mmh3.hash(token))
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+        """To emulate BM25 behaviour, we don't need to use weights in the query, and
+        it's enough to just hash the tokens and assign a weight of 1.0 to them.
+        """
+        if isinstance(query, str):
+            query = [query]
+
+        for text in query:
+            text = remove_non_alphanumeric(text)
+            tokens = self.tokenizer.tokenize(text)
+            stemmed_tokens = self._stem(tokens)
+            token_ids = np.array(
+                list(set(self.compute_token_id(token) for token in stemmed_tokens)),
+                dtype=np.int32,
+            )
+            values = np.ones_like(token_ids)
+            yield SparseEmbedding(indices=token_ids, values=values)
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["Bm25Worker"]:
+        return Bm25Worker
+
+
+class Bm25Worker(Worker):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ):
+        self.model = self.init_embedding(model_name, cache_dir, **kwargs)
+
+    @classmethod
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "Bm25Worker":
+        return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
+
+    def process(
+        self, items: Iterable[tuple[int, Any]]
+    ) -> Iterable[tuple[int, list[SparseEmbedding]]]:
+        for idx, batch in items:
+            onnx_output = self.model.raw_embed(batch)
+            yield idx, onnx_output
+
+    @staticmethod
+    def init_embedding(model_name: str, cache_dir: str, **kwargs: Any) -> Bm25:
+        return Bm25(model_name=model_name, cache_dir=cache_dir, **kwargs)

fastembed/sparse/bm42.py

@@ -0,0 +1,369 @@
+import math
+import string
+from pathlib import Path
+from typing import Any, Iterable, Sequence, Type
+
+import mmh3
+import numpy as np
+from py_rust_stemmers import SnowballStemmer
+
+from fastembed.common import OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
+from fastembed.common.utils import define_cache_dir
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
+from fastembed.common.model_description import SparseModelDescription, ModelSource
+
+supported_bm42_models: list[SparseModelDescription] = [
+    SparseModelDescription(
+        model="Qdrant/bm42-all-minilm-l6-v2-attentions",
+        vocab_size=30522,
+        description="Light sparse embedding model, which assigns an importance score to each token in the text",
+        license="apache-2.0",
+        size_in_GB=0.09,
+        sources=ModelSource(hf="Qdrant/all_miniLM_L6_v2_with_attentions"),
+        model_file="model.onnx",
+        additional_files=["stopwords.txt"],
+        requires_idf=True,
+    ),
+]
+
+
+_MODEL_TO_LANGUAGE = {
+    "Qdrant/bm42-all-minilm-l6-v2-attentions": "english",
+}
+MODEL_TO_LANGUAGE = {
+    model_name.lower(): language for model_name, language in _MODEL_TO_LANGUAGE.items()
+}
+
+
+def get_language_by_model_name(model_name: str) -> str:
+    return MODEL_TO_LANGUAGE[model_name.lower()]
+
+
+class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
+    """
+    Bm42 is an extension of BM25, which tries to better evaluate importance of tokens in the documents,
+    by extracting attention weights from the transformer model.
+
+    Traditional BM25 uses a count of tokens in the document to evaluate the importance of the token,
+    but this approach doesn't work well with short documents or chunks of text, as almost all tokens
+    there are unique.
+
+    BM42 addresses this issue by replacing the token count with the attention weights from the transformer model.
+    This allows sparse embeddings to work well with short documents, handle rare tokens and leverage traditional NLP
+    techniques like stemming and stopwords.
+
+    WARNING: This model is expected to be used with `modifier="idf"` in the sparse vector index of Qdrant.
+    """
+
+    ONNX_OUTPUT_NAMES = ["attention_6"]
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        alpha: float = 0.5,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The providers to use for onnxruntime.
+            alpha (float, optional): Parameter, that defines the importance of the token weight in the document
+                versus the importance of the token frequency in the corpus. Defaults to 0.5, based on empirical testing.
+                It is recommended to only change this parameter based on training data for a specific dataset.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: int | None = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+        self.invert_vocab: dict[int, str] = {}
+
+        self.special_tokens: set[str] = set()
+        self.special_tokens_ids: set[int] = set()
+        self.punctuation = set(string.punctuation)
+        self.stopwords = set(self._load_stopwords(self._model_dir))
+        self.stemmer = SnowballStemmer(get_language_by_model_name(self.model_name))
+        self.alpha = alpha
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+        for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
+            self.invert_vocab[idx] = token
+        self.special_tokens = set(self.special_token_to_id.keys())
+        self.special_tokens_ids = set(self.special_token_to_id.values())
+        self.stopwords = set(self._load_stopwords(self._model_dir))
+
+    def _filter_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
+        result: list[tuple[str, Any]] = []
+        for token, value in tokens:
+            if token in self.stopwords or token in self.punctuation:
+                continue
+            result.append((token, value))
+        return result
+
+    def _stem_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
+        result: list[tuple[str, Any]] = []
+        for token, value in tokens:
+            processed_token = self.stemmer.stem_word(token)
+            result.append((processed_token, value))
+        return result
+
+    @classmethod
+    def _aggregate_weights(
+        cls, tokens: list[tuple[str, list[int]]], weights: list[float]
+    ) -> list[tuple[str, float]]:
+        result: list[tuple[str, float]] = []
+        for token, idxs in tokens:
+            sum_weight = sum(weights[idx] for idx in idxs)
+            result.append((token, sum_weight))
+        return result
+
+    def _reconstruct_bpe(
+        self, bpe_tokens: Iterable[tuple[int, str]]
+    ) -> list[tuple[str, list[int]]]:
+        result: list[tuple[str, list[int]]] = []
+        acc: str = ""
+        acc_idx: list[int] = []
+
+        continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix  # type: ignore[union-attr]
+        continuing_subword_prefix_len = len(continuing_subword_prefix)
+
+        for idx, token in bpe_tokens:
+            if token in self.special_tokens:
+                continue
+
+            if token.startswith(continuing_subword_prefix):
+                acc += token[continuing_subword_prefix_len:]
+                acc_idx.append(idx)
+            else:
+                if acc:
+                    result.append((acc, acc_idx))
+                    acc_idx = []
+                acc = token
+                acc_idx.append(idx)
+
+        if acc:
+            result.append((acc, acc_idx))
+
+        return result
+
+    def _rescore_vector(self, vector: dict[str, float]) -> dict[int, float]:
+        """
+        Orders all tokens in the vector by their importance and generates a new score based on the importance order.
+        So that the scoring doesn't depend on absolute values assigned by the model, but on the relative importance.
+        """
+
+        new_vector: dict[int, float] = {}
+
+        for token, value in vector.items():
+            token_id = abs(mmh3.hash(token))
+            # Examples:
+            # Num 0: Log(1/1 + 1) = 0.6931471805599453
+            # Num 1: Log(1/2 + 1) = 0.4054651081081644
+            # Num 2: Log(1/3 + 1) = 0.28768207245178085
+            new_vector[token_id] = math.log(1.0 + value) ** self.alpha  # value
+
+        return new_vector
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
+        if output.input_ids is None:
+            raise ValueError("input_ids must be provided for document post-processing")
+
+        token_ids_batch = output.input_ids.astype(int)
+
+        # attention_value shape: (batch_size, num_heads, num_tokens, num_tokens)
+        pooled_attention = np.mean(output.model_output[:, :, 0], axis=1) * output.attention_mask
+
+        for document_token_ids, attention_value in zip(token_ids_batch, pooled_attention):
+            document_tokens_with_ids = (
+                (idx, self.invert_vocab[token_id])
+                for idx, token_id in enumerate(document_token_ids)
+            )
+
+            reconstructed = self._reconstruct_bpe(document_tokens_with_ids)
+
+            filtered = self._filter_pair_tokens(reconstructed)
+
+            stemmed = self._stem_pair_tokens(filtered)
+
+            weighted = self._aggregate_weights(stemmed, attention_value)
+
+            max_token_weight: dict[str, float] = {}
+
+            for token, weight in weighted:
+                max_token_weight[token] = max(max_token_weight.get(token, 0), weight)
+
+            rescored = self._rescore_vector(max_token_weight)
+
+            yield SparseEmbedding.from_dict(rescored)
+
+    @classmethod
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+        """
+        return supported_bm42_models
+
+    @classmethod
+    def _load_stopwords(cls, model_dir: Path) -> list[str]:
+        stopwords_path = model_dir / "stopwords.txt"
+        if not stopwords_path.exists():
+            return []
+
+        with open(stopwords_path, "r") as f:
+            return f.read().splitlines()
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> 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,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            alpha=self.alpha,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+        )
+
+    @classmethod
+    def _query_rehash(cls, tokens: Iterable[str]) -> dict[int, float]:
+        result: dict[int, float] = {}
+        for token in tokens:
+            token_id = abs(mmh3.hash(token))
+            result[token_id] = 1.0
+        return result
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+        """
+        To emulate BM25 behaviour, we don't need to use smart weights in the query, and
+        it's enough to just hash the tokens and assign a weight of 1.0 to them.
+        It is also faster, as we don't need to run the model for the query.
+        """
+        if isinstance(query, str):
+            query = [query]
+
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()
+
+        for text in query:
+            encoded = self.tokenizer.encode(text)  # type: ignore[union-attr]
+            document_tokens_with_ids = enumerate(encoded.tokens)
+            reconstructed = self._reconstruct_bpe(document_tokens_with_ids)
+            filtered = self._filter_pair_tokens(reconstructed)
+            stemmed = self._stem_pair_tokens(filtered)
+
+            yield SparseEmbedding.from_dict(self._query_rehash(token for token, _ in stemmed))
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
+        return Bm42TextEmbeddingWorker
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+        return self._token_count(texts, batch_size=batch_size, **kwargs)
+
+
+class Bm42TextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:
+        return Bm42(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            **kwargs,
+        )

fastembed/sparse/if_splade.py

@@ -0,0 +1,246 @@
+import json
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+
+from fastembed.common import OnnxProvider
+from fastembed.common.model_description import ModelSource, SparseModelDescription
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.preprocessor_utils import load_tokenizer
+from fastembed.common.types import Device
+from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
+
+IDF_FILE = "idf.json"
+
+supported_if_splade_models: list[SparseModelDescription] = [
+    SparseModelDescription(
+        model="opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte",
+        vocab_size=30522,
+        description="Inference-free SPLADE model. Documents are expanded with an ONNX encoder at index "
+        "time, queries are encoded with a tokenizer and an IDF lookup table only, "
+        "without any model inference.",
+        license="apache-2.0",
+        size_in_GB=0.55,
+        sources=ModelSource(hf="Qdrant/opensearch-neural-sparse-encoding-doc-v3-gte"),
+        model_file="model.onnx",
+        additional_files=[IDF_FILE],
+        requires_idf=None,
+    ),
+]
+
+
+class IfSplade(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
+    """Inference-free (asymmetric) SPLADE model.
+
+    Documents are encoded with a neural encoder which expands them into a sparse vocabulary-sized
+    vector, while queries are encoded by tokenizing the text and looking up a precomputed IDF
+    weight per token — no neural inference happens at query time.
+
+    Query and document embeddings are compared with a dot product.
+    Special tokens are excluded from both document and query embeddings.
+    """
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
+        # Max-pool token logits over the sequence, masking out the padding
+        pooled = np.max(
+            output.model_output * np.expand_dims(output.attention_mask, axis=-1), axis=1
+        )
+        # v3 models of the opensearch-neural-sparse family use a double log activation,
+        # log(1 + log(1 + relu(x))), to increase sparsity of document embeddings
+        scores = np.log1p(np.log1p(np.maximum(pooled, 0.0)))
+
+        if self.special_tokens_ids:
+            scores[:, list(self.special_tokens_ids)] = 0.0
+
+        for row_scores in scores:
+            indices = row_scores.nonzero()[0]
+            yield SparseEmbedding(values=row_scores[indices], indices=indices)
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        # unlike `OnnxTextModel._token_count`, does not require the onnx model to be loaded
+        token_num = 0
+        texts = [texts] if isinstance(texts, str) else texts
+        for batch in iter_batch(texts, batch_size):
+            for tokens in self.tokenizer.encode_batch(batch):  # type: ignore[union-attr]
+                token_num += sum(tokens.attention_mask)
+        return token_num
+
+    @classmethod
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+        """
+        return supported_if_splade_models
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: int | None = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+        # The tokenizer and the idf table are lightweight and are required for query embedding,
+        # which does not involve any model inference, so they are loaded eagerly, while
+        # `lazy_load` only defers the initialization of the onnx model
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=self._model_dir)
+        self.special_tokens_ids: set[int] = set(self.special_token_to_id.values())
+        self._token_id_to_idf = self._load_idf()
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+    def _load_idf(self) -> dict[int, float]:
+        with open(self._model_dir / IDF_FILE) as f:
+            token_to_idf: dict[str, float] = json.load(f)
+
+        vocab: dict[str, int] = self.tokenizer.get_vocab()  # type: ignore[union-attr]
+        return {vocab[token]: idf for token, idf in token_to_idf.items() if token in vocab}
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[SparseEmbedding]:
+        """
+        Encode a list of documents into list of embeddings.
+
+        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,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+        """
+        Encode a list of queries into list of sparse embeddings without any model inference.
+
+        A query is tokenized, and each unique token is assigned its IDF weight from
+        a precomputed lookup table shipped with the model. Special tokens are ignored.
+        """
+        if isinstance(query, str):
+            query = [query]
+
+        for text in query:
+            token_ids = set(self.tokenizer.encode(text).ids) - self.special_tokens_ids  # type: ignore[union-attr]
+            embedding = {
+                token_id: self._token_id_to_idf[token_id]
+                for token_id in sorted(token_ids)
+                if token_id in self._token_id_to_idf
+            }
+            yield SparseEmbedding.from_dict(embedding)
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
+        return IfSpladeEmbeddingWorker
+
+
+class IfSpladeEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> IfSplade:
+        return IfSplade(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/sparse/minicoil.py

@@ -0,0 +1,372 @@
+from pathlib import Path
+
+from typing import Any, Sequence, Iterable, Type
+
+import numpy as np
+from numpy.typing import NDArray
+from py_rust_stemmers import SnowballStemmer
+from tokenizers import Tokenizer
+
+from fastembed.common.model_description import SparseModelDescription, ModelSource
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
+from fastembed.common.utils import define_cache_dir
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.sparse.utils.minicoil_encoder import Encoder
+from fastembed.sparse.utils.sparse_vectors_converter import SparseVectorConverter, WordEmbedding
+from fastembed.sparse.utils.vocab_resolver import VocabResolver, VocabTokenizer
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
+
+
+MINICOIL_MODEL_FILE = "minicoil.triplet.model.npy"
+MINICOIL_VOCAB_FILE = "minicoil.triplet.model.vocab"
+STOPWORDS_FILE = "stopwords.txt"
+
+
+supported_minicoil_models: list[SparseModelDescription] = [
+    SparseModelDescription(
+        model="Qdrant/minicoil-v1",
+        vocab_size=19125,
+        description="Sparse embedding model, that resolves semantic meaning of the words, "
+        "while keeping exact keyword match behavior. "
+        "Based on jinaai/jina-embeddings-v2-small-en-tokens",
+        license="apache-2.0",
+        size_in_GB=0.09,
+        sources=ModelSource(hf="Qdrant/minicoil-v1"),
+        model_file="onnx/model.onnx",
+        additional_files=[
+            STOPWORDS_FILE,
+            MINICOIL_MODEL_FILE,
+            MINICOIL_VOCAB_FILE,
+        ],
+        requires_idf=True,
+    ),
+]
+
+_MODEL_TO_LANGUAGE = {
+    "Qdrant/minicoil-v1": "english",
+}
+MODEL_TO_LANGUAGE = {
+    model_name.lower(): language for model_name, language in _MODEL_TO_LANGUAGE.items()
+}
+
+
+def get_language_by_model_name(model_name: str) -> str:
+    return MODEL_TO_LANGUAGE[model_name.lower()]
+
+
+class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
+    """
+        MiniCOIL is a sparse embedding model, that resolves semantic meaning of the words,
+        while keeping exact keyword match behavior.
+
+        Each vocabulary token is converted into 4d component of a sparse vector, which is then weighted by the token frequency in the corpus.
+        If the token is not found in the corpus, it is treated exactly like in BM25.
+    `
+        The model is based on `jinaai/jina-embeddings-v2-small-en-tokens`
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        k: float = 1.2,
+        b: float = 0.75,
+        avg_len: float = 150.0,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The providers to use for onnxruntime.
+            k (float, optional): The k parameter in the BM25 formula. Defines the saturation of the term frequency.
+                I.e. defines how fast the moment when additional terms stop to increase the score. Defaults to 1.2.
+            b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.
+                Defaults to 0.75.
+            avg_len (float, optional): The average length of the documents in the corpus. Defaults to 150.0.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self.device_ids = device_ids
+        self.cuda = cuda
+        self.device_id = device_id
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        self.k = k
+        self.b = b
+        self.avg_len = avg_len
+
+        # Initialize class attributes
+        self.tokenizer: Tokenizer | None = None
+        self.invert_vocab: dict[int, str] = {}
+        self.special_tokens: set[str] = set()
+        self.special_tokens_ids: set[int] = set()
+        self.stopwords: set[str] = set()
+        self.vocab_resolver: VocabResolver | None = None
+        self.encoder: Encoder | None = None
+        self.output_dim: int | None = None
+        self.sparse_vector_converter: SparseVectorConverter | None = None
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+        assert self.tokenizer is not None
+
+        for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
+            self.invert_vocab[idx] = token
+        self.special_tokens = set(self.special_token_to_id.keys())
+        self.special_tokens_ids = set(self.special_token_to_id.values())
+        self.stopwords = set(self._load_stopwords(self._model_dir))
+
+        stemmer = SnowballStemmer(get_language_by_model_name(self.model_name))
+
+        self.vocab_resolver = VocabResolver(
+            tokenizer=VocabTokenizer(self.tokenizer),
+            stopwords=self.stopwords,
+            stemmer=stemmer,
+        )
+        self.vocab_resolver.load_json_vocab(str(self._model_dir / MINICOIL_VOCAB_FILE))
+
+        weights = np.load(str(self._model_dir / MINICOIL_MODEL_FILE), mmap_mode="r")
+        self.encoder = Encoder(weights)
+        self.output_dim = self.encoder.output_dim
+
+        self.sparse_vector_converter = SparseVectorConverter(
+            stopwords=self.stopwords,
+            stemmer=stemmer,
+            k=self.k,
+            b=self.b,
+            avg_len=self.avg_len,
+        )
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        return self._token_count(texts, batch_size=batch_size, **kwargs)
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> 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,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            k=self.k,
+            b=self.b,
+            avg_len=self.avg_len,
+            is_query=False,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+        """
+        Encode a list of queries into list of embeddings.
+        """
+        yield from self._embed_documents(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            documents=query,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            k=self.k,
+            b=self.b,
+            avg_len=self.avg_len,
+            is_query=True,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            **kwargs,
+        )
+
+    @classmethod
+    def _load_stopwords(cls, model_dir: Path) -> list[str]:
+        stopwords_path = model_dir / STOPWORDS_FILE
+        if not stopwords_path.exists():
+            return []
+
+        with open(stopwords_path, "r") as f:
+            return f.read().splitlines()
+
+    @classmethod
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+        """
+        return supported_minicoil_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, is_query: bool = False, **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
+        if output.input_ids is None:
+            raise ValueError("input_ids must be provided for document post-processing")
+
+        assert self.vocab_resolver is not None
+        assert self.encoder is not None
+        assert self.sparse_vector_converter is not None
+
+        # Size: (batch_size, sequence_length, hidden_size)
+        embeddings = output.model_output
+        # Size: (batch_size, sequence_length)
+        assert output.attention_mask is not None
+        masks = output.attention_mask
+
+        vocab_size = self.vocab_resolver.vocab_size()
+        embedding_size = self.encoder.output_dim
+
+        # For each document we only select those embeddings that are not masked out
+
+        for i in range(embeddings.shape[0]):
+            # Size: (sequence_length, hidden_size)
+            token_embeddings = embeddings[i, masks[i] == 1]
+
+            # Size: (sequence_length)
+            token_ids: NDArray[np.int64] = output.input_ids[i, masks[i] == 1]
+
+            word_ids_array, counts, oov, forms = self.vocab_resolver.resolve_tokens(token_ids)
+
+            # Size: (1, words)
+            word_ids_array_expanded: NDArray[np.int64] = np.expand_dims(word_ids_array, axis=0)
+
+            # Size: (1, words, embedding_size)
+            token_embeddings_array: NDArray[np.float32] = np.expand_dims(token_embeddings, axis=0)
+
+            assert word_ids_array_expanded.shape[1] == token_embeddings_array.shape[1]
+
+            # Size of word_ids_mapping: (unique_words, 2) - [vocab_id, batch_id]
+            # Size of embeddings: (unique_words, embedding_size)
+            ids_mapping, minicoil_embeddings = self.encoder.forward(
+                word_ids_array_expanded, token_embeddings_array
+            )
+
+            # Size of counts: (unique_words)
+            words_ids: list[int] = ids_mapping[:, 0].tolist()  # type: ignore[assignment]
+
+            sentence_result: dict[str, WordEmbedding] = {}
+
+            words = [self.vocab_resolver.lookup_word(word_id) for word_id in words_ids]
+
+            for word, word_id, emb in zip(words, words_ids, minicoil_embeddings.tolist()):  # type: ignore[arg-type]
+                if word_id == 0:
+                    continue
+
+                sentence_result[word] = WordEmbedding(
+                    word=word,
+                    forms=forms[word],
+                    count=int(counts[word_id]),
+                    word_id=int(word_id),
+                    embedding=emb,  # type: ignore[arg-type]
+                )
+
+            for oov_word, count in oov.items():
+                # {
+                #     "word": oov_word,
+                #     "forms": [oov_word],
+                #     "count": int(count),
+                #     "word_id": -1,
+                #     "embedding": [1]
+                # }
+                sentence_result[oov_word] = WordEmbedding(
+                    word=oov_word, forms=[oov_word], count=int(count), word_id=-1, embedding=[1]
+                )
+
+            if not is_query:
+                yield self.sparse_vector_converter.embedding_to_vector(
+                    sentence_result, vocab_size=vocab_size, embedding_size=embedding_size
+                )
+            else:
+                yield self.sparse_vector_converter.embedding_to_vector_query(
+                    sentence_result, vocab_size=vocab_size, embedding_size=embedding_size
+                )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["MiniCoilTextEmbeddingWorker"]:
+        return MiniCoilTextEmbeddingWorker
+
+
+class MiniCoilTextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> MiniCOIL:
+        return MiniCOIL(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/sparse/sparse_embedding_base.py

@@ -0,0 +1,90 @@
+from dataclasses import dataclass
+from typing import Iterable, Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+from fastembed.common.model_description import SparseModelDescription
+from fastembed.common.types import NumpyArray
+from fastembed.common.model_management import ModelManagement
+
+
+@dataclass
+class SparseEmbedding:
+    values: NumpyArray
+    indices: NDArray[np.int64] | NDArray[np.int32]
+
+    def as_object(self) -> dict[str, NumpyArray]:
+        return {
+            "values": self.values,
+            "indices": self.indices,
+        }
+
+    def as_dict(self) -> dict[int, float]:
+        return {int(i): float(v) for i, v in zip(self.indices, self.values)}  # type: ignore
+
+    @classmethod
+    def from_dict(cls, data: dict[int, float]) -> "SparseEmbedding":
+        if len(data) == 0:
+            return cls(values=np.array([]), indices=np.array([]))
+        indices, values = zip(*data.items())
+        return cls(values=np.array(values), indices=np.array(indices))
+
+
+class SparseTextEmbeddingBase(ModelManagement[SparseModelDescription]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        **kwargs: Any,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> Iterable[SparseEmbedding]:
+        raise NotImplementedError()
+
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+        """
+        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[SparseEmbedding]: The sparse embeddings.
+        """
+
+        # This is model-specific, so that different models can have specialized implementations
+        yield from self.embed(texts, **kwargs)
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+        """
+        Embeds queries
+
+        Args:
+            query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
+
+        Returns:
+            Iterable[SparseEmbedding]: The sparse embeddings.
+        """
+
+        # This is model-specific, so that different models can have specialized implementations
+        if isinstance(query, str):
+            yield from self.embed([query], **kwargs)
+        else:
+            yield from self.embed(query, **kwargs)
+
+    def token_count(self, texts: str | Iterable[str], **kwargs: Any) -> int:
+        """Returns the number of tokens in the texts."""
+        raise NotImplementedError("Subclasses must implement this method")

fastembed/sparse/sparse_text_embedding.py

@@ -0,0 +1,150 @@
+from typing import Any, Iterable, Sequence, Type
+from dataclasses import asdict
+
+from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
+from fastembed.sparse.bm25 import Bm25
+from fastembed.sparse.bm42 import Bm42
+from fastembed.sparse.if_splade import IfSplade
+from fastembed.sparse.minicoil import MiniCOIL
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.sparse.splade_pp import SpladePP
+import warnings
+from fastembed.common.model_description import SparseModelDescription
+
+
+class SparseTextEmbedding(SparseTextEmbeddingBase):
+    EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [
+        SpladePP,
+        Bm42,
+        Bm25,
+        MiniCOIL,
+        IfSplade,
+    ]
+
+    @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",
+                        "license": "apache-2.0",
+                        "size_in_GB": 0.532,
+                        "sources": {
+                            "hf": "qdrant/SPLADE_PP_en_v1",
+                        },
+                    }
+                ]
+                ```
+        """
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
+        result: list[SparseModelDescription] = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding._list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        if model_name.lower() == "prithvida/Splade_PP_en_v1".lower():
+            warnings.warn(
+                "The right spelling is prithivida/Splade_PP_en_v1. "
+                "Support of this name will be removed soon, please fix the model_name",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            model_name = "prithivida/Splade_PP_en_v1"
+
+        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=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **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: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> 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)
+
+    def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
+        """
+        Embeds queries
+
+        Args:
+            query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
+
+        Returns:
+            Iterable[SparseEmbedding]: The sparse embeddings.
+        """
+        yield from self.model.query_embed(query, **kwargs)
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        """Returns the number of tokens in the texts.
+
+        Args:
+            texts (str | Iterable[str]): The list of texts to embed.
+            batch_size (int): Batch size for encoding
+
+        Returns:
+            int: Sum of number of tokens in the texts.
+        """
+        return self.model.token_count(texts, batch_size=batch_size, **kwargs)

fastembed/sparse/splade_pp.py

@@ -0,0 +1,196 @@
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+from fastembed.common import OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import Device
+from fastembed.common.utils import define_cache_dir
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
+from fastembed.common.model_description import SparseModelDescription, ModelSource
+
+supported_splade_models: list[SparseModelDescription] = [
+    SparseModelDescription(
+        model="prithivida/Splade_PP_en_v1",
+        vocab_size=30522,
+        description="Independent Implementation of SPLADE++ Model for English.",
+        license="apache-2.0",
+        size_in_GB=0.532,
+        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
+        model_file="model.onnx",
+    ),
+    SparseModelDescription(
+        model="prithvida/Splade_PP_en_v1",
+        vocab_size=30522,
+        description="Independent Implementation of SPLADE++ Model for English.",
+        license="apache-2.0",
+        size_in_GB=0.532,
+        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
+        model_file="model.onnx",
+    ),
+]
+
+
+class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[SparseEmbedding]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
+        relu_log = np.log(1 + np.maximum(output.model_output, 0))
+
+        weighted_log = relu_log * np.expand_dims(output.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)
+
+    def token_count(
+        self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        return self._token_count(texts, batch_size=batch_size, **kwargs)
+
+    @classmethod
+    def _list_supported_models(cls) -> list[SparseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information.
+        """
+        return supported_splade_models
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        """
+        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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
+            lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
+                Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+            specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+        self.providers = providers
+        self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        self.device_id: int | None = None
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = str(define_cache_dir(cache_dir))
+
+        self._specific_model_path = specific_model_path
+        self._model_dir = self.download_model(
+            self.model_description,
+            self.cache_dir,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+        )
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description.model_file,
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
+        )
+
+    def embed(
+        self,
+        documents: str | Iterable[str],
+        batch_size: int = 256,
+        parallel: int | None = None,
+        **kwargs: Any,
+    ) -> 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,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
+            **kwargs,
+        )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
+        return SpladePPEmbeddingWorker
+
+
+class SpladePPEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
+        return SpladePP(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/sparse/utils/minicoil_encoder.py

@@ -0,0 +1,146 @@
+"""
+Pure numpy implementation of encoder model for a single word.
+
+This model is not trainable, and should only be used for inference.
+"""
+
+import numpy as np
+from fastembed.common.types import NumpyArray
+
+
+class Encoder:
+    """
+    Encoder(768, 4, 10000)
+
+    Will look like this:
+
+
+                                         Per-word
+                                         Encoder Matrix
+     ┌─────────────────────┐
+     │ Token Embedding(768)├──────┐      (10k, 768, 4)
+     └─────────────────────┘      │         ┌─────────┐
+                                  │         │         │
+     ┌─────────────────────┐      │       ┌─┴───────┐ │
+     │                     │      │       │         │ │
+     └─────────────────────┘      │     ┌─┴───────┐ │ │      ┌─────────┐
+                                  └────►│         │ │ ├─────►│Tanh     │
+     ┌─────────────────────┐            │         │ │ │      └─────────┘
+     │                     │            │         │ ├─┘
+     └─────────────────────┘            │         ├─┘
+                                        │         │
+     ┌─────────────────────┐            └─────────┘
+     │                     │
+     └─────────────────────┘
+
+     Final linear transformation is accompanied by a non-linear activation function: Tanh.
+
+     Tanh is used to ensure that the output is in the range [-1, 1].
+     It would be easier to visually interpret the output of the model, assuming that each dimension
+     would need to encode a type of semantic cluster.
+    """
+
+    def __init__(
+        self,
+        weights: NumpyArray,
+    ):
+        self.weights = weights
+        self.vocab_size, self.input_dim, self.output_dim = weights.shape
+
+        self.encoder_weights: NumpyArray = weights
+
+        # Activation function
+        self.activation = np.tanh
+
+    @staticmethod
+    def convert_vocab_ids(vocab_ids: NumpyArray) -> NumpyArray:
+        """
+        Convert vocab_ids of shape (batch_size, seq_len) into (batch_size, seq_len, 2)
+        by appending batch_id alongside each vocab_id.
+        """
+        batch_size, seq_len = vocab_ids.shape
+        batch_ids = np.arange(batch_size, dtype=vocab_ids.dtype).reshape(batch_size, 1)
+        batch_ids = np.repeat(batch_ids, seq_len, axis=1)
+        # Stack vocab_ids and batch_ids along the last dimension
+        combined: NumpyArray = np.stack((vocab_ids, batch_ids), axis=2).astype(np.int32)
+        return combined
+
+    @classmethod
+    def avg_by_vocab_ids(
+        cls, vocab_ids: NumpyArray, embeddings: NumpyArray
+    ) -> tuple[NumpyArray, NumpyArray]:
+        """
+        Takes:
+            vocab_ids: (batch_size, seq_len) int array
+            embeddings: (batch_size, seq_len, input_dim) float array
+
+        Returns:
+            unique_flattened_vocab_ids: (total_unique, 2) array of [vocab_id, batch_id]
+            unique_flattened_embeddings: (total_unique, input_dim) averaged embeddings
+        """
+        input_dim = embeddings.shape[2]
+
+        # Flatten vocab_ids and embeddings
+        # flattened_vocab_ids: (batch_size*seq_len, 2)
+        flattened_vocab_ids = cls.convert_vocab_ids(vocab_ids).reshape(-1, 2)
+
+        # flattened_embeddings: (batch_size*seq_len, input_dim)
+        flattened_embeddings = embeddings.reshape(-1, input_dim)
+
+        # Find unique (vocab_id, batch_id) pairs
+        unique_flattened_vocab_ids, inverse_indices = np.unique(
+            flattened_vocab_ids, axis=0, return_inverse=True
+        )
+
+        # Prepare arrays to accumulate sums
+        unique_count = unique_flattened_vocab_ids.shape[0]
+        unique_flattened_embeddings = np.zeros((unique_count, input_dim), dtype=np.float32)
+        unique_flattened_count = np.zeros(unique_count, dtype=np.int32)
+
+        # Use np.add.at to accumulate sums based on inverse indices
+        np.add.at(unique_flattened_embeddings, inverse_indices, flattened_embeddings)
+        np.add.at(unique_flattened_count, inverse_indices, 1)
+
+        # Compute averages
+        unique_flattened_embeddings /= unique_flattened_count[:, None]
+
+        return unique_flattened_vocab_ids.astype(np.int32), unique_flattened_embeddings.astype(
+            np.float32
+        )
+
+    def forward(
+        self, vocab_ids: NumpyArray, embeddings: NumpyArray
+    ) -> tuple[NumpyArray, NumpyArray]:
+        """
+        Args:
+            vocab_ids: (batch_size, seq_len) int array
+            embeddings: (batch_size, seq_len, input_dim) float array
+
+        Returns:
+            unique_flattened_vocab_ids_and_batch_ids: (total_unique, 2)
+            unique_flattened_encoded: (total_unique, output_dim)
+        """
+        # Average embeddings for duplicate vocab_ids
+        unique_flattened_vocab_ids_and_batch_ids, unique_flattened_embeddings = (
+            self.avg_by_vocab_ids(vocab_ids, embeddings)
+        )
+
+        # Select the encoder weights for each unique vocab_id
+        unique_flattened_vocab_ids = unique_flattened_vocab_ids_and_batch_ids[:, 0].astype(
+            np.int32
+        )
+
+        # unique_encoder_weights: (total_unique, input_dim, output_dim)
+        unique_encoder_weights = self.encoder_weights[unique_flattened_vocab_ids]
+
+        # Compute linear transform: (total_unique, output_dim)
+        # Using Einstein summation for matrix multiplication:
+        # 'bi,bio->bo' means: for each "b" (batch element), multiply embeddings (b,i) by weights (b,i,o) -> (b,o)
+        unique_flattened_encoded = np.einsum(
+            "bi,bio->bo", unique_flattened_embeddings, unique_encoder_weights
+        )
+
+        # Apply Tanh activation and ensure float32 type
+        unique_flattened_encoded = self.activation(unique_flattened_encoded).astype(np.float32)
+
+        return unique_flattened_vocab_ids_and_batch_ids.astype(np.int32), unique_flattened_encoded

fastembed/sparse/utils/sparse_vectors_converter.py

@@ -0,0 +1,244 @@
+import copy
+from dataclasses import dataclass
+
+import mmh3
+import numpy as np
+from py_rust_stemmers import SnowballStemmer
+
+from fastembed.common.utils import get_all_punctuation, remove_non_alphanumeric
+from fastembed.sparse.sparse_embedding_base import SparseEmbedding
+
+GAP = 32000
+INT32_MAX = 2**31 - 1
+
+
+@dataclass
+class WordEmbedding:
+    word: str
+    forms: list[str]
+    count: int
+    word_id: int
+    embedding: list[float]
+
+
+class SparseVectorConverter:
+    def __init__(
+        self,
+        stopwords: set[str],
+        stemmer: SnowballStemmer,
+        k: float = 1.2,
+        b: float = 0.75,
+        avg_len: float = 150.0,
+    ):
+        punctuation = set(get_all_punctuation())
+        special_tokens = {"[CLS]", "[SEP]", "[PAD]", "[UNK]", "[MASK]"}
+
+        self.stemmer = stemmer
+        self.unwanted_tokens = punctuation | special_tokens | stopwords
+
+        self.k = k
+        self.b = b
+        self.avg_len = avg_len
+
+    @classmethod
+    def unkn_word_token_id(
+        cls, word: str, shift: int
+    ) -> int:  # 2-3 words can collide in 1 index with this mapping, not considering mm3 collisions
+        token_hash = abs(mmh3.hash(word))
+
+        range_size = INT32_MAX - shift
+        remapped_hash = shift + (token_hash % range_size)
+
+        return remapped_hash
+
+    def bm25_tf(self, num_occurrences: int, sentence_len: int) -> float:
+        res = num_occurrences * (self.k + 1)
+        res /= num_occurrences + self.k * (1 - self.b + self.b * sentence_len / self.avg_len)
+        return res
+
+    @classmethod
+    def normalize_vector(cls, vector: list[float]) -> list[float]:
+        norm = sum([x**2 for x in vector]) ** 0.5
+        if norm < 1e-8:
+            return vector
+        return [x / norm for x in vector]
+
+    def clean_words(
+        self, sentence_embedding: dict[str, WordEmbedding], token_max_length: int = 40
+    ) -> dict[str, WordEmbedding]:
+        """
+        Clean miniCOIL-produced sentence_embedding, as unknown to the miniCOIL's stemmer tokens should fully resemble
+        our BM25 token representation.
+
+        sentence_embedding = {"9°": {"word": "9°", "word_id": -1, "count": 2, "embedding": [1], "forms": ["9°"]},
+                "9": {"word": "9", "word_id": -1, "count": 2, "embedding": [1], "forms": ["9"]},
+                "bat": {"word": "bat", "word_id": 2, "count": 3, "embedding": [0.2, 0.1, -0.2, -0.2], "forms": ["bats", "bat"]},
+                "9°9": {"word": "9°9", "word_id": -1, "count": 1, "embedding": [1], "forms": ["9°9"]},
+                "screech": {"word": "screech", "word_id": -1, "count": 1, "embedding": [1], "forms": ["screech"]},
+                "screeched": {"word": "screeched", "word_id": -1, "count": 1, "embedding": [1], "forms": ["screeched"]}
+                }
+        cleaned_embedding_ground_truth = {
+                "9": {"word": "9", "word_id": -1, "count": 6, "embedding": [1], "forms": ["9°", "9", "9°9", "9°9"]},
+                "bat": {"word": "bat", "word_id": 2, "count": 3, "embedding": [0.2, 0.1, -0.2, -0.2], "forms": ["bats", "bat"]},
+                "screech": {"word": "screech", "word_id": -1, "count": 2, "embedding": [1], "forms": ["screech", "screeched"]}
+                }
+        """
+
+        new_sentence_embedding: dict[str, WordEmbedding] = {}
+
+        for word, embedding in sentence_embedding.items():
+            # embedding = {
+            #     "word": "vector",
+            #     "forms": ["vector", "vectors"],
+            #     "count": 2,
+            #     "word_id": 1231,
+            #     "embedding": [0.1, 0.2, 0.3, 0.4]
+            # }
+            if embedding.word_id > 0:
+                # Known word, no need to clean
+                new_sentence_embedding[word] = embedding
+            else:
+                # Unknown word
+                if word in self.unwanted_tokens:
+                    continue
+
+                # Example complex word split:
+                # word = `word^vec`
+                word_cleaned = remove_non_alphanumeric(word).strip()
+                # word_cleaned = `word vec`
+
+                if len(word_cleaned) > 0:
+                    # Subwords: ['word', 'vec']
+                    for subword in word_cleaned.split():
+                        stemmed_subword: str = self.stemmer.stem_word(subword)
+                        if (
+                            len(stemmed_subword) <= token_max_length
+                            and stemmed_subword not in self.unwanted_tokens
+                        ):
+                            if stemmed_subword not in new_sentence_embedding:
+                                new_sentence_embedding[stemmed_subword] = copy.deepcopy(embedding)
+                                new_sentence_embedding[stemmed_subword].word = stemmed_subword
+                            else:
+                                new_sentence_embedding[stemmed_subword].count += embedding.count
+                                new_sentence_embedding[stemmed_subword].forms += embedding.forms
+
+        return new_sentence_embedding
+
+    def embedding_to_vector(
+        self,
+        sentence_embedding: dict[str, WordEmbedding],
+        embedding_size: int,
+        vocab_size: int,
+    ) -> SparseEmbedding:
+        """
+        Convert miniCOIL sentence embedding to Qdrant sparse vector
+
+        Example input:
+
+        ```
+        {
+            "vector": WordEmbedding({ // Vocabulary word, encoded with miniCOIL normally
+                "word": "vector",
+                "forms": ["vector", "vectors"],
+                "count": 2,
+                "word_id": 1231,
+                "embedding": [0.1, 0.2, 0.3, 0.4]
+            }),
+            "axiotic": WordEmbedding({ // Out-of-vocabulary word, fallback to BM25
+                "word": "axiotic",
+                "forms": ["axiotics"],
+                "count": 1,
+                "word_id": -1,
+            })
+        }
+        ```
+
+        """
+
+        indices: list[int] = []
+        values: list[float] = []
+
+        # Example:
+        # vocab_size = 10000
+        # embedding_size = 4
+        # GAP = 32000
+        #
+        # We want to start random words section from the bucket, that is guaranteed to not
+        # include any vocab words.
+        # We need (vocab_size * embedding_size) slots for vocab words.
+        # Therefore we need (vocab_size * embedding_size) // GAP + 1 buckets for vocab words.
+        # Therefore, we can start random words from bucket (vocab_size * embedding_size) // GAP + 1 + 1
+
+        # ID at which the scope of OOV words starts
+        unknown_words_shift = ((vocab_size * embedding_size) // GAP + 2) * GAP
+        sentence_embedding_cleaned = self.clean_words(sentence_embedding)
+
+        # Calculate sentence length after cleaning
+        sentence_len = 0
+        for embedding in sentence_embedding_cleaned.values():
+            sentence_len += embedding.count
+
+        for embedding in sentence_embedding_cleaned.values():
+            word_id = embedding.word_id
+            num_occurrences = embedding.count
+            tf = self.bm25_tf(num_occurrences, sentence_len)
+            if (
+                word_id > 0
+            ):  # miniCOIL starts with ID 1, we generally won't have word_id == 0 (UNK), as we don't add
+                # these words to sentence_embedding
+                embedding_values = embedding.embedding
+                normalized_embedding = self.normalize_vector(embedding_values)
+
+                for val_id, value in enumerate(normalized_embedding):
+                    indices.append(
+                        word_id * embedding_size + val_id
+                    )  # since miniCOIL IDs start with 1
+                    values.append(value * tf)
+            else:
+                indices.append(self.unkn_word_token_id(embedding.word, unknown_words_shift))
+                values.append(tf)
+
+        return SparseEmbedding(
+            indices=np.array(indices, dtype=np.int32),
+            values=np.array(values, dtype=np.float32),
+        )
+
+    def embedding_to_vector_query(
+        self,
+        sentence_embedding: dict[str, WordEmbedding],
+        embedding_size: int,
+        vocab_size: int,
+    ) -> SparseEmbedding:
+        """
+        Same as `embedding_to_vector`, but no TF
+        """
+
+        indices: list[int] = []
+        values: list[float] = []
+
+        # ID at which the scope of OOV words starts
+        unknown_words_shift = ((vocab_size * embedding_size) // GAP + 2) * GAP
+
+        sentence_embedding_cleaned = self.clean_words(sentence_embedding)
+
+        for embedding in sentence_embedding_cleaned.values():
+            word_id = embedding.word_id
+            tf = 1.0
+
+            if word_id >= 0:  # miniCOIL starts with ID 1
+                embedding_values = embedding.embedding
+                normalized_embedding = self.normalize_vector(embedding_values)
+
+                for val_id, value in enumerate(normalized_embedding):
+                    indices.append(
+                        word_id * embedding_size + val_id
+                    )  # since miniCOIL IDs start with 1
+                    values.append(value * tf)
+            else:
+                indices.append(self.unkn_word_token_id(embedding.word, unknown_words_shift))
+                values.append(tf)
+
+        return SparseEmbedding(
+            indices=np.array(indices, dtype=np.int32),
+            values=np.array(values, dtype=np.float32),
+        )

fastembed/sparse/utils/tokenizer.py

@@ -0,0 +1,120 @@
+# This code is a modified copy of the `NLTKWordTokenizer` class from `NLTK` library.
+
+import re
+
+
+class SimpleTokenizer:
+    @staticmethod
+    def tokenize(text: str) -> list[str]:
+        text = re.sub(r"[^\w]", " ", text.lower())
+        text = re.sub(r"\s+", " ", text)
+
+        return text.strip().split()
+
+
+class WordTokenizer:
+    """The tokenizer is "destructive" such that the regexes applied will munge the
+    input string to a state beyond re-construction.
+    """
+
+    # Starting quotes.
+    STARTING_QUOTES = [
+        (re.compile("([«“‘„]|[`]+)", re.U), r" \1 "),
+        (re.compile(r"^\""), r"``"),
+        (re.compile(r"(``)"), r" \1 "),
+        (re.compile(r"([ \(\[{<])(\"|\'{2})"), r"\1 `` "),
+        (re.compile(r"(?i)(\')(?!re|ve|ll|m|t|s|d|n)(\w)\b", re.U), r"\1 \2"),
+    ]
+
+    # Ending quotes.
+    ENDING_QUOTES = [
+        (re.compile("([»”’])", re.U), r" \1 "),
+        (re.compile(r"''"), " '' "),
+        (re.compile(r'"'), " '' "),
+        (re.compile(r"([^' ])('[sS]|'[mM]|'[dD]|') "), r"\1 \2 "),
+        (re.compile(r"([^' ])('ll|'LL|'re|'RE|'ve|'VE|n't|N'T) "), r"\1 \2 "),
+    ]
+
+    # Punctuation.
+    PUNCTUATION = [
+        (re.compile(r'([^\.])(\.)([\]\)}>"\'' "»”’ " r"]*)\s*$", re.U), r"\1 \2 \3 "),
+        (re.compile(r"([:,])([^\d])"), r" \1 \2"),
+        (re.compile(r"([:,])$"), r" \1 "),
+        (
+            re.compile(r"\.{2,}", re.U),
+            r" \g<0> ",
+        ),
+        (re.compile(r"[;@#$%&]"), r" \g<0> "),
+        (
+            re.compile(r'([^\.])(\.)([\]\)}>"\']*)\s*$'),
+            r"\1 \2\3 ",
+        ),  # Handles the final period.
+        (re.compile(r"[?!]"), r" \g<0> "),
+        (re.compile(r"([^'])' "), r"\1 ' "),
+        (
+            re.compile(r"[*]", re.U),
+            r" \g<0> ",
+        ),
+    ]
+
+    # Pads parentheses
+    PARENS_BRACKETS = (re.compile(r"[\]\[\(\)\{\}\<\>]"), r" \g<0> ")
+    DOUBLE_DASHES = (re.compile(r"--"), r" -- ")
+
+    # List of contractions adapted from Robert MacIntyre's tokenizer.
+    CONTRACTIONS2 = [
+        re.compile(pattern)
+        for pattern in (
+            r"(?i)\b(can)(?#X)(not)\b",
+            r"(?i)\b(d)(?#X)('ye)\b",
+            r"(?i)\b(gim)(?#X)(me)\b",
+            r"(?i)\b(gon)(?#X)(na)\b",
+            r"(?i)\b(got)(?#X)(ta)\b",
+            r"(?i)\b(lem)(?#X)(me)\b",
+            r"(?i)\b(more)(?#X)('n)\b",
+            r"(?i)\b(wan)(?#X)(na)(?=\s)",
+        )
+    ]
+    CONTRACTIONS3 = [
+        re.compile(pattern) for pattern in (r"(?i) ('t)(?#X)(is)\b", r"(?i) ('t)(?#X)(was)\b")
+    ]
+
+    @classmethod
+    def tokenize(cls, text: str) -> list[str]:
+        """Return a tokenized copy of `text`.
+
+        >>> s = '''Good muffins cost $3.88 (roughly 3,36 euros)\nin New York.'''
+        >>> WordTokenizer().tokenize(s)
+        ['Good', 'muffins', 'cost', '$', '3.88', '(', 'roughly', '3,36', 'euros', ')', 'in', 'New', 'York', '.']
+
+        Args:
+            text: The text to be tokenized.
+
+        Returns:
+            A list of tokens.
+        """
+        for regexp, substitution in cls.STARTING_QUOTES:
+            text = regexp.sub(substitution, text)
+
+        for regexp, substitution in cls.PUNCTUATION:
+            text = regexp.sub(substitution, text)
+
+        # Handles parentheses.
+        regexp, substitution = cls.PARENS_BRACKETS
+        text = regexp.sub(substitution, text)
+
+        # Handles double dash.
+        regexp, substitution = cls.DOUBLE_DASHES
+        text = regexp.sub(substitution, text)
+
+        # add extra space to make things easier
+        text = " " + text + " "
+
+        for regexp, substitution in cls.ENDING_QUOTES:
+            text = regexp.sub(substitution, text)
+
+        for regexp in cls.CONTRACTIONS2:
+            text = regexp.sub(r" \1 \2 ", text)
+        for regexp in cls.CONTRACTIONS3:
+            text = regexp.sub(r" \1 \2 ", text)
+        return text.split()

fastembed/sparse/utils/vocab_resolver.py

@@ -0,0 +1,202 @@
+from collections import defaultdict
+from typing import Iterable
+
+from py_rust_stemmers import SnowballStemmer
+import numpy as np
+from tokenizers import Tokenizer
+from numpy.typing import NDArray
+
+from fastembed.common.types import NumpyArray
+
+
+class VocabTokenizerBase:
+    def tokenize(self, sentence: str) -> NumpyArray:
+        raise NotImplementedError()
+
+    def convert_ids_to_tokens(self, token_ids: NumpyArray) -> list[str]:
+        raise NotImplementedError()
+
+
+class VocabTokenizer(VocabTokenizerBase):
+    def __init__(self, tokenizer: Tokenizer):
+        self.tokenizer = tokenizer
+
+    def tokenize(self, sentence: str) -> NumpyArray:
+        return np.array(self.tokenizer.encode(sentence).ids)
+
+    def convert_ids_to_tokens(self, token_ids: NumpyArray) -> list[str]:
+        return [self.tokenizer.id_to_token(token_id) for token_id in token_ids]
+
+
+class VocabResolver:
+    def __init__(self, tokenizer: VocabTokenizerBase, stopwords: set[str], stemmer: SnowballStemmer):
+        # Word to id mapping
+        self.vocab: dict[str, int] = {}
+        # Id to word mapping
+        self.words: list[str] = []
+        # Lemma to word mapping
+        self.stem_mapping: dict[str, str] = {}
+        self.tokenizer: VocabTokenizerBase = tokenizer
+        self.stemmer = stemmer
+        self.stopwords: set[str] = stopwords
+
+    def tokenize(self, sentence: str) -> NumpyArray:
+        return self.tokenizer.tokenize(sentence)
+
+    def lookup_word(self, word_id: int) -> str:
+        if word_id == 0:
+            return "UNK"
+        return self.words[word_id - 1]
+
+    def convert_ids_to_tokens(self, token_ids: NumpyArray) -> list[str]:
+        return self.tokenizer.convert_ids_to_tokens(token_ids)
+
+    def vocab_size(self) -> int:
+        # We need +1 for UNK token
+        return len(self.vocab) + 1
+
+    def save_vocab(self, path: str) -> None:
+        with open(path, "w") as f:
+            for word in self.words:
+                f.write(word + "\n")
+
+    def save_json_vocab(self, path: str) -> None:
+        import json
+
+        with open(path, "w") as f:
+            json.dump({"vocab": self.words, "stem_mapping": self.stem_mapping}, f, indent=2)
+
+    def load_json_vocab(self, path: str) -> None:
+        import json
+
+        with open(path, "r") as f:
+            data = json.load(f)
+            self.words = data["vocab"]
+            self.vocab = {word: idx + 1 for idx, word in enumerate(self.words)}
+            self.stem_mapping = data["stem_mapping"]
+
+    def add_word(self, word: str) -> None:
+        if word not in self.vocab:
+            self.vocab[word] = len(self.vocab) + 1
+            self.words.append(word)
+            stem = self.stemmer.stem_word(word)
+            if stem not in self.stem_mapping:
+                self.stem_mapping[stem] = word
+            else:
+                existing_word = self.stem_mapping[stem]
+                if len(existing_word) > len(word):
+                    # Prefer shorter words for the same stem
+                    # Example: "swim" is preferred over "swimming"
+                    self.stem_mapping[stem] = word
+
+    def load_vocab(self, path: str) -> None:
+        with open(path, "r") as f:
+            for line in f:
+                self.add_word(line.strip())
+
+    @classmethod
+    def _reconstruct_bpe(
+        cls, bpe_tokens: Iterable[tuple[int, str]]
+    ) -> list[tuple[str, list[int]]]:
+        result: list[tuple[str, list[int]]] = []
+        acc: str = ""
+        acc_idx: list[int] = []
+
+        continuing_subword_prefix = "##"
+        continuing_subword_prefix_len = len(continuing_subword_prefix)
+
+        for idx, token in bpe_tokens:
+            if token.startswith(continuing_subword_prefix):
+                acc += token[continuing_subword_prefix_len:]
+                acc_idx.append(idx)
+            else:
+                if acc:
+                    result.append((acc, acc_idx))
+                    acc_idx = []
+                acc = token
+                acc_idx.append(idx)
+
+        if acc:
+            result.append((acc, acc_idx))
+        return result
+
+    def resolve_tokens(
+        self, token_ids: NDArray[np.int64]
+    ) -> tuple[NDArray[np.int64], dict[int, int], dict[str, int], dict[str, list[str]]]:
+        """
+        Mark known tokens (including composed tokens) with vocab ids.
+
+        Args:
+            token_ids: (seq_len) - list of ids of tokens
+                Example:
+                    [
+                        101,  3897, 19332, 12718, 23348,
+                        1010,  1996,  7151,  2296, 4845,
+                        2359,  2005,  4234,  1010,  4332,
+                        2871,  3191,  2062, 102
+                    ]
+
+            returns:
+                - token_ids with vocab ids
+                    [
+                        0,  151, 151, 0, 0,
+                        912,  0,  0,  0, 332,
+                        332,  332,  0,  7121,  191,
+                        0,  0,  332, 0
+                    ]
+                - counts of each token
+                    {
+                        151: 1,
+                        332: 3,
+                        7121: 1,
+                        191: 1,
+                        912: 1
+                    }
+                - oov counts of each token
+                    {
+                        "the": 1,
+                        "a": 1,
+                        "[CLS]": 1,
+                        "[SEP]": 1,
+                        ...
+                    }
+                - forms of each token
+                    {
+                        "hello": ["hello"],
+                        "world": ["worlds", "world", "worlding"],
+                    }
+
+        """
+        tokens = self.convert_ids_to_tokens(token_ids)
+        tokens_mapping = self._reconstruct_bpe(enumerate(tokens))
+
+        counts: dict[int, int] = defaultdict(int)
+        oov_count: dict[str, int] = defaultdict(int)
+
+        forms: dict[str, list[str]] = defaultdict(list)
+
+        for token, mapped_token_ids in tokens_mapping:
+            vocab_id = 0
+            if token in self.stopwords:
+                vocab_id = 0
+            elif token in self.vocab:
+                vocab_id = self.vocab[token]
+                forms[token].append(token)
+            elif token in self.stem_mapping:
+                vocab_id = self.vocab[self.stem_mapping[token]]
+                forms[self.stem_mapping[token]].append(token)
+            else:
+                stem = self.stemmer.stem_word(token)
+                if stem in self.stem_mapping:
+                    vocab_id = self.vocab[self.stem_mapping[stem]]
+                    forms[self.stem_mapping[stem]].append(token)
+
+            for token_id in mapped_token_ids:
+                token_ids[token_id] = vocab_id
+
+            if vocab_id == 0:
+                oov_count[token] += 1
+            else:
+                counts[vocab_id] += 1
+        return token_ids, counts, oov_count, forms
+

fastembed/text/__init__.py

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

fastembed/text/builtin_sentence_embedding.py

@@ -0,0 +1,69 @@
+from typing import Any, Iterable, Type
+
+
+from fastembed.common.types import NumpyArray
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+
+supported_builtin_sentence_embedding_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="google/embeddinggemma-300m",
+        dim=768,
+        description=(
+            "Text embeddings, Unimodal (text), multilingual, 2048 input tokens truncation, "
+            "Prefixes for queries/documents: `task: search result | query: {content}` for query, "
+            "`title: {title | 'none'} | text: {content}` for documents, 2025 year."
+        ),
+        license="gemma",
+        size_in_GB=1.24,
+        sources=ModelSource(
+            hf="onnx-community/embeddinggemma-300m-ONNX",
+        ),
+        model_file="onnx/model.onnx",
+        additional_files=["onnx/model.onnx_data"],
+    ),
+]
+
+
+class BuiltinSentenceEmbedding(OnnxTextEmbedding):
+    """Builtin Sentence Embedding uses built-in pooling and normalization of underlying onnx models"""
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return BuiltinSentenceEmbeddingWorker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_builtin_sentence_embedding_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        return output.model_output
+
+    def _run_model(
+        self, onnx_input: dict[str, Any], onnx_output_names: list[str] | None = None
+    ) -> NumpyArray:
+        return self.model.run(onnx_output_names, onnx_input)[1]  # type: ignore[union-attr]
+
+
+class BuiltinSentenceEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxTextEmbedding:
+        return BuiltinSentenceEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/clip_embedding.py

@@ -0,0 +1,56 @@
+from typing import Any, Iterable, Type
+
+from fastembed.common.types import NumpyArray
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_clip_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="Qdrant/clip-ViT-B-32-text",
+        dim=512,
+        description=(
+            "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2021 year"
+        ),
+        license="mit",
+        size_in_GB=0.25,
+        sources=ModelSource(hf="Qdrant/clip-ViT-B-32-text"),
+        model_file="model.onnx",
+    ),
+]
+
+
+class CLIPOnnxEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return CLIPEmbeddingWorker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_clip_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        return output.model_output
+
+
+class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxTextEmbedding:
+        return CLIPOnnxEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/custom_text_embedding.py

@@ -0,0 +1,144 @@
+from typing import Sequence, Any, Iterable, Type
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import NDArray
+
+from fastembed.common import OnnxProvider
+from fastembed.common.model_description import (
+    PoolingType,
+    DenseModelDescription,
+)
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.types import NumpyArray, Device
+from fastembed.common.utils import normalize, mean_pooling, last_token_pooling
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+
+@dataclass(frozen=True)
+class PostprocessingConfig:
+    pooling: PoolingType
+    normalization: bool
+
+
+class CustomTextEmbedding(OnnxTextEmbedding):
+    SUPPORTED_MODELS: list[DenseModelDescription] = []
+    POSTPROCESSING_MAPPING: dict[str, PostprocessingConfig] = {}
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        lazy_load: bool = False,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_ids=device_ids,
+            lazy_load=lazy_load,
+            device_id=device_id,
+            specific_model_path=specific_model_path,
+            **kwargs,
+        )
+        postprocessing_config = self.POSTPROCESSING_MAPPING[self.model_description.model]
+        self._pooling = postprocessing_config.pooling
+        self._normalization = postprocessing_config.normalization
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        return cls.SUPPORTED_MODELS
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker[NumpyArray]"]:
+        return CustomTextEmbeddingWorker
+
+    def _get_worker_init_kwargs(self) -> dict[str, Any]:
+        return {
+            "model_description": self.model_description,
+            "postprocessing_config": self.POSTPROCESSING_MAPPING[self.model_description.model],
+        }
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        return self._normalize(self._pool(output.model_output, output.attention_mask))
+
+    def _pool(
+        self, embeddings: NumpyArray, attention_mask: NDArray[np.int64] | None = None
+    ) -> NumpyArray:
+        if self._pooling == PoolingType.CLS:
+            return embeddings[:, 0]
+
+        if self._pooling == PoolingType.MEAN:
+            if attention_mask is None:
+                raise ValueError("attention_mask must be provided for mean pooling")
+            return mean_pooling(embeddings, attention_mask)
+
+        if self._pooling == PoolingType.LAST_TOKEN:
+            if attention_mask is None:
+                raise ValueError("attention_mask must be provided for last token pooling")
+            return last_token_pooling(embeddings, attention_mask)
+
+        if self._pooling == PoolingType.DISABLED:
+            return embeddings
+
+        raise ValueError(
+            f"Unsupported pooling type {self._pooling}. "
+            f"Supported types are: {PoolingType.CLS}, {PoolingType.MEAN}, "
+            f"{PoolingType.LAST_TOKEN}, {PoolingType.DISABLED}."
+        )
+
+    def _normalize(self, embeddings: NumpyArray) -> NumpyArray:
+        return normalize(embeddings) if self._normalization else embeddings
+
+    @classmethod
+    def add_model(
+        cls,
+        model_description: DenseModelDescription,
+        pooling: PoolingType,
+        normalization: bool,
+    ) -> None:
+        cls.SUPPORTED_MODELS.append(model_description)
+        cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig(
+            pooling=pooling, normalization=normalization
+        )
+
+
+class CustomTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        model_description: DenseModelDescription | None = None,
+        postprocessing_config: PostprocessingConfig | None = None,
+        **kwargs: Any,
+    ) -> CustomTextEmbedding:
+        if model_description is None or postprocessing_config is None:
+            raise ValueError(
+                "`model_description` and `postprocessing_config` are required to initialize a "
+                "custom model in a worker process, they are provided by "
+                "`CustomTextEmbedding._get_worker_init_kwargs`"
+            )
+        # custom models live in a class-level registry, which spawned workers don't inherit
+        CustomTextEmbedding.add_model(
+            model_description,
+            pooling=postprocessing_config.pooling,
+            normalization=postprocessing_config.normalization,
+        )
+        return CustomTextEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )