Use downloaded int8 embeddings instead of an index in this repository

app.py

@@ -1,52 +1,71 @@
 import time
 import gradio as gr
-from datasets import load_dataset
+from datasets import load_dataset, load_from_disk
 import pandas as pd
 from sentence_transformers import SentenceTransformer
 from sentence_transformers.quantization import quantize_embeddings
 import faiss
-from usearch.index import Index
+import numpy as np
 
-# Load titles and texts
-title_text_dataset = load_dataset("mixedbread-ai/wikipedia-data-en-2023-11", split="train", num_proc=4).select_columns(["title", "text"])
+# Load titles, texts, and int8 embeddings in a lazy Dataset, allowing us to efficiently access specific rows on demand
+# Note that we never actually use the int8 embeddings for search directly, they are only used for rescoring after the binary search
+title_text_int8_dataset = load_dataset("sentence-transformers/wikipedia-mxbai-embed-int8-index", split="train").select_columns(["title", "text", "embedding"])
+# title_text_int8_dataset = load_from_disk("wikipedia-mxbai-embed-int8-index").select_columns(["url", "title", "text", "embedding"])
 
-# 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)
+# Load the binary indices
 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")
+binary_ivf_index: faiss.IndexBinaryIVF = faiss.read_index_binary("wikipedia_ubinary_ivf_faiss_50m.index")
 
 # Load the SentenceTransformer model for embedding the queries
-model = SentenceTransformer(
-    "mixedbread-ai/mxbai-embed-large-v1",
-    prompts={
-        "retrieval": "Represent this sentence for searching relevant passages: ",
-    },
-    default_prompt_name="retrieval",
-)
-
-
-def search(query, top_k: int = 100, rescore_multiplier: int = 1, use_approx: bool = False):
+model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
+if model.device.type == "cuda":
+    model.bfloat16()
+
+warmup_queries = [
+    "What is the capital of France?",
+    "Who is the president of the United States?",
+    "What is the largest mammal?",
+    "How to bake a chocolate cake?",
+    "What is the theory of relativity?",
+]
+model.encode(warmup_queries)
+
+
+def search(
+    query,
+    top_k: int = 20,
+    rescore_multiplier: int = 4,
+    use_approx: bool = False,
+    display_score: bool = True,
+    display_binary_rank: bool = False,
+):
     # 1. Embed the query as float32
     start_time = time.time()
-    query_embedding = model.encode(query)
+    query_embedding = model.encode_query(query)
     embed_time = time.time() - start_time
 
     # 2. Quantize the query to ubinary
     start_time = time.time()
-    query_embedding_ubinary = quantize_embeddings(query_embedding.reshape(1, -1), "ubinary")
+    query_embedding_ubinary = quantize_embeddings(
+        query_embedding.reshape(1, -1), "ubinary"
+    )
     quantize_time = time.time() - start_time
 
     # 3. Search the binary index (either exact or approximate)
-    index = binary_ivf if use_approx else binary_index
+    index = binary_ivf_index if use_approx else binary_index
     start_time = time.time()
-    _scores, binary_ids = 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
 
     # 4. Load the corresponding int8 embeddings
     start_time = time.time()
-    int8_embeddings = int8_view[binary_ids].astype(int)
-    load_time = time.time() - start_time
+    int8_embeddings = np.array(
+        title_text_int8_dataset[binary_ids]["embedding"], dtype=np.int8
+    )
+    load_int8_time = time.time() - start_time
 
     # 5. Rescore the top_k * rescore_multiplier using the float32 query embedding and the int8 document embeddings
     start_time = time.time()
@@ -58,22 +77,40 @@ def search(query, top_k: int = 100, rescore_multiplier: int = 1, use_approx: boo
     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}
-    )
     sort_time = time.time() - start_time
 
+    # 7. Load titles and texts for the top_k results
+    start_time = time.time()
+    top_k_titles = title_text_int8_dataset[top_k_indices]["title"]
+    top_k_urls = title_text_int8_dataset[top_k_indices]["url"]
+    top_k_texts = title_text_int8_dataset[top_k_indices]["text"]
+    top_k_titles = [f"[{title}]({url})" for title, url in zip(top_k_titles, top_k_urls)]
+    load_text_time = time.time() - start_time
+
+    rank = np.arange(1, top_k + 1)
+    data = {
+        "Score": [f"{score:.2f}" for score in top_k_scores],
+        "#": rank,
+        "Binary #": indices + 1,
+        "Title": top_k_titles,
+        "Text": top_k_texts,
+    }
+    if not display_score:
+        del data["Score"]
+    if not display_binary_rank:
+        del data["Binary #"]
+        del data["#"]
+    df = pd.DataFrame(data)
+
     return df, {
         "Embed Time": f"{embed_time:.4f} s",
         "Quantize Time": f"{quantize_time:.4f} s",
         "Search Time": f"{search_time:.4f} s",
-        "Load Time": f"{load_time:.4f} s",
+        "Load int8 Time": f"{load_int8_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",
+        "Load Text Time": f"{load_text_time:.4f} s",
+        "Total Retrieval Time": f"{quantize_time + search_time + load_int8_time + rescore_time + sort_time + load_text_time:.4f} s",
     }
 
 
@@ -81,17 +118,18 @@ with gr.Blocks(title="Quantized Retrieval") as demo:
     gr.Markdown(
         """
 ## Quantized Retrieval - Binary Search with Scalar (int8) Rescoring
-This demo showcases retrieval using [quantized embeddings](https://huggingface.co/blog/embedding-quantization) on a CPU. The corpus consists of 41 million texts from Wikipedia articles.
+This demo showcases retrieval using [quantized embeddings](https://huggingface.co/blog/embedding-quantization) on a CPU. The corpus consists of [41 million texts](https://huggingface.co/datasets/sentence-transformers/wikipedia-mxbai-embed-int8-index) from Wikipedia articles.
 
 <details><summary>Click to learn about the retrieval process</summary>
 
 Details:
 1. The query is embedded using the [`mixedbread-ai/mxbai-embed-large-v1`](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1) SentenceTransformer model.
 2. The query is quantized to binary using the `quantize_embeddings` function from the SentenceTransformers library.
-3. A binary index (41M binary embeddings; 5.2GB of memory/disk space) is searched using the quantized query for the top 40 documents.
-4. The top 40 documents are loaded on the fly from an int8 index on disk (41M int8 embeddings; 0 bytes of memory, 47.5GB of disk space).
-5. The top 40 documents are rescored using the float32 query and the int8 embeddings to get the top 10 documents.
-6. The top 10 documents are sorted by score and displayed.
+3. A binary index (41M binary embeddings; 5.2GB of memory/disk space) is searched using the quantized query for the top 80 documents.
+4. The top 80 documents are loaded on the fly from an int8 index on disk (41M int8 embeddings; 0 bytes of memory, 47.5GB of disk space).
+5. The top 80 documents are rescored using the float32 query and the int8 embeddings to get the top 20 documents.
+6. The top 20 documents are sorted by score.
+7. The titles and texts of the top 20 documents are loaded on the fly from disk and displayed.
 
 This process is designed to be memory efficient and fast, with the binary index being small enough to fit in memory and the int8 index being loaded as a view to save memory. 
 In total, this process requires keeping 1) the model in memory, 2) the binary index in memory, and 3) the int8 index on disk. With a dimensionality of 1024, 
@@ -103,13 +141,13 @@ Additionally, the binary index is much faster (up to 32x) to search than the flo
 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.
+- The approximate search index (a binary Inverted File Index (IVF)) is in beta and has not been trained with a lot of data.
 
 </details>
 """
     )
     with gr.Row():
-        with gr.Column(scale=75):
+        with gr.Column(scale=60):
             query = gr.Textbox(
                 label="Query for Wikipedia articles",
                 placeholder="Enter a query to search for relevant texts from Wikipedia.",
@@ -118,7 +156,16 @@ Notes:
             use_approx = gr.Radio(
                 choices=[("Exact Search", False), ("Approximate Search", True)],
                 value=True,
-                label="Search Index",
+                label="Search Settings",
+            )
+        with gr.Column(scale=15):
+            display_score = gr.Checkbox(
+                label="Display Score",
+                value=True,
+            )
+            display_binary_rank = gr.Checkbox(
+                label='Display Binary Rank',
+                value=False,
             )
 
     with gr.Row():
@@ -126,8 +173,8 @@ Notes:
             top_k = gr.Slider(
                 minimum=10,
                 maximum=1000,
-                step=5,
-                value=100,
+                step=1,
+                value=20,
                 label="Number of documents to retrieve",
                 info="Number of documents to retrieve from the binary search",
             )
@@ -136,7 +183,7 @@ Notes:
                 minimum=1,
                 maximum=10,
                 step=1,
-                value=1,
+                value=4,
                 label="Rescore multiplier",
                 info="Search for `rescore_multiplier` as many documents to rescore",
             )
@@ -145,12 +192,47 @@ Notes:
 
     with gr.Row():
         with gr.Column(scale=4):
-            output = gr.Dataframe(headers=["Score", "Title", "Text"])
+            output = gr.Dataframe(
+                headers=["Score", "#", "Binary #", "Title", "Text"],
+                datatype="markdown",
+            )
         with gr.Column(scale=1):
             json = gr.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])
+    examples = gr.Examples(
+        examples=[
+            "What is the coldest metal to the touch?",
+            "Who won the FIFA World Cup in 2018?",
+            "How to make a paper airplane?",
+            "Who was the first woman to cross the Pacific ocean by plane?",
+        ],
+        fn=search,
+        inputs=[query],
+        outputs=[output, json],
+        cache_examples=False,
+        run_on_click=True,
+    )
+
+    query.submit(
+        search,
+        inputs=[query, top_k, rescore_multiplier, use_approx, display_score, display_binary_rank],
+        outputs=[output, json],
+    )
+    search_button.click(
+        search,
+        inputs=[query, top_k, rescore_multiplier, use_approx, display_score, display_binary_rank],
+        outputs=[output, json],
+    )
+    display_score.change(
+        search,
+        inputs=[query, top_k, rescore_multiplier, use_approx, display_score, display_binary_rank],
+        outputs=[output, json],
+    )
+    display_binary_rank.change(
+        search,
+        inputs=[query, top_k, rescore_multiplier, use_approx, display_score, display_binary_rank],
+        outputs=[output, json],
+    )
 
 demo.queue()
 demo.launch()

requirements.txt

@@ -1,7 +1,6 @@
-git+https://github.com/tomaarsen/sentence-transformers@feat/quantization
+sentence-transformers==5.2.0
 datasets
 pandas
 huggingface_hub>=0.24.0
 
-usearch
 faiss-cpu