Spaces:
Runtime error
Runtime error
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"]) |