app.py

import time
import gradio as gr
from datasets import load_dataset
import pandas as pd
from sentence_transformers import SentenceTransformer
from sentence_transformers.quantization import quantize_embeddings
import faiss
from usearch.index import Index

# Load titles and texts
wikipedia_dataset = load_dataset("bourdoiscatie/wikipedia_fr_2022_250K", split="train", num_proc=4).select_columns(["title", "text", "wiki_id"])

def add_link(example):
    example["title"] = '['+example["title"]+']('+'https://fr.wikipedia.org/wiki?curid='+str(example["wiki_id"])+')'
    return example
wikipedia_dataset = wikipedia_dataset.map(add_link)
    
# 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_fr_2022_250K_int8_usearch.index", view=True)
binary_index: faiss.IndexBinaryFlat = faiss.read_index_binary("wikipedia_fr_2022_250K_ubinary_faiss.index")
binary_ivf: faiss.IndexBinaryIVF = faiss.read_index_binary("wikipedia_fr_2022_250K_ubinary_ivf_faiss.index")

# Load the SentenceTransformer model for embedding the queries
model = SentenceTransformer("OrdalieTech/Solon-embeddings-large-0.1")


def search(query, top_k: int = 20, rescore_multiplier: int = 1, use_approx: bool = False):
    # 1. Embed the query as float32
    start_time = time.time()
    query_embedding = model.encode(query, prompt="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")
    quantize_time = time.time() - start_time

    # 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 = 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

    # 5. Rescore the top_k * rescore_multiplier using the float32 query embedding and the int8 document embeddings
    start_time = time.time()
    scores = query_embedding @ int8_embeddings.T
    rescore_time = time.time() - start_time

    # 6. Sort the scores and return the top_k
    start_time = time.time()
    indices = scores.argsort()[::-1][:top_k]
    top_k_indices = binary_ids[indices]
    top_k_scores = scores[indices]
    top_k_titles, top_k_texts = zip(*[(wikipedia_dataset[idx]["title"], wikipedia_dataset[idx]["text"]) for idx in top_k_indices.tolist()])
    df = pd.DataFrame({"Score_paragraphe": [round(value, 2) for value in top_k_scores], "Titre": top_k_titles, "Texte": top_k_texts})
    score_sum = df.groupby('Titre')['Score_paragraphe'].sum().reset_index()
    df = pd.merge(df, score_sum, on='Titre', how='left')
    df.rename(columns={'Score_paragraphe_y': 'Score_article'}, inplace=True)
    df.rename(columns={'Score_paragraphe_x': 'Score_paragraphe'}, inplace=True)
    df = df[["Score_article", "Score_paragraphe", "Titre", "Texte"]]
    df = df.sort_values('Score_article', ascending=False)
    # df = df.groupby('Titre')[['Score', 'Texte']].agg({'Score': 'sum', 'Texte': '\n\n'.join}).reset_index().sort_values('Score', ascending=False)
    sort_time = time.time() - start_time

    return df, {
        "Temps pour enchâsser la requête ": f"{embed_time:.4f} s",
        "Temps pour la quantisation ": f"{quantize_time:.4f} s",
        "Temps pour effectuer la recherche ": f"{search_time:.4f} s",
        "Temps de chargement ": f"{load_time:.4f} s",
        "Temps de rescorage ": f"{rescore_time:.4f} s",
        "Temps pour trier les résustats ": f"{sort_time:.4f} s",
        "Temps total pour la recherche ": f"{quantize_time + search_time + load_time + rescore_time + sort_time:.4f} s",
    }


with gr.Blocks(title="Requêter Wikipedia en temps réel 🔍") as demo:
    
    gr.Markdown(
        """
## Requêter Wikipedia en temps réel 🔍

Ce démonstrateur permet de requêter un corpus composé des 250K paragraphes les plus consultés du Wikipédia francophone.    
Les résultats sont renvoyés en temps réel via un pipeline tournant sur un CPU 🚀  
Nous nous sommes grandement inspirés du Space [quantized-retrieval](https://huggingface.co/spaces/sentence-transformers/quantized-retrieval) conçu par [Tom Aarsen](https://huggingface.co/tomaarsen) 🤗    
Si vous voulez en savoir plus sur le processus complet derrière ce démonstrateur, n'hésitez pas à déplier les liens ci-dessous.

<details><summary>1. Détails sur les données</summary>
Le corpus utilisé correspond au 250 000 premières lignes du jeu de données <a href="https://hf.co/datasets/Cohere/wikipedia-22-12-fr-embeddings"><i>wikipedia-22-12-fr-embeddings</i></a> mis en ligne par Cohere.  
Comme son nom l'indique il s'agit d'un jeu de données datant de décembre 2022. Cette information est à prendre en compte lorsque vous effectuez votre requête.  
De même il s'agit ici d'un sous-ensemble du jeu de données total, à savoir les 250 000 paragraphes les plus consultés à cette date-là.  
Ainsi, si vous effectuez une recherche pointue sur un sujet peu consulté, ce démonstrateur ne reverra probablement rien de pertinent.  
A noter également que Cohere a effectué un prétraitement sur les données ce qui a conduit à la suppression de dates par exemple.  
Ce jeu de données n'est donc pas optimal. L'idée était de pouvoir proposer quelque chose en peu de temps. 
Dans un deuxième temps, ce démonstrateur sera étendu à l'ensemble du jeu de données <i>wikipedia-22-12-fr-embeddings</i> (soit 13M de paragraphes).
Il n'est pas exclus d'ensuite utiliser une version plus récente de Wikipedia (on peut penser par exemple à <a href="https://hf.co/datasets/wikimedia/wikipedia"><i>wikimedia/wikipedia</i></a>
</details>

<details><summary>2. Détails le pipeline</summary>
1. La requête est enchâssée en float32 à l'aide du modèle <a href="https://hf.co/OrdalieTech/Solon-embeddings-large-0.1">Solon-embeddings-large-0.1</a> d'Ordalie.
2. La requête est quantizée en binaire à l'aide de la fonction `quantize_embeddings` de la bibliothèque <a href="https://sbert.net/">SentenceTransformers</a>.
3. Un index binaire (250K <i>embeddings</i> binaires pesant 32MB de mémoire/espace disque) est requêté (en binaire si l'option approximative est sélectionnée, en int8 si l'option exacte est sélectionnée).
4. Les <i>n</i> textes demandés par l'utilisateur jugés les plus pertinents sont chargés à la volée à partir d'un index int8 sur disque (250K <i>embeddings</i> int8 ; 0 bytes de mémoire, 293MB d'espace disque).
5. Les <i>n</i> textes sont rescorés en utilisant la requête en float32 et les enchâssements en int8.
6. Les <i>n</i> premiers textes sont triés par score et affichés. Le "Score_paragraphe" correspond au score individuel de chaque paragraphe d'être pertinant vis-à-vis de la requête. Le "Score_article" correspond à la somme de tous les scores individuels des paragraphes issus d'un même article Wikipedia. L'objectif est alors de mettre en avant l'article source plutôt qu'un bout de texte le composant.

Ce processus est conçu pour être rapide et efficace en termes de mémoire : l'index binaire étant suffisamment petit pour tenir dans la mémoire et l'index int8 étant chargé en tant que vue pour économiser de la mémoire. 
Au total, ce processus nécessite de conserver 1) le modèle en mémoire, 2) l'index binaire en mémoire et 3) l'index int8 sur le disque.  
Avec une dimension de 1024, nous avons besoin de `1024 / 8 * num_docs` octets pour l'index binaire et de `1024 * num_docs` octets pour l'index int8.

C'est nettement moins cher que de faire le même processus avec des enchâssements en float32 qui nécessiterait `4 * 1024 * num_docs` octets de mémoire/espace disque pour l'index float32, soit 32x plus de mémoire et 4x plus d'espace disque.
De plus, l'index binaire est beaucoup plus rapide (jusqu'à 32x) à rechercher que l'index float32, tandis que le rescorage est également extrêmement efficace.
En conclusion, ce processus permet une recherche rapide, évolutive, peu coûteuse et efficace en termes de mémoire.
</details>
"""
    )
    with gr.Row():
        with gr.Column(scale=75):
            query = gr.Textbox(
                label="Requêter le Wikipédia francophone",
                placeholder="Saisissez une requête pour rechercher des textes pertinents dans Wikipédia.",
            )
        with gr.Column(scale=25):
            use_approx = gr.Radio(
                choices=[("Exacte", False), ("Approximative", True)],
                value=True,
                label="Type de recherche",
            )

    with gr.Row():
        with gr.Column(scale=2):
            top_k = gr.Slider(
                minimum=3,
                maximum=40,
                step=1,
                value=15,
                label="Nombre de documents à rechercher",
                info="Recherche effectué via un bi-encodeur binaire",
            )
        with gr.Column(scale=2):
            rescore_multiplier = gr.Slider(
                minimum=1,
                maximum=10,
                step=1,
                value=1,
                label="Coefficient de rescorage",
                info="Reranking via le coefficient",
            )

    search_button = gr.Button(value="Search")

    output = gr.Dataframe(headers=["Score_article", "Score_paragraphe", "Titre", "Texte"], datatype="markdown")
    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])
    gr.Image("/file=catie(2).png", height=250,width=80, show_download_button=False)

demo.queue()
demo.launch(allowed_paths=["catie(2).png"])