{ "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': ,\n 'optimizer_status': ,\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': ,\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": "
\n\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
candidatesrescoreaccuracytime
01False0.900.221826
11True0.910.134167
23False0.880.115299
33True0.970.209320
45False0.840.154485
55True0.910.124424
63False0.990.121695
73True0.960.123257
89False0.940.119629
99True0.980.119372
1015False0.900.121621
1115True0.970.125466
1210False0.930.135910
1310True0.950.138135
1430False0.940.177928
1530True0.980.254588
1650False0.940.268659
1750True0.960.269792
1820False0.960.249941
1920True0.960.247138
2060False0.970.251301
2160True0.980.256504
22100False0.980.270049
23100True0.970.248972
2450False0.970.306356
2550True0.980.257544
26150False0.980.238811
27150True0.990.263939
28250False0.990.256558
29250True1.000.335823
\n
", "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 }