Update app.py

app.py

@@ -8,8 +8,13 @@ import faiss
 from usearch.index import Index
 
 # Load titles and texts
-title_text_dataset = load_dataset("bourdoiscatie/wikipedia_fr_2022_250K", split="train", num_proc=4).select_columns(["title", "text"])
+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")
@@ -52,8 +57,15 @@ def search(query, top_k: int = 20, rescore_multiplier: int = 1, use_approx: bool
     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], "Titre": top_k_titles, "Texte": top_k_texts})
+    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_ind": [round(value, 2) for value in top_k_scores], "Titre": top_k_titles, "Texte": top_k_texts})
+    score_sum = df.groupby('Titre')['Score_ind'].sum().reset_index()
+    df = pd.merge(df, score_sum, on='Titre', how='left')
+    df.rename(columns={'Score_ind_y': 'Score_sum'}, inplace=True)
+    df.rename(columns={'Score_ind_x': 'Score_ind'}, inplace=True)
+    df = df[["Score_sum", "Score_ind", "Titre", "Texte"]]
+    df = df.sort_values('Score_sum', 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, {
@@ -67,24 +79,34 @@ def search(query, top_k: int = 20, rescore_multiplier: int = 1, use_approx: bool
     }
 
 
-with gr.Blocks(title="Requêter Wikipedia en temps réel") as demo:
+with gr.Blocks(title="Requêter Wikipedia en temps réel 🔍") as demo:
     gr.Markdown(
         """
-## Requêter Wikipedia en temps réel
-
-Effectuer une requête dans un corpus composé de 250K paragraphes tirés d'articles de Wikipédia.    
-Les résultats sont renvoyés en temps réel via une architecture tournant sur un CPU 🚀
-
-
-<details><summary>Détails du processus</summary>
+## 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 [wikipedia-22-12-fr-embeddings](https://huggingface.co/datasets/Cohere/wikipedia-22-12-fr-embeddings) 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 *wikipedia-22-12-fr-embeddings* (soit 13M de paragraphes).
+Il n'est pas exclus d'ensuite utiliser une version plus récente de Wikipedia (on peut penser par exemple à [wikimedia/wikipedia](https://huggingface.co/datasets/wikimedia/wikipedia) de Wikimedia) mais qui demandera d'effectuer plusieurs nettoyages.
+</details>
 
-Détails :
-1. La requête est enchâssée en float32 à l'aide du modèle [`OrdalieTech/Solon-embeddings-large-0.1`](https://hf.co/OrdalieTech/Solon-embeddings-large-0.1).
+<details><summary>2. Détails le pipeline</summary>
+1. La requête est enchâssée en float32 à l'aide du modèle [`OrdalieTech/Solon-embeddings-large-0.1`](https://hf.co/OrdalieTech/Solon-embeddings-large-0.1) d'Ordalie.
 2. La requête est quantizée en binaire à l'aide de la fonction `quantize_embeddings` de la bibliothèque [SentenceTransformers](https://sbert.net/).
-3. Un index binaire (XXM d'enchâssements binaires pesant XXGB de mémoire/espace disque) est requêté (en binaire si l'option approximatie sélectionnée, en int8 si l'option exacte est sélectionnée).
-4. Les *n* textes demandés par l'utilisateur jugés les plus pertinents sont chargés à la volée à partir d'un index int8 sur disque (XXM d'enchâssements int8 ; 0 bytes de mémoire, XXGB d'espace disque).
+3. Un index binaire (250K *embeddings* 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 *n* 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 *embeddings* int8 ; 0 bytes de mémoire, 293MB d'espace disque).
 5. Les *n* textes sont rescorés en utilisant la requête en float32 et les enchâssements en int8.
-6. Les *n* premiers textes sont triés par score et affichés.
+6. Les *n* premiers textes sont triés par score et affichés. Le "Score_ind" correspond au score individuel de chaque paragraphe d'être pertinant vis-à-vis de la requête. Le "Score_sum" 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.  
@@ -93,30 +115,29 @@ Avec une dimension de 1024, nous avons besoin de `1024 / 8 * num_docs` octets po
 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="Recherche d'articles dans le Wikipédia francophone",
+                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=[("Recherche exacte", False), ("Recherche approximative", True)],
+                choices=[("Exacte", False), ("Approximative", True)],
                 value=True,
-                label="Index de recherche",
+                label="Type de recherche",
             )
 
     with gr.Row():
         with gr.Column(scale=2):
             top_k = gr.Slider(
-                minimum=10,
-                maximum=200,
-                step=5,
-                value=20,
+                minimum=3,
+                maximum=40,
+                step=1,
+                value=15,
                 label="Nombre de documents à rechercher",
                 info="Recherche effectué via un bi-encodeur binaire",
             )
@@ -127,16 +148,13 @@ En conclusion, ce processus permet une recherche rapide, évolutive, peu coûteu
                 step=1,
                 value=1,
                 label="Coefficient de rescorage",
-                info="Reranking via le coefficient`",
+                info="Reranking via le coefficient",
             )
 
     search_button = gr.Button(value="Search")
 
-    with gr.Row():
-        with gr.Column(scale=4):
-            output = gr.Dataframe(headers=["Score", "Titre", "Texte"])
-        with gr.Column(scale=1):
-            json = gr.JSON()
+    output = gr.Dataframe(headers=["Score_sum", "Score_ind", "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])