Update app.py

app.py

@@ -2,89 +2,284 @@ import time
 import gradio as gr
 from datasets import load_dataset
 import pandas as pd
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, SparseEncoder
 from sentence_transformers.quantization import quantize_embeddings
 import faiss
 from usearch.index import Index
+import torch
+import numpy as np
+from collections import defaultdict
+from scipy import stats
 
-# Load titles and texts
-wikipedia_dataset = load_dataset("bourdoiscatie/wikipedia_fr_2022_250K", split="train", num_proc=4).select_columns(["title", "text", "wiki_id"])
 
+
+# Load titles, texts and pre-computed sparse embeddings
+wikipedia_dataset = load_dataset("CATIE-AQ/wikipedia_fr_2022_250K", split="train", num_proc=4).select_columns(["title", "text", "wiki_id", "sparse_emb"])
+
+
+# A function to make the titles of Wikipedia articles clickable in the final dataframe, so that you can consult the article on the website
 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.
+
+
+# Load the int8 and binary indices for dense retrieval
 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")
+
+# Load models
+dense_model = SentenceTransformer("OrdalieTech/Solon-embeddings-large-0.1")
+sparse_model = SparseEncoder("CATIE-AQ/SPLADE_camembert-base_STS")
+
+
+# Fusion methods
+def reciprocal_rank_fusion(dense_results, sparse_results, k=20):
+    """
+    Perform Reciprocal Rank Fusion to combine dense and sparse retrieval results
+    
+    Args:
+        dense_results: List of (doc_id, score) from dense retrieval
+        sparse_results: List of (doc_id, score) from sparse retrieval  
+        k: RRF parameter (default 20)
+    
+    Returns:
+        List of (doc_id, rrf_score) sorted by RRF score
+    """
+    rrf_scores = defaultdict(float)
+    
+    # Add scores from dense retrieval
+    for rank, (doc_id, _) in enumerate(dense_results, 1):
+        rrf_scores[doc_id] += 1 / (k + rank)
+        
+    # Add scores from sparse retrieval
+    for rank, (doc_id, _) in enumerate(sparse_results, 1):
+        rrf_scores[doc_id] += 1 / (k + rank)
+        
+    # Sort by RRF score
+    sorted_results = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
+    return sorted_results
+
+    
+def normalized_score_fusion(dense_results, sparse_results, dense_weight=0.5, sparse_weight=0.5):
+    """
+    Perform Normalized Score Fusion (NSF) with z-score normalization
+    
+    Args:
+        dense_results: List of (doc_id, score) from dense retrieval
+        sparse_results: List of (doc_id, score) from sparse retrieval
+        dense_weight: Weight for dense scores (default 0.5)
+        sparse_weight: Weight for sparse scores (default 0.5)
+    
+    Returns:
+        List of (doc_id, normalized_score) sorted by normalized score
+    """
+    # Extract scores for normalization
+    dense_scores = np.array([score for _, score in dense_results])
+    sparse_scores = np.array([score for _, score in sparse_results])
+    
+    # Z-score normalization (handle edge cases)
+    if len(dense_scores) > 1 and np.std(dense_scores) > 1e-10:
+        dense_scores_norm = stats.zscore(dense_scores)
+    else:
+        dense_scores_norm = np.zeros_like(dense_scores)
+    
+    if len(sparse_scores) > 1 and np.std(sparse_scores) > 1e-10:
+        sparse_scores_norm = stats.zscore(sparse_scores)
+    else:
+        sparse_scores_norm = np.zeros_like(sparse_scores)
+    
+    # Create dictionaries for normalized scores
+    dense_norm_dict = {doc_id: score for (doc_id, _), score in zip(dense_results, dense_scores_norm)}
+    sparse_norm_dict = {doc_id: score for (doc_id, _), score in zip(sparse_results, sparse_scores_norm)}
+    
+    # Combine all unique document IDs
+    all_doc_ids = set()
+    all_doc_ids.update(doc_id for doc_id, _ in dense_results)
+    all_doc_ids.update(doc_id for doc_id, _ in sparse_results)
+    
+    # Calculate weighted normalized scores
+    nsf_scores = {}
+    for doc_id in all_doc_ids:
+        dense_norm_score = dense_norm_dict.get(doc_id, 0.0)
+        sparse_norm_score = sparse_norm_dict.get(doc_id, 0.0)
+        
+        # Weighted combination of normalized scores
+        nsf_scores[doc_id] = (dense_weight * dense_norm_score + 
+                              sparse_weight * sparse_norm_score)
+    
+    # Sort by NSF score
+    sorted_results = sorted(nsf_scores.items(), key=lambda x: x[1], reverse=True)
+    return sorted_results
+
+
+# Search part
+## Currentl to try to speed up things, I keep only non-zero values for the sparse search
+## Need to optimise the sparse research part (about 25-30s currently...)
+## Tom must surely have some tips on this point
+sparse_index = {}
+for i, sparse_emd in enumerate(wikipedia_dataset["sparse_emb"]):
+    # Convert sparse_emd to a NumPy array
+    sparse_emd = np.array(sparse_emd)
+
+    # Only store non-zero values and their indices
+    non_zero_indices = np.nonzero(sparse_emd)[0]
+    non_zero_values = sparse_emd[non_zero_indices]
+    sparse_index[i] = (non_zero_indices, non_zero_values)
+    
+
+def sparse_search(query, top_k=20):
+    """
+    Perform sparse retrieval using SPLADE representations with a dictionary-based sparse index.
+    """
+    # Encode the query, the output is a sparse torch.Tensor
+    query_sparse_vector = sparse_model.encode(query, convert_to_numpy=False)
+    
+    # Convert the sparse tensor to a dense format before using np.nonzero
+    query_sparse_vector = query_sparse_vector.to_dense().numpy()
+    
+    # Get non-zero values and indices for the query
+    query_non_zero_indices = np.nonzero(query_sparse_vector)[0]
+    query_non_zero_values = query_sparse_vector[query_non_zero_indices]
+    
+    # Compute dot product similarity efficiently
+    scores = defaultdict(float)
+    for doc_id, (doc_indices, doc_values) in sparse_index.items():
+        # Find the intersection of non-zero indices
+        common_indices = np.intersect1d(query_non_zero_indices, doc_indices)
+        # Compute dot product for common indices only
+        for idx in common_indices:
+            query_val = query_non_zero_values[np.where(query_non_zero_indices == idx)[0][0]]
+            doc_val = doc_values[np.where(doc_indices == idx)[0][0]]
+            scores[doc_id] += query_val * doc_val
+            
+    # Sort and get top_k
+    sorted_scores = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+    return sorted_scores[:top_k]
 
 
-def search(query, top_k: int = 20, rescore_multiplier: int = 1, use_approx: bool = False):
-    # 1. Embed the query as float32
+# Final search combining dense and sparse results
+def search(query, top_k: int = 20, rescore_multiplier: int = 1, use_approx: bool = False, 
+           fusion_method: str = "rrf", rrf_k: int = 20, dense_weight: float = 0.5, sparse_weight: float = 0.5):
+    total_start_time = time.time()
+    
+    # 1. Dense retrieval pipeline (existing code)
     start_time = time.time()
-    query_embedding = model.encode(query, prompt="query: ")
+    query_embedding = dense_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)
+    _scores, binary_ids = index.search(query_embedding_ubinary, top_k * rescore_multiplier * 2)  # Get more for fusion
     binary_ids = binary_ids[0]
-    search_time = time.time() - start_time
+    dense_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
+    # Prepare dense results
+    dense_results = [(binary_ids[i], scores[i]) for i in range(len(binary_ids))]
+    dense_results.sort(key=lambda x: x[1], reverse=True)
+    
+    timing_info = {
+        "Temps pour créer l'embedding de la requête (dense)": f"{embed_time:.4f} s",
+        "Temps pour la quantification": f"{quantize_time:.4f} s",
+        "Temps pour effectuer la recherche dense": f"{dense_search_time:.4f} s",
+        "Temps de chargement": f"{load_time:.4f} s",
+        "Temps de rescorage": f"{rescore_time:.4f} s",
+    }
+
+    
+    if fusion_method != "dense_only":
+        # 2. Sparse retrieval pipeline
+        start_time = time.time()
+        sparse_results = sparse_search(query, top_k * rescore_multiplier * 2)
+        sparse_search_time = time.time() - start_time
+        
+        # 3. Apply selected fusion method
+        start_time = time.time()
+        if fusion_method == "rrf":
+            fusion_results = reciprocal_rank_fusion(dense_results, sparse_results, k=rrf_k)
+            fusion_method_name = f"RRF (k={rrf_k})"
+        elif fusion_method == "nsf":
+            fusion_results = normalized_score_fusion(dense_results, sparse_results, 
+                                                     dense_weight=dense_weight, sparse_weight=sparse_weight)
+            fusion_method_name = f"NSF Z-score (dense={dense_weight:.1f}, sparse={sparse_weight:.1f})"
+        else:
+            fusion_results = dense_results  # fallback
+            fusion_method_name = "Dense uniquement (fallback)"
+        
+        fusion_time = time.time() - start_time
+        
+        # Use fusion results
+        final_results = fusion_results[:top_k * rescore_multiplier]
+        final_doc_ids = [doc_id for doc_id, _ in final_results]
+        final_scores = [score for _, score in final_results]
+        
+        timing_info.update({
+            # Add time to create query embedding with sparse model
+            # Do quantification for sparse model as dense one?
+            "Temps pour la recherche sparse": f"{sparse_search_time:.4f} s",
+            "Temps pour la fusion": f"{fusion_time:.4f} s",
+        })
+        timing_info["Méthode de fusion utilisée"] = fusion_method_name
+    else:
+        # Use only dense results
+        final_doc_ids = [doc_id for doc_id, _ in dense_results[:top_k * rescore_multiplier]]
+        final_scores = [score for _, score in dense_results[:top_k * rescore_multiplier]]
+        timing_info["Méthode de fusion utilisée"] = "Dense uniquement"
+    
+    # 4. Prepare final results
     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)
+    try:
+        top_k_titles, top_k_texts = zip(*[(wikipedia_dataset[int(idx)]["title"], wikipedia_dataset[int(idx)]["text"]) 
+                                             for idx in final_doc_ids[:top_k]])
+        
+        # Create DataFrame with results
+        df = pd.DataFrame({
+            "Score_paragraphe": [round(float(score), 4) for score in final_scores[:top_k]], 
+            "Titre": top_k_titles, 
+            "Texte": top_k_texts
+        })
+        
+        # Calculate article scores (sum of paragraph scores)
+        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', 'Score_paragraphe_x': 'Score_paragraphe'}, inplace=True)
+        df = df[["Score_article", "Score_paragraphe", "Titre", "Texte"]]
+        df = df.sort_values('Score_article', ascending=False)
+        
+    except Exception as e:
+        print(f"Error creating results DataFrame: {e}")
+        df = pd.DataFrame({"Error": [f"No results found or error processing results: {e}"]})
+    
     sort_time = time.time() - start_time
+    total_time = time.time() - total_start_time
+    
+    timing_info.update({
+        "Temps pour afficher les résultats": f"{sort_time:.4f} s",
+        "Temps total": f"{total_time:.4f} s",
+    })
 
-    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",
-    }
+    return df, timing_info
 
 
-with gr.Blocks(title="Requêter Wikipedia en temps réel 🔍") as demo:
+with gr.Blocks(title="Requêter Wikipedia avec Fusion Hybride 🔍") 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) 🤗    
@@ -102,23 +297,11 @@ Il n'est pas exclus d'ensuite utiliser une version plus récente de Wikipedia (o
 </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.
+A écrire quand ça sera terminé.
 </details>
 """
     )
+    
     with gr.Row():
         with gr.Column(scale=75):
             query = gr.Textbox(
@@ -129,7 +312,7 @@ En conclusion, ce processus permet une recherche rapide, évolutive, peu coûteu
             use_approx = gr.Radio(
                 choices=[("Exacte", False), ("Approximative", True)],
                 value=True,
-                label="Type de recherche",
+                label="Type de recherche dense",
             )
 
     with gr.Row():
@@ -139,8 +322,7 @@ En conclusion, ce processus permet une recherche rapide, évolutive, peu coûteu
                 maximum=40,
                 step=1,
                 value=15,
-                label="Nombre de documents à rechercher",
-                info="Recherche effectué via un bi-encodeur binaire",
+                label="Nombre de documents à retourner",
             )
         with gr.Column(scale=2):
             rescore_multiplier = gr.Slider(
@@ -149,17 +331,72 @@ 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="Augmente le nombre de candidats avant fusion",
+            )
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            fusion_method = gr.Radio(
+                choices=[
+                    ("Dense uniquement", "dense_only"),
+                    ("Fusion RRF", "rrf"),  
+                    ("Fusion NSF Z-score", "nsf")
+                ],
+                value="rrf",
+                label="Méthode de fusion",
+                info="Choisissez comment combiner les résultats des modèles dense et sparse"
             )
+        with gr.Column(scale=2):
+            rrf_k = gr.Slider(
+                minimum=1,
+                maximum=200,
+                step=1,
+                value=60,
+                label="Paramètre k pour RRF",
+                info="Plus k est élevé, moins les rangs ont d'importance",
+                visible=True
+            )
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            dense_weight = gr.Slider(
+                minimum=0.0,
+                maximum=1.0,
+                step=0.1,
+                value=0.5,
+                label="Poids Dense (NSF)",
+                info="Importance des résultats du modèle dense",
+                visible=False
+            )
+        with gr.Column(scale=2):
+            sparse_weight = gr.Slider(
+                minimum=0.0,
+                maximum=1.0,
+                step=0.1,
+                value=0.5,
+                label="Poids Sparse (NSF)",
+                info="Importance des résultats du modèle sparse",
+                visible=False
+            )
+    
+    # JavaScript to show/hide parameters based on fusion method
+    fusion_method.change(
+        fn=lambda method: (
+            gr.update(visible=(method == "rrf")),  # rrf_k
+            gr.update(visible=(method == "nsf")),  # dense_weight  
+            gr.update(visible=(method == "nsf"))    # sparse_weight
+        ),
+        inputs=[fusion_method],
+        outputs=[rrf_k, dense_weight, sparse_weight]
+    )
 
-    search_button = gr.Button(value="Search")
+    search_button = gr.Button(value="Rechercher", variant="primary")
 
     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)
+    json = gr.JSON(label="Informations de performance")
 
+    query.submit(search, inputs=[query, top_k, rescore_multiplier, use_approx, fusion_method, rrf_k, dense_weight, sparse_weight], outputs=[output, json])
+    search_button.click(search, inputs=[query, top_k, rescore_multiplier, use_approx, fusion_method, rrf_k, dense_weight, sparse_weight], outputs=[output, json])
+    
 demo.queue()
-demo.launch(allowed_paths=["catie(2).png"])
+demo.launch()