chore: Pin numpy <2 (#278)

* MiniLM fix

* Added MiniLM to text embedding
Fixed MiniLM source destination
Black + isort for repo

* Fixed model all-MiniLM-L6-v2 description
Recomputed canonical vector for all-MiniLM-L6-v2 in test

---------

Co-authored-by: d.rudenko <dimitriyrudenk@gmail.com>

.github/workflows/python-tests.yml

@@ -1,4 +1,5 @@
 name: Tests
+run-name: Tests (gpu)
 
 on:
   push:
@@ -23,8 +24,6 @@ jobs:
           - '3.12.x'
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
 
     runs-on: ${{ matrix.os }}
 

README.md

@@ -18,11 +18,9 @@ To install the FastEmbed library, pip works best. You can install it with or wit
 
 ```bash
 pip install fastembed
-```
 
-### ⚡️ With GPU
+# or with GPU support
 
-```bash
 pip install fastembed-gpu
 ```
 
@@ -48,7 +46,99 @@ embeddings_list = list(embedding_model.embed(documents))
 len(embeddings_list[0]) # Vector of 384 dimensions
 ```
 
-### ⚡️ FastEmbed on a GPU
+Fastembed supports a variety of models for different tasks and modalities.
+The list of all the available models can be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/)
+### 🎒 Dense text embeddings
+
+```python
+from fastembed import TextEmbedding
+
+model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
+#   array([-0.1019,  0.0635, -0.0332,  0.0522, ...], dtype=float32)
+# ]
+
+```
+
+
+
+### 🔱 Sparse text embeddings
+
+* SPLADE++
+
+```python
+from fastembed import SparseTextEmbedding
+
+model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
+#   SparseEmbedding(indices=[ 38,  12,  91, ... ], values=[0.11, 0.22, 0.39, ...])
+# ]
+```
+
+<!--
+* BM42 - ([link](ToDo))
+
+```
+from fastembed import SparseTextEmbedding
+
+model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
+#   SparseEmbedding(indices=[ 38,  12,  91, ... ], values=[0.11, 0.22, 0.39, ...])
+# ]
+```
+-->
+
+### 🦥 Late interaction models (aka ColBERT)
+
+
+```python
+from fastembed import LateInteractionTextEmbedding
+
+model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   array([
+#       [-0.1115,  0.0097,  0.0052,  0.0195, ...],
+#       [-0.1019,  0.0635, -0.0332,  0.0522, ...],
+#   ]),
+#   array([
+#       [-0.9019,  0.0335, -0.0032,  0.0991, ...],
+#       [-0.2115,  0.8097,  0.1052,  0.0195, ...],
+#   ]),  
+# ]
+```
+
+### 🖼️ Image embeddings
+
+```python
+from fastembed import ImageEmbedding
+
+images = [
+    "./path/to/image1.jpg",
+    "./path/to/image2.jpg",
+]
+
+model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
+embeddings = list(embedding_model.embed(images))
+
+# [
+#   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
+#   array([-0.1019,  0.0635, -0.0332,  0.0522, ...], dtype=float32)
+# ]
+```
+
+
+## ⚡️ FastEmbed on a GPU
 
 FastEmbed supports running on GPU devices.
 It requires installation of the `fastembed-gpu` package.

docs/examples/Hybrid_Search.ipynb

@@ -328,7 +328,9 @@
    ],
    "source": [
     "source_df = dataset.to_pandas()\n",
-    "df = source_df.drop_duplicates(subset=[\"product_text\", \"product_title\", \"product_bullet_point\", \"product_brand\"])\n",
+    "df = source_df.drop_duplicates(\n",
+    "    subset=[\"product_text\", \"product_title\", \"product_bullet_point\", \"product_brand\"]\n",
+    ")\n",
     "df = df.dropna(subset=[\"product_text\", \"product_title\", \"product_bullet_point\", \"product_brand\"])\n",
     "df.head()"
    ]
@@ -367,7 +369,9 @@
    },
    "outputs": [],
    "source": [
-    "df[\"combined_text\"] = df[\"product_title\"] + \"\\n\" + df[\"product_text\"] + \"\\n\" + df[\"product_bullet_point\"]"
+    "df[\"combined_text\"] = (\n",
+    "    df[\"product_title\"] + \"\\n\" + df[\"product_text\"] + \"\\n\" + df[\"product_bullet_point\"]\n",
+    ")"
    ]
   },
   {
@@ -489,7 +493,9 @@
     "    return list(sparse_model.embed(texts, batch_size=32))\n",
     "\n",
     "\n",
-    "sparse_embedding: List[SparseEmbedding] = make_sparse_embedding([\"Fastembed is a great library for text embeddings!\"])\n",
+    "sparse_embedding: List[SparseEmbedding] = make_sparse_embedding(\n",
+    "    [\"Fastembed is a great library for text embeddings!\"]\n",
+    ")\n",
     "sparse_embedding"
    ]
   },
@@ -628,7 +634,9 @@
     "        token_weight_dict[token] = weight\n",
     "\n",
     "    # Sort the dictionary by weights\n",
-    "    token_weight_dict = dict(sorted(token_weight_dict.items(), key=lambda item: item[1], reverse=True))\n",
+    "    token_weight_dict = dict(\n",
+    "        sorted(token_weight_dict.items(), key=lambda item: item[1], reverse=True)\n",
+    "    )\n",
     "    return token_weight_dict\n",
     "\n",
     "\n",
@@ -870,14 +878,21 @@
     "    dense_vectors = df[\"dense_embedding\"].tolist()\n",
     "    rows = df.to_dict(orient=\"records\")\n",
     "    points = []\n",
-    "    for idx, (text, sparse_vector, dense_vector) in enumerate(zip(product_texts, sparse_vectors, dense_vectors)):\n",
-    "        sparse_vector = SparseVector(indices=sparse_vector.indices.tolist(), values=sparse_vector.values.tolist())\n",
+    "    for idx, (text, sparse_vector, dense_vector) in enumerate(\n",
+    "        zip(product_texts, sparse_vectors, dense_vectors)\n",
+    "    ):\n",
+    "        sparse_vector = SparseVector(\n",
+    "            indices=sparse_vector.indices.tolist(), values=sparse_vector.values.tolist()\n",
+    "        )\n",
     "        point = PointStruct(\n",
     "            id=idx,\n",
-    "            payload={\"text\": text, \"product_id\": rows[idx][\"product_id\"]},  # Add any additional payload if necessary\n",
+    "            payload={\n",
+    "                \"text\": text,\n",
+    "                \"product_id\": rows[idx][\"product_id\"],\n",
+    "            },  # Add any additional payload if necessary\n",
     "            vector={\n",
     "                \"text-sparse\": sparse_vector,\n",
-    "                \"text-dense\": dense_vector,\n",
+    "                \"text-dense\": dense_vector.tolist(),\n",
     "            },\n",
     "        )\n",
     "        points.append(point)\n",
@@ -936,7 +951,7 @@
     "            SearchRequest(\n",
     "                vector=NamedVector(\n",
     "                    name=\"text-dense\",\n",
-    "                    vector=query_dense_vector[0],\n",
+    "                    vector=query_dense_vector[0].tolist(),\n",
     "                ),\n",
     "                limit=10,\n",
     "                with_payload=True,\n",
@@ -1133,8 +1148,12 @@
     }
    ],
    "source": [
-    "def find_point_by_id(client: QdrantClient, collection_name: str, rrf_rank_list: List[Tuple[int, float]]):\n",
-    "    return client.retrieve(collection_name=collection_name, ids=[item[0] for item in rrf_rank_list])\n",
+    "def find_point_by_id(\n",
+    "    client: QdrantClient, collection_name: str, rrf_rank_list: List[Tuple[int, float]]\n",
+    "):\n",
+    "    return client.retrieve(\n",
+    "        collection_name=collection_name, ids=[item[0] for item in rrf_rank_list]\n",
+    "    )\n",
     "\n",
     "\n",
     "find_point_by_id(client, collection_name, rrf_rank_list)"

docs/experimental/Binary Quantization from Scratch.ipynb

@@ -14,23 +14,33 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
+   "execution_count": 1,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:00:06.460001Z",
+     "start_time": "2024-06-06T17:00:04.214098Z"
+    }
+   },
    "outputs": [],
    "source": [
-    "!pip install matplotlib tqdm pandas numpy --quiet"
+    "!pip install matplotlib tqdm pandas numpy datasets --quiet --upgrade"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 2,
    "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:00:07.041784Z",
+     "start_time": "2024-06-06T17:00:06.461658Z"
+    },
     "id": "WBVTItUX4yyr"
    },
    "outputs": [],
    "source": [
     "import numpy as np\n",
     "import pandas as pd\n",
+    "from datasets import load_dataset\n",
     "from tqdm import tqdm"
    ]
   },
@@ -52,8 +62,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 3,
    "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:09.343230Z",
+     "start_time": "2024-06-06T17:00:07.042526Z"
+    },
     "colab": {
      "base_uri": "https://localhost:8080/",
      "height": 250
@@ -61,58 +75,24 @@
     "id": "REJpFqkG7EG2",
     "outputId": "7a43c0ae-fbcc-45fe-fd58-bfe691297b22"
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|██████████| 26/26 [00:10<00:00,  2.45it/s]\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "(1000000, 1536)"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "def get_openai_vectors(force_download: bool = False):\n",
-    "    res = []\n",
-    "    for i in tqdm(range(26)):\n",
-    "        if force_download:\n",
-    "            !wget https://huggingface.co/api/datasets/KShivendu/dbpedia-entities-openai-1M/parquet/KShivendu--dbpedia-entities-openai-1M/train/{i}.parquet\n",
-    "        df = pd.read_parquet(f\"{i}.parquet\", engine=\"pyarrow\")\n",
-    "        res.append(np.stack(df.openai))\n",
-    "        del df\n",
-    "\n",
-    "    openai_vectors = np.concatenate(res)\n",
-    "    del res\n",
-    "    return openai_vectors\n",
-    "\n",
-    "\n",
-    "openai_vectors = get_openai_vectors(force_download=False)\n",
-    "openai_vectors.shape"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## ㆓ Binary Conversion\n",
-    "\n",
-    "Here, we will use 0 as the threshold for the binary conversion. All values greater than 0 will be set to 1, and others will remain 0. This is a simple and effective way to convert continuous values into binary values for OpenAI embeddings."
+    "# Download from Huggingface Hub\n",
+    "ds = load_dataset(\n",
+    "    \"Qdrant/dbpedia-entities-openai3-text-embedding-3-large-3072-100K\", split=\"train\"\n",
+    ")\n",
+    "openai_vectors = np.array(ds[\"text-embedding-3-large-3072-embedding\"])\n",
+    "del ds"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 4,
    "metadata": {
-    "id": "0JM2-Bj2Jkab"
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:10.900963Z",
+     "start_time": "2024-06-06T17:01:09.344842Z"
+    }
    },
    "outputs": [],
    "source": [
@@ -120,6 +100,30 @@
     "openai_bin[openai_vectors > 0] = 1"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:10.906827Z",
+     "start_time": "2024-06-06T17:01:10.901820Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "3072"
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "n_dim = openai_vectors.shape[1]\n",
+    "n_dim"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -131,8 +135,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 6,
    "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:10.909730Z",
+     "start_time": "2024-06-06T17:01:10.908166Z"
+    },
     "id": "FqshI-GlIERd"
    },
    "outputs": [],
@@ -141,7 +149,7 @@
     "    scores = np.dot(openai_vectors, openai_vectors[idx])\n",
     "    dot_results = np.argsort(scores)[-limit:][::-1]\n",
     "\n",
-    "    bin_scores = 1536 - np.logical_xor(openai_bin, openai_bin[idx]).sum(axis=1)\n",
+    "    bin_scores = n_dim - np.logical_xor(openai_bin, openai_bin[idx]).sum(axis=1)\n",
     "    bin_results = np.argsort(bin_scores)[-(limit * oversampling) :][::-1]\n",
     "\n",
     "    return len(set(dot_results).intersection(set(bin_results))) / limit"
@@ -156,8 +164,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 7,
    "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:25.206592Z",
+     "start_time": "2024-06-06T17:01:10.911971Z"
+    },
     "colab": {
      "base_uri": "https://localhost:8080/"
     },
@@ -169,110 +181,128 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "  0%|          | 0/4 [00:00<?, ?it/s]"
+      "  0%|          | 0/4 [00:00<?, ?it/s]\n",
+      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      " 50%|█████     | 1/2 [00:02<00:02,  2.05s/it]\u001b[A"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 1, 'limit': 10, 'recall': 0.8}\n"
+      "{'sampling_rate': 1, 'limit': 3, 'mean_acc': 0.9}\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 2/2 [00:33<00:00, 16.98s/it]\n",
-      " 25%|██▌       | 1/4 [00:33<01:41, 33.96s/it]"
+      "\n",
+      "100%|██████████| 2/2 [00:04<00:00,  2.02s/it]\u001b[A\n",
+      " 25%|██▌       | 1/4 [00:04<00:12,  4.05s/it]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 1, 'limit': 100, 'recall': 0.708}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": []
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'sampling_rate': 2, 'limit': 10, 'recall': 0.95}\n"
+      "{'sampling_rate': 1, 'limit': 10, 'mean_acc': 0.8300000000000001}\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 2/2 [00:32<00:00, 16.38s/it]\n",
-      " 50%|█████     | 2/4 [01:06<01:06, 33.26s/it]"
+      "\n",
+      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      " 50%|█████     | 1/2 [00:01<00:01,  1.72s/it]\u001b[A"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 2, 'limit': 100, 'recall': 0.877}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": []
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'sampling_rate': 3, 'limit': 10, 'recall': 0.96}\n"
+      "{'sampling_rate': 2, 'limit': 3, 'mean_acc': 1.0}\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 2/2 [00:32<00:00, 16.49s/it]\n",
-      " 75%|███████▌  | 3/4 [01:39<00:33, 33.13s/it]"
+      "\n",
+      "100%|██████████| 2/2 [00:03<00:00,  1.76s/it]\u001b[A\n",
+      " 50%|█████     | 2/4 [00:07<00:07,  3.75s/it]"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 3, 'limit': 100, 'recall': 0.937}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": []
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'sampling_rate': 5, 'limit': 10, 'recall': 0.9800000000000001}\n"
+      "{'sampling_rate': 2, 'limit': 10, 'mean_acc': 0.9700000000000001}\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 2/2 [00:32<00:00, 16.47s/it]\n",
-      "100%|██████████| 4/4 [02:12<00:00, 33.17s/it]"
+      "\n",
+      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      " 50%|█████     | 1/2 [00:01<00:01,  1.72s/it]\u001b[A"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'sampling_rate': 5, 'limit': 100, 'recall': 0.977}\n"
+      "{'sampling_rate': 3, 'limit': 3, 'mean_acc': 1.0}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 2/2 [00:03<00:00,  1.69s/it]\u001b[A\n",
+      " 75%|███████▌  | 3/4 [00:10<00:03,  3.58s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 3, 'limit': 10, 'mean_acc': 0.9800000000000001}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      " 50%|█████     | 1/2 [00:01<00:01,  1.68s/it]\u001b[A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 5, 'limit': 3, 'mean_acc': 1.0}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 2/2 [00:03<00:00,  1.65s/it]\u001b[A\n",
+      "100%|██████████| 4/4 [00:14<00:00,  3.57s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'sampling_rate': 5, 'limit': 10, 'mean_acc': 0.99}\n"
      ]
     },
     {
@@ -285,119 +315,53 @@
    ],
    "source": [
     "number_of_samples = 10\n",
-    "limits = [10, 100]\n",
+    "limits = [3, 10]\n",
     "sampling_rate = [1, 2, 3, 5]\n",
     "results = []\n",
     "\n",
     "\n",
     "def mean_accuracy(number_of_samples, limit, sampling_rate):\n",
-    "    return np.mean([accuracy(i, limit=limit, oversampling=sampling_rate) for i in range(number_of_samples)])\n",
+    "    return np.mean(\n",
+    "        [accuracy(i, limit=limit, oversampling=sampling_rate) for i in range(number_of_samples)]\n",
+    "    )\n",
     "\n",
     "\n",
     "for i in tqdm(sampling_rate):\n",
     "    for j in tqdm(limits):\n",
-    "        result = {\"sampling_rate\": i, \"limit\": j, \"recall\": mean_accuracy(number_of_samples, j, i)}\n",
+    "        result = {\n",
+    "            \"sampling_rate\": i,\n",
+    "            \"limit\": j,\n",
+    "            \"mean_acc\": mean_accuracy(number_of_samples, j, i),\n",
+    "        }\n",
     "        print(result)\n",
     "        results.append(result)"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 19,
+   "cell_type": "markdown",
    "metadata": {},
+   "source": [
+    "## ㆓ Binary Conversion\n",
+    "\n",
+    "Here, we will use 0 as the threshold for the binary conversion. All values greater than 0 will be set to 1, and others will remain 0. This is a simple and effective way to convert continuous values into binary values for OpenAI embeddings."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-06-06T17:01:25.247495Z",
+     "start_time": "2024-06-06T17:01:25.213508Z"
+    }
+   },
    "outputs": [
     {
      "data": {
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>sampling_rate</th>\n",
-       "      <th>limit</th>\n",
-       "      <th>recall</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>1</td>\n",
-       "      <td>10</td>\n",
-       "      <td>0.800</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>1</td>\n",
-       "      <td>100</td>\n",
-       "      <td>0.708</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>2</td>\n",
-       "      <td>10</td>\n",
-       "      <td>0.950</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>2</td>\n",
-       "      <td>100</td>\n",
-       "      <td>0.877</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>3</td>\n",
-       "      <td>10</td>\n",
-       "      <td>0.960</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>5</th>\n",
-       "      <td>3</td>\n",
-       "      <td>100</td>\n",
-       "      <td>0.937</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6</th>\n",
-       "      <td>5</td>\n",
-       "      <td>10</td>\n",
-       "      <td>0.980</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>7</th>\n",
-       "      <td>5</td>\n",
-       "      <td>100</td>\n",
-       "      <td>0.977</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "   sampling_rate  limit  recall\n",
-       "0              1     10   0.800\n",
-       "1              1    100   0.708\n",
-       "2              2     10   0.950\n",
-       "3              2    100   0.877\n",
-       "4              3     10   0.960\n",
-       "5              3    100   0.937\n",
-       "6              5     10   0.980\n",
-       "7              5    100   0.977"
-      ]
+      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>sampling_rate</th>\n      <th>limit</th>\n      <th>mean_acc</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>1</td>\n      <td>3</td>\n      <td>0.90</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>1</td>\n      <td>10</td>\n      <td>0.83</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>2</td>\n      <td>3</td>\n      <td>1.00</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>2</td>\n      <td>10</td>\n      <td>0.97</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>3</td>\n      <td>3</td>\n      <td>1.00</td>\n    </tr>\n    <tr>\n      <th>5</th>\n      <td>3</td>\n      <td>10</td>\n      <td>0.98</td>\n    </tr>\n    <tr>\n      <th>6</th>\n      <td>5</td>\n      <td>3</td>\n      <td>1.00</td>\n    </tr>\n    <tr>\n      <th>7</th>\n      <td>5</td>\n      <td>10</td>\n      <td>0.99</td>\n    </tr>\n  </tbody>\n</table>\n</div>",
+      "text/plain": "   sampling_rate  limit  mean_acc\n0              1      3      0.90\n1              1     10      0.83\n2              2      3      1.00\n3              2     10      0.97\n4              3      3      1.00\n5              3     10      0.98\n6              5      3      1.00\n7              5     10      0.99"
      },
-     "execution_count": 19,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -408,22 +372,13 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "| sampling_rate | limit | accuracy |\n",
-    "|---------------|-------|----------|\n",
-    "| 1             | 10    | 0.800    |\n",
-    "| 1             | 100   | 0.708    |\n",
-    "| 2             | 10    | 0.950    |\n",
-    "| 2             | 100   | 0.877    |\n",
-    "| 4             | 10    | 0.970    |\n",
-    "| 4             | 100   | 0.956    |\n",
-    "| 8             | 10    | 0.990    |\n",
-    "| 8             | 100   | 0.990    |\n",
-    "| 16            | 10    | 1.000    |\n",
-    "| 16            | 100   | 0.998    |"
-   ]
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -432,7 +387,8 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
    "name": "python3"
   },
   "language_info": {
@@ -445,7 +401,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.17"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

experiments/attention_export.py

@@ -12,4 +12,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 #     print("Model already exported")
 # except FileNotFoundError:
 print(f"Exporting model to {output_dir}")
-main_export(model_id, output=output_dir, no_post_process=True, model_kwargs=model_kwargs)
+main_export(
+    model_id, output=output_dir, no_post_process=True, model_kwargs=model_kwargs
+)

experiments/try_attention_export.py

@@ -17,10 +17,14 @@ input_ids = tokenizer_output["input_ids"]
 attention_mask = tokenizer_output["attention_mask"]
 print(attention_mask)
 # Prepare the input
-input_ids = np.array(input_ids).astype(np.int64)  # Replace your_input_ids with actual input data
+input_ids = np.array(input_ids).astype(
+    np.int64
+)  # Replace your_input_ids with actual input data
 
 # Run the ONNX model
-outputs = ort_session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
+outputs = ort_session.run(
+    None, {"input_ids": input_ids, "attention_mask": attention_mask}
+)
 
 # Get the attention weights
 attentions = outputs[-1]

fastembed/__init__.py

@@ -1,9 +1,9 @@
 import importlib.metadata
 
 from fastembed.image import ImageEmbedding
-from fastembed.text import TextEmbedding
-from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
 from fastembed.late_interaction import LateInteractionTextEmbedding
+from fastembed.sparse import SparseEmbedding, SparseTextEmbedding
+from fastembed.text import TextEmbedding
 
 try:
     version = importlib.metadata.version("fastembed")

fastembed/common/__init__.py

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

fastembed/common/model_management.py

@@ -2,13 +2,13 @@ import os
 import shutil
 import tarfile
 from pathlib import Path
-from typing import List, Optional, Dict, Any
+from typing import Any, Dict, List, Optional
 
 import requests
 from huggingface_hub import snapshot_download
 from huggingface_hub.utils import RepositoryNotFoundError
-from tqdm import tqdm
 from loguru import logger
+from tqdm import tqdm
 
 
 class ModelManagement:
@@ -42,7 +42,9 @@ class ModelManagement:
         raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
 
     @classmethod
-    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
+    def download_file_from_gcs(
+        cls, url: str, output_path: str, show_progress: bool = True
+    ) -> str:
         """
         Downloads a file from Google Cloud Storage.
 
@@ -71,12 +73,17 @@ class ModelManagement:
 
         # Warn if the total size is zero
         if total_size_in_bytes == 0:
-            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
+            print(
+                f"Warning: Content-length header is missing or zero in the response from {url}."
+            )
 
         show_progress = total_size_in_bytes and show_progress
 
         with tqdm(
-            total=total_size_in_bytes, unit="iB", unit_scale=True, disable=not show_progress
+            total=total_size_in_bytes,
+            unit="iB",
+            unit_scale=True,
+            disable=not show_progress,
         ) as progress_bar:
             with open(output_path, "wb") as file:
                 for chunk in response.iter_content(chunk_size=1024):
@@ -156,7 +163,9 @@ class ModelManagement:
         return cache_dir
 
     @classmethod
-    def retrieve_model_gcs(cls, model_name: str, source_url: str, cache_dir: str) -> Path:
+    def retrieve_model_gcs(
+        cls, model_name: str, source_url: str, cache_dir: str
+    ) -> Path:
         fast_model_name = f"fast-{model_name.split('/')[-1]}"
 
         cache_tmp_dir = Path(cache_dir) / "tmp"
@@ -182,8 +191,12 @@ class ModelManagement:
             output_path=str(model_tar_gz),
         )
 
-        cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
-        assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
+        cls.decompress_to_cache(
+            targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir)
+        )
+        assert (
+            model_tmp_dir.exists()
+        ), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
 
         model_tar_gz.unlink()
         # Rename from tmp to final name is atomic

fastembed/common/onnx_model.py

@@ -1,7 +1,17 @@
+import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, Generic, Iterable, Optional, Tuple, Type, TypeVar, Sequence
-import warnings
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+)
 
 import numpy as np
 import onnxruntime as ort
@@ -9,7 +19,6 @@ import onnxruntime as ort
 from fastembed.common.types import OnnxProvider
 from fastembed.parallel_processor import Worker
 
-
 # Holds type of the embedding result
 T = TypeVar("T")
 
@@ -51,7 +60,9 @@ class OnnxModel(Generic[T]):
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
 
-        onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
+        onnx_providers = (
+            ["CPUExecutionProvider"] if providers is None else list(providers)
+        )
         available_providers = ort.get_available_providers()
         requested_provider_names = []
         for provider in onnx_providers:
@@ -92,6 +103,7 @@ class EmbeddingWorker(Worker):
         self,
         model_name: str,
         cache_dir: str,
+        **kwargs,
     ) -> OnnxModel:
         raise NotImplementedError()
 
@@ -99,15 +111,13 @@ class EmbeddingWorker(Worker):
         self,
         model_name: str,
         cache_dir: str,
+        **kwargs,
     ):
-        self.model = self.init_embedding(model_name, cache_dir)
+        self.model = self.init_embedding(model_name, cache_dir, **kwargs)
 
     @classmethod
     def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
-        return cls(
-            model_name=model_name,
-            cache_dir=cache_dir,
-        )
+        return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
     def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
         raise NotImplementedError("Subclasses must implement this method")

fastembed/common/preprocessor_utils.py

@@ -2,7 +2,7 @@ import json
 from pathlib import Path
 from typing import Tuple
 
-from tokenizers import Tokenizer, AddedToken
+from tokenizers import AddedToken, Tokenizer
 
 from fastembed.image.transform.operators import Compose
 
@@ -40,7 +40,9 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, d
     tokens_map = load_special_tokens(model_dir)
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
-    tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
+    tokenizer.enable_truncation(
+        max_length=min(tokenizer_config["model_max_length"], max_length)
+    )
     tokenizer.enable_padding(
         pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
     )

fastembed/common/types.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import Union, Iterable, Tuple, Dict, Any
+from typing import Any, Dict, Iterable, Tuple, Union
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias

fastembed/common/utils.py

@@ -2,7 +2,7 @@ import os
 import tempfile
 from itertools import islice
 from pathlib import Path
-from typing import Union, Iterable, Generator, Optional
+from typing import Generator, Iterable, Optional, Union
 
 import numpy as np
 

fastembed/image/__init__.py

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

fastembed/image/image_embedding.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterable, List, Optional, Type, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
 
 import numpy as np
 
@@ -51,9 +51,16 @@ class ImageEmbedding(ImageEmbeddingBase):
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+            if any(
+                model_name.lower() == model["model"].lower()
+                for model in supported_models
+            ):
                 self.model = EMBEDDING_MODEL_TYPE(
-                    model_name, cache_dir, threads, providers=providers, **kwargs
+                    model_name,
+                    cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    **kwargs,
                 )
                 return
 

fastembed/image/image_embedding_base.py

@@ -31,6 +31,11 @@ class ImageEmbeddingBase(ModelManagement):
 
         Args:
             images - The list of image paths to preprocess and embed.
+            batch_size: Batch size for encoding
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:

fastembed/image/onnx_embedding.py

@@ -1,12 +1,12 @@
-from typing import Dict, Optional, Iterable, Type, List, Any, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
 
 import numpy as np
 
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize, define_cache_dir
 from fastembed.common import ImageInput, OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import define_cache_dir, normalize
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
-from fastembed.image.onnx_image_model import OnnxImageModel, ImageEmbeddingWorker
+from fastembed.image.onnx_image_model import ImageEmbeddingWorker, OnnxImageModel
 
 supported_onnx_models = [
     {
@@ -56,9 +56,9 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
         model_description = self._get_model_description(model_name)
-        cache_dir = define_cache_dir(cache_dir)
+        self.cache_dir = define_cache_dir(cache_dir)
         model_dir = self.download_model(
-            model_description, cache_dir, local_files_only=self._local_files_only
+            model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         self.load_onnx_model(
@@ -106,6 +106,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
             images=images,
             batch_size=batch_size,
             parallel=parallel,
+            **kwargs,
         )
 
     @classmethod
@@ -126,9 +127,5 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
 
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
-    def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
-    ) -> OnnxImageEmbedding:
-        return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
+        return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

fastembed/image/onnx_image_model.py

@@ -1,15 +1,15 @@
-import os
 import contextlib
+import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type
 
-from PIL import Image
 import numpy as np
+from PIL import Image
 
+from fastembed.common import ImageInput, OnnxProvider, PathInput
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
-from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T, OnnxOutputContext
-from fastembed.common import PathInput, ImageInput, OnnxProvider
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
@@ -44,7 +44,10 @@ class OnnxImageModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
     ) -> None:
         super().load_onnx_model(
-            model_dir=model_dir, model_file=model_file, threads=threads, providers=providers
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
         )
         self.processor = load_preprocessor(model_dir=model_dir)
 
@@ -59,9 +62,7 @@ class OnnxImageModel(OnnxModel[T]):
         onnx_input = self._preprocess_onnx_input(onnx_input)
         model_output = self.model.run(None, onnx_input)
         embeddings = model_output[0].reshape(len(images), -1)
-        return OnnxOutputContext(
-            model_output=embeddings
-        )
+        return OnnxOutputContext(model_output=embeddings)
 
     def _embed_images(
         self,
@@ -70,6 +71,7 @@ class OnnxImageModel(OnnxModel[T]):
         images: ImageInput,
         batch_size: int = 256,
         parallel: Optional[int] = None,
+        **kwargs,
     ) -> Iterable[T]:
         is_small = False
 
@@ -88,11 +90,10 @@ class OnnxImageModel(OnnxModel[T]):
             for batch in iter_batch(images, batch_size):
                 yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            params = {
-                "model_name": model_name,
-                "cache_dir": cache_dir,
-            }
+            start_method = (
+                "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            )
+            params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
             pool = ParallelWorkerPool(
                 parallel, self._get_worker_class(), start_method=start_method
             )

fastembed/image/transform/functional.py

@@ -1,8 +1,7 @@
-from typing import Union, Tuple, Sized
-
-from PIL import Image
+from typing import Sized, Tuple, Union
 
 import numpy as np
+from PIL import Image
 
 
 def convert_to_rgb(image: Image.Image) -> Image.Image:

fastembed/image/transform/operators.py

@@ -1,14 +1,15 @@
-from typing import List, Tuple, Union, Any, Dict
+from typing import Any, Dict, List, Tuple, Union
 
 import numpy as np
 from PIL import Image
+
 from fastembed.image.transform.functional import (
     center_crop,
-    normalize,
-    resize,
     convert_to_rgb,
+    normalize,
+    pil2ndarray,
     rescale,
-    pil2ndarray
+    resize,
 )
 
 
@@ -49,7 +50,9 @@ class Resize(Transform):
         self.resample = resample
 
     def __call__(self, images: List[Image.Image]) -> List[Image.Image]:
-        return [resize(image, size=self.size, resample=self.resample) for image in images]
+        return [
+            resize(image, size=self.size, resample=self.resample) for image in images
+        ]
 
 
 class Rescale(Transform):
@@ -59,15 +62,21 @@ class Rescale(Transform):
     def __call__(self, images: List[np.ndarray]) -> List[np.ndarray]:
         return [rescale(image, scale=self.scale) for image in images]
 
+
 class PILtoNDarray(Transform):
-    def __call__(self, images: List[Union[Image.Image, np.ndarray]]) -> List[np.ndarray]:
+    def __call__(
+        self, images: List[Union[Image.Image, np.ndarray]]
+    ) -> List[np.ndarray]:
         return [pil2ndarray(image) for image in images]
 
+
 class Compose:
     def __init__(self, transforms: List[Transform]):
         self.transforms = transforms
 
-    def __call__(self, images: Union[List[Image.Image], List[np.ndarray]]) -> Union[List[np.ndarray], List[Image.Image]]:
+    def __call__(
+        self, images: Union[List[Image.Image], List[np.ndarray]]
+    ) -> Union[List[np.ndarray], List[Image.Image]]:
         for transform in self.transforms:
             images = transform(images)
         return images
@@ -75,25 +84,25 @@ class Compose:
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> "Compose":
         """Creates processor from a config dict.
-            Args:
-                config (Dict[str, Any]): Configuration dictionary.
+        Args:
+            config (Dict[str, Any]): Configuration dictionary.
 
-                    Valid keys:
-                        - do_resize
-                        - size
-                        - do_center_crop
-                        - crop_size
-                        - do_rescale
-                        - rescale_factor
-                        - do_normalize
-                        - image_mean
-                        - image_std
-                    Valid size keys (nested):
-                        - {"height", "width"}
-                        - {"shortest_edge"}
+                Valid keys:
+                    - do_resize
+                    - size
+                    - do_center_crop
+                    - crop_size
+                    - do_rescale
+                    - rescale_factor
+                    - do_normalize
+                    - image_mean
+                    - image_std
+                Valid size keys (nested):
+                    - {"height", "width"}
+                    - {"shortest_edge"}
 
-            Returns:
-                Compose: Image processor.
+        Returns:
+            Compose: Image processor.
         """
         transforms = []
         cls._get_convert_to_rgb(transforms, config)
@@ -110,8 +119,8 @@ class Compose:
 
     @staticmethod
     def _get_resize(transforms: List[Transform], config: Dict[str, Any]):
-        mode = config.get('image_processor_type', 'CLIPImageProcessor')
-        if mode == 'CLIPImageProcessor':
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+        if mode == "CLIPImageProcessor":
             if config.get("do_resize", False):
                 size = config["size"]
                 if "shortest_edge" in size:
@@ -119,25 +128,44 @@ class Compose:
                 elif "height" in size and "width" in size:
                     size = (size["height"], size["width"])
                 else:
-                    raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.")
-                transforms.append(Resize(size=size, resample=config.get("resample", Image.Resampling.BICUBIC)))
-        elif mode == 'ConvNextFeatureExtractor':
-            if 'size' in config and "shortest_edge" not in config['size']:
-                raise ValueError(f"Size dictionary must contain 'shortest_edge' key. Got {config['size'].keys()}")
-            shortest_edge = config['size']["shortest_edge"]
+                    raise ValueError(
+                        "Size must contain either 'shortest_edge' or 'height' and 'width'."
+                    )
+                transforms.append(
+                    Resize(
+                        size=size,
+                        resample=config.get("resample", Image.Resampling.BICUBIC),
+                    )
+                )
+        elif mode == "ConvNextFeatureExtractor":
+            if "size" in config and "shortest_edge" not in config["size"]:
+                raise ValueError(
+                    f"Size dictionary must contain 'shortest_edge' key. Got {config['size'].keys()}"
+                )
+            shortest_edge = config["size"]["shortest_edge"]
             crop_pct = config.get("crop_pct", 0.875)
             if shortest_edge < 384:
                 # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
                 resize_shortest_edge = int(shortest_edge / crop_pct)
-                transforms.append(Resize(size=resize_shortest_edge, resample=config.get("resample", Image.Resampling.BICUBIC)))
+                transforms.append(
+                    Resize(
+                        size=resize_shortest_edge,
+                        resample=config.get("resample", Image.Resampling.BICUBIC),
+                    )
+                )
                 transforms.append(CenterCrop(size=(shortest_edge, shortest_edge)))
             else:
-                transforms.append(Resize(size=(shortest_edge, shortest_edge), resample=config.get("resample", Image.Resampling.BICUBIC)))
+                transforms.append(
+                    Resize(
+                        size=(shortest_edge, shortest_edge),
+                        resample=config.get("resample", Image.Resampling.BICUBIC),
+                    )
+                )
 
     @staticmethod
     def _get_center_crop(transforms: List[Transform], config: Dict[str, Any]):
-        mode = config.get('image_processor_type', 'CLIPImageProcessor')
-        if mode == 'CLIPImageProcessor':
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+        if mode == "CLIPImageProcessor":
             if config.get("do_center_crop", False):
                 crop_size = config["crop_size"]
                 if isinstance(crop_size, int):
@@ -147,7 +175,7 @@ class Compose:
                 else:
                     raise ValueError(f"Invalid crop size: {crop_size}")
                 transforms.append(CenterCrop(size=crop_size))
-        elif mode == 'ConvNextFeatureExtractor':
+        elif mode == "ConvNextFeatureExtractor":
             pass
         else:
             raise ValueError(f"Preprocessor {mode} is not supported")
@@ -165,4 +193,6 @@ class Compose:
     @staticmethod
     def _get_normalize(transforms: List[Transform], config: Dict[str, Any]):
         if config.get("do_normalize", False):
-            transforms.append(Normalize(mean=config["image_mean"], std=config["image_std"]))
+            transforms.append(
+                Normalize(mean=config["image_mean"], std=config["image_std"])
+            )

fastembed/late_interaction/__init__.py

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

fastembed/late_interaction/colbert.py

@@ -1,5 +1,5 @@
-from typing import Any, Dict, Iterable, List, Optional, Union, Type, Sequence
 import string
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
 
 import numpy as np
 from tokenizers import Encoding
@@ -12,7 +12,6 @@ from fastembed.late_interaction.late_interaction_embedding_base import (
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 
-
 supported_colbert_models = [
     {
         "model": "colbert-ir/colbertv2.0",
@@ -76,7 +75,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             if self.tokenizer.padding:
                 prev_padding = self.tokenizer.padding
             self.tokenizer.enable_padding(
-                pad_token=self.MASK_TOKEN, pad_id=self.mask_token_id, length=self.MIN_QUERY_LENGTH
+                pad_token=self.MASK_TOKEN,
+                pad_id=self.mask_token_id,
+                length=self.MIN_QUERY_LENGTH,
             )
             encoded = self.tokenizer.encode_batch(query)
             if prev_padding is None:
@@ -123,10 +124,10 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
         model_description = self._get_model_description(model_name)
-        cache_dir = define_cache_dir(cache_dir)
+        self.cache_dir = define_cache_dir(cache_dir)
 
         model_dir = self.download_model(
-            model_description, cache_dir, local_files_only=self._local_files_only
+            model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         self.load_onnx_model(
@@ -171,6 +172,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
+            **kwargs,
         )
 
     def query_embed(self, query: Union[str, List[str]], **kwargs) -> np.ndarray:
@@ -188,9 +190,5 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
 
 
 class ColbertEmbeddingWorker(TextEmbeddingWorker):
-    def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
-    ) -> Colbert:
-        return Colbert(model_name=model_name, cache_dir=cache_dir, threads=1)
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Colbert:
+        return Colbert(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

fastembed/late_interaction/late_interaction_embedding_base.py

@@ -42,7 +42,9 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs
+    ) -> Iterable[np.ndarray]:
         """
         Embeds queries
 

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -1,12 +1,12 @@
-from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
 
 import numpy as np
 
 from fastembed.common import OnnxProvider
+from fastembed.late_interaction.colbert import Colbert
 from fastembed.late_interaction.late_interaction_embedding_base import (
     LateInteractionTextEmbeddingBase,
 )
-from fastembed.late_interaction.colbert import Colbert
 
 
 class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
@@ -54,7 +54,10 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+            if any(
+                model_name.lower() == model["model"].lower()
+                for model in supported_models
+            ):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name, cache_dir, threads, providers=providers, **kwargs
                 )
@@ -89,7 +92,9 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs
+    ) -> Iterable[np.ndarray]:
         """
         Embeds queries
 

fastembed/parallel_processor.py

@@ -83,7 +83,9 @@ def _worker(
 
 
 class ParallelWorkerPool:
-    def __init__(self, num_workers: int, worker: Type[Worker], start_method: Optional[str] = None):
+    def __init__(
+        self, num_workers: int, worker: Type[Worker], start_method: Optional[str] = None
+    ):
         self.worker_class = worker
         self.num_workers = num_workers
         self.input_queue: Optional[Queue] = None
@@ -118,7 +120,9 @@ class ParallelWorkerPool:
             process.start()
             self.processes.append(process)
 
-    def ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Any]:
+    def ordered_map(
+        self, stream: Iterable[Any], *args: Any, **kwargs: Any
+    ) -> Iterable[Any]:
         buffer = defaultdict(Any)
         next_expected = 0
 

fastembed/sparse/bm25.py

@@ -0,0 +1,284 @@
+import os
+import string
+from collections import defaultdict
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union
+
+import mmh3
+import numpy as np
+from snowballstemmer import stemmer as get_stemmer
+
+from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.parallel_processor import ParallelWorkerPool, Worker
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.sparse.utils.tokenizer import WordTokenizer
+
+supported_bm25_models = [
+    {
+        "model": "Qdrant/bm25",
+        "description": "BM25 as sparse embeddings meant to be used with Qdrant",
+        "size_in_GB": 0.01,
+        "sources": {
+            "hf": "Qdrant/bm25",
+        },
+        "model_file": "mock.file",  # bm25 does not require a model, so we just use a mock
+        "additional_files": ["stopwords.txt"],
+    },
+]
+
+MODEL_TO_LANGUAGE = {
+    "Qdrant/bm25": "english",
+}
+
+
+class Bm25(SparseTextEmbeddingBase):
+    """Implements traditional BM25 in a form of sparse embeddings.
+    Uses a count of tokens in the document to evaluate the importance of the token.
+
+    WARNING: This model is expected to be used with `modifier="idf"` in the sparse vector index of Qdrant.
+
+    BM25 formula:
+
+    score(q, d) = SUM[ IDF(q_i) * (f(q_i, d) * (k + 1)) / (f(q_i, d) + k * (1 - b + b * (|d| / avg_len))) ],
+
+    where IDF is the inverse document frequency, computed on Qdrant's side
+    f(q_i, d) is the term frequency of the token q_i in the document d
+    k, b, avg_len are hyperparameters, described below.
+
+    Args:
+        model_name (str): The name of the model to use.
+        cache_dir (str, optional): The path to the cache directory.
+            Can be set using the `FASTEMBED_CACHE_PATH` env variable.
+            Defaults to `fastembed_cache` in the system's temp directory.
+        k (float, optional): The k parameter in the BM25 formula. Defines the saturation of the term frequency.
+            I.e. defines how fast the moment when additional terms stop to increase the score. Defaults to 1.2.
+        b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.
+            Defaults to 0.75.
+        avg_len (float, optional): The average length of the documents in the corpus. Defaults to 256.0.
+    Raises:
+        ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        k: float = 1.2,
+        b: float = 0.75,
+        avg_len: float = 256.0,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, **kwargs)
+
+        self.k = k
+        self.b = b
+        self.avg_len = avg_len
+
+        model_description = self._get_model_description(model_name)
+        self.cache_dir = define_cache_dir(cache_dir)
+
+        model_dir = self.download_model(
+            model_description, self.cache_dir, local_files_only=self._local_files_only
+        )
+
+        self.punctuation = set(string.punctuation)
+        self.stopwords = set(self._load_stopwords(model_dir))
+        self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
+        self.tokenizer = WordTokenizer
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_bm25_models
+
+    @classmethod
+    def _load_stopwords(cls, model_dir: Path) -> List[str]:
+        stopwords_path = model_dir / "stopwords.txt"
+        if not stopwords_path.exists():
+            return []
+
+        with open(stopwords_path, "r") as f:
+            return f.read().splitlines()
+
+    def _embed_documents(
+        self,
+        model_name: str,
+        cache_dir: str,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+    ) -> Iterable[SparseEmbedding]:
+        is_small = False
+
+        if isinstance(documents, str):
+            documents = [documents]
+            is_small = True
+
+        if isinstance(documents, list):
+            if len(documents) < batch_size:
+                is_small = True
+
+        if parallel == 0:
+            parallel = os.cpu_count()
+
+        if parallel is None or is_small:
+            for batch in iter_batch(documents, batch_size):
+                yield from self.raw_embed(batch)
+        else:
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "k": self.k,
+                "b": self.b,
+                "avg_len": self.avg_len,
+            }
+            pool = ParallelWorkerPool(
+                parallel, self._get_worker_class(), start_method=start_method
+            )
+            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
+                for record in batch:
+                    yield record
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[SparseEmbedding]:
+        """
+        Encode a list of documents into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+
+        Args:
+            documents: Iterator of documents or single document to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self._embed_documents(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            documents=documents,
+            batch_size=batch_size,
+            parallel=parallel,
+        )
+
+    def _stem(self, tokens: List[str]) -> List[str]:
+        stemmed_tokens = []
+        for token in tokens:
+            if token in self.punctuation:
+                continue
+
+            if token in self.stopwords:
+                continue
+
+            stemmed_token = self.stemmer.stemWord(token)
+
+            if stemmed_token:
+                stemmed_tokens.append(stemmed_token)
+        return stemmed_tokens
+
+    def raw_embed(
+        self,
+        documents: List[str],
+    ) -> List[SparseEmbedding]:
+        embeddings = []
+        for document in documents:
+            tokens = self.tokenizer.tokenize(document)
+            stemmed_tokens = self._stem(tokens)
+            token_id2value = self._term_frequency(stemmed_tokens)
+            embeddings.append(SparseEmbedding.from_dict(token_id2value))
+        return embeddings
+
+    def _term_frequency(self, tokens: List[str]) -> Dict[int, float]:
+        """Calculate the term frequency part of the BM25 formula.
+
+        (
+            f(q_i, d) * (k + 1)
+        ) / (
+            f(q_i, d) + k * (1 - b + b * (|d| / avg_len))
+        )
+
+        Args:
+            tokens (List[str]): The list of tokens in the document.
+
+        Returns:
+            Dict[int, float]: The token_id to term frequency mapping.
+        """
+        tf_map = {}
+        counter = defaultdict(int)
+        for stemmed_token in tokens:
+            counter[stemmed_token] += 1
+
+        doc_len = len(tokens)
+        for stemmed_token in counter:
+            token_id = self.compute_token_id(stemmed_token)
+            num_occurrences = counter[stemmed_token]
+            tf_map[token_id] = num_occurrences * (self.k + 1)
+            tf_map[token_id] /= num_occurrences + self.k * (
+                1 - self.b + self.b * doc_len / self.avg_len
+            )
+        return tf_map
+
+    @classmethod
+    def compute_token_id(cls, token: str) -> int:
+        return abs(mmh3.hash(token))
+
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
+        """To emulate BM25 behaviour, we don't need to use weights in the query, and
+        it's enough to just hash the tokens and assign a weight of 1.0 to them.
+        """
+        if isinstance(query, str):
+            query = [query]
+
+        for text in query:
+            tokens = self.tokenizer.tokenize(text)
+            stemmed_tokens = self._stem(tokens)
+            token_ids = np.array(
+                [self.compute_token_id(token) for token in stemmed_tokens],
+                dtype=np.float32,
+            )
+            values = np.ones_like(token_ids)
+            yield SparseEmbedding(indices=token_ids, values=values)
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["Bm25Worker"]:
+        return Bm25Worker
+
+
+class Bm25Worker(Worker):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
+    ):
+        self.model = self.init_embedding(model_name, cache_dir, **kwargs)
+
+    @classmethod
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "Bm25Worker":
+        return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
+
+    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+        for idx, batch in items:
+            onnx_output = self.model.raw_embed(batch)
+            yield idx, onnx_output
+
+    @staticmethod
+    def init_embedding(model_name: str, cache_dir: str, **kwargs) -> Bm25:
+        return Bm25(model_name=model_name, cache_dir=cache_dir, **kwargs)

fastembed/sparse/bm42.py

@@ -1,16 +1,19 @@
 import math
 import string
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Union
 
-import numpy as np
 import mmh3
+import numpy as np
 from snowballstemmer import stemmer as get_stemmer
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir
-from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 
 supported_bm42_models = [
@@ -47,17 +50,17 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
     WARNING: This model is expected to be used with `modifier="idf"` in the sparse vector index of Qdrant.
     """
-    
+
     ONNX_OUTPUT_NAMES = ["attention_6"]
 
     def __init__(
-            self,
-            model_name: str,
-            cache_dir: Optional[str] = None,
-            threads: Optional[int] = None,
-            providers: Optional[Sequence[OnnxProvider]] = None,
-            alpha: float = 0.5,
-            **kwargs,
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        alpha: float = 0.5,
+        **kwargs,
     ):
         """
         Args:
@@ -78,10 +81,10 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
         model_description = self._get_model_description(model_name)
-        cache_dir = define_cache_dir(cache_dir)
+        self.cache_dir = define_cache_dir(cache_dir)
 
         model_dir = self.download_model(
-            model_description, cache_dir, local_files_only=self._local_files_only
+            model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         self.load_onnx_model(
@@ -119,7 +122,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         return result
 
     @classmethod
-    def _aggregate_weights(cls, tokens: List[Tuple[str, List[int]]], weights: List[float]) -> List[Tuple[str, float]]:
+    def _aggregate_weights(
+        cls, tokens: List[Tuple[str, List[int]]], weights: List[float]
+    ) -> List[Tuple[str, float]]:
         result = []
         for token, idxs in tokens:
             sum_weight = sum(weights[idx] for idx in idxs)
@@ -127,7 +132,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         return result
 
     def _reconstruct_bpe(
-            self, bpe_tokens: Iterable[Tuple[int, str]]
+        self, bpe_tokens: Iterable[Tuple[int, str]]
     ) -> List[Tuple[str, List[int]]]:
         result = []
         acc = ""
@@ -169,7 +174,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             # Num 0: Log(1/1 + 1) = 0.6931471805599453
             # Num 1: Log(1/2 + 1) = 0.4054651081081644
             # Num 2: Log(1/3 + 1) = 0.28768207245178085
-            new_vector[token_id] = math.log(1. + value) ** self.alpha  # value
+            new_vector[token_id] = math.log(1.0 + value) ** self.alpha  # value
 
         return new_vector
 
@@ -221,11 +226,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             return f.read().splitlines()
 
     def embed(
-            self,
-            documents: Union[str, Iterable[str]],
-            batch_size: int = 256,
-            parallel: Optional[int] = None,
-            **kwargs,
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
     ) -> Iterable[SparseEmbedding]:
         """
         Encode a list of documents into list of embeddings.
@@ -248,6 +253,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
+            alpha=self.alpha,
         )
 
     @classmethod
@@ -278,4 +284,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
-        return TextEmbeddingWorker
+        return Bm42TextEmbeddingWorker
+
+
+class Bm42TextEmbeddingWorker(TextEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Bm42:
+        return Bm42(model_name=model_name, cache_dir=cache_dir, **kwargs)

fastembed/sparse/sparse_embedding_base.py

@@ -28,11 +28,11 @@ class SparseEmbedding:
 
 class SparseTextEmbeddingBase(ModelManagement):
     def __init__(
-            self,
-            model_name: str,
-            cache_dir: Optional[str] = None,
-            threads: Optional[int] = None,
-            **kwargs,
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
@@ -40,15 +40,17 @@ class SparseTextEmbeddingBase(ModelManagement):
         self._local_files_only = kwargs.pop("local_files_only", False)
 
     def embed(
-            self,
-            documents: Union[str, Iterable[str]],
-            batch_size: int = 256,
-            parallel: Optional[int] = None,
-            **kwargs,
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
     ) -> Iterable[SparseEmbedding]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[SparseEmbedding]:
+    def passage_embed(
+        self, texts: Iterable[str], **kwargs
+    ) -> Iterable[SparseEmbedding]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -63,7 +65,9 @@ class SparseTextEmbeddingBase(ModelManagement):
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs
+    ) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 

fastembed/sparse/sparse_text_embedding.py

@@ -1,16 +1,17 @@
-from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
 
 from fastembed.common import OnnxProvider
+from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.bm42 import Bm42
-from fastembed.sparse.sparse_embedding_base import SparseTextEmbeddingBase, SparseEmbedding
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
 from fastembed.sparse.splade_pp import SpladePP
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
-    EMBEDDINGS_REGISTRY: List[Type[SparseTextEmbeddingBase]] = [
-        SpladePP,
-        Bm42,
-    ]
+    EMBEDDINGS_REGISTRY: List[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25]
 
     @classmethod
     def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -52,9 +53,16 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+            if any(
+                model_name.lower() == model["model"].lower()
+                for model in supported_models
+            ):
                 self.model = EMBEDDING_MODEL_TYPE(
-                    model_name, cache_dir, threads, providers=providers, **kwargs
+                    model_name,
+                    cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    **kwargs,
                 )
                 return
 
@@ -87,7 +95,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs
+    ) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 

fastembed/sparse/splade_pp.py

@@ -1,14 +1,16 @@
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
 
 import numpy as np
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir
-from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 
-
 supported_splade_models = [
     {
         "model": "prithvida/Splade_PP_en_v1",
@@ -34,9 +36,7 @@ supported_splade_models = [
 
 
 class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext
-    ) -> Iterable[SparseEmbedding]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
         relu_log = np.log(1 + np.maximum(output.model_output, 0))
 
         weighted_log = relu_log * np.expand_dims(output.attention_mask, axis=-1)
@@ -82,10 +82,10 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
         model_description = self._get_model_description(model_name)
-        cache_dir = define_cache_dir(cache_dir)
+        self.cache_dir = define_cache_dir(cache_dir)
 
         model_dir = self.download_model(
-            model_description, cache_dir, local_files_only=self._local_files_only
+            model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         self.load_onnx_model(
@@ -131,9 +131,5 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
 
 class SpladePPEmbeddingWorker(TextEmbeddingWorker):
-    def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
-    ) -> SpladePP:
-        return SpladePP(model_name=model_name, cache_dir=cache_dir, threads=1)
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> SpladePP:
+        return SpladePP(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

fastembed/sparse/utils/tokenizer.py

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

fastembed/text/clip_embedding.py

@@ -1,4 +1,4 @@
-from typing import Type, List, Dict, Any, Tuple, Iterable
+from typing import Any, Dict, Iterable, List, Type
 
 import numpy as np
 
@@ -42,8 +42,8 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
 
 class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
     def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
+        self, model_name: str, cache_dir: str, **kwargs
     ) -> OnnxTextEmbedding:
-        return CLIPOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)
+        return CLIPOnnxEmbedding(
+            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+        )

fastembed/text/e5_onnx_embedding.py

@@ -1,4 +1,4 @@
-from typing import Type, List, Dict, Any
+from typing import Any, Dict, List, Type
 
 import numpy as np
 
@@ -45,7 +45,9 @@ class E5OnnxEmbedding(OnnxTextEmbedding):
         """
         return supported_multilingual_e5_models
 
-    def _preprocess_onnx_input(self, onnx_input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+    def _preprocess_onnx_input(
+        self, onnx_input: Dict[str, np.ndarray], **kwargs
+    ) -> Dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -55,8 +57,8 @@ class E5OnnxEmbedding(OnnxTextEmbedding):
 
 class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
     def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
+        self, model_name: str, cache_dir: str, **kwargs
     ) -> E5OnnxEmbedding:
-        return E5OnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)
+        return E5OnnxEmbedding(
+            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+        )

fastembed/text/jina_onnx_embedding.py

@@ -1,4 +1,4 @@
-from typing import Type, List, Dict, Any, Tuple, Iterable
+from typing import Any, Dict, Iterable, List, Type
 
 import numpy as np
 
@@ -24,6 +24,14 @@ supported_jina_models = [
         "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
         "model_file": "onnx/model.onnx",
     },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-de",
+        "dim": 768,
+        "description": "German embedding model supporting 8192 sequence length",
+        "size_in_GB": 0.32,
+        "sources": {"hf": "jinaai/jina-embeddings-v2-base-de"},
+        "model_file": "onnx/model_fp16.onnx",
+    },
 ]
 
 
@@ -61,8 +69,8 @@ class JinaOnnxEmbedding(OnnxTextEmbedding):
 
 class JinaEmbeddingWorker(OnnxTextEmbeddingWorker):
     def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
+        self, model_name: str, cache_dir: str, **kwargs
     ) -> OnnxTextEmbedding:
-        return JinaOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)
+        return JinaOnnxEmbedding(
+            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+        )

fastembed/text/mini_lm_embedding.py

@@ -0,0 +1,58 @@
+from typing import Any, Dict, Iterable, List, Type
+
+import numpy as np
+
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import normalize
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+supported_mini_lm_models = [
+    {
+        "model": "sentence-transformers/all-MiniLM-L6-v2",
+        "dim": 384,
+        "description": "Sentence Transformer model, MiniLM-L6-v2",
+        "size_in_GB": 0.09,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
+            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
+        },
+        "model_file": "model.onnx",
+    }
+]
+
+
+class MiniLMOnnxEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+        return MiniLMEmbeddingWorker
+
+    @classmethod
+    def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
+        token_embeddings = model_output
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
+        input_mask_expanded = input_mask_expanded.astype(float)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.sum(input_mask_expanded, axis=1)
+        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+        return pooled_embeddings
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_mini_lm_models
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+        embeddings = output.model_output
+        attn_mask = output.attention_mask
+        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+
+
+class MiniLMEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxTextEmbedding:
+        return MiniLMOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

fastembed/text/onnx_embedding.py

@@ -1,14 +1,13 @@
-from typing import Dict, Optional, Union, Iterable, Type, List, Any, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
 
 import numpy as np
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize, define_cache_dir
-from fastembed.text.onnx_text_model import TextEmbeddingWorker, OnnxTextModel
+from fastembed.common.utils import define_cache_dir, normalize
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 from fastembed.text.text_embedding_base import TextEmbeddingBase
 
-
 supported_onnx_models = [
     {
         "model": "BAAI/bge-base-en",
@@ -71,17 +70,6 @@ supported_onnx_models = [
         },
         "model_file": "model_optimized.onnx",
     },
-    {
-        "model": "sentence-transformers/all-MiniLM-L6-v2",
-        "dim": 384,
-        "description": "Sentence Transformer model, MiniLM-L6-v2",
-        "size_in_GB": 0.09,
-        "sources": {
-            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
-            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
-        },
-        "model_file": "model.onnx",
-    },
     {
         "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
         "dim": 384,
@@ -231,9 +219,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
         model_description = self._get_model_description(model_name)
-        cache_dir = define_cache_dir(cache_dir)
+        self.cache_dir = define_cache_dir(cache_dir)
         model_dir = self.download_model(
-            model_description, cache_dir, local_files_only=self._local_files_only
+            model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         self.load_onnx_model(
@@ -271,6 +259,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
+            **kwargs,
         )
 
     @classmethod
@@ -295,5 +284,6 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
         self,
         model_name: str,
         cache_dir: str,
+        **kwargs,
     ) -> OnnxTextEmbedding:
-        return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)
+        return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

fastembed/text/onnx_text_model.py

@@ -1,14 +1,14 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Union
 
 import numpy as np
 from tokenizers import Encoding
 
 from fastembed.common import OnnxProvider
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T, OnnxOutputContext
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
@@ -44,7 +44,10 @@ class OnnxTextModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
     ) -> None:
         super().load_onnx_model(
-            model_dir=model_dir, model_file=model_file, threads=threads, providers=providers
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
         )
         self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
 
@@ -86,6 +89,7 @@ class OnnxTextModel(OnnxModel[T]):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
+        **kwargs,
     ) -> Iterable[T]:
         is_small = False
 
@@ -104,11 +108,10 @@ class OnnxTextModel(OnnxModel[T]):
             for batch in iter_batch(documents, batch_size):
                 yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            params = {
-                "model_name": model_name,
-                "cache_dir": cache_dir,
-            }
+            start_method = (
+                "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            )
+            params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
             pool = ParallelWorkerPool(
                 parallel, self._get_worker_class(), start_method=start_method
             )

fastembed/text/text_embedding.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterable, List, Optional, Type, Union, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
 
 import numpy as np
 
@@ -6,6 +6,7 @@ from fastembed.common import OnnxProvider
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
 from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
 from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
+from fastembed.text.mini_lm_embedding import MiniLMOnnxEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
 
@@ -16,6 +17,7 @@ class TextEmbedding(TextEmbeddingBase):
         E5OnnxEmbedding,
         JinaOnnxEmbedding,
         CLIPOnnxEmbedding,
+        MiniLMOnnxEmbedding,
     ]
 
     @classmethod
@@ -59,9 +61,16 @@ class TextEmbedding(TextEmbeddingBase):
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+            if any(
+                model_name.lower() == model["model"].lower()
+                for model in supported_models
+            ):
                 self.model = EMBEDDING_MODEL_TYPE(
-                    model_name, cache_dir, threads, providers=providers, **kwargs
+                    model_name,
+                    cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    **kwargs,
                 )
                 return
 

fastembed/text/text_embedding_base.py

@@ -42,7 +42,9 @@ class TextEmbeddingBase(ModelManagement):
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
+    def query_embed(
+        self, query: Union[str, Iterable[str]], **kwargs
+    ) -> Iterable[np.ndarray]:
         """
         Embeds queries
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
-name = "fastembed"
-version = "0.3.0"
+name = "fastembed-gpu"
+version = "0.3.1"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -12,16 +12,15 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
 
 [tool.poetry.dependencies]
 python = ">=3.8.0,<3.13"
-onnx = "^1.15.0"
-onnxruntime = "^1.17.0"
+onnxruntime-gpu = "^1.17.0"
 tqdm = "^4.66"
 requests = "^2.31"
 tokenizers = ">=0.15,<1.0"
 huggingface-hub = ">=0.20,<1.0"
 loguru = "^0.7.2"
 numpy = [
-    { version = ">=1.21", python = "<3.12" },
-    { version = ">=1.26", python = ">=3.12" }
+    { version = ">=1.21, <2", python = "<3.12" },
+    { version = ">=1.26, <2", python = ">=3.12" }
 ]
 pillow = "^10.3.0"
 snowballstemmer = "^2.2.0"

tests/profiling.py

@@ -57,7 +57,9 @@ class HF:
         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")
+        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)
@@ -84,7 +86,9 @@ embedding_model = DefaultEmbedding()
 
 
 # %%
-def calculate_time_stats(embed_func: Callable, documents: list, k: int) -> Tuple[float, float, float]:
+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
@@ -101,13 +105,17 @@ def calculate_time_stats(embed_func: Callable, documents: list, k: int) -> Tuple
 # %%
 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)
+fst_stats = calculate_time_stats(
+    lambda x: list(embedding_model.embed(x)), documents, k=2
+)
 print(f"FastEmbed (Average, Max, Min): {fst_stats}")
 
 
 # %%
 def plot_character_per_second_comparison(
-    hf_stats: Tuple[float, float, float], fst_stats: Tuple[float, float, float], documents: list
+    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)

tests/test_attention_embeddings.py

@@ -1,10 +1,14 @@
 import numpy as np
+import pytest
 
 from fastembed import SparseTextEmbedding
 
 
-def test_attention_embeddings():
-    model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
+@pytest.mark.parametrize(
+    "model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]
+)
+def test_attention_embeddings(model_name):
+    model = SparseTextEmbedding(model_name=model_name)
 
     output = list(
         model.query_embed(
@@ -56,3 +60,25 @@ def test_attention_embeddings():
     for result in output:
         assert len(result.indices) == len(result.values)
         assert len(result.indices) == 2
+
+
+@pytest.mark.parametrize(
+    "model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]
+)
+def test_parallel_processing(model_name):
+    model = SparseTextEmbedding(model_name=model_name)
+
+    docs = ["hello world", "attention embedding"] * 100
+    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+
+    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+
+    embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
+
+    assert len(embeddings) == len(docs)
+
+    for emb_1, emb_2, emb_3 in zip(embeddings, embeddings_2, embeddings_3):
+        assert np.allclose(emb_1.indices, emb_2.indices)
+        assert np.allclose(emb_1.indices, emb_3.indices)
+        assert np.allclose(emb_1.values, emb_2.values)
+        assert np.allclose(emb_1.values, emb_3.values)

tests/test_late_interaction_embeddings.py

@@ -1,7 +1,8 @@
 import numpy as np
 
-from fastembed.late_interaction.late_interaction_text_embedding import LateInteractionTextEmbedding
-
+from fastembed.late_interaction.late_interaction_text_embedding import (
+    LateInteractionTextEmbedding,
+)
 
 # vectors are abridged and rounded for brevity
 CANONICAL_COLUMN_VALUES = {
@@ -76,7 +77,7 @@ def test_single_embedding():
 
     for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
         print("evaluating", model_name)
-        model = LateInteractionTextEmbedding(model_name=model_name, cache_dir="colbert-cache")
+        model = LateInteractionTextEmbedding(model_name=model_name)
         result = next(iter(model.embed(docs_to_embed, batch_size=6)))
         token_num, abridged_dim = expected_result.shape
         assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4)
@@ -87,16 +88,14 @@ def test_single_embedding_query():
 
     for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
         print("evaluating", model_name)
-        model = LateInteractionTextEmbedding(model_name=model_name, cache_dir="colbert-cache")
+        model = LateInteractionTextEmbedding(model_name=model_name)
         result = next(iter(model.query_embed(queries_to_embed)))
         token_num, abridged_dim = expected_result.shape
         assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4)
 
 
 def test_parallel_processing():
-    model = LateInteractionTextEmbedding(
-        model_name="colbert-ir/colbertv2.0", cache_dir="colbert-cache"
-    )
+    model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
     token_dim = 128
     docs = ["hello world", "flag embedding"] * 100
     embeddings = list(model.embed(docs, batch_size=10, parallel=2))

tests/test_sparse_embeddings.py

@@ -93,5 +93,9 @@ def test_parallel_processing():
             == sparse_embedding_duo.indices.tolist()
             == sparse_embedding_all.indices.tolist()
         )
-        assert np.allclose(sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3)
-        assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3)
+        assert np.allclose(
+            sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3
+        )
+        assert np.allclose(
+            sparse_embedding.values, sparse_embedding_all.values, atol=1e-3
+        )

tests/test_text_onnx_embeddings.py

@@ -26,17 +26,30 @@ CANONICAL_VECTOR_VALUES = {
     "BAAI/bge-large-en-v1.5-quantized": np.array(
         [0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]
     ),
-    "sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259, 0.0058, 0.0114, 0.0380, -0.0233]),
+    "sentence-transformers/all-MiniLM-L6-v2": np.array(
+        [-0.034478, 0.03102, 0.00673, 0.02611, -0.039362]
+    ),
     "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array(
         [0.0094, 0.0184, 0.0328, 0.0072, -0.0351]
     ),
-    "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
+    "intfloat/multilingual-e5-large": np.array(
+        [0.0098, 0.0045, 0.0066, -0.0354, 0.0070]
+    ),
     "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array(
         [-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
     ),
-    "jinaai/jina-embeddings-v2-small-en": np.array([-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]),
-    "jinaai/jina-embeddings-v2-base-en": np.array([-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]),
-    "nomic-ai/nomic-embed-text-v1": np.array([0.0061, 0.0103, -0.0296, -0.0242, -0.0170]),
+    "jinaai/jina-embeddings-v2-small-en": np.array(
+        [-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]
+    ),
+    "jinaai/jina-embeddings-v2-base-en": np.array(
+        [-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]
+    ),
+    "jinaai/jina-embeddings-v2-base-de": np.array(
+      [-0.0085, 0.0417, 0.0342, 0.0309, -0.0149]
+    ),
+    "nomic-ai/nomic-embed-text-v1": np.array(
+        [0.0061, 0.0103, -0.0296, -0.0242, -0.0170]
+    ),
     "nomic-ai/nomic-embed-text-v1.5": np.array(
         [-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
     ),
@@ -49,13 +62,21 @@ CANONICAL_VECTOR_VALUES = {
     "mixedbread-ai/mxbai-embed-large-v1": np.array(
         [0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]
     ),
-    "snowflake/snowflake-arctic-embed-xs": np.array([0.0092, 0.0619, 0.0196, 0.009, -0.0114]),
-    "snowflake/snowflake-arctic-embed-s": np.array([-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]),
-    "snowflake/snowflake-arctic-embed-m": np.array([-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]),
+    "snowflake/snowflake-arctic-embed-xs": np.array(
+        [0.0092, 0.0619, 0.0196, 0.009, -0.0114]
+    ),
+    "snowflake/snowflake-arctic-embed-s": np.array(
+        [-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]
+    ),
+    "snowflake/snowflake-arctic-embed-m": np.array(
+        [-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]
+    ),
     "snowflake/snowflake-arctic-embed-m-long": np.array(
         [0.0080, -0.0266, -0.0335, 0.0282, 0.0143]
     ),
-    "snowflake/snowflake-arctic-embed-l": np.array([0.0189, -0.0673, 0.0183, 0.0124, 0.0146]),
+    "snowflake/snowflake-arctic-embed-l": np.array(
+        [0.0189, -0.0673, 0.0183, 0.0124, 0.0146]
+    ),
     "Qdrant/clip-ViT-B-32-text": np.array([0.0083, 0.0103, -0.0138, 0.0199, -0.0069]),
 }