Merge pull request #15 from qdrant/v0.0.5

bump version 0.0.5

.github/workflows/python-tests.yml

@@ -20,6 +20,8 @@ jobs:
           - '3.11.x'
         os:
           - ubuntu-latest
+          - macos-latest
+          - windows-latest
 
     runs-on: ${{ matrix.os }}
 

README.md

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

docs/Getting Started.ipynb

@@ -16,12 +16,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "ada95c6a",
    "metadata": {},
    "outputs": [],
    "source": [
-    "!pip install fastembed --upgrade # Install fastembed"
+    "!pip install fastembed --upgrade --quiet # Install fastembed "
    ]
   },
   {
@@ -34,10 +34,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "id": "b61c6552",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([384])\n"
+     ]
+    }
+   ],
    "source": [
     "from typing import List\n",
     "import numpy as np\n",
@@ -51,10 +66,7 @@
     "]\n",
     "# Initialize the DefaultEmbedding class with the desired parameters\n",
     "embedding_model = DefaultEmbedding(model_name=\"BAAI/bge-small-en\", max_length=512)\n",
-    "embeddings: List[np.ndarray] = list(\n",
-    "    embedding_model.embed(documents)\n",
-    ")  # notice that we are casting the generator to a list\n",
-    "\n",
+    "embeddings: List[np.ndarray] = embedding_model.embed(documents)\n",
     "print(embeddings[0].shape)"
    ]
   },
@@ -78,7 +90,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "c0a6f634",
    "metadata": {},
    "outputs": [],
@@ -107,7 +119,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "id": "145a56ce",
    "metadata": {},
    "outputs": [],
@@ -139,7 +151,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "id": "272c8915",
    "metadata": {},
    "outputs": [],
@@ -159,14 +171,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "id": "8013eee9",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
+     ]
+    }
+   ],
    "source": [
-    "embeddings: List[np.ndarray] = list(\n",
-    "    embedding_model.embed(documents)\n",
-    ")  # notice that we are casting the generator to a list"
+    "embeddings: List[np.ndarray] = embedding_model.embed(documents)"
    ]
   },
   {
@@ -179,10 +197,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "id": "0d8c8e08",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([384])\n"
+     ]
+    }
+   ],
    "source": [
     "print(embeddings[0].shape)  # (384,) or similar output"
    ]

docs/examples/Supported_Models.ipynb

@@ -0,0 +1,115 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>model</th>\n",
+       "      <th>dim</th>\n",
+       "      <th>description</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>BAAI/bge-small-en</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Fast and Default English model</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>BAAI/bge-base-en</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Base English model</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Sentence Transformer model, MiniLM-L6-v2</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>intfloat/multilingual-e5-large</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Multilingual model, e5-large. Recommend using this model for non-English languages. Recommend using this via Torch implementation of FastEmbed</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                    model   dim  \\\n",
+       "0                       BAAI/bge-small-en   384   \n",
+       "1                        BAAI/bge-base-en   768   \n",
+       "2  sentence-transformers/all-MiniLM-L6-v2   384   \n",
+       "3          intfloat/multilingual-e5-large  1024   \n",
+       "\n",
+       "                                                                                                                                      description  \n",
+       "0                                                                                                                  Fast and Default English model  \n",
+       "1                                                                                                                              Base English model  \n",
+       "2                                                                                                        Sentence Transformer model, MiniLM-L6-v2  \n",
+       "3  Multilingual model, e5-large. Recommend using this model for non-English languages. Recommend using this via Torch implementation of FastEmbed  "
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "from fastembed.embedding import Embedding\n",
+    "import pandas as pd\n",
+    "pd.set_option('display.max_colwidth', None)\n",
+    "pd.DataFrame(Embedding.list_supported_models())"
+   ]
+  }
+ ],
+ "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

@@ -5,9 +5,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# [Experimental] Usage With Qdrant\n",
-    "\n",
-    "> **Note:** This notebook is experimental and is subject to change. For working with this, use the dev branch of QdrantClient.\n",
+    "# Usage With Qdrant\n",
     "\n",
     "This notebook demonstrates how to use FastEmbed and Qdrant to perform vector search and retrieval. Qdrant is an open-source vector similarity search engine that is used to store, organize, and query collections of high-dimensional vectors. \n",
     "\n",
@@ -28,13 +26,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 1,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[33mDEPRECATION: pytorch-lightning 1.6.5 has a non-standard dependency specifier torch>=1.8.*. pip 23.3 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
+      "\u001b[0m"
+     ]
+    }
+   ],
    "source": [
-    "# !pip install fastembed --quiet --upgrade\n",
-    "\n",
-    "# !pip install git+https://github.com/qdrant/qdrant_client.git@dev"
+    "!pip install 'qdrant-client[fastembed]' --quiet --upgrade"
    ]
   },
   {
@@ -46,7 +51,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -68,7 +73,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -105,25 +110,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "['901ff9d7a90a4e56afe655d1de8e7c06',\n",
-       " '7bae5aa894164398a9be68d47d72ff7a',\n",
-       " 'd940bba124e24469ae3166c3110c62f8',\n",
-       " 'c8dcf956a4f444c6bfe69a00cbeb85ce',\n",
-       " 'd2aeb24c51f549c5b21851be05cb4048',\n",
-       " '6d4c672bafef4db68ae72c8986ee8a72',\n",
-       " 'ff198613768a4361a6d8230b40ee7f78',\n",
-       " 'a133ad72876b48d48e716533c2d78cf0',\n",
-       " 'bdf2b016136e4f57b7ac4b65534031b8',\n",
-       " '5f4f279f304c47839e1a74b0f28de136']"
+       "['77e1e4724dd243b08608f57d5692f6aa',\n",
+       " '74841e5dc3594646bda2c6a6d2795dbd',\n",
+       " '6ef39a9445604d0da84d04f760cd7cf7',\n",
+       " 'e659503d3b3748ef90f23c778274835b',\n",
+       " 'b999675068cd413f93faa0cc890c3819',\n",
+       " '8e452f2935cf4e4b80d8eea68c2aad58',\n",
+       " '28ed4fd4592c48c9a0519618d51bb86e',\n",
+       " '59378c784c5f49109bef65fdc4061334',\n",
+       " 'a78c9b598f7942749156334283a6f24f',\n",
+       " 'f72bb24701c64fabb0182c9e757b581b']"
       ]
      },
-     "execution_count": 17,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -150,38 +162,57 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[42, 2]"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Prepare your documents, metadata, and IDs\n",
+    "docs = [\"Qdrant has Langchain integrations\", \"Qdrant also has Llama Index integrations\"]\n",
+    "metadata = [\n",
+    "    {\"source\": \"Langchain-docs\"},\n",
+    "    {\"source\": \"Linkedin-docs\"},\n",
+    "]\n",
+    "ids = [42, 2]\n",
+    "\n",
+    "# Use the new add method\n",
+    "client.add(\n",
+    "    collection_name=\"demo_collection\",\n",
+    "    documents=docs,\n",
+    "    metadata=metadata,\n",
+    "    ids=ids\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Top 5 results:\n",
-      "Rank 1: Maharana Pratap was a Rajput warrior king from Mewar. Score: 0.77\n",
-      "Rank 2: Maharana Pratap is considered a symbol of Rajput resistance against foreign rule. Score: 0.77\n",
-      "Rank 3: His legacy is celebrated in Rajasthan through festivals and monuments. Score: 0.69\n",
-      "Rank 4: He had 11 wives and 17 sons, including Amar Singh I who succeeded him as ruler of Mewar. Score: 0.68\n",
-      "Rank 5: He fought against the Mughal Empire led by Akbar. Score: 0.67\n"
+      "[QueryResponse(id='42', embedding=None, metadata={'document': 'Qdrant has Langchain integrations', 'source': 'Langchain-docs'}, document='Qdrant has Langchain integrations', score=0.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"
      ]
     }
    ],
    "source": [
-    "from qdrant_client.qdrant_fastembed import QueryResponse\n",
-    "\n",
-    "\n",
-    "def print_top_k_results(results: List[QueryResponse], k: int = 5):\n",
-    "    print(f\"Top {k} results:\")\n",
-    "    for i, result in enumerate(results[:k]):\n",
-    "        print(f\"Rank {i + 1}: {result.document}. Score: {result.score:.2f}\")\n",
-    "\n",
-    "\n",
-    "query_text = \"Who is Maharana Pratap?\"\n",
-    "results = client.query(\n",
-    "    collection_name=\"test_collection\", query_text=query_text, limit=7\n",
-    ")  # Returns limit most relevant documents\n",
-    "\n",
-    "print_top_k_results(results)"
+    "search_result = client.query(\n",
+    "    collection_name=\"demo_collection\",\n",
+    "    query_text=[\"This is a query document\"]\n",
+    ")\n",
+    "print(search_result)"
    ]
   },
   {

docs/index.md

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

fastembed/embedding.py

@@ -1,24 +1,26 @@
+import json
 import os
 import shutil
 import tarfile
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Iterable, List
+from typing import Dict, Iterable, List, Union
 
+import onnxruntime as ort
 import numpy as np
 import requests
-from optimum.onnxruntime import ORTModelForFeatureExtraction
+from tokenizers import Tokenizer, AddedToken
 from tqdm import tqdm
-from transformers import AutoTokenizer
 
 
-def normalize(input_array, p=2.0, dim=1, eps=1e-12):
+def normalize(input_array, p=2, dim=1, eps=1e-12):
     # Calculate the Lp norm along the specified dimension
     norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
-    norm = np.maximum(norm, eps) # Avoid division by zero
+    norm = np.maximum(norm, eps)  # Avoid division by zero
     normalized_array = input_array / norm
     return normalized_array
 
+
 class Embedding(ABC):
     """
     Abstract class for embeddings.
@@ -43,6 +45,34 @@ class Embedding(ABC):
     def embed(self, texts: List[str]) -> List[np.ndarray]:
         raise NotImplementedError
 
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Union[str, int]]]:
+        """
+        Lists the supported models.
+        """
+        return [
+            {
+                "model": "BAAI/bge-small-en",
+                "dim": 384,
+                "description": "Fast and Default English model",
+            },
+            {
+                "model": "BAAI/bge-base-en",
+                "dim": 768,
+                "description": "Base English model",
+            },
+            {
+                "model": "sentence-transformers/all-MiniLM-L6-v2",
+                "dim": 384,
+                "description": "Sentence Transformer model, MiniLM-L6-v2",
+            },
+            {
+                "model": "intfloat/multilingual-e5-large",
+                "dim": 1024,
+                "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
+            },
+        ]
+
     @classmethod
     def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
         """
@@ -159,7 +189,7 @@ class Embedding(ABC):
         try:
             self.download_file_from_gcs(
                 f"https://storage.googleapis.com/qdrant-fastembed/{fast_model_name}.tar.gz",
-            output_path=str(model_tar_gz),
+                output_path=str(model_tar_gz),
             )
         except PermissionError:
             simple_model_name = model_name.replace("/", "-")
@@ -219,6 +249,40 @@ class FlagEmbedding(Embedding):
         Embedding (_type_): _description_
     """
 
+    @classmethod
+    def load_tokenizer(cls, model_dir: Path, max_length: int = 512) -> Tokenizer:
+        config_path = model_dir / "config.json"
+        if not config_path.exists():
+            raise ValueError(f"Could not find config.json in {model_dir}")
+
+        tokenizer_path = model_dir / "tokenizer.json"
+        if not tokenizer_path.exists():
+            raise ValueError(f"Could not find tokenizer.json in {model_dir}")
+
+        tokenizer_config_path = model_dir / "tokenizer_config.json"
+        if not tokenizer_config_path.exists():
+            raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
+
+        tokens_map_path = model_dir / "special_tokens_map.json"
+        if not tokens_map_path.exists():
+            raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
+
+        config = json.load(open(str(config_path)))
+        tokenizer_config = json.load(open(str(tokenizer_config_path)))
+        tokens_map = json.load(open(str(tokens_map_path)))
+
+        tokenizer = Tokenizer.from_file(str(tokenizer_path))
+        tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
+        tokenizer.enable_padding(pad_id=config["pad_token_id"], pad_token=tokenizer_config["pad_token"])
+
+        for token in tokens_map.values():
+            if isinstance(token, str):
+                tokenizer.add_special_tokens([token])
+            elif isinstance(token, dict):
+                tokenizer.add_special_tokens([AddedToken(**token)])
+
+        return tokenizer
+
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en",
@@ -229,31 +293,63 @@ class FlagEmbedding(Embedding):
         Args:
             model_name (str): The name of the model to use.
             max_length (int, optional): The maximum number of tokens. Defaults to 512. Unknown behavior for values > 512.
+            cache_dir (str, optional): The path to the cache directory. Defaults to `local_cache` in the current directory.
 
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
         """
+        self.model_name = model_name
+
         if cache_dir is None:
             cache_dir = Path(".").resolve() / "local_cache"
             cache_dir.mkdir(parents=True, exist_ok=True)
+
         model_dir = self.retrieve_model(model_name, cache_dir)
-        if not (model_dir / "tokenizer.json").exists():
-            raise ValueError(f"Could not find tokenizer.json in {model_dir}")
-        if not (model_dir / "model.onnx").exists():
+
+        model_path = model_dir / "model.onnx"
+        optimized_model_path = model_dir / "model_optimized.onnx"
+
+        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+        onnx_providers = ["CPUExecutionProvider"]
+
+        if not model_path.exists():
             # Rename file model_optimized.onnx to model.onnx if it exists
-            if (model_dir / "model_optimized.onnx").exists():
-                (model_dir / "model_optimized.onnx").rename(model_dir / "model.onnx")
+            if optimized_model_path.exists():
+                optimized_model_path.rename(model_path)
             else:
                 raise ValueError(f"Could not find model.onnx in {model_dir}")
 
-        self.tokenizer = AutoTokenizer.from_pretrained(str(model_dir))
-        self.model = ORTModelForFeatureExtraction.from_pretrained(str(model_dir))
+        # Hacky support for multilingual model
+        self.exclude_token_type_ids = False
+        if model_name == "intfloat/multilingual-e5-large":
+            self.exclude_token_type_ids = True
 
-    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)
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        self.tokenizer = self.load_tokenizer(model_dir, max_length=max_length)
+
+        self.model = ort.InferenceSession(str(model_path), providers=onnx_providers, sess_options=so)
+
+    def onnx_embed(self, documents: List[str]) -> np.ndarray:
+        encoded = self.tokenizer.encode_batch(documents)
+        input_ids = np.array([e.ids for e in encoded])
+        attention_mask = np.array([e.attention_mask for e in encoded])
+
+        onnx_input = {
+            "input_ids": np.array(input_ids, dtype=np.int64),
+            "attention_mask": np.array(attention_mask, dtype=np.int64),
+        }
+
+        if not self.exclude_token_type_ids:
+            onnx_input["token_type_ids"] = np.array(
+                [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
+            )
+
+        model_output = self.model.run(None, onnx_input)
+        last_hidden_state = model_output[0][:, 0]
+        embeddings = normalize(last_hidden_state).astype(np.float32)
+        return embeddings
 
     def embed(self, documents: List[str], batch_size: int = 256) -> Iterable[np.ndarray]:
         """
@@ -267,13 +363,16 @@ class FlagEmbedding(Embedding):
         Returns:
             List of embeddings, one per document
         """
+        if type(documents) == str:
+            documents = [documents]
         # TODO: Replace loop with parallelized batching
         if len(documents) >= batch_size:
             for i in range(0, len(documents), batch_size):
                 batch = documents[i : i + batch_size]
-                return self.onnx_embed(batch)
+                yield from self.onnx_embed(batch)
         else:
-            return self.onnx_embed(documents)
+            vectors = self.onnx_embed(documents)
+            yield from vectors
 
 
 class DefaultEmbedding(FlagEmbedding):
@@ -287,12 +386,9 @@ class DefaultEmbedding(FlagEmbedding):
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en",
-        onnx_providers: List[str] = None,
         max_length: int = 512,
         cache_dir: str = None,
     ):
-        # if onnx_providers is None:
-        #     onnx_providers = [ONNXProviders.CPU]
         super().__init__(model_name, max_length=max_length, cache_dir=cache_dir)
 
 
@@ -306,4 +402,3 @@ class OpenAIEmbedding(Embedding):
         # Use your OpenAI model to embed the texts
         # return self.model.embed(texts)
         raise NotImplementedError
-        raise NotImplementedError

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.0.4"
+version = "0.0.5"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -12,18 +12,16 @@ 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"
+onnx = "^1.11"
+onnxruntime = "^1.15"
+tqdm = "^4.65"
+requests = "^2.31"
+tokenizers = "^0.13"
 
 [tool.poetry.dev-dependencies]
 ruff = "^0.0.277"
 isort = "^5.12.0"
 black = "^23.7.0"
-onnx = "^1.11.0"
 notebook = ">=7.0.2"
 mkdocs-material = "^9.1.21"
 mkdocstrings = "^0.22.0"
@@ -32,10 +30,6 @@ cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
 pytest = "^7.4.0"
 
-[tool.poetry.dependencies.onnxruntime-silicon]
-version = "^1.15.0"
-markers = "sys_platform == 'darwin'"  # This makes it macOS specific
-
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

tests/test_onnx_embeddings.py

@@ -1,13 +1,35 @@
 import numpy as np
 
-from fastembed.embedding import DefaultEmbedding
+from fastembed.embedding import DefaultEmbedding, Embedding
+
+
+CANONICAL_VECTOR_VALUES = {
+    "BAAI/bge-small-en": np.array([-0.0232, -0.0255,  0.0174, -0.0639, -0.0006]),
+    "BAAI/bge-base-en": np.array([0.0115,  0.0372,  0.0295,  0.0121,  0.0346]),
+    "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]),
+}
 
 
 def test_default_embedding():
+    for model_desc in Embedding.list_supported_models():
+        dim = model_desc["dim"]
+        model = DefaultEmbedding(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)
+
+
+def test_batch_embedding():
     model = DefaultEmbedding()
 
-    docs = ["hello world", "flag embedding"]
-    embeddings = np.array(model.embed(docs))
-    assert embeddings.shape == (2, 384)
+    docs = ["hello world", "flag embedding"] * 100
+    embeddings = list(model.embed(docs, batch_size=10))
+    embeddings = np.stack(embeddings, axis=0)
 
-test_default_embedding()
+    assert embeddings.shape == (200, 384)