Add dynamic quantization from @Xenova

* check for existing files

* (chore: model_management.py):  Add comment

---------

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

README.md

@@ -1,6 +1,6 @@
 # ⚡️ What is FastEmbed?
 
-FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a Github issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
+FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
 The default 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/).
 
@@ -23,7 +23,7 @@ To install the FastEmbed library, pip works:
 pip install fastembed
 ```
 
-## 📖 Usage
+## 📖 Quickstart
 
 ```python
 from fastembed import TextEmbedding
@@ -48,16 +48,14 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```
 
-Might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
+You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
 from qdrant_client import QdrantClient
 
 # Initialize the client
 client = QdrantClient("localhost", port=6333) # For production
-# OR if you just want to try it out quickly:
-# client = QdrantClient(":memory:")
-# client = QdrantClient(path="path/to/db")
+# client = QdrantClient(":memory:") # For small experiments
 
 # Prepare your documents, metadata, and IDs
 docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
@@ -67,7 +65,12 @@ metadata = [
 ]
 ids = [42, 2]
 
-# Use the new add method
+# If you want to change the model:
+# client.set_model("sentence-transformers/all-MiniLM-L6-v2")
+# List of supported models: https://qdrant.github.io/fastembed/examples/Supported_Models
+
+# Use the new add() instead of upsert()
+# This internally calls embed() of the configured embedding model
 client.add(
     collection_name="demo_collection",
     documents=docs,

docs/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

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -40,6 +40,7 @@
        "      <th>dim</th>\n",
        "      <th>description</th>\n",
        "      <th>size_in_GB</th>\n",
+       "      <th>sources</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
@@ -49,6 +50,7 @@
        "      <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",
@@ -56,123 +58,196 @@
        "      <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>3</th>\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>4</th>\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>5</th>\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",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6</th>\n",
-       "      <td>intfloat/multilingual-e5-large</td>\n",
-       "      <td>1024</td>\n",
-       "      <td>Multilingual model, e5-large. Recommend using this model for non-English languages</td>\n",
-       "      <td>2.24</td>\n",
+       "      <td>{'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>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",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>8</th>\n",
-       "      <td>jinaai/jina-embeddings-v2-small-en</td>\n",
-       "      <td>512</td>\n",
-       "      <td>English embedding model supporting 8192 sequence length</td>\n",
-       "      <td>0.13</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>9</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>xenova/multilingual-e5-large</td>\n",
+       "      <td>thenlper/gte-large</td>\n",
        "      <td>1024</td>\n",
-       "      <td>Multilingual model. Recommended for non-English languages</td>\n",
-       "      <td>2.24</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>xenova/paraphrase-multilingual-mpnet-base-v2</td>\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  1024   \n",
-       "3                              BAAI/bge-small-en   384   \n",
-       "4                         BAAI/bge-small-en-v1.5   384   \n",
-       "5                         BAAI/bge-small-zh-v1.5   512   \n",
-       "6                 intfloat/multilingual-e5-large  1024   \n",
-       "7              jinaai/jina-embeddings-v2-base-en   768   \n",
-       "8             jinaai/jina-embeddings-v2-small-en   512   \n",
-       "9         sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "10                  xenova/multilingual-e5-large  1024   \n",
-       "11  xenova/paraphrase-multilingual-mpnet-base-v2   768   \n",
+       "                                                          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                                                                   Fast English model   \n",
-       "4                                                       Fast and Default English model   \n",
-       "5                                                   Fast and recommended Chinese model   \n",
-       "6   Multilingual model, e5-large. Recommend using this model for non-English languages   \n",
-       "7                              English embedding model supporting 8192 sequence length   \n",
-       "8                              English embedding model supporting 8192 sequence length   \n",
-       "9                                             Sentence Transformer model, MiniLM-L6-v2   \n",
-       "10                           Multilingual model. Recommended for non-English languages   \n",
-       "11            Sentence-transformers model for tasks like clustering or semantic search   \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         0.20  \n",
-       "4         0.13  \n",
-       "5         0.10  \n",
-       "6         2.24  \n",
-       "7         0.55  \n",
-       "8         0.13  \n",
-       "9         0.09  \n",
-       "10        2.24  \n",
-       "11        1.11  "
+       "    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,
@@ -181,11 +256,11 @@
     }
    ],
    "source": [
-    "from fastembed.embedding import Embedding\n",
+    "from fastembed import TextEmbedding\n",
     "import pandas as pd\n",
     "\n",
     "pd.set_option(\"display.max_colwidth\", None)\n",
-    "pd.DataFrame(Embedding.list_supported_models())"
+    "pd.DataFrame(TextEmbedding.list_supported_models())"
    ]
   }
  ],
@@ -205,7 +280,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.7"
+   "version": "3.10.13"
   },
   "orig_nbformat": 4
  },

docs/examples/Usage_With_Qdrant.ipynb

@@ -102,19 +102,26 @@
    "execution_count": 4,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 77.7M/77.7M [00:05<00:00, 14.6MiB/s]\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "['6e8fcf7e0ecc407b9b6bb011d169f629',\n",
-       " 'c9d26e7e0ea741b2b1082d097796b28b',\n",
-       " 'cf05747e7eb34d2490b1df1f8be94049',\n",
-       " '208c197266d547a880dfb65e46738b19',\n",
-       " '27bd985c5d6f49d68fc2cf73dac74199',\n",
-       " 'c5e929c8837f4370818c97f63996f8ef',\n",
-       " 'c12213c6cdac470aa2471f2d30dc4041',\n",
-       " '974e64a7d8624f6e9824fa7b9c94f99d',\n",
-       " '0129fae193c740eba092512d8e53ab4a',\n",
-       " '492cad6e741e4aeebb196bd818a97d17']"
+       "['4fa8b10c78da4b18ba0830ba8a57367a',\n",
+       " '2eae04b515ee4e9185a9a0e6be812bba',\n",
+       " 'c6039f88486f47f1835ae3b069c5823c',\n",
+       " 'c2c8c51e305144d1917b373125fb4d95',\n",
+       " '79fd23b9ec0648cdab38d1947c6b933e',\n",
+       " '036aa200d8c3492b8a438e4f825f5e7f',\n",
+       " 'c35c77f3ea37460a9a13723fb77b7367',\n",
+       " '6ebccbca571b40d0ab6e83e5e0f2f562',\n",
+       " '38048c2ccc1d4962a4f8f1bd89c8357a',\n",
+       " 'c6b09308360140c7b4f106af3658a31e']"
       ]
      },
      "execution_count": 4,
@@ -187,12 +194,12 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[QueryResponse(id='42', embedding=None, metadata={'document': 'Qdrant has Langchain integrations', 'source': 'Langchain-docs'}, document='Qdrant has Langchain integrations', score=0.8496814051311954), QueryResponse(id='2', embedding=None, metadata={'document': 'Qdrant also has Llama Index integrations', 'source': 'Linkedin-docs'}, document='Qdrant also has Llama Index integrations', score=0.8478494193031256)]\n"
+      "[QueryResponse(id=42, embedding=None, metadata={'document': 'Qdrant has Langchain integrations', 'source': 'Langchain-docs'}, document='Qdrant has Langchain integrations', score=0.8276550115796268), QueryResponse(id=2, embedding=None, metadata={'document': 'Qdrant also has Llama Index integrations', 'source': 'Linkedin-docs'}, document='Qdrant also has Llama Index integrations', score=0.8265536935180283)]\n"
      ]
     }
    ],
    "source": [
-    "search_result = client.query(collection_name=\"demo_collection\", query_text=[\"This is a query document\"])\n",
+    "search_result = client.query(collection_name=\"demo_collection\", query_text=\"This is a query document\")\n",
     "print(search_result)"
    ]
   },
@@ -226,7 +233,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.11.5"
   },
   "orig_nbformat": 4
  },

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

@@ -131,7 +131,8 @@ class ModelManagement:
         model_tmp_dir = cache_tmp_dir / fast_model_name
         model_dir = Path(cache_dir) / fast_model_name
 
-        if model_dir.exists():
+        # 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():

fastembed/common/models.py

@@ -33,7 +33,7 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
     tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
-    tokenizer.enable_padding(pad_id=config["pad_token_id"], pad_token=tokenizer_config["pad_token"])
+    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):

fastembed/common/utils.py

@@ -24,10 +24,10 @@ def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
     """
     if cache_dir is None:
         default_cache_dir = os.path.join(tempfile.gettempdir(), "fastembed_cache")
-        cache_dir = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
+        cache_path = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
     else:
-        cache_dir = Path(cache_dir)
+        cache_path = Path(cache_dir)
 
-    cache_dir.mkdir(parents=True, exist_ok=True)
+    cache_path.mkdir(parents=True, exist_ok=True)
 
-    return cache_dir
+    return cache_path

fastembed/embedding.py

@@ -4,7 +4,7 @@ from loguru import logger
 
 from fastembed.text.text_embedding import TextEmbedding
 
-logger.warning("DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated." " Use TextEmbedding instead.")
+logger.warning("DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated." "Use from fastembed import TextEmbedding instead.")
 
 DefaultEmbedding = TextEmbedding
 FlagEmbedding = TextEmbedding

fastembed/models.json

@@ -1,133 +0,0 @@
-[
-    {
-        "model": "BAAI/bge-base-en",
-        "dim": 768,
-        "description": "Base English model",
-        "size_in_GB": 0.5,
-        "hf_sources": [],
-        "compressed_url_sources": [
-            "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,
-        "hf_sources": [
-            "qdrant/bge-base-en-v1.5-onnx-q"
-        ],
-        "compressed_url_sources": [
-            "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz"
-        ]
-    },
-    {
-        "model": "BAAI/bge-large-en-v1.5",
-        "dim": 1024,
-        "description": "Large English model, v1.5",
-        "size_in_GB": 1.34,
-        "hf_sources": [
-            "qdrant/bge-large-en-v1.5-onnx",
-            "qdrant/bge-large-en-v1.5-onnx-q"
-        ],
-        "compressed_url_sources": []
-    },
-    {
-        "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,
-        "hf_sources": [
-            "qdrant/bge-small-en-v1.5-onnx-q"
-        ],
-        "compressed_url_sources": [
-            "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz"
-        ]
-    },
-    {
-        "model": "BAAI/bge-small-zh-v1.5",
-        "dim": 512,
-        "description": "Fast and recommended Chinese model",
-        "size_in_GB": 0.1,
-        "hf_sources": [],
-        "compressed_url_sources": [
-            "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz"
-        ]
-    },
-    {
-        "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,
-        "hf_sources": [
-            "qdrant/multilingual-e5-large-onnx"
-        ],
-        "compressed_url_sources": [
-            "https://storage.googleapis.com/qdrant-fastembed/intfloat-multilingual-e5-large.tar.gz"
-        ]
-    },
-    {
-        "model": "jinaai/jina-embeddings-v2-base-en",
-        "dim": 768,
-        "description": "English embedding model supporting 8192 sequence length",
-        "size_in_GB": 0.55,
-        "hf_sources": [
-            "xenova/jina-embeddings-v2-base-en"
-        ],
-        "compressed_url_sources": []
-    },
-    {
-        "model": "jinaai/jina-embeddings-v2-small-en",
-        "dim": 512,
-        "description": " English embedding model supporting 8192 sequence length",
-        "size_in_GB": 0.13,
-        "hf_sources": [
-            "xenova/jina-embeddings-v2-small-en"
-        ],
-        "compressed_url_sources": []
-    },
-    {
-        "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"
-        ]
-    },
-    {
-        "model": "xenova/multilingual-e5-large",
-        "dim": 1024,
-        "description": "Multilingual model. Recommended for non-English languages",
-        "size_in_GB": 2.24,
-        "hf_sources": [
-            "xenova/multilingual-e5-large"
-        ],
-        "compressed_url_sources": []
-    },
-    {
-        "model": "xenova/paraphrase-multilingual-mpnet-base-v2",
-        "dim": 768,
-        "description": "Sentence-transformers model for tasks like clustering or semantic search",
-        "size_in_GB": 1.11,
-        "hf_sources": [
-            "xenova/paraphrase-multilingual-mpnet-base-v2"
-        ],
-        "compressed_url_sources": []
-    }
-]

fastembed/text/e5_onnx_embedding.py

@@ -22,8 +22,8 @@ supported_multilingual_e5_models = [
         "size_in_GB": 1.11,
         "sources": {
             "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
-        }
-    }
+        },
+    },
 ]
 
 
@@ -51,8 +51,8 @@ class E5OnnxEmbedding(OnnxTextEmbedding):
 
 class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
     def init_embedding(
-            self,
-            model_name: str,
-            cache_dir: str,
+        self,
+        model_name: str,
+        cache_dir: str,
     ) -> E5OnnxEmbedding:
         return E5OnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)

fastembed/text/onnx_embedding.py

@@ -1,6 +1,6 @@
 import os
 from multiprocessing import get_all_start_methods
-from typing import List, Dict, Any, Tuple, Union, Iterable, Type
+from typing import List, Dict, Any, Optional, Tuple, Union, Iterable, Type
 
 import numpy as np
 import onnxruntime as ort
@@ -98,6 +98,42 @@ supported_onnx_models = [
             "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,
@@ -149,8 +185,8 @@ class OnnxTextEmbedding(TextEmbeddingBase):
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
-        cache_dir: str = None,
-        threads: int = None,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -193,7 +229,7 @@ class OnnxTextEmbedding(TextEmbeddingBase):
         self,
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int = None,
+        parallel: Optional[int] = None,
         **kwargs,
     ) -> Iterable[np.ndarray]:
         """

fastembed/text/text_embedding.py

@@ -53,24 +53,22 @@ class TextEmbedding(TextEmbeddingBase):
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
-        self.model = None
         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)
-                break
+                return
 
-        if self.model is None:
-            raise ValueError(
-                f"Model {model_name} is not supported in TextEmbedding."
-                "Please check the supported models using `TextEmbedding.list_supported_models()`"
-            )
+        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: int = None,
+        parallel: Optional[int] = None,
         **kwargs,
     ) -> Iterable[np.ndarray]:
         """

fastembed/text/text_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union, Iterable, List, Dict, Any
+from typing import Any, Dict, Iterable, List, Optional, Union
 
 import numpy as np
 
@@ -19,7 +19,7 @@ class TextEmbeddingBase(ModelManagement):
         self,
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
-        parallel: int = None,
+        parallel: Optional[int] = None,
         **kwargs,
     ) -> Iterable[np.ndarray]:
         raise NotImplementedError()
@@ -39,17 +39,19 @@ class TextEmbeddingBase(ModelManagement):
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: str, **kwargs) -> np.ndarray:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
         """
-        Embeds a query
+        Embeds queries
 
         Args:
-            query (str): The query to search for.
+            query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            np.ndarray: The embeddings.
+            Iterable[np.ndarray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
-        query_embedding = list(self.embed([query], **kwargs))[0]
-        return query_embedding
+        if isinstance(query, str):
+            yield from self.embed([query], **kwargs)
+        if isinstance(query, Iterable):
+            yield from self.embed(query, **kwargs)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.2.0"
+version = "0.2.2"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -26,14 +26,15 @@ numpy = [
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.2"
-ruff = "^0.1.13"
+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"
-pre-commit = { version = "^3.6.0", python = ">=3.9,<3.12" }
+pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
+click = "^8.1.7"
 
 
 [build-system]

tests/test_onnx_embeddings.py

@@ -1,76 +0,0 @@
-import os
-
-import numpy as np
-import pytest
-
-from fastembed.embedding import DefaultEmbedding, JinaEmbedding
-
-CANONICAL_VECTOR_VALUES = {
-    "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
-    "BAAI/bge-small-en-v1.5": np.array([0.01522374, -0.02271799, 0.00860278, -0.07424029, 0.00386434]),
-    "BAAI/bge-small-zh-v1.5": np.array([-0.01023294, 0.07634465, 0.0691722, -0.04458365, -0.03160762]),
-    "BAAI/bge-base-en": np.array([0.0115, 0.0372, 0.0295, 0.0121, 0.0346]),
-    "BAAI/bge-base-en-v1.5": np.array([0.01129394, 0.05493144, 0.02615099, 0.00328772, 0.02996045]),
-    "BAAI/bge-large-en-v1.5": 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]),
-    "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
-    "xenova/multilingual-e5-large": np.array([0.00975464, 0.00446568, 0.00655449, -0.0354155, 0.00702112]),
-    "xenova/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]),
-}
-
-
-@pytest.mark.parametrize("embedding_class", [DefaultEmbedding, JinaEmbedding])
-def test_embedding(embedding_class):
-    is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
-
-    for model_desc in embedding_class.list_supported_models():
-        if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
-            continue
-
-        if model_desc["model"] not in CANONICAL_VECTOR_VALUES:
-            continue
-
-        dim = model_desc["dim"]
-        model = embedding_class(model_name=model_desc["model"])
-
-        docs = ["hello world", "flag embedding"]
-        embeddings = list(model.embed(docs))
-        embeddings = np.stack(embeddings, axis=0)
-        assert embeddings.shape == (2, dim)
-
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
-        assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3), model_desc["model"]
-
-
-@pytest.mark.parametrize("n_dims,embedding_class", [(384, DefaultEmbedding), (768, JinaEmbedding)])
-def test_batch_embedding(n_dims, embedding_class):
-    model = embedding_class()
-
-    docs = ["hello world", "flag embedding"] * 100
-    embeddings = list(model.embed(docs, batch_size=10))
-    embeddings = np.stack(embeddings, axis=0)
-
-    assert embeddings.shape == (200, n_dims)
-
-
-@pytest.mark.parametrize("n_dims,embedding_class", [(384, DefaultEmbedding), (768, JinaEmbedding)])
-def test_parallel_processing(n_dims, embedding_class):
-    model = embedding_class()
-
-    docs = ["hello world", "flag embedding"] * 100
-    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
-    embeddings = np.stack(embeddings, axis=0)
-
-    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
-    embeddings_2 = np.stack(embeddings_2, axis=0)
-
-    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
-    embeddings_3 = np.stack(embeddings_3, axis=0)
-
-    assert embeddings.shape == (200, n_dims)
-    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
-    assert np.allclose(embeddings, embeddings_3, atol=1e-3)

tests/test_text_onnx_embeddings.py

@@ -14,12 +14,18 @@ CANONICAL_VECTOR_VALUES = {
     "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]),
 }