Files
fastembed/docs/qdrant/Binary_Quantization_with_Qdrant.ipynb

378 lines
19 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Binary Quantization with Qdrant & OpenAI Embedding\n",
"\n",
"---\n",
"In the world of large-scale data retrieval and processing, efficiency is crucial. With the exponential growth of data, the ability to retrieve information quickly and accurately can significantly affect system performance. This blog post explores a technique known as binary quantization applied to OpenAI embeddings, demonstrating how it can enhance **retrieval latency by 20x** or more.\n",
"\n",
"## What Are OpenAI Embeddings?\n",
"OpenAI embeddings are numerical representations of textual information. They transform text into a vector space where semantically similar texts are mapped close together. This mathematical representation enables computers to understand and process human language more effectively.\n",
"\n",
"## Binary Quantization\n",
"Binary quantization is a method which converts continuous numerical values into binary values (0 or 1). It simplifies the data structure, allowing faster computations. Here's a brief overview of the binary quantization process applied to OpenAI embeddings:\n",
"\n",
"1. **Load Embeddings**: OpenAI embeddings are loaded from parquet files.\n",
"2. **Binary Transformation**: The continuous valued vectors are converted into binary form. Here, values greater than 0 are set to 1, and others remain 0.\n",
"3. **Comparison & Retrieval**: Binary vectors are used for comparison using logical XOR operations and other efficient algorithms.\n",
"\n",
"Binary Quantization is a promising approach to improve retrieval speeds and reduce memory footprint of vector search engines. In this notebook we will show how to use Qdrant to perform binary quantization of vectors and perform fast similarity search on the resulting index.\n",
"\n",
"## Table of Contents\n",
"1. Imports\n",
"2. Download and Slice Dataset\n",
"3. Create Qdrant Collection\n",
"4. Indexing\n",
"5. Search\n",
"\n",
"## 1. Imports"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:33:03.013948Z",
"start_time": "2024-04-01T16:33:01.019043Z"
}
},
"outputs": [],
"source": [
"!pip install qdrant-client pandas dataset --quiet --upgrade"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:33:03.914729Z",
"start_time": "2024-04-01T16:33:03.015394Z"
}
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/joein/work/qdrant/fastembed/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n"
]
}
],
"source": [
"import os\n",
"import random\n",
"import time\n",
"\n",
"import numpy as np\n",
"import pandas as pd\n",
"from qdrant_client import QdrantClient, models\n",
"\n",
"random.seed(37)\n",
"np.random.seed(37)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Download and Slice Dataset\n",
"\n",
"We will be using the [dbpedia-entities](https://huggingface.co/datasets/Qdrant/dbpedia-entities-openai3-text-embedding-3-small-1536-100K) dataset from the [HuggingFace Datasets](https://huggingface.co/datasets) library. This contains 100K vectors of 1536 dimensions each"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:33:09.085853Z",
"start_time": "2024-04-01T16:33:03.912688Z"
}
},
"outputs": [
{
"data": {
"text/plain": "100000"
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import datasets\n",
"\n",
"dataset = datasets.load_dataset(\n",
" \"Qdrant/dbpedia-entities-openai3-text-embedding-3-small-1536-100K\", split=\"train\"\n",
")\n",
"len(dataset)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:33:09.176212Z",
"start_time": "2024-04-01T16:33:09.084550Z"
}
},
"outputs": [
{
"data": {
"text/plain": "True"
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"client = QdrantClient(\n",
" prefer_grpc=True,\n",
")\n",
"\n",
"collection_name = \"binary-quantization\"\n",
"client.recreate_collection(\n",
" collection_name=collection_name,\n",
" vectors_config=models.VectorParams(\n",
" size=1536,\n",
" distance=models.Distance.DOT,\n",
" on_disk=True,\n",
" ),\n",
" quantization_config=models.BinaryQuantization(\n",
" binary=models.BinaryQuantizationConfig(always_ram=True),\n",
" ),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.383986Z",
"start_time": "2024-04-01T16:33:09.175725Z"
}
},
"outputs": [],
"source": [
"def iter_dataset(dataset):\n",
" for point in dataset:\n",
" yield point[\"openai\"], {\"text\": point[\"text\"]}\n",
"\n",
"\n",
"vectors, payload = zip(*iter_dataset(dataset))\n",
"client.upload_collection(\n",
" collection_name=collection_name,\n",
" vectors=vectors,\n",
" payload=payload,\n",
" parallel=max(1, (os.cpu_count() // 2)),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.390886Z",
"start_time": "2024-04-01T16:34:13.385961Z"
}
},
"outputs": [
{
"data": {
"text/plain": "{'status': <CollectionStatus.YELLOW: 'yellow'>,\n 'optimizer_status': <OptimizersStatusOneOf.OK: 'ok'>,\n 'vectors_count': 116640,\n 'indexed_vectors_count': 43520,\n 'points_count': 116640,\n 'segments_count': 6,\n 'config': {'params': {'vectors': {'size': 1536,\n 'distance': <Distance.DOT: 'Dot'>,\n 'hnsw_config': None,\n 'quantization_config': None,\n 'on_disk': True},\n 'shard_number': 1,\n 'sharding_method': None,\n 'replication_factor': 1,\n 'write_consistency_factor': 1,\n 'read_fan_out_factor': None,\n 'on_disk_payload': True,\n 'sparse_vectors': None},\n 'hnsw_config': {'m': 16,\n 'ef_construct': 100,\n 'full_scan_threshold': 10000,\n 'max_indexing_threads': 0,\n 'on_disk': False,\n 'payload_m': None},\n 'optimizer_config': {'deleted_threshold': 0.2,\n 'vacuum_min_vector_number': 1000,\n 'default_segment_number': 0,\n 'max_segment_size': None,\n 'memmap_threshold': None,\n 'indexing_threshold': 20000,\n 'flush_interval_sec': 5,\n 'max_optimization_threads': None},\n 'wal_config': {'wal_capacity_mb': 32, 'wal_segments_ahead': 0},\n 'quantization_config': {'binary': {'always_ram': True}}},\n 'payload_schema': {}}"
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"collection_info = client.get_collection(collection_name=f\"{collection_name}\")\n",
"collection_info.dict()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Oversampling vs Recall\n",
"\n",
"### Preparing a query dataset\n",
"\n",
"For the purpose of this illustration, we'll take a few vectors which we know are already in the index and query them. We should get the same vectors back as results from the Qdrant index. "
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.453626Z",
"start_time": "2024-04-01T16:34:13.391567Z"
}
},
"outputs": [
{
"data": {
"text/plain": "[89391,\n 79659,\n 12006,\n 80978,\n 87219,\n 97885,\n 83155,\n 67504,\n 4645,\n 82711,\n 48395,\n 57375,\n 69208,\n 14136,\n 89515,\n 59880,\n 78730,\n 36952,\n 49620,\n 96486,\n 55473,\n 58179,\n 18926,\n 6489,\n 11931,\n 54146,\n 9850,\n 71259,\n 37825,\n 47331,\n 84964,\n 92399,\n 56669,\n 77042,\n 73744,\n 47993,\n 83780,\n 92429,\n 75114,\n 4463,\n 69030,\n 81185,\n 27950,\n 66217,\n 54652,\n 8260,\n 1151,\n 993,\n 85954,\n 66863,\n 47303,\n 8992,\n 92688,\n 76030,\n 29472,\n 3077,\n 42454,\n 46120,\n 69140,\n 20877,\n 2844,\n 95423,\n 1770,\n 28568,\n 96448,\n 94227,\n 40837,\n 91684,\n 29785,\n 66936,\n 85121,\n 39546,\n 81910,\n 5514,\n 37068,\n 35731,\n 93990,\n 26685,\n 63076,\n 18762,\n 27922,\n 34916,\n 80976,\n 83189,\n 6328,\n 57508,\n 58860,\n 13758,\n 72976,\n 85030,\n 332,\n 34963,\n 85009,\n 31344,\n 11560,\n 58108,\n 85163,\n 17064,\n 44712,\n 45962]"
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"query_indices = random.sample(range(len(dataset)), 100)\n",
"query_dataset = dataset[query_indices]\n",
"query_indices"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.453928Z",
"start_time": "2024-04-01T16:34:13.452405Z"
}
},
"outputs": [],
"source": [
"## Add Gaussian noise to any vector\n",
"\n",
"\n",
"def add_noise(vector, noise=0.05):\n",
" return vector + noise * np.random.randn(*vector.shape)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.457839Z",
"start_time": "2024-04-01T16:34:13.455431Z"
}
},
"outputs": [],
"source": [
"def correct(results, text):\n",
" return text in [x.payload[\"text\"] for x in results]\n",
"\n",
"\n",
"def count_correct(query_dataset, limit=1, oversampling=1, rescore=False):\n",
" correct_results = 0\n",
" for query_vector, text in zip(query_dataset[\"openai\"], query_dataset[\"text\"]):\n",
" results = client.search(\n",
" collection_name=collection_name,\n",
" query_vector=add_noise(np.array(query_vector)),\n",
" limit=limit,\n",
" search_params=models.SearchParams(\n",
" quantization=models.QuantizationSearchParams(\n",
" rescore=rescore,\n",
" oversampling=oversampling,\n",
" )\n",
" ),\n",
" )\n",
" correct_results += correct(results, text)\n",
" return correct_results"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:40:48.347002Z",
"start_time": "2024-04-01T16:40:42.228551Z"
},
"collapsed": false
},
"outputs": [],
"source": [
"limit_grid = [1, 3, 10, 20, 50]\n",
"oversampling_grid = [1.0, 3.0, 5.0]\n",
"rescore_grid = [False, True]\n",
"results = []\n",
"\n",
"for limit in limit_grid:\n",
" for oversampling in oversampling_grid:\n",
" for rescore in rescore_grid:\n",
" start = time.perf_counter()\n",
" correct_results = count_correct(\n",
" query_dataset, limit=limit, oversampling=oversampling, rescore=rescore\n",
" )\n",
" end = time.perf_counter()\n",
" results.append(\n",
" {\n",
" \"limit\": limit,\n",
" \"oversampling\": oversampling,\n",
" \"candidates\": int(oversampling * limit),\n",
" \"rescore\": rescore,\n",
" \"accuracy\": correct_results / 100,\n",
" \"total queries\": len(query_dataset[\"text\"]),\n",
" \"time\": end - start,\n",
" }\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:41:55.445405Z",
"start_time": "2024-04-01T16:41:55.442687Z"
}
},
"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>candidates</th>\n <th>rescore</th>\n <th>accuracy</th>\n <th>time</th>\n </tr>\n </thead>\n <tbody>\n <tr>\n <th>0</th>\n <td>1</td>\n <td>False</td>\n <td>0.90</td>\n <td>0.221826</td>\n </tr>\n <tr>\n <th>1</th>\n <td>1</td>\n <td>True</td>\n <td>0.91</td>\n <td>0.134167</td>\n </tr>\n <tr>\n <th>2</th>\n <td>3</td>\n <td>False</td>\n <td>0.88</td>\n <td>0.115299</td>\n </tr>\n <tr>\n <th>3</th>\n <td>3</td>\n <td>True</td>\n <td>0.97</td>\n <td>0.209320</td>\n </tr>\n <tr>\n <th>4</th>\n <td>5</td>\n <td>False</td>\n <td>0.84</td>\n <td>0.154485</td>\n </tr>\n <tr>\n <th>5</th>\n <td>5</td>\n <td>True</td>\n <td>0.91</td>\n <td>0.124424</td>\n </tr>\n <tr>\n <th>6</th>\n <td>3</td>\n <td>False</td>\n <td>0.99</td>\n <td>0.121695</td>\n </tr>\n <tr>\n <th>7</th>\n <td>3</td>\n <td>True</td>\n <td>0.96</td>\n <td>0.123257</td>\n </tr>\n <tr>\n <th>8</th>\n <td>9</td>\n <td>False</td>\n <td>0.94</td>\n <td>0.119629</td>\n </tr>\n <tr>\n <th>9</th>\n <td>9</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.119372</td>\n </tr>\n <tr>\n <th>10</th>\n <td>15</td>\n <td>False</td>\n <td>0.90</td>\n <td>0.121621</td>\n </tr>\n <tr>\n <th>11</th>\n <td>15</td>\n <td>True</td>\n <td>0.97</td>\n <td>0.125466</td>\n </tr>\n <tr>\n <th>12</th>\n <td>10</td>\n <td>False</td>\n <td>0.93</td>\n <td>0.135910</td>\n </tr>\n <tr>\n <th>13</th>\n <td>10</td>\n <td>True</td>\n <td>0.95</td>\n <td>0.138135</td>\n </tr>\n <tr>\n <th>14</th>\n <td>30</td>\n <td>False</td>\n <td>0.94</td>\n <td>0.177928</td>\n </tr>\n <tr>\n <th>15</th>\n <td>30</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.254588</td>\n </tr>\n <tr>\n <th>16</th>\n <td>50</td>\n <td>False</td>\n <td>0.94</td>\n <td>0.268659</td>\n </tr>\n <tr>\n <th>17</th>\n <td>50</td>\n <td>True</td>\n <td>0.96</td>\n <td>0.269792</td>\n </tr>\n <tr>\n <th>18</th>\n <td>20</td>\n <td>False</td>\n <td>0.96</td>\n <td>0.249941</td>\n </tr>\n <tr>\n <th>19</th>\n <td>20</td>\n <td>True</td>\n <td>0.96</td>\n <td>0.247138</td>\n </tr>\n <tr>\n <th>20</th>\n <td>60</td>\n <td>False</td>\n <td>0.97</td>\n <td>0.251301</td>\n </tr>\n <tr>\n <th>21</th>\n <td>60</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.256504</td>\n </tr>\n <tr>\n <th>22</th>\n <td>100</td>\n <td>False</td>\n <td>0.98</td>\n <td>0.270049</td>\n </tr>\n <tr>\n <th>23</th>\n <td>100</td>\n <td>True</td>\n <td>0.97</td>\n <td>0.248972</td>\n </tr>\n <tr>\n <th>24</th>\n <td>50</td>\n <td>False</td>\n <td>0.97</td>\n <td>0.306356</td>\n </tr>\n <tr>\n <th>25</th>\n <td>50</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.257544</td>\n </tr>\n <tr>\n <th>26</th>\n <td>150</td>\n <td>False</td>\n <td>0.98</td>\n <td>0.238811</td>\n </tr>\n <tr>\n <th>27</th>\n <td>150</td>\n <td>True</td>\n <td>0.99</td>\n <td>0.263939</td>\n </tr>\n <tr>\n <th>28</th>\n <td>250</td>\n <td>False</td>\n <td>0.99</td>\n <td>0.256558</td>\n </tr>\n <tr>\n <th>29</th>\n <td>250</td>\n <td>True</td>\n <td>1.00</td>\n <td>0.335823</td>\n </tr>\n </tbody>\n</table>\n</div>",
"text/plain": " candidates rescore accuracy time\n0 1 False 0.90 0.221826\n1 1 True 0.91 0.134167\n2 3 False 0.88 0.115299\n3 3 True 0.97 0.209320\n4 5 False 0.84 0.154485\n5 5 True 0.91 0.124424\n6 3 False 0.99 0.121695\n7 3 True 0.96 0.123257\n8 9 False 0.94 0.119629\n9 9 True 0.98 0.119372\n10 15 False 0.90 0.121621\n11 15 True 0.97 0.125466\n12 10 False 0.93 0.135910\n13 10 True 0.95 0.138135\n14 30 False 0.94 0.177928\n15 30 True 0.98 0.254588\n16 50 False 0.94 0.268659\n17 50 True 0.96 0.269792\n18 20 False 0.96 0.249941\n19 20 True 0.96 0.247138\n20 60 False 0.97 0.251301\n21 60 True 0.98 0.256504\n22 100 False 0.98 0.270049\n23 100 True 0.97 0.248972\n24 50 False 0.97 0.306356\n25 50 True 0.98 0.257544\n26 150 False 0.98 0.238811\n27 150 True 0.99 0.263939\n28 250 False 0.99 0.256558\n29 250 True 1.00 0.335823"
},
"execution_count": 22,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"df = pd.DataFrame(results)\n",
"df[[\"candidates\", \"rescore\", \"accuracy\", \"time\"]]\n",
"# df.to_csv(\"candidates-rescore-time.csv\", index=False)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.13"
}
},
"nbformat": 4,
"nbformat_minor": 4
}