Add dynamic quantization from @Xenova

* check for existing files

* (chore: model_management.py):  Add comment

---------

Co-authored-by: Nirant <NirantK@users.noreply.github.com>

.github/workflows/ci.yml

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

.github/workflows/python-tests.yml

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

.pre-commit-config.yaml

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

README.md

@@ -1,61 +1,90 @@
 # ⚡️ What is FastEmbed?
 
-FastEmbed is an easy to use -- lightweight, fast, Python library built for retrieval embedding generation. 
+FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
-The default embedding supports "query" and "passage" prefixes for the input text. The default model is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval%20with%20FastEmbed/)
+The default text embedding (`TextEmbedding`) model is Flag Embedding, the top model in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
 
-1. Light
+1. Light & Fast
     - Quantized model weights
-    - ONNX Runtime for inference
-    - No hidden dependencies on PyTorch or TensorFlow via Huggingface Transformers
+    - ONNX Runtime, no PyTorch dependency
+    - CPU-first design
+    - Data-parallelism for encoding of large datasets
 
 2. Accuracy/Recall
     - Better than OpenAI Ada-002
     - Default is Flag Embedding, which is top of the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard
-
-3. Fast
-    - About 2x faster than Huggingface (PyTorch) transformers on single queries
-    - Lot faster for batches!
-    - ONNX Runtime allows you to use dedicated runtimes for even higher throughput and lower latency 
+    - List of [supported models](https://qdrant.github.io/fastembed/examples/Supported_Models/) - including multilingual models
 
 ## 🚀 Installation
 
-To install the FastEmbed library, pip works: 
+To install the FastEmbed library, pip works:
 
 ```bash
 pip install fastembed
 ```
 
-## 📖 Usage
+## 📖 Quickstart
 
 ```python
-from fastembed.embedding import FlagEmbedding as Embedding
+from fastembed import TextEmbedding
+from typing import List
+import numpy as np
 
 documents: List[str] = [
     "passage: Hello, World!",
     "query: Hello, World!", # these are two different embedding
     "passage: This is an example passage.",
-    # You can leave out the prefix but it's recommended
-    "fastembed is supported by and maintained by Qdrant." 
+    "fastembed is supported by and maintained by Qdrant." # You can leave out the prefix but it's recommended
 ]
-embedding_model = Embedding(model_name="BAAI/bge-base-en", max_length=512) 
-embeddings: List[np.ndarray] = list(embedding_model.embed(documents))
+embedding_model = TextEmbedding(model_name="BAAI/bge-base-en")
+embeddings: List[np.ndarray] = list(embedding_model.embed(documents)) # Note the list() call - this is a generator
 ```
 
-### Why fast?
+## Usage with Qdrant
 
-It's important we justify the "fast" in FastEmbed. FastEmbed is fast because:
+Installation with Qdrant Client in Python:
 
-1. Quantized model weights
-2. ONNX Runtime which allows for fast inference on CPU and other dedicated runtimes
+```bash
+pip install qdrant-client[fastembed]
+```
 
-### Why light?
-1. No hidden dependencies on PyTorch or TensorFlow via Huggingface Transformers
-2. We do use the tokenizer from Huggingface Transformers, but it's a light dependency
+You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
-### Why accurate?
-1. Better than OpenAI Ada-002
-2. Top of the Embedding leaderboards e.g. [MTEB](https://huggingface.co/spaces/mteb/leaderboard)
+```python
+from qdrant_client import QdrantClient
+
+# Initialize the client
+client = QdrantClient("localhost", port=6333) # For production
+# client = QdrantClient(":memory:") # For small experiments
+
+# Prepare your documents, metadata, and IDs
+docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
+metadata = [
+    {"source": "Langchain-docs"},
+    {"source": "Llama-index-docs"},
+]
+ids = [42, 2]
+
+# If you want to change the model:
+# client.set_model("sentence-transformers/all-MiniLM-L6-v2")
+# List of supported models: https://qdrant.github.io/fastembed/examples/Supported_Models
+
+# Use the new add() instead of upsert()
+# This internally calls embed() of the configured embedding model
+client.add(
+    collection_name="demo_collection",
+    documents=docs,
+    metadata=metadata,
+    ids=ids
+)
+
+search_result = client.query(
+    collection_name="demo_collection",
+    query_text="This is a query document"
+)
+print(search_result)
+```
 
 #### Similar Work
-Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.
+
+Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.

docs/Getting Started.ipynb

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

docs/examples/Retrieval_with_FastEmbed.ipynb

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

docs/examples/Supported_Models.ipynb

@@ -0,0 +1,289 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>model</th>\n",
+       "      <th>dim</th>\n",
+       "      <th>description</th>\n",
+       "      <th>size_in_GB</th>\n",
+       "      <th>sources</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>BAAI/bge-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Base English model</td>\n",
+       "      <td>0.50</td>\n",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>BAAI/bge-base-en-v1.5</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Base English model, v1.5</td>\n",
+       "      <td>0.44</td>\n",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz', 'hf': 'qdrant/bge-base-en-v1.5-onnx-q'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>BAAI/bge-large-en-v1.5-quantized</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Large English model, v1.5</td>\n",
+       "      <td>1.34</td>\n",
+       "      <td>{'hf': 'qdrant/bge-large-en-v1.5-onnx-q'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>BAAI/bge-large-en-v1.5</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Large English model, v1.5</td>\n",
+       "      <td>1.34</td>\n",
+       "      <td>{'hf': 'qdrant/bge-large-en-v1.5-onnx'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>BAAI/bge-small-en</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Fast English model</td>\n",
+       "      <td>0.20</td>\n",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>BAAI/bge-small-en-v1.5</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Fast and Default English model</td>\n",
+       "      <td>0.13</td>\n",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz', 'hf': 'qdrant/bge-small-en-v1.5-onnx-q'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
+       "      <td>BAAI/bge-small-zh-v1.5</td>\n",
+       "      <td>512</td>\n",
+       "      <td>Fast and recommended Chinese model</td>\n",
+       "      <td>0.10</td>\n",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</th>\n",
+       "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Sentence Transformer model, MiniLM-L6-v2</td>\n",
+       "      <td>0.09</td>\n",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz', 'hf': 'qdrant/all-MiniLM-L6-v2-onnx'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>8</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
+       "      <td>768</td>\n",
+       "      <td>8192 context length english model</td>\n",
+       "      <td>0.54</td>\n",
+       "      <td>{'hf': 'nomic-ai/nomic-embed-text-v1'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>9</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
+       "      <td>768</td>\n",
+       "      <td>8192 context length english model</td>\n",
+       "      <td>0.54</td>\n",
+       "      <td>{'hf': 'nomic-ai/nomic-embed-text-v1.5'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>10</th>\n",
+       "      <td>thenlper/gte-large</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Large general text embeddings model</td>\n",
+       "      <td>1.34</td>\n",
+       "      <td>{'hf': 'qdrant/gte-large-onnx'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>11</th>\n",
+       "      <td>intfloat/multilingual-e5-large</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Multilingual model, e5-large. Recommend using this model for non-English languages</td>\n",
+       "      <td>2.24</td>\n",
+       "      <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz', 'hf': 'qdrant/multilingual-e5-large-onnx'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>12</th>\n",
+       "      <td>sentence-transformers/paraphrase-multilingual-mpnet-base-v2</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Sentence-transformers model for tasks like clustering or semantic search</td>\n",
+       "      <td>1.11</td>\n",
+       "      <td>{'hf': 'xenova/paraphrase-multilingual-mpnet-base-v2'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>13</th>\n",
+       "      <td>sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2</td>\n",
+       "      <td>0.46</td>\n",
+       "      <td>{'hf': 'qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>14</th>\n",
+       "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>English embedding model supporting 8192 sequence length</td>\n",
+       "      <td>0.55</td>\n",
+       "      <td>{'hf': 'xenova/jina-embeddings-v2-base-en'}</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>15</th>\n",
+       "      <td>jinaai/jina-embeddings-v2-small-en</td>\n",
+       "      <td>512</td>\n",
+       "      <td>English embedding model supporting 8192 sequence length</td>\n",
+       "      <td>0.13</td>\n",
+       "      <td>{'hf': 'xenova/jina-embeddings-v2-small-en'}</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                                          model   dim  \\\n",
+       "0                                              BAAI/bge-base-en   768   \n",
+       "1                                         BAAI/bge-base-en-v1.5   768   \n",
+       "2                              BAAI/bge-large-en-v1.5-quantized  1024   \n",
+       "3                                        BAAI/bge-large-en-v1.5  1024   \n",
+       "4                                             BAAI/bge-small-en   384   \n",
+       "5                                        BAAI/bge-small-en-v1.5   384   \n",
+       "6                                        BAAI/bge-small-zh-v1.5   512   \n",
+       "7                        sentence-transformers/all-MiniLM-L6-v2   384   \n",
+       "8                                  nomic-ai/nomic-embed-text-v1   768   \n",
+       "9                                nomic-ai/nomic-embed-text-v1.5   768   \n",
+       "10                                           thenlper/gte-large  1024   \n",
+       "11                               intfloat/multilingual-e5-large  1024   \n",
+       "12  sentence-transformers/paraphrase-multilingual-mpnet-base-v2   768   \n",
+       "13  sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2   384   \n",
+       "14                            jinaai/jina-embeddings-v2-base-en   768   \n",
+       "15                           jinaai/jina-embeddings-v2-small-en   512   \n",
+       "\n",
+       "                                                                           description  \\\n",
+       "0                                                                   Base English model   \n",
+       "1                                                             Base English model, v1.5   \n",
+       "2                                                            Large English model, v1.5   \n",
+       "3                                                            Large English model, v1.5   \n",
+       "4                                                                   Fast English model   \n",
+       "5                                                       Fast and Default English model   \n",
+       "6                                                   Fast and recommended Chinese model   \n",
+       "7                                             Sentence Transformer model, MiniLM-L6-v2   \n",
+       "8                                                    8192 context length english model   \n",
+       "9                                                    8192 context length english model   \n",
+       "10                                                 Large general text embeddings model   \n",
+       "11  Multilingual model, e5-large. Recommend using this model for non-English languages   \n",
+       "12            Sentence-transformers model for tasks like clustering or semantic search   \n",
+       "13                   Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2   \n",
+       "14                             English embedding model supporting 8192 sequence length   \n",
+       "15                             English embedding model supporting 8192 sequence length   \n",
+       "\n",
+       "    size_in_GB  \\\n",
+       "0         0.50   \n",
+       "1         0.44   \n",
+       "2         1.34   \n",
+       "3         1.34   \n",
+       "4         0.20   \n",
+       "5         0.13   \n",
+       "6         0.10   \n",
+       "7         0.09   \n",
+       "8         0.54   \n",
+       "9         0.54   \n",
+       "10        1.34   \n",
+       "11        2.24   \n",
+       "12        1.11   \n",
+       "13        0.46   \n",
+       "14        0.55   \n",
+       "15        0.13   \n",
+       "\n",
+       "                                                                                                                                           sources  \n",
+       "0                                                               {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz'}  \n",
+       "1                  {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz', 'hf': 'qdrant/bge-base-en-v1.5-onnx-q'}  \n",
+       "2                                                                                                        {'hf': 'qdrant/bge-large-en-v1.5-onnx-q'}  \n",
+       "3                                                                                                          {'hf': 'qdrant/bge-large-en-v1.5-onnx'}  \n",
+       "4                                                              {'url': 'https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz'}  \n",
+       "5                {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz', 'hf': 'qdrant/bge-small-en-v1.5-onnx-q'}  \n",
+       "6                                                         {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz'}  \n",
+       "7   {'url': 'https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz', 'hf': 'qdrant/all-MiniLM-L6-v2-onnx'}  \n",
+       "8                                                                                                           {'hf': 'nomic-ai/nomic-embed-text-v1'}  \n",
+       "9                                                                                                         {'hf': 'nomic-ai/nomic-embed-text-v1.5'}  \n",
+       "10                                                                                                                 {'hf': 'qdrant/gte-large-onnx'}  \n",
+       "11         {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz', 'hf': 'qdrant/multilingual-e5-large-onnx'}  \n",
+       "12                                                                                          {'hf': 'xenova/paraphrase-multilingual-mpnet-base-v2'}  \n",
+       "13                                                                                   {'hf': 'qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q'}  \n",
+       "14                                                                                                     {'hf': 'xenova/jina-embeddings-v2-base-en'}  \n",
+       "15                                                                                                    {'hf': 'xenova/jina-embeddings-v2-small-en'}  "
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from fastembed import TextEmbedding\n",
+    "import pandas as pd\n",
+    "\n",
+    "pd.set_option(\"display.max_colwidth\", None)\n",
+    "pd.DataFrame(TextEmbedding.list_supported_models())"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "fst",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.13"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/examples/Usage_With_Qdrant.ipynb

@@ -5,9 +5,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# [Experimental] Usage With Qdrant\n",
-    "\n",
-    "> **Note:** This notebook is experimental and is subject to change. For working with this, use the dev branch of QdrantClient.\n",
+    "# Usage With Qdrant\n",
     "\n",
     "This notebook demonstrates how to use FastEmbed and Qdrant to perform vector search and retrieval. Qdrant is an open-source vector similarity search engine that is used to store, organize, and query collections of high-dimensional vectors. \n",
     "\n",
@@ -28,13 +26,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# !pip install fastembed --quiet --upgrade\n",
-    "\n",
-    "# !pip install git+https://github.com/qdrant/qdrant_client.git@dev"
+    "!pip install 'qdrant-client[fastembed]' --quiet --upgrade"
    ]
   },
   {
@@ -46,13 +42,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
     "from typing import List\n",
-    "import numpy as np\n",
-    "from fastembed.embedding import FlagEmbedding as Embedding\n",
     "from qdrant_client import QdrantClient"
    ]
   },
@@ -68,7 +62,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -105,25 +99,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 77.7M/77.7M [00:05<00:00, 14.6MiB/s]\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "['901ff9d7a90a4e56afe655d1de8e7c06',\n",
-       " '7bae5aa894164398a9be68d47d72ff7a',\n",
-       " 'd940bba124e24469ae3166c3110c62f8',\n",
-       " 'c8dcf956a4f444c6bfe69a00cbeb85ce',\n",
-       " 'd2aeb24c51f549c5b21851be05cb4048',\n",
-       " '6d4c672bafef4db68ae72c8986ee8a72',\n",
-       " 'ff198613768a4361a6d8230b40ee7f78',\n",
-       " 'a133ad72876b48d48e716533c2d78cf0',\n",
-       " 'bdf2b016136e4f57b7ac4b65534031b8',\n",
-       " '5f4f279f304c47839e1a74b0f28de136']"
+       "['4fa8b10c78da4b18ba0830ba8a57367a',\n",
+       " '2eae04b515ee4e9185a9a0e6be812bba',\n",
+       " 'c6039f88486f47f1835ae3b069c5823c',\n",
+       " 'c2c8c51e305144d1917b373125fb4d95',\n",
+       " '79fd23b9ec0648cdab38d1947c6b933e',\n",
+       " '036aa200d8c3492b8a438e4f825f5e7f',\n",
+       " 'c35c77f3ea37460a9a13723fb77b7367',\n",
+       " '6ebccbca571b40d0ab6e83e5e0f2f562',\n",
+       " '38048c2ccc1d4962a4f8f1bd89c8357a',\n",
+       " 'c6b09308360140c7b4f106af3658a31e']"
       ]
      },
-     "execution_count": 17,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -150,38 +151,56 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[42, 2]"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Prepare your documents, metadata, and IDs\n",
+    "docs = [\"Qdrant has Langchain integrations\", \"Qdrant also has Llama Index integrations\"]\n",
+    "metadata = [\n",
+    "    {\"source\": \"Langchain-docs\"},\n",
+    "    {\"source\": \"Linkedin-docs\"},\n",
+    "]\n",
+    "ids = [42, 2]\n",
+    "\n",
+    "# Use the new add method\n",
+    "client.add(collection_name=\"demo_collection\", documents=docs, metadata=metadata, ids=ids)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Behind the scenes, Qdrant Client uses the FastEmbed library to make a passage embedding and then uses the Qdrant API to upsert the documents with metadata, put together as a Points into the collection."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Top 5 results:\n",
-      "Rank 1: Maharana Pratap was a Rajput warrior king from Mewar. Score: 0.77\n",
-      "Rank 2: Maharana Pratap is considered a symbol of Rajput resistance against foreign rule. Score: 0.77\n",
-      "Rank 3: His legacy is celebrated in Rajasthan through festivals and monuments. Score: 0.69\n",
-      "Rank 4: He had 11 wives and 17 sons, including Amar Singh I who succeeded him as ruler of Mewar. Score: 0.68\n",
-      "Rank 5: He fought against the Mughal Empire led by Akbar. Score: 0.67\n"
+      "[QueryResponse(id=42, embedding=None, metadata={'document': 'Qdrant has Langchain integrations', 'source': 'Langchain-docs'}, document='Qdrant has Langchain integrations', score=0.8276550115796268), QueryResponse(id=2, embedding=None, metadata={'document': 'Qdrant also has Llama Index integrations', 'source': 'Linkedin-docs'}, document='Qdrant also has Llama Index integrations', score=0.8265536935180283)]\n"
      ]
     }
    ],
    "source": [
-    "from qdrant_client.qdrant_fastembed import QueryResponse\n",
-    "\n",
-    "\n",
-    "def print_top_k_results(results: List[QueryResponse], k: int = 5):\n",
-    "    print(f\"Top {k} results:\")\n",
-    "    for i, result in enumerate(results[:k]):\n",
-    "        print(f\"Rank {i + 1}: {result.document}. Score: {result.score:.2f}\")\n",
-    "\n",
-    "\n",
-    "query_text = \"Who is Maharana Pratap?\"\n",
-    "results = client.query(\n",
-    "    collection_name=\"test_collection\", query_text=query_text, limit=7\n",
-    ")  # Returns limit most relevant documents\n",
-    "\n",
-    "print_top_k_results(results)"
+    "search_result = client.query(collection_name=\"demo_collection\", query_text=\"This is a query document\")\n",
+    "print(search_result)"
    ]
   },
   {
@@ -214,7 +233,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.11.5"
   },
   "orig_nbformat": 4
  },

docs/experimental/Binary Quantization from Scratch.ipynb

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

docs/experimental/Binary Quantization with Qdrant.ipynb

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

docs/index.md

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

docs/overrides/main.html

@@ -5,7 +5,7 @@
     <a href="{{ page.nb_url }}" title="Download Notebook" class="md-content__button md-icon jp-DownloadNB">
         {% include ".icons/material/download.svg" %}
     </a>
-{% endif %} 
+{% endif %}
 
 {{ super() }}
 
@@ -20,7 +20,8 @@
         </span>
         <strong>Qdrant Discord server</strong>
     </a>
-    to get help and share your work! Or check out <a rel="me" href="https://login.cloud.qdrant.io/">Qdrant Cloud</a> to
+    to get help and share your work! Or check out <a rel="me"
+        href="https://cloud.qdrant.io?utm_source=twitter&utm_medium=website&utm_campaign=fastembed">Qdrant Cloud</a> to
     get started with vector search!
 </div>
-{% endblock %}
+{% endblock %}

experiments/01_ONNX_Port.ipynb

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

experiments/dynamic_quantization_export.py

@@ -0,0 +1,473 @@
+"""
+This script is used to convert HuggingFace models to ONNX format and optionally quantize the model using dynamic quantization. 
+This is courtesy of Joshua aka @Xenova
+"""
+
+import json
+import os
+import shutil
+from dataclasses import dataclass, field
+from typing import Optional, Set
+
+import onnx
+from onnxruntime.quantization import QuantType, quantize_dynamic
+from optimum.exporters.onnx import export_models, main_export
+from optimum.exporters.tasks import TasksManager
+from tqdm import tqdm
+from transformers import AutoConfig, AutoTokenizer, HfArgumentParser
+
+DEFAULT_QUANTIZE_PARAMS = {
+    "per_channel": True,
+    "reduce_range": True,
+}
+
+MODEL_SPECIFIC_QUANTIZE_PARAMS = {
+    # Decoder-only models
+    "codegen": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gpt2": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gpt_bigcode": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gptj": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gpt-neo": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "gpt-neox": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "mpt": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "bloom": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "llama": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "opt": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "mistral": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "falcon": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "phi": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "qwen2": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    # Encoder-decoder models
+    "whisper": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    "vision-encoder-decoder": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+    # Encoder-only models
+    "owlv2": {
+        "per_channel": False,
+        "reduce_range": False,
+    },
+}
+
+MODELS_WITHOUT_TOKENIZERS = [
+    "wav2vec2",
+    "wav2vec2-bert",
+    "wavlm",
+    "hubert",
+]
+
+
+@dataclass
+class ConversionArguments:
+    """
+    Arguments used for converting HuggingFace models to onnx.
+    """
+
+    model_id: str = field(metadata={"help": "Model identifier"})
+    tokenizer_id: str = field(default=None, metadata={"help": "Tokenizer identifier (if different to `model_id`)"})
+    quantize: bool = field(default=False, metadata={"help": "Whether to quantize the model."})
+    output_parent_dir: str = field(
+        default="./models/", metadata={"help": "Path where the converted model will be saved to."}
+    )
+
+    task: Optional[str] = field(
+        default="auto",
+        metadata={
+            "help": (
+                "The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:"
+                f" {str(TasksManager.get_all_tasks())}. For decoder models, use `xxx-with-past` to export the model using past key values in the decoder."
+            )
+        },
+    )
+
+    opset: int = field(
+        default=None,
+        metadata={
+            "help": (
+                "If specified, ONNX opset version to export the model with. Otherwise, the default opset will be used."
+            )
+        },
+    )
+
+    device: str = field(default="cpu", metadata={"help": "The device to use to do the export."})
+    skip_validation: bool = field(default=False, metadata={"help": "Whether to skip validation of the converted model"})
+
+    per_channel: bool = field(default=None, metadata={"help": "Whether to quantize weights per channel"})
+    reduce_range: bool = field(
+        default=None,
+        metadata={
+            "help": "Whether to quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode"
+        },
+    )
+
+    output_attentions: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to output attentions from the model. NOTE: This is only supported for whisper models right now."
+        },
+    )
+
+    split_modalities: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to split multimodal models. NOTE: This is only supported for CLIP models right now."
+        },
+    )
+
+    trust_remote_code: bool = field(
+        default=False,
+        metadata={
+            "help": "Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories"
+            "you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository."
+        },
+    )
+
+    custom_onnx_configs: str = field(
+        default=None,
+        metadata={
+            "help": "Experimental usage: override the default ONNX config used for the given model. This argument may be useful for advanced users "
+            "that desire a finer-grained control on the export."
+        },
+    )
+
+
+def get_operators(model: onnx.ModelProto) -> Set[str]:
+    operators = set()
+
+    def traverse_graph(graph):
+        for node in graph.node:
+            operators.add(node.op_type)
+            for attr in node.attribute:
+                if attr.type == onnx.AttributeProto.GRAPH:
+                    subgraph = attr.g
+                    traverse_graph(subgraph)
+
+    traverse_graph(model.graph)
+    return operators
+
+
+def quantize(model_names_or_paths, **quantize_kwargs):
+    """
+    Quantize the weights of the model from float32 to int8 to allow very efficient inference on modern CPU
+
+    Uses unsigned ints for activation values, signed ints for weights, per
+    https://onnxruntime.ai/docs/performance/quantization.html#data-type-selection
+    it is faster on most CPU architectures
+    Args:
+        onnx_model_path: Path to location the exported ONNX model is stored
+    Returns: The Path generated for the quantized
+    """
+
+    quantize_config = dict(**quantize_kwargs, per_model_config={})
+
+    for model in tqdm(model_names_or_paths, desc="Quantizing"):
+        directory_path = os.path.dirname(model)
+        file_name_without_extension = os.path.splitext(os.path.basename(model))[0]
+
+        # NOTE:
+        # As of 2023/04/20, the current latest version of onnxruntime-web is 1.14.0, and does not support INT8 weights for Conv layers.
+        # For this reason, we choose model weight types to ensure compatibility with onnxruntime-web.
+        #
+        # As per docs, signed weight type (QInt8) is faster on most CPUs, so, we use that unless the model contains a Conv layer.
+        # For more information, see:
+        #  - https://github.com/microsoft/onnxruntime/issues/3130#issuecomment-1105200621
+        #  - https://github.com/microsoft/onnxruntime/issues/2339
+
+        loaded_model = onnx.load_model(model)
+        op_types = get_operators(loaded_model)
+        weight_type = QuantType.QUInt8 if "Conv" in op_types else QuantType.QInt8
+
+        quantize_dynamic(
+            model_input=model,
+            model_output=os.path.join(directory_path, f"{file_name_without_extension}_quantized.onnx"),
+            weight_type=weight_type,
+            # TODO allow user to specify these
+            # op_types_to_quantize=['MatMul', 'Add', 'Conv'],
+            extra_options=dict(EnableSubgraph=True),
+            **quantize_kwargs,
+        )
+
+        quantize_config["per_model_config"][file_name_without_extension] = dict(
+            op_types=list(op_types),
+            weight_type=str(weight_type),
+        )
+
+    # Save quantization config
+    with open(os.path.join(directory_path, "quantize_config.json"), "w") as fp:
+        json.dump(quantize_config, fp, indent=4)
+
+
+def main():
+    parser = HfArgumentParser((ConversionArguments,))
+    (conv_args,) = parser.parse_args_into_dataclasses()
+
+    model_id = conv_args.model_id
+    tokenizer_id = conv_args.tokenizer_id or model_id
+
+    output_model_folder = os.path.join(conv_args.output_parent_dir, model_id)
+
+    # Create output folder
+    os.makedirs(output_model_folder, exist_ok=True)
+
+    from_pretrained_kwargs = dict(
+        trust_remote_code=conv_args.trust_remote_code,
+    )
+
+    # Saving the model config
+    config = AutoConfig.from_pretrained(model_id, **from_pretrained_kwargs)
+
+    custom_kwargs = {}
+    if conv_args.custom_onnx_configs is not None:
+        if conv_args.task == "auto":
+            raise Exception("`--task` must be set when exporting with `--custom_onnx_configs`")
+        custom_onnx_configs = json.loads(conv_args.custom_onnx_configs)
+
+        for key in custom_onnx_configs:
+            onnx_configs = TasksManager._SUPPORTED_MODEL_TYPE[custom_onnx_configs[key]]["onnx"]
+            mapping = onnx_configs[conv_args.task]
+            custom_onnx_configs[key] = mapping.func(config, **mapping.keywords)
+
+        custom_kwargs["custom_onnx_configs"] = custom_onnx_configs
+
+    tokenizer = None
+    try:
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, **from_pretrained_kwargs)
+
+        # To avoid inserting all chat templates into tokenizers.js, we save the chat template
+        # to the tokenizer_config.json file, and load it when the tokenizer is loaded.
+        if getattr(tokenizer, "chat_template", None) is None and getattr(tokenizer, "use_default_system_prompt", False):
+            # No chat template specified, and we use the default
+            setattr(tokenizer, "chat_template", tokenizer.default_chat_template)
+
+    except KeyError:
+        pass  # No Tokenizer
+
+    except Exception as e:
+        if config.model_type not in MODELS_WITHOUT_TOKENIZERS:
+            raise e
+
+    core_export_kwargs = dict(
+        opset=conv_args.opset,
+        device=conv_args.device,
+        trust_remote_code=conv_args.trust_remote_code,
+        **custom_kwargs,
+    )
+
+    export_kwargs = dict(
+        model_name_or_path=model_id,
+        output=output_model_folder,
+        task=conv_args.task,
+        do_validation=not conv_args.skip_validation,
+        library_name="transformers",
+        **core_export_kwargs,
+    )
+
+    # Handle special cases
+    if config.model_type == "marian":
+        from .extra.marian import generate_tokenizer_json
+
+        tokenizer_json = generate_tokenizer_json(model_id, tokenizer)
+
+        with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
+            json.dump(tokenizer_json, fp, indent=4)
+
+    elif config.model_type == "esm":
+        from .extra.esm import generate_fast_tokenizer
+
+        fast_tokenizer = generate_fast_tokenizer(tokenizer)
+        fast_tokenizer.save(os.path.join(output_model_folder, "tokenizer.json"))
+
+    elif config.model_type == "whisper":
+        if conv_args.output_attentions:
+            from .extra.whisper import get_main_export_kwargs
+
+            export_kwargs.update(**get_main_export_kwargs(config, "automatic-speech-recognition"))
+
+    elif config.model_type in ("wav2vec2", "wav2vec2-bert", "hubert"):
+        if tokenizer is not None:
+            from .extra.wav2vec2 import generate_tokenizer_json
+
+            tokenizer_json = generate_tokenizer_json(tokenizer)
+
+            with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
+                json.dump(tokenizer_json, fp, indent=4)
+
+    elif config.model_type == "vits":
+        if tokenizer is not None:
+            from .extra.vits import generate_tokenizer_json
+
+            tokenizer_json = generate_tokenizer_json(tokenizer)
+
+            with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
+                json.dump(tokenizer_json, fp, indent=4)
+
+    elif config.model_type == "speecht5":
+        # TODO allow user to specify vocoder path
+        export_kwargs["model_kwargs"] = {"vocoder": "microsoft/speecht5_hifigan"}
+
+        if tokenizer is not None:
+            from .extra.speecht5 import generate_tokenizer_json
+
+            tokenizer_json = generate_tokenizer_json(tokenizer)
+
+            with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
+                json.dump(tokenizer_json, fp, indent=4)
+
+    elif config.model_type in ("owlvit", "owlv2"):
+        # Override default batch size to 1, needed because non-maximum suppression is performed for exporting.
+        # For more information, see https://github.com/huggingface/optimum/blob/e3b7efb1257c011db907ef40ab340e795cc5684c/optimum/exporters/onnx/model_configs.py#L1028-L1032
+        export_kwargs["batch_size"] = 1
+
+    else:
+        pass  # TODO
+
+    # Step 1. convert huggingface model to onnx
+    if not conv_args.split_modalities:
+        main_export(**export_kwargs)
+    else:
+        custom_export_kwargs = dict(
+            output_dir=output_model_folder,
+            **core_export_kwargs,
+        )
+
+        if config.model_type == "clip":
+            # Handle special case for exporting text and vision models separately
+            from transformers.models.clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection
+
+            from .extra.clip import CLIPTextModelWithProjectionOnnxConfig, CLIPVisionModelWithProjectionOnnxConfig
+
+            text_model = CLIPTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
+            vision_model = CLIPVisionModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
+
+            export_models(
+                models_and_onnx_configs={
+                    "text_model": (text_model, CLIPTextModelWithProjectionOnnxConfig(text_model.config)),
+                    "vision_model": (vision_model, CLIPVisionModelWithProjectionOnnxConfig(vision_model.config)),
+                },
+                **custom_export_kwargs,
+            )
+
+        elif config.model_type == "siglip":
+            # Handle special case for exporting text and vision models separately
+            from transformers.models.siglip import SiglipTextModel, SiglipVisionModel
+
+            from .extra.siglip import SiglipTextModelOnnxConfig, SiglipVisionModelOnnxConfig
+
+            text_model = SiglipTextModel.from_pretrained(model_id, **from_pretrained_kwargs)
+            vision_model = SiglipVisionModel.from_pretrained(model_id, **from_pretrained_kwargs)
+
+            export_models(
+                models_and_onnx_configs={
+                    "text_model": (text_model, SiglipTextModelOnnxConfig(text_model.config)),
+                    "vision_model": (vision_model, SiglipVisionModelOnnxConfig(vision_model.config)),
+                },
+                **custom_export_kwargs,
+            )
+
+        # TODO: Enable once https://github.com/huggingface/optimum/pull/1552 is merged
+        # elif config.model_type == 'clap':
+        #     # Handle special case for exporting text and audio models separately
+        #     from .extra.clap import ClapTextModelWithProjectionOnnxConfig, ClapAudioModelWithProjectionOnnxConfig
+        #     from transformers.models.clap import ClapTextModelWithProjection, ClapAudioModelWithProjection
+
+        #     text_model = ClapTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
+        #     audio_model = ClapAudioModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
+
+        #     export_models(
+        #         models_and_onnx_configs={
+        #             "text_model": (text_model, ClapTextModelWithProjectionOnnxConfig(text_model.config)),
+        #             "audio_model": (audio_model, ClapAudioModelWithProjectionOnnxConfig(audio_model.config)),
+        #         },
+        #         **custom_export_kwargs,
+        #     )
+
+        else:
+            raise Exception(f"Unable to export {config.model_type} model with `--split_modalities`.")
+
+    # Step 2. (optional, recommended) quantize the converted model for fast inference and to reduce model size.
+    if conv_args.quantize:
+        # Update quantize config with model specific defaults
+        quantize_config = MODEL_SPECIFIC_QUANTIZE_PARAMS.get(config.model_type, DEFAULT_QUANTIZE_PARAMS)
+
+        # Update if user specified values
+        if conv_args.per_channel is not None:
+            quantize_config["per_channel"] = conv_args.per_channel
+
+        if conv_args.reduce_range is not None:
+            quantize_config["reduce_range"] = conv_args.reduce_range
+
+        quantize(
+            [
+                os.path.join(output_model_folder, x)
+                for x in os.listdir(output_model_folder)
+                if x.endswith(".onnx") and not x.endswith("_quantized.onnx")
+            ],
+            **quantize_config,
+        )
+
+    # Step 3. Move .onnx files to the 'onnx' subfolder
+    os.makedirs(os.path.join(output_model_folder, "onnx"), exist_ok=True)
+    for file in os.listdir(output_model_folder):
+        if file.endswith((".onnx", ".onnx_data")):
+            shutil.move(os.path.join(output_model_folder, file), os.path.join(output_model_folder, "onnx", file))
+
+    # Step 4. Update the generation config if necessary
+    if config.model_type == "whisper":
+        from transformers import GenerationConfig
+
+        from .extra.whisper import get_alignment_heads
+
+        generation_config = GenerationConfig.from_pretrained(model_id, **from_pretrained_kwargs)
+        generation_config.alignment_heads = get_alignment_heads(config)
+        generation_config.save_pretrained(output_model_folder)
+
+
+if __name__ == "__main__":
+    main()

experiments/try_quantized_model.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+from typing import List
+
+import click
+import numpy as np
+import torch
+import torch.nn.functional as F
+from optimum.onnxruntime import ORTModelForFeatureExtraction
+from optimum.pipelines import pipeline
+from torch import Tensor
+from transformers import AutoModel, AutoTokenizer
+
+
+def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+
+def hf_embed(model_id: str, texts: List[str], tokenizer):
+    # Tokenize the input texts
+    model = AutoModel.from_pretrained(model_id)
+    model.eval()
+    encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
+
+    model_output = model(**encoded_input)
+    sentence_embeddings = model_output[0][:, 0]
+    sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+    return sentence_embeddings
+
+
+@click.command()
+@click.option("--model_id", help="model id from huggingface.co/models")
+@click.option("--model_dir", help="The person to greet.")
+def setup(model_id, model_dir):
+    text = "This is a test sentence"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    output_dir = Path(model_dir)
+    model = ORTModelForFeatureExtraction.from_pretrained(output_dir)
+    onnx_quant_embed = pipeline(
+        "feature-extraction", model=model, accelerator="ort", tokenizer=tokenizer, return_tensors=True
+    )
+    quant_embeddings = onnx_quant_embed([text])
+    quant_embeddings = F.normalize(quant_embeddings[0][:,0], p=2, dim=1)
+    quant_embeddings = quant_embeddings.detach().numpy()
+    print(quant_embeddings.shape)
+
+    torch_embeddings = hf_embed(model_id, texts=[text], tokenizer=tokenizer)
+    torch_embeddings = F.normalize(torch_embeddings, p=2, dim=1)
+    torch_embeddings = torch_embeddings.detach().numpy()
+    print(torch_embeddings.shape)
+    assert quant_embeddings.shape == torch_embeddings.shape
+    print(np.allclose(quant_embeddings, torch_embeddings, atol=1e-5))
+
+
+if __name__ == "__main__":
+    setup()

fastembed/__init__.py

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

fastembed/common/model_management.py

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

fastembed/common/models.py

@@ -0,0 +1,52 @@
+import json
+from pathlib import Path
+
+import numpy as np
+from tokenizers import Tokenizer, AddedToken
+
+
+def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
+    config_path = model_dir / "config.json"
+    if not config_path.exists():
+        raise ValueError(f"Could not find config.json in {model_dir}")
+
+    tokenizer_path = model_dir / "tokenizer.json"
+    if not tokenizer_path.exists():
+        raise ValueError(f"Could not find tokenizer.json in {model_dir}")
+
+    tokenizer_config_path = model_dir / "tokenizer_config.json"
+    if not tokenizer_config_path.exists():
+        raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
+
+    tokens_map_path = model_dir / "special_tokens_map.json"
+    if not tokens_map_path.exists():
+        raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
+
+    with open(str(config_path)) as config_file:
+        config = json.load(config_file)
+
+    with open(str(tokenizer_config_path)) as tokenizer_config_file:
+        tokenizer_config = json.load(tokenizer_config_file)
+
+    with open(str(tokens_map_path)) as tokens_map_file:
+        tokens_map = json.load(tokens_map_file)
+
+    tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
+    tokenizer.enable_padding(pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"])
+
+    for token in tokens_map.values():
+        if isinstance(token, str):
+            tokenizer.add_special_tokens([token])
+        elif isinstance(token, dict):
+            tokenizer.add_special_tokens([AddedToken(**token)])
+
+    return tokenizer
+
+
+def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
+    # Calculate the Lp norm along the specified dimension
+    norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
+    norm = np.maximum(norm, eps)  # Avoid division by zero
+    normalized_array = input_array / norm
+    return normalized_array

fastembed/common/utils.py

@@ -0,0 +1,33 @@
+import os
+import tempfile
+from itertools import islice
+from pathlib import Path
+from typing import Union, Iterable, Generator, Optional
+
+
+def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
+    """
+    >>> list(iter_batch([1,2,3,4,5], 3))
+    [[1, 2, 3], [4, 5]]
+    """
+    source_iter = iter(iterable)
+    while source_iter:
+        b = list(islice(source_iter, size))
+        if len(b) == 0:
+            break
+        yield b
+
+
+def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
+    """
+    Define the cache directory for fastembed
+    """
+    if cache_dir is None:
+        default_cache_dir = os.path.join(tempfile.gettempdir(), "fastembed_cache")
+        cache_path = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
+    else:
+        cache_path = Path(cache_dir)
+
+    cache_path.mkdir(parents=True, exist_ok=True)
+
+    return cache_path

fastembed/embedding.py

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

fastembed/parallel_processor.py

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

fastembed/text/e5_onnx_embedding.py

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

fastembed/text/jina_onnx_embedding.py

@@ -0,0 +1,62 @@
+from typing import Type, List, Dict, Any, Tuple
+
+import numpy as np
+
+from fastembed.common.models import normalize
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, EmbeddingWorker, OnnxTextEmbeddingWorker
+
+supported_jina_models = [
+    {
+        "model": "jinaai/jina-embeddings-v2-base-en",
+        "dim": 768,
+        "description": "English embedding model supporting 8192 sequence length",
+        "size_in_GB": 0.55,
+        "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
+    },
+    {
+        "model": "jinaai/jina-embeddings-v2-small-en",
+        "dim": 512,
+        "description": "English embedding model supporting 8192 sequence length",
+        "size_in_GB": 0.13,
+        "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
+    },
+]
+
+
+class JinaOnnxEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[EmbeddingWorker]:
+        return JinaEmbeddingWorker
+
+    @classmethod
+    def mean_pooling(cls, model_output, attention_mask) -> np.ndarray:
+        token_embeddings = model_output
+        input_mask_expanded = (np.expand_dims(attention_mask, axis=-1)).astype(float)
+
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        mask_sum = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
+
+        return sum_embeddings / mask_sum
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_jina_models
+
+    @classmethod
+    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> np.ndarray:
+        embeddings, attn_mask = output
+        return normalize(cls.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+
+
+class JinaEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> OnnxTextEmbedding:
+        return JinaOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)

fastembed/text/onnx_embedding.py

@@ -0,0 +1,343 @@
+import os
+from multiprocessing import get_all_start_methods
+from typing import List, Dict, Any, Optional, Tuple, Union, Iterable, Type
+
+import numpy as np
+import onnxruntime as ort
+
+from fastembed.common.model_management import locate_model_file
+from fastembed.common.models import load_tokenizer, normalize
+from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.parallel_processor import ParallelWorkerPool, Worker
+from fastembed.text.text_embedding_base import TextEmbeddingBase
+
+supported_onnx_models = [
+    {
+        "model": "BAAI/bge-base-en",
+        "dim": 768,
+        "description": "Base English model",
+        "size_in_GB": 0.5,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
+        },
+    },
+    {
+        "model": "BAAI/bge-base-en-v1.5",
+        "dim": 768,
+        "description": "Base English model, v1.5",
+        "size_in_GB": 0.44,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
+            "hf": "qdrant/bge-base-en-v1.5-onnx-q",
+        },
+    },
+    {
+        "model": "BAAI/bge-large-en-v1.5-quantized",
+        "dim": 1024,
+        "description": "Large English model, v1.5",
+        "size_in_GB": 1.34,
+        "sources": {
+            "hf": "qdrant/bge-large-en-v1.5-onnx-q",
+        },
+    },
+    {
+        "model": "BAAI/bge-large-en-v1.5",
+        "dim": 1024,
+        "description": "Large English model, v1.5",
+        "size_in_GB": 1.34,
+        "sources": {
+            "hf": "qdrant/bge-large-en-v1.5-onnx",
+        },
+    },
+    {
+        "model": "BAAI/bge-small-en",
+        "dim": 384,
+        "description": "Fast English model",
+        "size_in_GB": 0.2,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
+        },
+    },
+    # {
+    #     "model": "BAAI/bge-small-en",
+    #     "dim": 384,
+    #     "description": "Fast English model",
+    #     "size_in_GB": 0.2,
+    #     "hf_sources": [],
+    #     "compressed_url_sources": [
+    #         "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en.tar.gz",
+    #         "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz"
+    #     ]
+    # },
+    {
+        "model": "BAAI/bge-small-en-v1.5",
+        "dim": 384,
+        "description": "Fast and Default English model",
+        "size_in_GB": 0.13,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz",
+            "hf": "qdrant/bge-small-en-v1.5-onnx-q",
+        },
+    },
+    {
+        "model": "BAAI/bge-small-zh-v1.5",
+        "dim": 512,
+        "description": "Fast and recommended Chinese model",
+        "size_in_GB": 0.1,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
+        },
+    },
+    {  # todo: it is not a flag embedding
+        "model": "sentence-transformers/all-MiniLM-L6-v2",
+        "dim": 384,
+        "description": "Sentence Transformer model, MiniLM-L6-v2",
+        "size_in_GB": 0.09,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
+            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
+        },
+    },
+    {
+        "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+        "dim": 384,
+        "description": "Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2",
+        "size_in_GB": 0.46,
+        "sources": {
+            "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
+        },
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1",
+        },
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+    },
+    {
+        "model": "thenlper/gte-large",
+        "dim": 1024,
+        "description": "Large general text embeddings model",
+        "size_in_GB": 1.34,
+        "sources": {
+            "hf": "qdrant/gte-large-onnx",
+        },
+    },
+    # {
+    #     "model": "sentence-transformers/all-MiniLM-L6-v2",
+    #     "dim": 384,
+    #     "description": "Sentence Transformer model, MiniLM-L6-v2",
+    #     "size_in_GB": 0.09,
+    #     "hf_sources": [
+    #         "qdrant/all-MiniLM-L6-v2-onnx"
+    #     ],
+    #     "compressed_url_sources": [
+    #         "https://storage.googleapis.com/qdrant-fastembed/fast-all-MiniLM-L6-v2.tar.gz",
+    #         "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz"
+    #     ]
+    # }
+]
+
+
+class OnnxTextEmbedding(TextEmbeddingBase):
+    """Implementation of the Flag Embedding model."""
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_onnx_models
+
+    @classmethod
+    def _get_model_description(cls, model_name: str) -> Dict[str, Any]:
+        """
+        Gets the model description from the model_name.
+
+        Args:
+            model_name (str): The name of the model.
+
+        raises:
+            ValueError: If the model_name is not supported.
+
+        Returns:
+            Dict[str, Any]: The model description.
+        """
+        for model in cls.list_supported_models():
+            if model_name == model["model"]:
+                return model
+
+        raise ValueError(f"Model {model_name} is not supported in FlagEmbedding.")
+
+    def __init__(
+        self,
+        model_name: str = "BAAI/bge-small-en-v1.5",
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            cache_dir (str, optional): The path to the cache directory.
+                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
+                                       Defaults to `fastembed_cache` in the system's temp directory.
+            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        self.model_name = model_name
+        self._model_description = self._get_model_description(model_name)
+
+        self._cache_dir = define_cache_dir(cache_dir)
+        self._model_dir = self.download_model(self._model_description, self._cache_dir)
+        self._max_length = 512
+
+        model_path = locate_model_file(self._model_dir, ["model.onnx", "model_optimized.onnx"])
+
+        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+        onnx_providers = ["CPUExecutionProvider"]
+
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        if self.threads is not None:
+            so.intra_op_num_threads = self.threads
+            so.inter_op_num_threads = self.threads
+
+        self.tokenizer = load_tokenizer(model_dir=self._model_dir, max_length=self._max_length)
+        self.model = ort.InferenceSession(str(model_path), providers=onnx_providers, sess_options=so)
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        Encode a list of documents into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+
+        Args:
+            documents: Iterator of documents or single document to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        is_small = False
+
+        if isinstance(documents, str):
+            documents = [documents]
+            is_small = True
+
+        if isinstance(documents, list):
+            if len(documents) < batch_size:
+                is_small = True
+
+        if parallel == 0:
+            parallel = os.cpu_count()
+
+        if parallel is None or is_small:
+            for batch in iter_batch(documents, batch_size):
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
+        else:
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": self.model_name,
+                "cache_dir": str(self._cache_dir),
+            }
+            pool = ParallelWorkerPool(parallel, self._get_worker_class(), start_method=start_method)
+            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
+                yield from self._post_process_onnx_output(batch)
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
+        return OnnxTextEmbeddingWorker
+
+    def _preprocess_onnx_input(self, onnx_input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    @classmethod
+    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]):
+        embeddings, _ = output
+        return normalize(embeddings[:, 0]).astype(np.float32)
+
+    def onnx_embed(self, documents: List[str]) -> Tuple[np.ndarray, np.ndarray]:
+        encoded = self.tokenizer.encode_batch(documents)
+        input_ids = np.array([e.ids for e in encoded])
+        attention_mask = np.array([e.attention_mask for e in encoded])
+
+        onnx_input = {
+            "input_ids": np.array(input_ids, dtype=np.int64),
+            "attention_mask": np.array(attention_mask, dtype=np.int64),
+            "token_type_ids": np.array([np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64),
+        }
+
+        onnx_input = self._preprocess_onnx_input(onnx_input)
+
+        model_output = self.model.run(None, onnx_input)
+        embeddings = model_output[0]
+        return embeddings, attention_mask
+
+
+class EmbeddingWorker(Worker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> OnnxTextEmbedding:
+        raise NotImplementedError()
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ):
+        self.model = self.init_embedding(model_name, cache_dir)
+
+    @classmethod
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
+        return cls(
+            model_name=model_name,
+            cache_dir=cache_dir,
+        )
+
+    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+        for idx, batch in items:
+            embeddings, attn_mask = self.model.onnx_embed(batch)
+            yield idx, (embeddings, attn_mask)
+
+
+class OnnxTextEmbeddingWorker(EmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> OnnxTextEmbedding:
+        return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)

fastembed/text/text_embedding.py

@@ -0,0 +1,89 @@
+from typing import Optional, Union, Iterable, List, Dict, Any, Type
+
+import numpy as np
+
+from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
+from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.text_embedding_base import TextEmbeddingBase
+
+
+class TextEmbedding(TextEmbeddingBase):
+    EMBEDDINGS_REGISTRY: List[Type[TextEmbeddingBase]] = [
+        OnnxTextEmbedding,
+        E5OnnxEmbedding,
+        JinaOnnxEmbedding,
+    ]
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """
+        Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+
+            Example:
+                ```
+                [
+                    {
+                        "model": "intfloat/multilingual-e5-large",
+                        "dim": 1024,
+                        "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
+                        "size_in_GB": 2.24,
+                        "sources": {
+                            "gcp": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
+                            "hf": "qdrant/multilingual-e5-large-onnx",
+                        }
+                    }
+                ]
+                ```
+        """
+        result = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding.list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str = "BAAI/bge-small-en-v1.5",
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        for embedding in self.EMBEDDINGS_REGISTRY:
+            supported_models = embedding.list_supported_models()
+            if any(model_name == model["model"] for model in supported_models):
+                self.model = embedding(model_name, cache_dir, threads, **kwargs)
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in TextEmbedding."
+            "Please check the supported models using `TextEmbedding.list_supported_models()`"
+        )
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        Encode a list of documents into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+
+        Args:
+            documents: Iterator of documents or single document to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self.model.embed(documents, batch_size, parallel, **kwargs)

fastembed/text/text_embedding_base.py

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

mkdocs.yml

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

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.0.4"
+version = "0.2.2"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -11,40 +11,35 @@ repository = "https://github.com/qdrant/fastembed"
 keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-transformers"]
 
 [tool.poetry.dependencies]
-python = ">=3.8.0,<3.12"
-onnxruntime = "^1.15.1"
-torch = ">=2.0.0, !=2.0.1"
-optimum = ">1.12.0"
-tqdm = "^4.65.0"
-requests = "^2.31.0"
-tokenizers = "^0.13.3"
+python = ">=3.8.0,<3.13"
+onnx = "^1.15.0"
+onnxruntime = "^1.17.0"
+tqdm = "^4.66"
+requests = "^2.31"
+tokenizers = "^0.15.1"
+huggingface-hub = "^0.20"
+loguru = "^0.7.2"
+numpy = [
+    { version = ">=1.21", python = "<3.12" },
+    { version = ">=1.26", python = ">=3.12" }
+]
 
-[tool.poetry.dev-dependencies]
-ruff = "^0.0.277"
-isort = "^5.12.0"
-black = "^23.7.0"
-onnx = "^1.11.0"
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.4.2"
+ruff = "^0.2.2"
 notebook = ">=7.0.2"
-mkdocs-material = "^9.1.21"
-mkdocstrings = "^0.22.0"
-pillow = "^10.0.0"
+mkdocs-material = "^9.5.10"
+mkdocstrings = "^0.24.0"
+pillow = "^10.2.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
-pytest = "^7.4.0"
+pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
+click = "^8.1.7"
 
-[tool.poetry.dependencies.onnxruntime-silicon]
-version = "^1.15.0"
-markers = "sys_platform == 'darwin'"  # This makes it macOS specific
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
-[tool.black]
-line-length = 120
-
-[tool.isort]
-profile = "black"
-
 [tool.ruff]
 line-length = 120

tests/profiling.py

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

tests/test_onnx_embeddings.py

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

tests/test_text_onnx_embeddings.py

@@ -0,0 +1,82 @@
+import os
+
+import numpy as np
+import pytest
+
+from fastembed.text.text_embedding import TextEmbedding
+
+CANONICAL_VECTOR_VALUES = {
+    "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
+    "BAAI/bge-small-en-v1.5": np.array([0.01522374, -0.02271799, 0.00860278, -0.07424029, 0.00386434]),
+    "BAAI/bge-small-zh-v1.5": np.array([-0.01023294, 0.07634465, 0.0691722, -0.04458365, -0.03160762]),
+    "BAAI/bge-base-en": np.array([0.0115, 0.0372, 0.0295, 0.0121, 0.0346]),
+    "BAAI/bge-base-en-v1.5": np.array([0.01129394, 0.05493144, 0.02615099, 0.00328772, 0.02996045]),
+    "BAAI/bge-large-en-v1.5": np.array([0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]),
+    "BAAI/bge-large-en-v1.5-quantized": np.array([0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]),
+    "sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259, 0.0058, 0.0114, 0.0380, -0.0233]),
+    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array([0.0094, 0.0184, 0.0328, 0.0072, -0.0351]),
+    "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
+    "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array(
+        [-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
+    ),
+    "jinaai/jina-embeddings-v2-small-en": np.array([-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]),
+    "jinaai/jina-embeddings-v2-base-en": np.array([-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]),
+    "nomic-ai/nomic-embed-text-v1": np.array([0.0061, 0.0103, -0.0296, -0.0242, -0.0170]),
+    "nomic-ai/nomic-embed-text-v1.5": np.array(
+        [-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
+    ),
+    "thenlper/gte-large": np.array([-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]),
+}
+
+
+def test_embedding():
+    is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
+
+    for model_desc in TextEmbedding.list_supported_models():
+        if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
+            continue
+
+        dim = model_desc["dim"]
+        model = TextEmbedding(model_name=model_desc["model"])
+
+        docs = ["hello world", "flag embedding"]
+        embeddings = list(model.embed(docs))
+        embeddings = np.stack(embeddings, axis=0)
+        assert embeddings.shape == (2, dim)
+
+        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
+        assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3), model_desc["model"]
+
+
+@pytest.mark.parametrize(
+    "n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")]
+)
+def test_batch_embedding(n_dims, model_name):
+    model = TextEmbedding(model_name=model_name)
+
+    docs = ["hello world", "flag embedding"] * 100
+    embeddings = list(model.embed(docs, batch_size=10))
+    embeddings = np.stack(embeddings, axis=0)
+
+    assert embeddings.shape == (200, n_dims)
+
+
+@pytest.mark.parametrize(
+    "n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")]
+)
+def test_parallel_processing(n_dims, model_name):
+    model = TextEmbedding(model_name=model_name)
+
+    docs = ["hello world", "flag embedding"] * 100
+    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+    embeddings = np.stack(embeddings, axis=0)
+
+    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+    embeddings_2 = np.stack(embeddings_2, axis=0)
+
+    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+    embeddings_3 = np.stack(embeddings_3, axis=0)
+
+    assert embeddings.shape == (200, n_dims)
+    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
+    assert np.allclose(embeddings, embeddings_3, atol=1e-3)