wehub/flagopen--flagembedding
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:39:21 +08:00
commit bc34f6df14
1149 changed files with 328099 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/tutorial/3_Indexing/3.1.1.ipynb
+411
View File
@@ -0,0 +1,411 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Indexing Using Faiss"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In practical cases, datasets contain thousands or millions of rows. Looping through the whole corpus to find the best answer to a query is very time and space consuming. In this tutorial, we'll introduce how to use indexing to make our retrieval fast and neat."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 0: Setup"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Install the dependencies in the environment."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%pip install -U FlagEmbedding"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### faiss-gpu on Linux (x86_64)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Faiss maintain the latest updates on conda. So if you have GPUs on Linux x86_64, create a conda virtual environment and run:\n",
"\n",
"```conda install -c pytorch -c nvidia faiss-gpu=1.8.0```\n",
"\n",
"and make sure you select that conda env as the kernel for this notebook."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### faiss-cpu\n",
"\n",
"Otherwise it's simple, just run the following cell to install `faiss-cpu`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%pip install -U faiss-cpu"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 1: Dataset"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Below is a super tiny courpus with only 10 sentences, which will be the dataset we use.\n",
"\n",
"Each sentence is a concise discription of a famous people in specific domain."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"corpus = [\n",
" \"Michael Jackson was a legendary pop icon known for his record-breaking music and dance innovations.\",\n",
" \"Fei-Fei Li is a professor in Stanford University, revolutionized computer vision with the ImageNet project.\",\n",
" \"Brad Pitt is a versatile actor and producer known for his roles in films like 'Fight Club' and 'Once Upon a Time in Hollywood.'\",\n",
" \"Geoffrey Hinton, as a foundational figure in AI, received Turing Award for his contribution in deep learning.\",\n",
" \"Eminem is a renowned rapper and one of the best-selling music artists of all time.\",\n",
" \"Taylor Swift is a Grammy-winning singer-songwriter known for her narrative-driven music.\",\n",
" \"Sam Altman leads OpenAI as its CEO, with astonishing works of GPT series and pursuing safe and beneficial AI.\",\n",
" \"Morgan Freeman is an acclaimed actor famous for his distinctive voice and diverse roles.\",\n",
" \"Andrew Ng spread AI knowledge globally via public courses on Coursera and Stanford University.\",\n",
" \"Robert Downey Jr. is an iconic actor best known for playing Iron Man in the Marvel Cinematic Universe.\",\n",
"]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"And a few queries (add your own queries and check the result!): "
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"queries = [\n",
" \"Who is Robert Downey Jr.?\",\n",
" \"An expert of neural network\",\n",
" \"A famous female singer\",\n",
"]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 2: Text Embedding"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here, for the sake of speed, we just embed the first 500 docs in the corpus."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"shape of the corpus embeddings: (10, 768)\n",
"data type of the embeddings: float32\n"
]
}
],
"source": [
"from FlagEmbedding import FlagModel\n",
"\n",
"# get the BGE embedding model\n",
"model = FlagModel('BAAI/bge-base-en-v1.5',\n",
" query_instruction_for_retrieval=\"Represent this sentence for searching relevant passages:\",\n",
" use_fp16=True)\n",
"\n",
"# get the embedding of the corpus\n",
"corpus_embeddings = model.encode(corpus)\n",
"\n",
"print(\"shape of the corpus embeddings:\", corpus_embeddings.shape)\n",
"print(\"data type of the embeddings: \", corpus_embeddings.dtype)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Faiss only accepts float32 inputs.\n",
"\n",
"So make sure the dtype of corpus_embeddings is float32 before adding them to the index."
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"corpus_embeddings = corpus_embeddings.astype(np.float32)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 3: Indexing"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this step, we build an index and add the embedding vectors to it."
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
"import faiss\n",
"\n",
"# get the length of our embedding vectors, vectors by bge-base-en-v1.5 have length 768\n",
"dim = corpus_embeddings.shape[-1]\n",
"\n",
"# create the faiss index and store the corpus embeddings into the vector space\n",
"index = faiss.index_factory(dim, 'Flat', faiss.METRIC_INNER_PRODUCT)\n",
"\n",
"# if you installed faiss-gpu, uncomment the following lines to make the index on your GPUs.\n",
"\n",
"# co = faiss.GpuMultipleClonerOptions()\n",
"# index = faiss.index_cpu_to_all_gpus(index, co)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"No need to train if we use \"Flat\" quantizer and METRIC_INNER_PRODUCT as metric. Some other indices that using quantization might need training."
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"True\n",
"total number of vectors: 10\n"
]
}
],
"source": [
"# check if the index is trained\n",
"print(index.is_trained) \n",
"# index.train(corpus_embeddings)\n",
"\n",
"# add all the vectors to the index\n",
"index.add(corpus_embeddings)\n",
"\n",
"print(f\"total number of vectors: {index.ntotal}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Step 3.5 (Optional): Saving Faiss index"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Once you have your index with the embedding vectors, you can save it locally for future usage."
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [],
"source": [
"# change the path to where you want to save the index\n",
"path = \"./index.bin\"\n",
"faiss.write_index(index, path)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If you already have stored index in your local directory, you can load it by:"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [],
"source": [
"index = faiss.read_index(\"./index.bin\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 4: Find answers to the query"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"First, get the embeddings of all the queries:"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [],
"source": [
"query_embeddings = model.encode_queries(queries)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then, use the Faiss index to do a knn search in the vector space:"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[[0.6686779 0.37858668 0.3767978 ]\n",
" [0.6062041 0.59364545 0.527691 ]\n",
" [0.5409331 0.5097007 0.42427146]]\n",
"[[9 7 2]\n",
" [3 1 8]\n",
" [5 0 4]]\n"
]
}
],
"source": [
"dists, ids = index.search(query_embeddings, k=3)\n",
"print(dists)\n",
"print(ids)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now let's see the result:"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"query:\tWho is Robert Downey Jr.?\n",
"answer:\tRobert Downey Jr. is an iconic actor best known for playing Iron Man in the Marvel Cinematic Universe.\n",
"\n",
"query:\tAn expert of neural network\n",
"answer:\tGeoffrey Hinton, as a foundational figure in AI, received Turing Award for his contribution in deep learning.\n",
"\n",
"query:\tA famous female singer\n",
"answer:\tTaylor Swift is a Grammy-winning singer-songwriter known for her narrative-driven music.\n",
"\n"
]
}
],
"source": [
"for i, q in enumerate(queries):\n",
" print(f\"query:\\t{q}\\nanswer:\\t{corpus[ids[i][0]]}\\n\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "base",
"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": 2
}
docs/source/tutorial/3_Indexing/3.1.2.ipynb
+373
View File
@@ -0,0 +1,373 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Faiss GPU"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In the last tutorial, we went through the basics of indexing using faiss-cpu. While for the use cases in research and industry. The size of dataset for indexing will be extremely large, the frequency of searching might also be very high. In this tutorial we'll see how to combine Faiss and GPU almost seamlessly."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Installation"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Faiss maintain the latest updates on conda. And its gpu version only supports Linux x86_64\n",
"\n",
"create a conda virtual environment and run:\n",
"\n",
"```conda install -c pytorch -c nvidia faiss-gpu=1.8.0```\n",
"\n",
"make sure you select that conda env as the kernel for this notebook. After installation, restart the kernal."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If your system does not satisfy the requirement, install faiss-cpu and just skip the steps with gpu related codes."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Data Preparation"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"First let's create two datasets with \"fake embeddings\" of corpus and queries:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import faiss\n",
"import numpy as np\n",
"\n",
"dim = 768\n",
"corpus_size = 1000\n",
"# np.random.seed(111)\n",
"\n",
"corpus = np.random.random((corpus_size, dim)).astype('float32')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. Create Index on CPU"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Option 1:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Faiss provides a great amount of choices of indexes by initializing directly:"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# first build a flat index (on CPU)\n",
"index = faiss.IndexFlatIP(dim)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Option 2:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Besides the basic index class, we can also use the index_factory function to produce composite Faiss index."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"index = faiss.index_factory(dim, \"Flat\", faiss.METRIC_L2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. Build GPU Index and Search"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"All the GPU indexes are built with `StandardGpuResources` object. It contains all the needed resources for each GPU in use. By default it will allocate 18% of the total VRAM as a temporary scratch space."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `GpuClonerOptions` and `GpuMultipleClonerOptions` objects are optional when creating index from cpu to gpu. They are used to adjust the way the GPUs stores the objects."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Single GPU:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"# use a single GPU\n",
"rs = faiss.StandardGpuResources()\n",
"co = faiss.GpuClonerOptions()\n",
"\n",
"# then make it to gpu index\n",
"index_gpu = faiss.index_cpu_to_gpu(provider=rs, device=0, index=index, options=co)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 5.31 ms, sys: 6.26 ms, total: 11.6 ms\n",
"Wall time: 8.94 ms\n"
]
}
],
"source": [
"%%time\n",
"index_gpu.add(corpus)\n",
"D, I = index_gpu.search(corpus, 4)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### All Available GPUs"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If your system contains multiple GPUs, Faiss provides the option to deploy al available GPUs. You can control their usages through `GpuMultipleClonerOptions`, e.g. whether to shard or replicate the index acrross GPUs."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"# cloner options for multiple GPUs\n",
"co = faiss.GpuMultipleClonerOptions()\n",
"\n",
"index_gpu = faiss.index_cpu_to_all_gpus(index=index, co=co)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 29.8 ms, sys: 26.8 ms, total: 56.6 ms\n",
"Wall time: 33.9 ms\n"
]
}
],
"source": [
"%%time\n",
"index_gpu.add(corpus)\n",
"D, I = index_gpu.search(corpus, 4)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Multiple GPUs"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"There's also option that use multiple GPUs but not all:"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"ngpu = 4\n",
"resources = [faiss.StandardGpuResources() for _ in range(ngpu)]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Create vectors for the GpuResources and divices, then pass them to the index_cpu_to_gpu_multiple() function."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
"vres = faiss.GpuResourcesVector()\n",
"vdev = faiss.Int32Vector()\n",
"for i, res in zip(range(ngpu), resources):\n",
" vdev.push_back(i)\n",
" vres.push_back(res)\n",
"index_gpu = faiss.index_cpu_to_gpu_multiple(vres, vdev, index)"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 3.49 ms, sys: 13.4 ms, total: 16.9 ms\n",
"Wall time: 9.03 ms\n"
]
}
],
"source": [
"%%time\n",
"index_gpu.add(corpus)\n",
"D, I = index_gpu.search(corpus, 4)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. Results"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"All the three approaches should lead to identical result. Now let's do a quick sanity check:"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"# The nearest neighbor of each vector in the corpus is itself\n",
"assert np.all(corpus[:] == corpus[I[:, 0]])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"And the corresponding distance should be 0."
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[[ 0. 111.30057 113.2251 113.342316]\n",
" [ 0. 111.158875 111.742325 112.09038 ]\n",
" [ 0. 116.44429 116.849915 117.30502 ]]\n"
]
}
],
"source": [
"print(D[:3])"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "faiss",
"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.11.9"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
docs/source/tutorial/3_Indexing/3.1.3.ipynb
+417
View File
@@ -0,0 +1,417 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Faiss Indexes"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This tutorial will go through several widely used indexes in Faiss that fits different requirements, and how to use them."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Preparation"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For CPU usage, use:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"%pip install faiss-cpu"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For GPU on Linux x86_64 system, use Conda:\n",
"\n",
"```conda install -c pytorch -c nvidia faiss-gpu=1.8.0```"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import faiss\n",
"import numpy as np\n",
"\n",
"np.random.seed(768)\n",
"\n",
"data = np.random.random((1000, 128))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. `IndexFlat*`"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Flat index is the very fundamental index structure. It does not do any preprocess for the incoming vectors. All the vectors are stored directly without compression or quantization. Thus no training is need for flat indexes.\n",
"\n",
"When searching, Flat index will decode all the vectors sequentially and compute the similarity score to the query vectors. Thus, Flat Index guarantees the global optimum of results."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Flat index family is small: just `IndexFlatL2` and `IndexFlatIP`, which are just different by the similarity metrics of Euclidean distance and inner product."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Usage:"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"d = 128 # dimension of the vector\n",
"k = 3 # number of nearest neighbors to search\n",
"\n",
"# just simply create the index and add all the data\n",
"index = faiss.IndexFlatL2(d)\n",
"index.add(data)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Sanity check:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"closest elements: [[ 0 471 188]]\n",
"distance: [[ 0. 16.257435 16.658928]]\n"
]
}
],
"source": [
"# search for the k nearest neighbor for the first element in data\n",
"D, I = index.search(data[:1], k)\n",
"\n",
"print(f\"closest elements: {I}\")\n",
"print(f\"distance: {D}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Flat Indexes guarantee the perfect quality but with terrible speed. It works well on small datasets or the cases that speed is not a crucial factor. \n",
"\n",
"But what about the cases that speed is important? There's no way to have it all. So we want some indexes that only sacrifice as small as possible quality to speed up. That's why approximate nearest-neighbors (ANN) algorithms are widely accepted. Now we will go through a few popular ANN methods used in vector searching."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. `IndexIVF*`"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Intro\n",
"\n",
"Inverted File Flat (IVF) Index is a widely accepted technique to speed up searching by using k-means or Voronoi diagram to create a number of cells (or say, clusters) in the whole space. Then when given a query, an amount of closest cells will be searched. After that, `k` closest elements to the query will be searched in those cells.\n",
"\n",
"- `quantizer` is another index/quantizer to assign vectors to inverted lists.\n",
"- `nlist` is the number of cells the space to be partitioned.\n",
"- `nprob` is the nuber of closest cells to visit for searching in query time."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Tradeoff\n",
"\n",
"Increasing `nlist` will shrink the size of each cell, which speed up the search process. But the smaller coverage will sacrifice accuracy and increase the possibility of the edge/surface problem discribed above.\n",
"\n",
"Increasing `nprob` will have a greater scope, preferring search quality by the tradeoff of slower speed."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Shortage\n",
"\n",
"There could be a problem when the query vector lands on the edge/surface of the cell. It is possible that the closest element falls into the neighbor cell, which may not be considered due to `nprob` is not large enough."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Example"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"nlist = 5\n",
"nprob = 2\n",
"\n",
"# the quantizer defines how to store and compare the vectors\n",
"quantizer = faiss.IndexFlatL2(d)\n",
"index = faiss.IndexIVFFlat(quantizer, d, nlist)\n",
"\n",
"# note different from flat index, IVF index first needs training to create the cells\n",
"index.train(data)\n",
"index.add(data)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"closest elements: [[ 0 471 188]]\n",
"distance: [[ 0. 16.257435 16.658928]]\n"
]
}
],
"source": [
"# set nprob before searching\n",
"index.nprobe = 8\n",
"D, I = index.search(data[:1], k)\n",
"\n",
"print(f\"closest elements: {I}\")\n",
"print(f\"distance: {D}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. `IndexHNSW*`"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Intro\n",
"\n",
"Hierarchical Navigable Small World (HNSW) indexing is a graph based method, which is an extension of navigable small world (NSW). It builds a multi-layered graph where nodes (vectors) are connected based on their proximity, forming \"small-world\" structures that allow efficient navigation through the space.\n",
"\n",
"- `M` is the number of neighbors each vector has in the graph.\n",
"- `efConstruction` is the number of entry points to explore when building the index.\n",
"- `efSearch` is the number of entry points to explore when searching."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Tradeoff\n",
"\n",
"Increasing `M` or `efSearch` will make greater fidelity with reasonable longer time. Larger `efConstruction` mainly increases the index construction time.\n",
"\n",
"HNSW has great searching quality and speed. But it is memory-consuming due to the graph structure. Scaling up `M` will cause a linear increase of memory usage.\n",
"\n",
"Note that HNSW index does not support vector's removal because removing nodes will distroy graph structure.\n",
"\n",
"Thus HNSW is a great index to choose when RAM is not a limiting factor."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Example"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"M = 32\n",
"ef_search = 16\n",
"ef_construction = 32\n",
"\n",
"index = faiss.IndexHNSWFlat(d, M)\n",
"# set the two parameters before adding data\n",
"index.hnsw.efConstruction = ef_construction\n",
"index.hnsw.efSearch = ef_search\n",
"\n",
"index.add(data)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"closest elements: [[ 0 471 188]]\n",
"distance: [[ 0. 16.257435 16.658928]]\n"
]
}
],
"source": [
"D, I = index.search(data[:1], k)\n",
"\n",
"print(f\"closest elements: {I}\")\n",
"print(f\"distance: {D}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 4. `IndexLSH`"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Intro\n",
"\n",
"Locality Sensitive Hashing (LSH) is an ANN method that hashing data points into buckets. While well known use cases of hash function such as dictionary/hashtabel are trying to avoid hashing collisions, LSH trys to maximize hashing collisions. Similar vectors will be grouped into same hash bucket.\n",
"\n",
"In Faiss, `IndexLSH` is a Flat index with binary codes. Vectors are hashed into binary codes and compared by Hamming distances.\n",
"\n",
"- `nbits` can be seen as the \"resolution\" of hashed vectors."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Tradeoff\n",
"\n",
"Increasing `nbits` can get higher fidelity with the cost of more memory and longer searching time.\n",
"\n",
"LSH suffers the curse of dimensionality when using a larger `d`. In order to get similar search quality, the `nbits` value needs to be scaled up to maintain the search quality."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Shortage\n",
"\n",
"LSH speeds up searching time with a reasonable sacrifice of quality. But that only applies to small dimension `d`. Even 128 is already too large for LSH. Thus for vectors generated by transformer based embedding models, LSH index is not a common choice."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Example"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"nbits = d * 8\n",
"\n",
"index = faiss.IndexLSH(d, nbits)\n",
"index.train(data)\n",
"index.add(data)"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"closest elements: [[ 0 471 392]]\n",
"distance: [[ 0. 197. 199.]]\n"
]
}
],
"source": [
"D, I = index.search(data[:1], k)\n",
"\n",
"print(f\"closest elements: {I}\")\n",
"print(f\"distance: {D}\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "faiss",
"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.11.9"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
docs/source/tutorial/3_Indexing/3.1.4.ipynb
+354
View File
@@ -0,0 +1,354 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Faiss Quantizers"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this notebook, we will introduce the quantizer object in Faiss and how to use them."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Preparation"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For CPU usage, run:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%pip install faiss-cpu"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For GPU on Linux x86_64 system, use Conda:\n",
"\n",
"```conda install -c pytorch -c nvidia faiss-gpu=1.8.0```"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import faiss\n",
"import numpy as np\n",
"\n",
"np.random.seed(768)\n",
"\n",
"data = np.random.random((1000, 128))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Scalar Quantizer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Normal data type of vector embeedings is usually 32 bit floats. Scalar quantization is transforming the 32 float representation to, for example, 8 bit interger. Thus with a 4x reduction in size. In this way, it can be seen as we distribute each dimension into 256 buckets."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"| Name | Class | Parameters |\n",
"|:------------:|:--------:|:-----------|\n",
"| `ScalarQuantizer` | Quantizer class | `d`: dimension of vectors<br>`qtype`: map dimension into $2^\\text{qtype}$ clusters |\n",
"| `IndexScalarQuantizer` | Flat index class | `d`: dimension of vectors<br>`qtype`: map dimension into $2^\\text{qtype}$ clusters<br>`metric`: similarity metric (L2 or IP) |\n",
"| `IndexIVFScalarQuantizer` | IVF index class | `d`: dimension of vectors<br>`nlist`: number of cells/clusters to partition the inverted file space<br>`qtype`: map dimension into $2^\\text{qtype}$ clusters<br>`metric`: similarity metric (L2 or IP)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Quantizer class objects are used to compress the data before adding into indexes. Flat index class objects and IVF index class objects can be used direct as and index. Quantization will be done automatically."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Scalar Quantizer"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[156 180 46 226 13 130 41 187 63 251 16 199 205 166 117 122 214 2\n",
" 206 137 71 186 20 131 59 57 68 114 35 45 28 210 27 93 74 245\n",
" 167 5 32 42 44 128 10 189 10 13 42 162 179 221 241 104 205 21\n",
" 70 87 52 219 172 138 193 0 228 175 144 34 59 88 170 1 233 220\n",
" 20 64 245 241 5 161 41 55 30 247 107 8 229 90 201 10 43 158\n",
" 238 184 187 114 232 90 116 205 14 214 135 158 237 192 205 141 232 176\n",
" 124 176 163 68 49 91 125 70 6 170 55 44 215 84 46 48 218 56\n",
" 107 176]\n"
]
}
],
"source": [
"d = 128\n",
"qtype = faiss.ScalarQuantizer.QT_8bit\n",
"\n",
"quantizer = faiss.ScalarQuantizer(d, qtype)\n",
"\n",
"quantizer.train(data)\n",
"new_data = quantizer.compute_codes(data)\n",
"\n",
"print(new_data[0])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Scalar Quantizer Index"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"d = 128\n",
"k = 3\n",
"qtype = faiss.ScalarQuantizer.QT_8bit\n",
"# nlist = 5\n",
"\n",
"index = faiss.IndexScalarQuantizer(d, qtype, faiss.METRIC_L2)\n",
"# index = faiss.IndexIVFScalarQuantizer(d, nlist, faiss.ScalarQuantizer.QT_8bit, faiss.METRIC_L2)\n",
"\n",
"index.train(data)\n",
"index.add(data)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"closest elements: [[ 0 471 188]]\n",
"distance: [[1.6511828e-04 1.6252808e+01 1.6658131e+01]]\n"
]
}
],
"source": [
"D, I = index.search(data[:1], k)\n",
"\n",
"print(f\"closest elements: {I}\")\n",
"print(f\"distance: {D}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Product Quantizer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When speed and memory are crucial factors in searching, product quantizer becomes a top choice. It is one of the effective quantizer on reducing memory size. \n",
"\n",
"The first step of PQ is dividing the original vectors with dimension `d` into smaller, low-dimensional sub-vectors with dimension `d/m`. Here `m` is the number of sub-vectors.\n",
"\n",
"Then clustering algorithms are used to create codebook of a fixed number of centroids.\n",
"\n",
"Next, each sub-vector of a vector is replaced by the index of the closest centroid from its corresponding codebook. Now each vector will be stored with only the indices instead of the full vector.\n",
"\n",
"When comuputing the distance between a query vector. Only the distances to the centroids in the codebooks are calculated, thus enable the quick approximate nearest neighbor searches."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"| Name | Class | Parameters |\n",
"|:------------:|:--------:|:-----------|\n",
"| `ProductQuantizer` | Quantizer class | `d`: dimension of vectors<br>`M`: number of sub-vectors that D % M == 0<br>`nbits`: number of bits per subquantizer, so each contain $2^\\text{nbits}$ centroids |\n",
"| `IndexPQ` | Flat index class | `d`: dimension of vectors<br>`M`: number of sub-vectors that D % M == 0<br>`nbits`: number of bits per subquantizer, so each contain $2^\\text{nbits}$ centroids<br>`metric`: similarity metric (L2 or IP) |\n",
"| `IndexIVFPQ` | IVF index class | `quantizer`: the quantizer used in computing distance phase.<br>`d`: dimension of vectors<br>`nlist`: number of cells/clusters to partition the inverted file space<br>`M`: number of sub-vectors that D % M == 0<br>`nbits`: number of bits per subquantizer, so each contain $2^\\text{nbits}$ centroids<br>`metric`: similarity metric (L2 or IP) |"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Product Quantizer"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"255\n",
"[[ 90 169 226 45]\n",
" [ 33 51 34 15]]\n"
]
}
],
"source": [
"d = 128\n",
"M = 8\n",
"nbits = 4\n",
"\n",
"quantizer = faiss.ProductQuantizer(d, M, nbits)\n",
"\n",
"quantizer.train(data)\n",
"new_data = quantizer.compute_codes(data)\n",
"\n",
"print(new_data.max())\n",
"print(new_data[:2])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Product Quantizer Index"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"index = faiss.IndexPQ(d, M, nbits, faiss.METRIC_L2)\n",
"\n",
"index.train(data)\n",
"index.add(data)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"closest elements: [[ 0 946 330]]\n",
"distance: [[ 8.823908 11.602461 11.746731]]\n"
]
}
],
"source": [
"D, I = index.search(data[:1], k)\n",
"\n",
"print(f\"closest elements: {I}\")\n",
"print(f\"distance: {D}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Product Quantizer IVF Index"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"nlist = 5\n",
"\n",
"quantizer = faiss.IndexFlat(d, faiss.METRIC_L2)\n",
"index = faiss.IndexIVFPQ(quantizer, d, nlist, M, nbits, faiss.METRIC_L2)\n",
"\n",
"index.train(data)\n",
"index.add(data)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"closest elements: [[ 0 899 521]]\n",
"distance: [[ 8.911423 12.088312 12.104569]]\n"
]
}
],
"source": [
"D, I = index.search(data[:1], k)\n",
"\n",
"print(f\"closest elements: {I}\")\n",
"print(f\"distance: {D}\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "base",
"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": 2
}
docs/source/tutorial/3_Indexing/3.1.5.ipynb
+624
View File
@@ -0,0 +1,624 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Choosing Index"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Give a great amount of indexes and quantizers, how to choose the one in the experiment/application? In this part, we will give a general suggestion on how to choose the one fits your need."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 0. Preparation"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Packages"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For CPU usage, run:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# %pip install -U faiss-cpu numpy h5py"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For GPU on Linux x86_64 system, use Conda:\n",
"\n",
"```conda install -c pytorch -c nvidia faiss-gpu=1.8.0```"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"from urllib.request import urlretrieve\n",
"import h5py\n",
"import faiss\n",
"import numpy as np"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Dataset"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this tutorial, we'll use [SIFT1M](http://corpus-texmex.irisa.fr/), a very popular dataset for ANN evaluation, as our dataset to demonstrate the comparison.\n",
"\n",
"Run the following cell to download the dataset or you can also manually download from the repo [ann-benchmarks](https://github.com/erikbern/ann-benchmarks?tab=readme-ov-file#data-sets))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data_url = \"http://ann-benchmarks.com/sift-128-euclidean.hdf5\"\n",
"destination = \"data.hdf5\"\n",
"urlretrieve(data_url, destination)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then load the data from the hdf5 file."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"(1000000, 128) float32\n",
"(10000, 128) float32\n"
]
}
],
"source": [
"with h5py.File('data.hdf5', 'r') as f:\n",
" corpus = f['train'][:]\n",
" query = f['test'][:]\n",
"\n",
"print(corpus.shape, corpus.dtype)\n",
"print(query.shape, corpus.dtype)"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"d = corpus[0].shape[0]\n",
"k = 100"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Helper function"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The following is a helper function for computing recall."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"# compute recall from the prediction results and ground truth\n",
"def compute_recall(res, truth):\n",
" recall = 0\n",
" for i in range(len(res)):\n",
" intersect = np.intersect1d(res[i], truth[i])\n",
" recall += len(intersect) / len(res[i])\n",
" recall /= len(res)\n",
"\n",
" return recall"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Flat Index"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Flat index use brute force to search neighbors for each query. It guarantees the optimal result with 100% recall. Thus we use the result from it as the ground truth."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 69.2 ms, sys: 80.6 ms, total: 150 ms\n",
"Wall time: 149 ms\n"
]
}
],
"source": [
"%%time\n",
"index = faiss.IndexFlatL2(d)\n",
"index.add(corpus)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 17min 30s, sys: 1.62 s, total: 17min 31s\n",
"Wall time: 2min 1s\n"
]
}
],
"source": [
"%%time\n",
"D, I_truth = index.search(query, k)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. IVF Index"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 10.6 s, sys: 831 ms, total: 11.4 s\n",
"Wall time: 419 ms\n"
]
}
],
"source": [
"%%time\n",
"nlist = 5\n",
"nprob = 3\n",
"\n",
"quantizer = faiss.IndexFlatL2(d)\n",
"index = faiss.IndexIVFFlat(quantizer, d, nlist)\n",
"index.nprobe = nprob\n",
"\n",
"index.train(corpus)\n",
"index.add(corpus)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 9min 15s, sys: 598 ms, total: 9min 16s\n",
"Wall time: 12.5 s\n"
]
}
],
"source": [
"%%time\n",
"D, I = index.search(query, k)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Recall: 0.9999189999999997\n"
]
}
],
"source": [
"recall = compute_recall(I, I_truth)\n",
"print(f\"Recall: {recall}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"From the test we can see that IVFFlatL2 has a pretty good promotion for the searching speed with a very tiny loss of recall."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. HNSW Index"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 11min 21s, sys: 595 ms, total: 11min 22s\n",
"Wall time: 17 s\n"
]
}
],
"source": [
"%%time\n",
"M = 64\n",
"ef_search = 32\n",
"ef_construction = 64\n",
"\n",
"index = faiss.IndexHNSWFlat(d, M)\n",
"# set the two parameters before adding data\n",
"index.hnsw.efConstruction = ef_construction\n",
"index.hnsw.efSearch = ef_search\n",
"\n",
"index.add(corpus)"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 5.14 s, sys: 3.94 ms, total: 5.14 s\n",
"Wall time: 110 ms\n"
]
}
],
"source": [
"%%time\n",
"D, I = index.search(query, k)"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Recall: 0.8963409999999716\n"
]
}
],
"source": [
"recall = compute_recall(I, I_truth)\n",
"print(f\"Recall: {recall}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"From the searching time of less than 1 second, we can see why HNSW is one of the best choice when looking for an extreme speed during searching phase. The reduction of recall is acceptable. But the longer time during creation of index and large memory footprint need to be considered."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. LSH"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 13.7 s, sys: 660 ms, total: 14.4 s\n",
"Wall time: 12.1 s\n"
]
}
],
"source": [
"%%time\n",
"nbits = d * 8\n",
"\n",
"index = faiss.IndexLSH(d, nbits)\n",
"index.train(corpus)\n",
"index.add(corpus)"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 3min 20s, sys: 84.2 ms, total: 3min 20s\n",
"Wall time: 5.64 s\n"
]
}
],
"source": [
"%%time\n",
"D, I = index.search(query, k)"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Recall: 0.5856720000000037\n"
]
}
],
"source": [
"recall = compute_recall(I, I_truth)\n",
"print(f\"Recall: {recall}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As we covered in the last notebook, LSH is not a good choice when the data dimension is large. Here 128 is already burdened for LSH. As we can see, even we choose a relatively small `nbits` of d * 8, the index creating time and search time are still pretty long. And the recall of about 58.6% is not satisfactory."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. Scalar Quantizer Index"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 550 ms, sys: 18 ms, total: 568 ms\n",
"Wall time: 87.4 ms\n"
]
}
],
"source": [
"%%time\n",
"qtype = faiss.ScalarQuantizer.QT_8bit\n",
"metric = faiss.METRIC_L2\n",
"\n",
"index = faiss.IndexScalarQuantizer(d, qtype, metric)\n",
"index.train(corpus)\n",
"index.add(corpus)"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 7min 36s, sys: 169 ms, total: 7min 36s\n",
"Wall time: 12.7 s\n"
]
}
],
"source": [
"%%time\n",
"D, I = index.search(query, k)"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Recall: 0.990444999999872\n"
]
}
],
"source": [
"recall = compute_recall(I, I_truth)\n",
"print(f\"Recall: {recall}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here scalar quantizer index's performance looks very similar to the Flat index. Because the elements of vectors in the SIFT dataset are integers in the range of [0, 218]. Thus the index does not lose to much information during scalar quantization. For the dataset with more complex distribution in float32. The difference will be more obvious."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 6. Product Quantizer Index"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 46.7 s, sys: 22.3 ms, total: 46.7 s\n",
"Wall time: 1.36 s\n"
]
}
],
"source": [
"%%time\n",
"M = 16\n",
"nbits = 8\n",
"metric = faiss.METRIC_L2\n",
"\n",
"index = faiss.IndexPQ(d, M, nbits, metric)\n",
"\n",
"index.train(corpus)\n",
"index.add(corpus)"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 1min 37s, sys: 106 ms, total: 1min 37s\n",
"Wall time: 2.8 s\n"
]
}
],
"source": [
"%%time\n",
"D, I = index.search(query, k)"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Recall: 0.630898999999999\n"
]
}
],
"source": [
"recall = compute_recall(I, I_truth)\n",
"print(f\"Recall: {recall}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Product quantizer index is not standout in any one of the aspect. But it somewhat balance the tradeoffs. It is widely used in real applications with the combination of other indexes such as IVF or HNSW."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "base",
"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": 2
}