Fix demo rescoring; add approximate index support; add sliders

app.py

@@ -1,4 +1,3 @@
-
 import time
 import gradio as gr
 from datasets import load_dataset
@@ -14,6 +13,7 @@ title_text_dataset = load_dataset("mixedbread-ai/wikipedia-data-en-2023-11", spl
 # Load the int8 and binary indices. Int8 is loaded as a view to save memory, as we never actually perform search with it.
 int8_view = Index.restore("wikipedia_int8_usearch_50m.index", view=True)
 binary_index: faiss.IndexBinaryFlat = faiss.read_index_binary("wikipedia_ubinary_faiss_50m.index")
+binary_ivf: faiss.IndexBinaryIVF = faiss.read_index_binary("wikipedia_ubinary_ivf_faiss_50m.index")
 
 # Load the SentenceTransformer model for embedding the queries
 model = SentenceTransformer(
@@ -25,7 +25,7 @@ model = SentenceTransformer(
 )
 
 
-def search(query, top_k: int = 10, rescore_multiplier: int = 4):
+def search(query, top_k: int = 100, rescore_multiplier: int = 1, use_approx: bool = False):
     # 1. Embed the query as float32
     start_time = time.time()
     query_embedding = model.encode(query)
@@ -36,9 +36,10 @@ def search(query, top_k: int = 10, rescore_multiplier: int = 4):
     query_embedding_ubinary = quantize_embeddings(query_embedding.reshape(1, -1), "ubinary")
     quantize_time = time.time() - start_time
 
-    # 3. Search the binary index
+    # 3. Search the binary index (either exact or approximate)
+    index = binary_ivf if use_approx else binary_index
     start_time = time.time()
-    _scores, binary_ids = binary_index.search(query_embedding_ubinary, top_k * rescore_multiplier)
+    _scores, binary_ids = index.search(query_embedding_ubinary, top_k * rescore_multiplier)
     binary_ids = binary_ids[0]
     search_time = time.time() - start_time
 
@@ -54,11 +55,15 @@ def search(query, top_k: int = 10, rescore_multiplier: int = 4):
 
     # 6. Sort the scores and return the top_k
     start_time = time.time()
-    indices = scores.argsort()[:top_k]
+    indices = scores.argsort()[::-1][:top_k]
     top_k_indices = binary_ids[indices]
     top_k_scores = scores[indices]
-    top_k_titles, top_k_texts = zip(*[(title_text_dataset[idx]["title"], title_text_dataset[idx]["text"]) for idx in top_k_indices.tolist()])
-    df = pd.DataFrame({"Score": [round(value, 2) for value in top_k_scores], "Title": top_k_titles, "Text": top_k_texts})
+    top_k_titles, top_k_texts = zip(
+        *[(title_text_dataset[idx]["title"], title_text_dataset[idx]["text"]) for idx in top_k_indices.tolist()]
+    )
+    df = pd.DataFrame(
+        {"Score": [round(value, 2) for value in top_k_scores], "Title": top_k_titles, "Text": top_k_texts}
+    )
     sort_time = time.time() - start_time
 
     return df, {
@@ -68,14 +73,15 @@ def search(query, top_k: int = 10, rescore_multiplier: int = 4):
         "Load Time": f"{load_time:.4f} s",
         "Rescore Time": f"{rescore_time:.4f} s",
         "Sort Time": f"{sort_time:.4f} s",
-        "Total Retrieval Time": f"{quantize_time + search_time + load_time + rescore_time + sort_time:.4f} s"
+        "Total Retrieval Time": f"{quantize_time + search_time + load_time + rescore_time + sort_time:.4f} s",
     }
 
+
 with gr.Blocks(title="Quantized Retrieval") as demo:
     gr.Markdown(
-"""
+        """
 ## Quantized Retrieval - Binary Search with Scalar (int8) Rescoring
-This demo showcases exact retrieval using [quantized embeddings](https://huggingface.co/blog/embedding-quantization). The corpus consists of 41 million texts from Wikipedia articles.
+This demo showcases retrieval using [quantized embeddings](https://huggingface.co/blog/embedding-quantization). The corpus consists of 41 million texts from Wikipedia articles.
 
 <details><summary>Click to learn about the retrieval process</summary>
 
@@ -94,11 +100,47 @@ we need `1024 / 8 * num_docs` bytes for the binary index and `1024 * num_docs` b
 This is notably cheaper than doing the same process with float32 embeddings, which would require `4 * 1024 * num_docs` bytes of memory/disk space for the float32 index, i.e. 32x as much memory and 4x as much disk space.
 Additionally, the binary index is much faster (up to 32x) to search than the float32 index, while the rescoring is also extremely efficient. In conclusion, this process allows for fast, scalable, cheap, and memory-efficient retrieval.
 
-Feel free to check out the [code for this demo](https://huggingface.co/spaces/tomaarsen/quantized_retrieval/blob/main/app.py) to learn more about how to apply this in practice.
+Feel free to check out the [code for this demo](https://huggingface.co/spaces/sentence-transformers/quantized-retrieval/blob/main/app.py) to learn more about how to apply this in practice.
+
+Notes:
+- The approximate search index (a binary Inverted File Index (IVF)) is in beta and has not been trained with a lot of data. A better IVF index will be released soon.
 
 </details>
-""")
-    query = gr.Textbox(label="Query for Wikipedia articles", placeholder="Enter a query to search for relevant texts from Wikipedia.")
+"""
+    )
+    with gr.Row():
+        with gr.Column(scale=75):
+            query = gr.Textbox(
+                label="Query for Wikipedia articles",
+                placeholder="Enter a query to search for relevant texts from Wikipedia.",
+            )
+        with gr.Column(scale=25):
+            use_approx = gr.Radio(
+                choices=[("Exact Search", False), ("Approximate Search", True)],
+                value=True,
+                label="Search Index",
+            )
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            top_k = gr.Slider(
+                minimum=10,
+                maximum=1000,
+                step=5,
+                value=100,
+                label="Number of documents to retrieve",
+                info="Number of documents to retrieve from the binary search",
+            )
+        with gr.Column(scale=2):
+            rescore_multiplier = gr.Slider(
+                minimum=1,
+                maximum=10,
+                step=1,
+                value=1,
+                label="Rescore multiplier",
+                info="Search for `rescore_multiplier` as many documents to rescore",
+            )
+
     search_button = gr.Button(value="Search")
 
     with gr.Row():
@@ -107,8 +149,8 @@ Feel free to check out the [code for this demo](https://huggingface.co/spaces/to
         with gr.Column(scale=1):
             json = gr.JSON()
 
-    query.submit(search, inputs=[query], outputs=[output, json])
-    search_button.click(search, inputs=[query], outputs=[output, json])
+    query.submit(search, inputs=[query, top_k, rescore_multiplier, use_approx], outputs=[output, json])
+    search_button.click(search, inputs=[query, top_k, rescore_multiplier, use_approx], outputs=[output, json])
 
 demo.queue()
-demo.launch(debug=True)
+demo.launch()