* fix(embedding.py): update NotImplementedError message to provide more specific information

* fix(embedding.py): update ValueError message to provide more specific reasons for the error

.gitignore

@@ -161,7 +161,18 @@ cython_debug/
 .DS_Store
 nbs/*.tar.gz
 *.tar.gz
-nbs/local_cache/*/*
-local_cache/*/*
 nbs/fast-*/*
-docs/local_cache/*/*
+local_cache/*/*
+*/local_cache/*/*
+*/*/local_cache/*/*
+docs/experimental/*.parquet
+docs/experimental/*.bin
+qdrant_storage/*
+fooling_around/fast-multilingual-e5-large/config.json
+fooling_around/fast-multilingual-e5-large/model_optimized.onnx
+fooling_around/fast-multilingual-e5-large/model_optimized.onnx.data
+fooling_around/fast-multilingual-e5-large/ort_config.json
+fooling_around/fast-multilingual-e5-large/sentencepiece.bpe.model
+fooling_around/fast-multilingual-e5-large/special_tokens_map.json
+fooling_around/fast-multilingual-e5-large/tokenizer_config.json
+fooling_around/fast-multilingual-e5-large/tokenizer.json

README.md

@@ -1,6 +1,8 @@
 # ⚡️ What is FastEmbed?
 
-FastEmbed is an easy to use -- lightweight, fast, Python library built for retrieval augmented generation. The default embedding supports "query" and "passage" prefixes for the input text.
+FastEmbed is an easy to use -- lightweight, fast, Python library built for retrieval embedding generation. 
+
+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/)
 
 1. Light
     - Quantized model weights
@@ -40,18 +42,20 @@ embedding_model = Embedding(model_name="BAAI/bge-base-en", max_length=512)
 embeddings: List[np.ndarray] = list(embedding_model.embed(documents))
 ```
 
-## 🚒 Under the hood
-
 ### Why fast?
 
 It's important we justify the "fast" in FastEmbed. FastEmbed is fast because:
 
 1. Quantized model weights
-2. ONNX Runtime which allows for inference on CPU, GPU, and other dedicated runtimes
+2. ONNX Runtime which allows for fast inference on CPU and other dedicated runtimes
 
 ### 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
 
 ### Why accurate?
 1. Better than OpenAI Ada-002
-2. Top of the Embedding leaderboards e.g. [MTEB](https://huggingface.co/spaces/mteb/leaderboard)
+2. Top of the Embedding leaderboards e.g. [MTEB](https://huggingface.co/spaces/mteb/leaderboard)
+
+#### Similar Work
+Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.

docs/Getting Started.ipynb

@@ -41,17 +41,19 @@
    "source": [
     "from typing import List\n",
     "import numpy as np\n",
-    "from fastembed.embedding import FlagEmbedding as Embedding\n",
+    "from fastembed.embedding import DefaultEmbedding\n",
     "\n",
-    "# Example list of documents \n",
+    "# Example list of documents\n",
     "documents: List[str] = [\n",
     "    \"Hello, World!\",\n",
     "    \"This is an example document.\",\n",
     "    \"fastembed is supported by and maintained by Qdrant.\",\n",
     "]\n",
     "# Initialize the DefaultEmbedding class with the desired parameters\n",
-    "embedding_model = Embedding(model_name=\"BAAI/bge-small-en\", max_length=512)\n",
-    "embeddings: List[np.ndarray] = list(embedding_model.embed(documents)) # notice that we are casting the generator to a list\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",
     "print(embeddings[0].shape)"
    ]
@@ -61,7 +63,7 @@
    "id": "8c49ae50",
    "metadata": {},
    "source": [
-    "## What is happening under the hood?"
+    "## Let's think step by step"
    ]
   },
   {
@@ -69,6 +71,8 @@
    "id": "92cf4b76",
    "metadata": {},
    "source": [
+    "### Setup\n",
+    "\n",
     "Importing the required classes and modules:"
    ]
   },
@@ -81,7 +85,7 @@
    "source": [
     "from typing import List\n",
     "import numpy as np\n",
-    "from fastembed.embedding import FlagEmbedding as Embedding"
+    "from fastembed.embedding import DefaultEmbedding as Embedding"
    ]
   },
   {
@@ -89,7 +93,7 @@
    "id": "3fd03a71",
    "metadata": {},
    "source": [
-    "Notice that we are using the FlagEmbedding -- which is state of the Art and beats OpenAI's Embedding by a large margin. \n",
+    "Notice that we are using the DefaultEmbedding -- which is a quantized, state of the Art Flag Embedding model which beats OpenAI's Embedding by a large margin. \n",
     "\n",
     "### Prepare your Documents\n",
     "You can define a list of documents that you'd like to embed. These can be sentences, paragraphs, or even entire documents. \n",
@@ -111,10 +115,10 @@
     "# Example list of documents\n",
     "documents: List[str] = [\n",
     "    \"passage: Hello, World!\",\n",
-    "    \"query: Hello, World!\", # these are two different embedding\n",
+    "    \"query: Hello, World!\",  # these are two different embedding\n",
     "    \"passage: This is an example passage.\",\n",
     "    # You can leave out the prefix but it's recommended\n",
-    "    \"fastembed is supported by and maintained by Qdrant.\" \n",
+    "    \"fastembed is supported by and maintained by Qdrant.\",\n",
     "]"
    ]
   },
@@ -126,15 +130,11 @@
     "### Load the Embedding Model Weights\n",
     "Next, initialize the Embedding class with the desired parameters. Here, \"BAAI/bge-small-en\" is the pre-trained model name, and max_length=512 is the maximum token length for each document.\n",
     "\n",
-    "#### 💡 Tip\n",
-    "A model from the same family is also our Default Model, so you can also initialize the DefaultEmbedding class without any parameters:\n",
-    "```python\n",
-    "from fastembed import DefaultEmbedding\n",
-    "embedding_model = DefaultEmbedding()\n",
-    "```\n",
     "This will download the model weights, decompress to directory `local_cache` and load them into the Embedding class.\n",
     "\n",
-    "#### Initialize FlagEmbedding"
+    "#### Initialize DefaultEmbedding\n",
+    "\n",
+    "We will initialize Flag Embeddings with the model name and the maximum token length. That is the DefaultEmbedding class with the model name \"BAAI/bge-small-en\" and max_length=512."
    ]
   },
   {
@@ -144,8 +144,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Initialize the FlagEmbedding class with the desired parameters\n",
-    "embedding_model = Embedding(model_name=\"BAAI/bge-small-en\", max_length=512)"
+    "embedding_model = DefaultEmbedding()"
    ]
   },
   {
@@ -165,7 +164,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "embeddings: List[np.ndarray] = list(embedding_model.embed(documents)) # notice that we are casting the generator to a list"
+    "embeddings: List[np.ndarray] = list(\n",
+    "    embedding_model.embed(documents)\n",
+    ")  # notice that we are casting the generator to a list"
    ]
   },
   {
@@ -183,7 +184,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(embeddings[0].shape) # (768,) or similar output"
+    "print(embeddings[0].shape)  # (384,) or similar output"
    ]
   }
  ],

docs/examples/Retrieval_with_FastEmbed.ipynb

@@ -0,0 +1,222 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ⚓️ Retrieval with FastEmbed\n",
+    "\n",
+    "This notebook demonstrates how to use FastEmbed to perform vector search and retrieval. It consists of the following sections:\n",
+    "\n",
+    "1. Setup: Installing the necessary packages.\n",
+    "2. Importing Libraries: Importing FastEmbed and other libraries.\n",
+    "3. Data Preparation: Example data and embedding generation.\n",
+    "4. Querying: Defining a function to search documents based on a query.\n",
+    "5. Running Queries: Running example queries.\n",
+    "\n",
+    "## Setup\n",
+    "\n",
+    "First, we need to install the dependencies. `fastembed` to create embeddings and perform retrieval."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# !pip install fastembed --quiet --upgrade"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Importing the necessary libraries:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import List\n",
+    "import numpy as np\n",
+    "from fastembed.embedding import FlagEmbedding as Embedding"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Data Preparation\n",
+    "We initialize the embedding model and generate embeddings for the documents.\n",
+    "\n",
+    "### 💡 Tip: Prefer using `query_embed` for queries and `passage_embed` for documents."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(384,) 10\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Example list of documents\n",
+    "documents: List[str] = [\n",
+    "    \"Maharana Pratap was a Rajput warrior king from Mewar\",\n",
+    "    \"He fought against the Mughal Empire led by Akbar\",\n",
+    "    \"The Battle of Haldighati in 1576 was his most famous battle\",\n",
+    "    \"He refused to submit to Akbar and continued guerrilla warfare\",\n",
+    "    \"His capital was Chittorgarh, which he lost to the Mughals\",\n",
+    "    \"He died in 1597 at the age of 57\",\n",
+    "    \"Maharana Pratap is considered a symbol of Rajput resistance against foreign rule\",\n",
+    "    \"His legacy is celebrated in Rajasthan through festivals and monuments\",\n",
+    "    \"He had 11 wives and 17 sons, including Amar Singh I who succeeded him as ruler of Mewar\",\n",
+    "    \"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",
+    "\n",
+    "# We'll use the passage_embed method to get the embeddings for the documents\n",
+    "embeddings: List[np.ndarray] = list(\n",
+    "    embedding_model.passage_embed(documents)\n",
+    ")  # notice that we are casting the generator to a list\n",
+    "\n",
+    "print(embeddings[0].shape, len(embeddings))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Querying\n",
+    "\n",
+    "We'll define a function to print the top k documents based on a query, and prepare a sample query."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "query = \"Who was Maharana Pratap?\"\n",
+    "query_embedding = list(embedding_model.query_embed(query))[0]\n",
+    "plain_query_embedding = list(embedding_model.embed(query))[0]\n",
+    "\n",
+    "\n",
+    "def print_top_k(query_embedding, embeddings, documents, k=5):\n",
+    "    # use numpy to calculate the cosine similarity between the query and the documents\n",
+    "    scores = np.dot(embeddings, query_embedding)\n",
+    "    # sort the scores in descending order\n",
+    "    sorted_scores = np.argsort(scores)[::-1]\n",
+    "    # print the top 5\n",
+    "    for i in range(k):\n",
+    "        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,
+   "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"
+     ]
+    }
+   ],
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The `query_embed` is specifically designed for queries, leading to more relevant and context-aware results. The retrieved documents tend to align closely with the query's intent.\n",
+    "\n",
+    "In contrast, `embed` is a more general-purpose representation that might not capture the nuances of the query as effectively. The retrieved documents using plain embeddings might be less relevant or ordered differently compared to the results obtained using query embeddings.\n",
+    "\n",
+    "Conclusion: Using query and passage embeddings leads to more relevant and context-aware results."
+   ]
+  }
+ ],
+ "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.9.17"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/examples/Usage_With_Qdrant.ipynb

@@ -0,0 +1,223 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "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",
+    "\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",
+    "We will use the Qdrant to add a collection of documents to the engine and then query the collection to retrieve the most relevant documents.\n",
+    "\n",
+    "It consists of the following sections:\n",
+    "\n",
+    "1. Setup: Installing necessary packages, including the Qdrant Client and FastEmbed.\n",
+    "2. Importing Libraries: Importing FastEmbed and other libraries\n",
+    "3. Data Preparation: Example data and embedding generation\n",
+    "4. Querying: Defining a function to search documents based on a query\n",
+    "5. Running Queries: Running example queries\n",
+    "\n",
+    "## Setup\n",
+    "\n",
+    "First, we need to install the dependencies. `fastembed` to create embeddings and perform retrieval, and `qdrant-client` to interact with the Qdrant database."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# !pip install fastembed --quiet --upgrade\n",
+    "\n",
+    "# !pip install git+https://github.com/qdrant/qdrant_client.git@dev"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Importing the necessary libraries:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "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"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Data Preparation\n",
+    "We initialize the embedding model and generate embeddings for the documents.\n",
+    "\n",
+    "### 💡 Tip: Prefer using `query_embed` for queries and `passage_embed` for documents."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Example list of documents\n",
+    "documents: List[str] = [\n",
+    "    \"Maharana Pratap was a Rajput warrior king from Mewar\",\n",
+    "    \"He fought against the Mughal Empire led by Akbar\",\n",
+    "    \"The Battle of Haldighati in 1576 was his most famous battle\",\n",
+    "    \"He refused to submit to Akbar and continued guerrilla warfare\",\n",
+    "    \"His capital was Chittorgarh, which he lost to the Mughals\",\n",
+    "    \"He died in 1597 at the age of 57\",\n",
+    "    \"Maharana Pratap is considered a symbol of Rajput resistance against foreign rule\",\n",
+    "    \"His legacy is celebrated in Rajasthan through festivals and monuments\",\n",
+    "    \"He had 11 wives and 17 sons, including Amar Singh I who succeeded him as ruler of Mewar\",\n",
+    "    \"His life has been depicted in various films, TV shows, and books\",\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This tutorial demonstrates how to utilize the QdrantClient to add documents to a collection and query the collection for relevant documents."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## ➕ Adding Documents\n",
+    "\n",
+    "The `add` creates a collection if it does not already exist. Now, we can add the documents to the collection:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['901ff9d7a90a4e56afe655d1de8e7c06',\n",
+       " '7bae5aa894164398a9be68d47d72ff7a',\n",
+       " 'd940bba124e24469ae3166c3110c62f8',\n",
+       " 'c8dcf956a4f444c6bfe69a00cbeb85ce',\n",
+       " 'd2aeb24c51f549c5b21851be05cb4048',\n",
+       " '6d4c672bafef4db68ae72c8986ee8a72',\n",
+       " 'ff198613768a4361a6d8230b40ee7f78',\n",
+       " 'a133ad72876b48d48e716533c2d78cf0',\n",
+       " 'bdf2b016136e4f57b7ac4b65534031b8',\n",
+       " '5f4f279f304c47839e1a74b0f28de136']"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "client = QdrantClient(\":memory:\")\n",
+    "client.add(collection_name=\"test_collection\", documents=documents)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "These are the ids of the documents we just added. We don't have a use for them in this tutorial, but they can be used to update or delete documents."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 📝 Running Queries\n",
+    "We'll define a function to print the top k documents based on a query, and prepare a sample query."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "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"
+     ]
+    }
+   ],
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 🎬 Conclusion\n",
+    "\n",
+    "This tutorial demonstrates the basics of working with the QdrantClient to add and query documents. By following this guide, you can easily integrate Qdrant into your projects for vector similarity search and retrieval.\n",
+    "\n",
+    "Remember to properly handle the closing of the client connection and further customization of the query parameters according to your specific needs.\n",
+    "\n",
+    "The official Qdrant Python client documentation can be found [here](https://github.com/qdrant/qdrant-client) for more details on customization and advanced features."
+   ]
+  }
+ ],
+ "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.9.17"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/experimental/Binary Quantization from Scratch.ipynb

@@ -0,0 +1,467 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Binary Quantization of OpenAI Embedding\n",
+    "---\n",
+    "\n",
+    "In the world of large-scale data retrieval and processing, efficiency is crucial. With the exponential growth of data, the ability to retrieve information quickly and accurately can significantly affect system performance. This blog post explores a technique known as binary quantization applied to OpenAI embeddings, demonstrating how it can enhance **retrieval latency by 20x** or more.\n",
+    "\n",
+    "## What Are OpenAI Embeddings?\n",
+    "OpenAI embeddings are numerical representations of textual information. They transform text into a vector space where semantically similar texts are mapped close together. This mathematical representation enables computers to understand and process human language more effectively.\n",
+    "\n",
+    "## Binary Quantization\n",
+    "Binary quantization is a method which converts continuous numerical values into binary values (0 or 1). It simplifies the data structure, allowing faster computations. Here's a brief overview of the binary quantization process applied to OpenAI embeddings:\n",
+    "\n",
+    "1. **Load Embeddings**: OpenAI embeddings are loaded from parquet files.\n",
+    "2. **Binary Transformation**: The continuous valued vectors are converted into binary form. Here, values greater than 0 are set to 1, and others remain 0.\n",
+    "3. **Comparison & Retrieval**: Binary vectors are used for comparison using logical XOR operations and other efficient algorithms."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup: Install Dependencies, Imports & Download Embeddings"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install matplotlib tqdm pandas numpy --quiet"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {
+    "id": "WBVTItUX4yyr"
+   },
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "from tqdm import tqdm"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 👨🏾‍💻 Code Walkthrough\n",
+    "Here's an explanation of the code structure provided:\n",
+    "\n",
+    "1. **Loading Data**: OpenAI embeddings are loaded from a parquet files (we can load upto 1M embedding) and concatenated into one array.\n",
+    "2. **Binary Conversion**: A new array with the same shape is initialized with zeros, and the positive values in the original vectors are set to 1.\n",
+    "3. **Accuracy Function**: The accuracy function compares original vectors with binary vectors for a given index, limit, and oversampling rate. The comparison is done using dot products and logical XOR, sorting the results, and measuring the intersection.\n",
+    "4. **Testing**: The accuracy is tested for different oversampling rates (1, 2, 4), revealing a correctness of ~0.96 for an oversampling of 4.\n",
+    "\n",
+    "\n",
+    "## 💿 Loading Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 250
+    },
+    "id": "REJpFqkG7EG2",
+    "outputId": "7a43c0ae-fbcc-45fe-fd58-bfe691297b22"
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 26/26 [00:10<00:00,  2.45it/s]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "(1000000, 1536)"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def get_openai_vectors(force_download: bool = False):\n",
+    "    res = []\n",
+    "    for i in tqdm(range(26)):\n",
+    "        if force_download:\n",
+    "            !wget https://huggingface.co/api/datasets/KShivendu/dbpedia-entities-openai-1M/parquet/KShivendu--dbpedia-entities-openai-1M/train/{i}.parquet\n",
+    "        df = pd.read_parquet(f\"{i}.parquet\", engine=\"pyarrow\")\n",
+    "        res.append(np.stack(df.openai))\n",
+    "        del df\n",
+    "\n",
+    "    openai_vectors = np.concatenate(res)\n",
+    "    del res\n",
+    "    return openai_vectors\n",
+    "\n",
+    "\n",
+    "openai_vectors = get_openai_vectors(force_download=False)\n",
+    "openai_vectors.shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## ㆓ Binary Conversion\n",
+    "\n",
+    "Here, we will use 0 as the threshold for the binary conversion. All values greater than 0 will be set to 1, and others will remain 0. This is a simple and effective way to convert continuous values into binary values for OpenAI embeddings."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {
+    "id": "0JM2-Bj2Jkab"
+   },
+   "outputs": [],
+   "source": [
+    "openai_bin = np.zeros_like(openai_vectors, dtype=np.int8)\n",
+    "openai_bin[openai_vectors > 0] = 1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 🎯 Accuracy Function\n",
+    "\n",
+    "We will use the accuracy function to compare the original vectors with the binary vectors for a given index, limit, and oversampling rate. The comparison is done using dot products and logical XOR, sorting the results, and measuring the intersection."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {
+    "id": "FqshI-GlIERd"
+   },
+   "outputs": [],
+   "source": [
+    "def accuracy(idx, limit: int, oversampling: int):\n",
+    "    scores = np.dot(openai_vectors, openai_vectors[idx])\n",
+    "    dot_results = np.argsort(scores)[-limit:][::-1]\n",
+    "\n",
+    "    bin_scores = 1536 - np.logical_xor(openai_bin, openai_bin[idx]).sum(axis=1)\n",
+    "    bin_results = np.argsort(bin_scores)[-(limit * oversampling) :][::-1]\n",
+    "\n",
+    "    return len(set(dot_results).intersection(set(bin_results))) / limit"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 📊 Results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "qtzUlq_sFTRf",
+    "outputId": "17fe04ea-4f73-4a57-990b-180f1c04b472"
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|          | 0/4 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 1, 'limit': 10, 'recall': 0.8}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 2/2 [00:33<00:00, 16.98s/it]\n",
+      " 25%|██▌       | 1/4 [00:33<01:41, 33.96s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 1, 'limit': 100, 'recall': 0.708}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": []
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 2, 'limit': 10, 'recall': 0.95}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 2/2 [00:32<00:00, 16.38s/it]\n",
+      " 50%|█████     | 2/4 [01:06<01:06, 33.26s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 2, 'limit': 100, 'recall': 0.877}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": []
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 3, 'limit': 10, 'recall': 0.96}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 2/2 [00:32<00:00, 16.49s/it]\n",
+      " 75%|███████▌  | 3/4 [01:39<00:33, 33.13s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 3, 'limit': 100, 'recall': 0.937}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": []
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 5, 'limit': 10, 'recall': 0.9800000000000001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 2/2 [00:32<00:00, 16.47s/it]\n",
+      "100%|██████████| 4/4 [02:12<00:00, 33.17s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 5, 'limit': 100, 'recall': 0.977}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "number_of_samples = 10\n",
+    "limits = [10, 100]\n",
+    "sampling_rate = [1, 2, 3, 5]\n",
+    "results = []\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",
+    "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",
+    "        print(result)\n",
+    "        results.append(result)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "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>sampling_rate</th>\n",
+       "      <th>limit</th>\n",
+       "      <th>recall</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>1</td>\n",
+       "      <td>10</td>\n",
+       "      <td>0.800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>1</td>\n",
+       "      <td>100</td>\n",
+       "      <td>0.708</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>2</td>\n",
+       "      <td>10</td>\n",
+       "      <td>0.950</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>2</td>\n",
+       "      <td>100</td>\n",
+       "      <td>0.877</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>3</td>\n",
+       "      <td>10</td>\n",
+       "      <td>0.960</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>3</td>\n",
+       "      <td>100</td>\n",
+       "      <td>0.937</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
+       "      <td>5</td>\n",
+       "      <td>10</td>\n",
+       "      <td>0.980</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</th>\n",
+       "      <td>5</td>\n",
+       "      <td>100</td>\n",
+       "      <td>0.977</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   sampling_rate  limit  recall\n",
+       "0              1     10   0.800\n",
+       "1              1    100   0.708\n",
+       "2              2     10   0.950\n",
+       "3              2    100   0.877\n",
+       "4              3     10   0.960\n",
+       "5              3    100   0.937\n",
+       "6              5     10   0.980\n",
+       "7              5    100   0.977"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "results = pd.DataFrame(results)\n",
+    "results"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "| sampling_rate | limit | accuracy |\n",
+    "|---------------|-------|----------|\n",
+    "| 1             | 10    | 0.800    |\n",
+    "| 1             | 100   | 0.708    |\n",
+    "| 2             | 10    | 0.950    |\n",
+    "| 2             | 100   | 0.877    |\n",
+    "| 4             | 10    | 0.970    |\n",
+    "| 4             | 100   | 0.956    |\n",
+    "| 8             | 10    | 0.990    |\n",
+    "| 8             | 100   | 0.990    |\n",
+    "| 16            | 10    | 1.000    |\n",
+    "| 16            | 100   | 0.998    |"
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "machine_shape": "hm",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.9.17"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

docs/index.md

@@ -1,8 +1,10 @@
 # ⚡️ What is FastEmbed?
 
-FastEmbed is an easy to use -- lightweight, fast, Python library built for retrieval augmented generation. The default embedding supports "query" and "passage" prefixes for the input text.
+FastEmbed is an easy to use -- lightweight, fast, Python library built for retrieval embedding generation. 
 
-## 🚀 Installation
+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. 
+
+Advanced user? Skip ahead to [Retrieval with FastEmbed](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/)
 
 To install the FastEmbed library, pip works: 
 
@@ -40,4 +42,4 @@ It's important we justify the "fast" in FastEmbed. FastEmbed is fast because:
 
 ### Why accurate?
 1. Better than OpenAI Ada-002
-2. Top of the Embedding leaderboards e.g. [MTEB](https://huggingface.co/spaces/mteb/leaderboard)
+2. Top of the Embedding leaderboards e.g. [MTEB](https://huggingface.co/spaces/mteb/leaderboard)

docs/overrides/main.html

@@ -14,11 +14,11 @@
 {% block announce %}
 <div style="text-align: center">
     If you're using FastEmbed from Qdrant, join the
-    <a rel="me" href="https://discord.gg/g6g4tagk">
+    <a rel="me" href="https://discord.gg/Qy6HCJK9Dc">
         <span class="twemoji mastodon">
             {% include ".icons/fontawesome/brands/discord.svg" %}
         </span>
-        <strong>Discord server</strong>
+        <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
     get started with vector search!

fastembed/embedding.py

@@ -6,47 +6,30 @@ from pathlib import Path
 from typing import Iterable, List
 
 import numpy as np
-import onnxruntime as ort
 import requests
-from tokenizers import Tokenizer
+from optimum.onnxruntime import ORTModelForFeatureExtraction
 from tqdm import tqdm
-
-# set the default logger to ERROR level with this
-ort.set_default_logger_severity(3)
+from transformers import AutoTokenizer
 
 
-# Use pytorches default epsilon for division by zero
-# https://pytorch.org/docs/stable/generated/torch.nn.functional.normalize.html
-def normalize(v):
-    norm = np.linalg.norm(v, axis=1)
-    norm[norm == 0] = 1e-12
-    return v / norm[:, np.newaxis]
-
-
-class ONNXProviders:
-    """List of Execution Providers: https://onnxruntime.ai/docs/execution-providers"""
-
-    CPU = "CPUExecutionProvider"
-    GPU = "CUDAExecutionProvider"
-    # GPU support is experimental, and can be improved: https://onnxruntime.ai/docs/api/python/api_summary.html#data-on-device
-    Metal = "CoreMLExecutionProvider"
-
+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 (): 
+        ABC ():
 
     Raises:
-        NotImplementedError: _description_
+        NotImplementedError: Raised when you call an abstract method that has not been implemented.
         PermissionError: _description_
-        ValueError: _description_
-        ValueError: _description_
-        ValueError: _description_
-        ValueError: _description_
-        ValueError: _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:
@@ -166,19 +149,25 @@ class Embedding(ABC):
 
         assert "/" in model_name, "model_name must be in the format <org>/<model> e.g. BAAI/bge-base-en"
 
-        model_name = model_name.split("/")[-1]
-
-        fast_model_name = f"fast-{model_name}"
+        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"
-        self.download_file_from_gcs(
-            f"https://storage.googleapis.com/qdrant-fastembed/{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}"
@@ -187,6 +176,40 @@ class Embedding(ABC):
 
         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):
     """
@@ -196,43 +219,41 @@ class FlagEmbedding(Embedding):
         Embedding (_type_): _description_
     """
 
-
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en",
-        onnx_providers=None,
         max_length: int = 512,
         cache_dir: str = None,
     ):
         """
         Args:
             model_name (str): The name of the model to use.
-            onnx_providers (List[str]): A list of ONNX providers to use.
-            max_length (int, optional): The maximum length of the input text. Defaults to 512.
+            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 onnx_providers is None:
-            onnx_providers = [ONNXProviders.CPU]
-
         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)
-        tokenizer_path = model_dir / "tokenizer.json"
-        if not tokenizer_path.exists():
+        if not (model_dir / "tokenizer.json").exists():
             raise ValueError(f"Could not find tokenizer.json in {model_dir}")
-        model_path = model_dir / "model_optimized.onnx"
-        if not model_path.exists():
-            raise ValueError(f"Could not find model_optimized.onnx 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 = Tokenizer.from_file(str(tokenizer_path))
-        self.tokenizer.enable_truncation(max_length=max_length)
-        self.tokenizer.enable_padding(pad_id=0, pad_token="[PAD]", length=max_length)
-        self.model = ort.InferenceSession(str(model_path), providers=onnx_providers)
+        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]:
         """
@@ -247,26 +268,12 @@ class FlagEmbedding(Embedding):
             List of embeddings, one per document
         """
         # TODO: Replace loop with parallelized batching
-        for i in range(0, len(documents), batch_size):
-            batch = documents[i : i + batch_size]
-            encoded = [self.tokenizer.encode(d) for d in batch]
-            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),
-            }
-            model_output = self.model.run(None, onnx_input)
-            last_hidden_state = model_output[0]
-            # Perform mean pooling with attention weighting
-            input_mask_expanded = np.broadcast_to(np.expand_dims(attention_mask, -1), last_hidden_state.shape)
-            embeddings = np.sum(last_hidden_state * input_mask_expanded, 1) / np.clip(
-                input_mask_expanded.sum(1), a_min=1e-9, a_max=None
-            )
-            # TODO: Should we normalize after all batches are done?
-            embeddings = normalize(embeddings).astype(np.float32)
-            yield from embeddings
+        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):
@@ -276,6 +283,7 @@ class DefaultEmbedding(FlagEmbedding):
     Args:
         FlagEmbedding (_type_): _description_
     """
+
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en",
@@ -283,9 +291,9 @@ class DefaultEmbedding(FlagEmbedding):
         max_length: int = 512,
         cache_dir: str = None,
     ):
-        if onnx_providers is None:
-            onnx_providers = [ONNXProviders.CPU]
-        super().__init__(model_name, onnx_providers, max_length, cache_dir)
+        # if onnx_providers is None:
+        #     onnx_providers = [ONNXProviders.CPU]
+        super().__init__(model_name, max_length=max_length, cache_dir=cache_dir)
 
 
 class OpenAIEmbedding(Embedding):
@@ -298,4 +306,4 @@ class OpenAIEmbedding(Embedding):
         # Use your OpenAI model to embed the texts
         # return self.model.embed(texts)
         raise NotImplementedError
-
+        raise NotImplementedError

fooling_around/01_ONNX_Port.ipynb

@@ -43,7 +43,7 @@
     "# model_id = \"sentence-transformers/all-MiniLM-L6-v2\"\n",
     "# model_id = \"thenlper/gte-base\"\n",
     "# model_id = \"intfloat/multilingual-e5-large\"\n",
-    "model_id = \"BAAI/bge-base-en\"\n",
+    "model_id = \"BAAI/bge-small-en-v1.5\"\n",
     "save_dir = f\"fast-{model_id.split('/')[1]}\"\n",
     "print(save_dir)"
    ]
@@ -64,11 +64,11 @@
     "input_texts = [\n",
     "    \"query: how much protein should a female eat\",\n",
     "    \"query: 南瓜的家常做法\",\n",
-    "    \"query: भारत का राष्ट्रीय खेल कौन-सा है?\", # Hindi text\n",
-    "    \"query: భారత్ దేశంలో రాష్ట్రపతి ఎవరు?\", # Telugu text\n",
-    "    \"query: இந்தியாவின் தேசிய கோப்பை எது?\", # Tamil text\n",
-    "    \"query: ಭಾರತದಲ್ಲಿ ರಾಷ್ಟ್ರಪತಿ ಯಾರು?\", # Kannada text\n",
-    "    \"query: ഇന്ത്യയുടെ രാഷ്ട്രീയ ഗാനം എന്താണ്?\", # Malayalam text\n",
+    "    \"query: भारत का राष्ट्रीय खेल कौन-सा है?\",  # Hindi text\n",
+    "    \"query: భారత్ దేశంలో రాష్ట్రపతి ఎవరు?\",  # Telugu text\n",
+    "    \"query: இந்தியாவின் தேசிய கோப்பை எது?\",  # Tamil text\n",
+    "    \"query: ಭಾರತದಲ್ಲಿ ರಾಷ್ಟ್ರಪತಿ ಯಾರು?\",  # Kannada text\n",
+    "    \"query: ഇന്ത്യയുടെ രാഷ്ട്രീയ ഗാനം എന്താണ്?\",  # Malayalam text\n",
     "]\n",
     "\n",
     "english_texts = [\n",
@@ -76,7 +76,7 @@
     "    \"Machine Learning: Turning data into knowledge, one algorithm at a time.\",\n",
     "    \"Python: The language that makes programming a piece of cake.\",\n",
     "    \"fastembed: Accelerating embeddings for lightning-fast similarity search.\",\n",
-    "    \"Qdrant: The ultimate tool for high-dimensional indexing and search.\"\n",
+    "    \"Qdrant: The ultimate tool for high-dimensional indexing and search.\",\n",
     "]"
    ]
   },
@@ -87,11 +87,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
     "def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:\n",
     "    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)\n",
     "    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]\n",
     "\n",
+    "\n",
     "def hf_embed(model_id: str, inputs: List[str]):\n",
     "    # Tokenize the input texts\n",
     "    batch_dict = hf_tokenizer(inputs, max_length=512, padding=True, truncation=True, return_tensors=\"pt\")\n",
@@ -101,9 +101,17 @@
     "\n",
     "    # normalize embeddings\n",
     "    embeddings = F.normalize(embeddings, p=2, dim=1)\n",
-    "    scores = (embeddings[:2] @ embeddings[2:].T) * 100\n",
-    "    # print(scores.tolist())\n",
-    "    return scores"
+    "    return embeddings.detach().numpy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "69bb4501",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hf_embed(inputs=english_texts, model_id=model_id)"
    ]
   },
   {
@@ -123,7 +131,7 @@
     "    p.unlink()\n",
     "\n",
     "# Load the optimization configuration detailing the optimization we wish to apply\n",
-    "optimization_config = AutoOptimizationConfig.O3()\n",
+    "optimization_config = AutoOptimizationConfig.O4()\n",
     "optimizer = ORTOptimizer.from_pretrained(model)\n",
     "\n",
     "optimizer.optimize(save_dir=save_dir, optimization_config=optimization_config, use_external_data_format=True)\n",
@@ -141,10 +149,30 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "onnx_quant_embed = pipeline(\"feature-extraction\", model=model, accelerator=\"ort\")\n",
+    "onnx_quant_embed = pipeline(\"feature-extraction\", model=model, accelerator=\"ort\", tokenizer=tokenizer,return_tensors=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "51fa5775",
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "df405d70",
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "\n",
-    "\n",
-    "def measure_pipeline_time(pipeline, input_texts: List[str], num_runs = 10, **kwargs: Any) -> Tuple[float, float]:\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",
     "    total_chars = sum(len(text) for text in input_texts)\n",
@@ -157,7 +185,7 @@
     "    mean_time = np.mean(times)\n",
     "    std_dev = np.std(times)\n",
     "    chars_per_second = total_chars / mean_time\n",
-    "    return mean_time, std_dev, chars_per_second\n"
+    "    return mean_time, std_dev, chars_per_second"
    ]
   },
   {
@@ -217,7 +245,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "id": "578b1d74",
    "metadata": {},
    "outputs": [],
@@ -226,6 +254,7 @@
     "from pathlib import Path\n",
     "import tarfile\n",
     "\n",
+    "save_dir = Path(\"../local_cache/fast-bge-small-en-v1.5\")\n",
     "\n",
     "def compress(directory_path):\n",
     "    directory_path = Path(directory_path)\n",
@@ -254,14 +283,30 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "id": "1dab9595",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/opt/homebrew/Caskroom/miniconda/base/envs/fst/lib/python3.9/site-packages/google/auth/_default.py:76: UserWarning: Your application has authenticated using end user credentials from Google Cloud SDK without a quota project. You might receive a \"quota exceeded\" or \"API not enabled\" error. See the following page for troubleshooting: https://cloud.google.com/docs/authentication/adc-troubleshooting/user-creds. \n",
+      "  warnings.warn(_CLOUD_SDK_CREDENTIALS_WARNING)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "File fast-bge-small-en-v1.5.tar.gz uploaded to qdrant-fastembed.\n"
+     ]
+    }
+   ],
    "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",
@@ -272,7 +317,7 @@
     "    print(f\"File {source_file_path} uploaded to {bucket_name}.\")\n",
     "\n",
     "\n",
-    "upload(\"qdrant-fastembed\", compressed_file_name)"
+    "upload(\"qdrant-fastembed\", source_file_path=compressed_file_name)"
    ]
   },
   {
@@ -304,7 +349,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.9.17"
   }
  },
  "nbformat": 4,

mkdocs.yml

@@ -8,7 +8,7 @@ remote_branch: gh-pages
 remote_name: origin
 
 copyright: |
-  Maintained by <a href="https://github.com/qdrant">Qdrant</a>.
+  Maintained by <a href="https://github.com/qdrant">Qdrant</a>. Originally created by <a href="https://nirantk.com/about">Nirant Kasliwal</a>.
 
 theme:
   name: material

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.0.3a0"
-description = "Fast, State of the Art Quantized Embedding Models"
+version = "0.0.4"
+description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
 readme = "README.md"
@@ -13,6 +13,8 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
 [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"
@@ -21,7 +23,6 @@ tokenizers = "^0.13.3"
 ruff = "^0.0.277"
 isort = "^5.12.0"
 black = "^23.7.0"
-optimum = "^1.11.1"
 onnx = "^1.11.0"
 notebook = ">=7.0.2"
 mkdocs-material = "^9.1.21"

tests/profiling.py

@@ -0,0 +1,145 @@
+# %% [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
+
+import numpy as np
+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.
+
+# %%
+documents: List[str] = [
+    "Chandrayaan-3 is India's third lunar mission",
+    "It aimed to land a rover on the Moon's surface - joining the US, China and Russia",
+    "The mission is a follow-up to Chandrayaan-2, which had partial success",
+    "Chandrayaan-3 will be launched by the Indian Space Research Organisation (ISRO)",
+    "The estimated cost of the mission is around $35 million",
+    "It will carry instruments to study the lunar surface and atmosphere",
+    "Chandrayaan-3 landed on the Moon's surface on 23rd August 2023",
+    "It consists of a lander named Vikram and a rover named Pragyan similar to Chandrayaan-2. Its propulsion module would act like an orbiter.",
+    "The propulsion module carries the lander and rover configuration until the spacecraft is in a 100-kilometre (62 mi) lunar orbit",
+    "The mission used GSLV Mk III rocket for its launch",
+    "Chandrayaan-3 was launched from the Satish Dhawan Space Centre in Sriharikota",
+    "Chandrayaan-3 was launched earlier in the year 2023",
+]
+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:
+    """
+    HuggingFace Transformer implementation of FlagEmbedding
+    Based on https://huggingface.co/BAAI/bge-base-en
+    """
+
+    def __init__(self, model_id: str):
+        self.model = AutoModel.from_pretrained(model_id)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+    def embed(self, texts: List[str]):
+        encoded_input = self.tokenizer(texts, max_length=512, padding=True, truncation=True, return_tensors="pt")
+        model_output = self.model(**encoded_input)
+        sentence_embeddings = model_output[0][:, 0]
+        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.
+
+# %%
+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)
+        end_time = time.time()
+
+        times.append(end_time - start_time)
+
+    # 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(
+    hf_stats: Tuple[float, float, float], fst_stats: Tuple[float, float, float], documents: list
+):
+    # Calculating total characters in documents
+    total_characters = sum(len(doc) for doc in documents)
+
+    # Calculating characters per second for each model
+    hf_chars_per_sec = total_characters / hf_stats[0]  # Mean time is at index 0
+    fst_chars_per_sec = total_characters / fst_stats[0]
+
+    # Plotting the bar chart
+    models = ["HF Embed (Torch)", "FastEmbed"]
+    chars_per_sec = [hf_chars_per_sec, fst_chars_per_sec]
+
+    bars = plt.bar(models, chars_per_sec, color=["#1f356c", "#dd1f4b"])
+    plt.ylabel("Characters per Second")
+    plt.title("Characters Processed per Second Comparison")
+
+    # Adding the number at the top of each bar
+    for bar, chars in zip(bars, chars_per_sec):
+        plt.text(
+            bar.get_x() + bar.get_width() / 2,
+            bar.get_height(),
+            f"{chars:.1f}",
+            ha="center",
+            va="bottom",
+            color="#1f356c",
+            fontsize=12,
+        )
+
+    plt.show()
+
+
+plot_character_per_second_comparison(hf_stats, fst_stats, documents)
+
+

tests/test_onnx_embeddigns.py

@@ -1,25 +0,0 @@
-from fastembed.embedding import DefaultEmbedding
-
-
-def test_onnx_inference():
-    model = DefaultEmbedding(
-        model_name="sentence-transformers/all-MiniLM-L6-v2"
-    )
-
-    docs = [
-        "hello world"
-    ]
-
-    expected_vector = [
-        -3.44772711e-02,  3.10231801e-02,  6.73499797e-03,  2.61089597e-02
-    ]
-
-    vector = list(model.embed(docs))[0]
-    assert len(vector) == 384
-
-    for i in range(len(expected_vector)):
-        assert abs(vector[i] - expected_vector[i]) < 1e-3
-
-
-
-

tests/test_onnx_embeddings.py

@@ -0,0 +1,13 @@
+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()