node-loki/quantized-retrieval

1import time2import gradio as gr3from datasets import load_dataset4import pandas as pd5from sentence_transformers import SentenceTransformer6from sentence_transformers.quantization import quantize_embeddings7import faiss8from usearch.index import Index9 10# Load titles and texts11title_text_dataset = load_dataset("mixedbread-ai/wikipedia-data-en-2023-11", split="train", num_proc=4).select_columns(["title", "text"])12 13# Load the int8 and binary indices. Int8 is loaded as a view to save memory, as we never actually perform search with it.14int8_view = Index.restore("wikipedia_int8_usearch_50m.index", view=True)15binary_index: faiss.IndexBinaryFlat = faiss.read_index_binary("wikipedia_ubinary_faiss_50m.index")16binary_ivf: faiss.IndexBinaryIVF = faiss.read_index_binary("wikipedia_ubinary_ivf_faiss_50m.index")17 18# Load the SentenceTransformer model for embedding the queries19model = SentenceTransformer(20    "mixedbread-ai/mxbai-embed-large-v1",21    prompts={22        "retrieval": "Represent this sentence for searching relevant passages: ",23    },24    default_prompt_name="retrieval",25)26 27 28def search(query, top_k: int = 100, rescore_multiplier: int = 1, use_approx: bool = False):29    # 1. Embed the query as float3230    start_time = time.time()31    query_embedding = model.encode(query)32    embed_time = time.time() - start_time33 34    # 2. Quantize the query to ubinary35    start_time = time.time()36    query_embedding_ubinary = quantize_embeddings(query_embedding.reshape(1, -1), "ubinary")37    quantize_time = time.time() - start_time38 39    # 3. Search the binary index (either exact or approximate)40    index = binary_ivf if use_approx else binary_index41    start_time = time.time()42    _scores, binary_ids = index.search(query_embedding_ubinary, top_k * rescore_multiplier)43    binary_ids = binary_ids[0]44    search_time = time.time() - start_time45 46    # 4. Load the corresponding int8 embeddings47    start_time = time.time()48    int8_embeddings = int8_view[binary_ids].astype(int)49    load_time = time.time() - start_time50 51    # 5. Rescore the top_k * rescore_multiplier using the float32 query embedding and the int8 document embeddings52    start_time = time.time()53    scores = query_embedding @ int8_embeddings.T54    rescore_time = time.time() - start_time55 56    # 6. Sort the scores and return the top_k57    start_time = time.time()58    indices = scores.argsort()[::-1][:top_k]59    top_k_indices = binary_ids[indices]60    top_k_scores = scores[indices]61    top_k_titles, top_k_texts = zip(62        *[(title_text_dataset[idx]["title"], title_text_dataset[idx]["text"]) for idx in top_k_indices.tolist()]63    )64    df = pd.DataFrame(65        {"Score": [round(value, 2) for value in top_k_scores], "Title": top_k_titles, "Text": top_k_texts}66    )67    sort_time = time.time() - start_time68 69    return df, {70        "Embed Time": f"{embed_time:.4f} s",71        "Quantize Time": f"{quantize_time:.4f} s",72        "Search Time": f"{search_time:.4f} s",73        "Load Time": f"{load_time:.4f} s",74        "Rescore Time": f"{rescore_time:.4f} s",75        "Sort Time": f"{sort_time:.4f} s",76        "Total Retrieval Time": f"{quantize_time + search_time + load_time + rescore_time + sort_time:.4f} s",77    }78 79 80with gr.Blocks(title="Quantized Retrieval") as demo:81    gr.Markdown(82        """83## Quantized Retrieval - Binary Search with Scalar (int8) Rescoring84This 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.85 86<details><summary>Click to learn about the retrieval process</summary>87 88Details:891. The query is embedded using the [`mixedbread-ai/mxbai-embed-large-v1`](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1) SentenceTransformer model.902. The query is quantized to binary using the `quantize_embeddings` function from the SentenceTransformers library.913. A binary index (41M binary embeddings; 5.2GB of memory/disk space) is searched using the quantized query for the top 40 documents.924. 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).935. The top 40 documents are rescored using the float32 query and the int8 embeddings to get the top 10 documents.946. The top 10 documents are sorted by score and displayed.95 96This 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. 97In 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, 98we need `1024 / 8 * num_docs` bytes for the binary index and `1024 * num_docs` bytes for the int8 index.99 100This 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.101Additionally, 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.102 103Feel 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.104 105Notes:106- 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.107 108</details>109"""110    )111    with gr.Row():112        with gr.Column(scale=75):113            query = gr.Textbox(114                label="Query for Wikipedia articles",115                placeholder="Enter a query to search for relevant texts from Wikipedia.",116            )117        with gr.Column(scale=25):118            use_approx = gr.Radio(119                choices=[("Exact Search", False), ("Approximate Search", True)],120                value=True,121                label="Search Index",122            )123 124    with gr.Row():125        with gr.Column(scale=2):126            top_k = gr.Slider(127                minimum=10,128                maximum=1000,129                step=5,130                value=100,131                label="Number of documents to retrieve",132                info="Number of documents to retrieve from the binary search",133            )134        with gr.Column(scale=2):135            rescore_multiplier = gr.Slider(136                minimum=1,137                maximum=10,138                step=1,139                value=1,140                label="Rescore multiplier",141                info="Search for `rescore_multiplier` as many documents to rescore",142            )143 144    search_button = gr.Button(value="Search")145 146    with gr.Row():147        with gr.Column(scale=4):148            output = gr.Dataframe(headers=["Score", "Title", "Text"])149        with gr.Column(scale=1):150            json = gr.JSON()151 152    query.submit(search, inputs=[query, top_k, rescore_multiplier, use_approx], outputs=[output, json])153    search_button.click(search, inputs=[query, top_k, rescore_multiplier, use_approx], outputs=[output, json])154 155demo.queue()156demo.launch()157