Updated app with code for deduplication

app.py

@@ -4,67 +4,74 @@ import numpy as np
 from model2vec import StaticModel
 from reach import Reach
 from difflib import ndiff
+import sys
+import tqdm
 
-def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024, progress=gr.Progress(track_tqdm=True)) -> tuple[np.ndarray, dict[int, int]]:
+# Load the model at startup
+model = StaticModel.from_pretrained("minishlab/M2V_base_output")
+
+# Load the default datasets at startup
+default_dataset1_name = "ag_news"
+default_dataset1_split = "train"
+default_dataset2_name = "ag_news"
+default_dataset2_split = "test"
+
+ds_default1 = load_dataset(default_dataset1_name, split=default_dataset1_split)
+ds_default2 = load_dataset(default_dataset2_name, split=default_dataset2_split)
+
+def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024, progress=None) -> tuple[np.ndarray, dict[int, int]]:
     """
     Deduplicate embeddings and return the deduplicated indices and a mapping of removed indices to their corresponding original indices.
     """
     reach = Reach(vectors=embedding_matrix, items=[str(i) for i in range(len(embedding_matrix))])
 
-    # Use a set for deduplicated indices and keep track of duplicates
-    deduplicated_indices = set(range(len(embedding_matrix)))  # Start with all indices as deduplicated
+    deduplicated_indices = set(range(len(embedding_matrix)))
     duplicate_to_original_mapping = {}
 
     results = reach.nearest_neighbor_threshold(
-        embedding_matrix, 
-        threshold=threshold, 
-        batch_size=batch_size, 
-        show_progressbar=True
+        embedding_matrix,
+        threshold=threshold,
+        batch_size=batch_size,
+        show_progressbar=False  # Disable internal progress bar
     )
 
     # Process duplicates
     for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates")):
         if i not in deduplicated_indices:
-            continue  # Skip already marked duplicates
+            continue
 
-        # Similar items are returned as (index, score), we are only interested in the index
         similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i]
 
-        # Mark similar documents as duplicates and map them to the original
         for sim_idx in similar_indices:
             if sim_idx in deduplicated_indices:
                 deduplicated_indices.remove(sim_idx)
-                duplicate_to_original_mapping[sim_idx] = i  # Map duplicate to original
+                duplicate_to_original_mapping[sim_idx] = i
 
     return np.array(list(deduplicated_indices)), duplicate_to_original_mapping
 
-def deduplicate_across_datasets(embedding_matrix_1: np.ndarray, embedding_matrix_2: np.ndarray, threshold: float, batch_size: int = 1024, progress=gr.Progress(track_tqdm=True)) -> tuple[list[int], dict[int, int]]:
+def deduplicate_across_datasets(embedding_matrix_1: np.ndarray, embedding_matrix_2: np.ndarray, threshold: float, batch_size: int = 1024, progress=None) -> tuple[list[int], dict[int, int]]:
     """
     Deduplicate embeddings across two datasets and return the indices of duplicates between them.
     """
     reach = Reach(vectors=embedding_matrix_1, items=[str(i) for i in range(len(embedding_matrix_1))])
 
-    # Keep track of duplicates in the second dataset
     duplicate_indices_in_test = []
     duplicate_to_original_mapping = {}
 
-    # Find nearest neighbors from the test set in the train set
     results = reach.nearest_neighbor_threshold(
-        embedding_matrix_2, 
-        threshold=threshold, 
-        batch_size=batch_size, 
-        show_progressbar=True
+        embedding_matrix_2,
+        threshold=threshold,
+        batch_size=batch_size,
+        show_progressbar=False  # Disable internal progress bar
     )
 
     # Process duplicates
     for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates across datasets")):
-        # Similar items are returned as (index, score), we are only interested in the index
-        similar_indices = [int(item[0]) for item in similar_items if item[1] >= threshold]  # Keep those above the threshold
+        similar_indices = [int(item[0]) for item in similar_items if item[1] >= threshold]
 
-        # If we find a similar item in the train set, mark it as a duplicate
         if similar_indices:
             duplicate_indices_in_test.append(i)
-            duplicate_to_original_mapping[i] = similar_indices[0]  # Map duplicate in test to original in train
+            duplicate_to_original_mapping[i] = similar_indices[0]
 
     return duplicate_indices_in_test, duplicate_to_original_mapping
 
@@ -83,85 +90,114 @@ def perform_deduplication(
     threshold=0.8,
     progress=gr.Progress(track_tqdm=True)
 ):
-    # Convert threshold to float
-    threshold = float(threshold)
-    
-    if deduplication_type == "Single dataset":
-        # Load the dataset
-        ds = load_dataset(dataset1_name, split=dataset1_split)
-        
-        # Extract texts
-        texts = [example[dataset1_text_column] for example in ds]
-        
-        # Compute embeddings
-        model = StaticModel.from_pretrained("minishlab/M2V_base_output")
-        embedding_matrix = model.encode(texts, show_progressbar=True)
-        
-        # Deduplicate
-        deduplicated_indices, duplicate_to_original_mapping = deduplicate(embedding_matrix, threshold, progress=progress)
-        
-        # Prepare the results
-        num_duplicates = len(duplicate_to_original_mapping)
-        num_total = len(texts)
-        num_deduplicated = len(deduplicated_indices)
-        
-        result_text = f"**Total documents:** {num_total}\n"
-        result_text += f"**Number of duplicates found:** {num_duplicates}\n"
-        result_text += f"**Number of unique documents after deduplication:** {num_deduplicated}\n\n"
-        
-        # Show deduplicated examples
-        result_text += "**Examples of duplicates found:**\n\n"
-        num_examples = min(5, num_duplicates)
-        for duplicate_idx, original_idx in list(duplicate_to_original_mapping.items())[:num_examples]:
-            original_text = texts[original_idx]
-            duplicate_text = texts[duplicate_idx]
-            differences = display_word_differences(original_text, duplicate_text)
-            result_text += f"**Original text:**\n{original_text}\n\n"
-            result_text += f"**Duplicate text:**\n{duplicate_text}\n\n"
-            result_text += f"**Differences:**\n{differences}\n"
-            result_text += "-" * 50 + "\n\n"
-        
-        return result_text
-        
-    elif deduplication_type == "Cross-dataset":
-        # Load datasets
-        ds1 = load_dataset(dataset1_name, split=dataset1_split)
-        ds2 = load_dataset(dataset2_name, split=dataset2_split)
-        
-        # Extract texts
-        texts1 = [example[dataset1_text_column] for example in ds1]
-        texts2 = [example[dataset2_text_column] for example in ds2]
-        
-        # Compute embeddings
-        model = StaticModel.from_pretrained("minishlab/M2V_base_output")
-        embedding_matrix1 = model.encode(texts1, show_progressbar=True)
-        embedding_matrix2 = model.encode(texts2, show_progressbar=True)
-        
-        # Deduplicate across datasets
-        duplicate_indices_in_ds2, duplicate_to_original_mapping = deduplicate_across_datasets(embedding_matrix1, embedding_matrix2, threshold, progress=progress)
-        
-        num_duplicates = len(duplicate_indices_in_ds2)
-        num_total_ds2 = len(texts2)
-        num_unique_ds2 = num_total_ds2 - num_duplicates
-        
-        result_text = f"**Total documents in {dataset2_name}/{dataset2_split}:** {num_total_ds2}\n"
-        result_text += f"**Number of duplicates found in {dataset2_name}/{dataset2_split}:** {num_duplicates}\n"
-        result_text += f"**Number of unique documents in {dataset2_name}/{dataset2_split} after deduplication:** {num_unique_ds2}\n\n"
-        
-        # Show deduplicated examples
-        result_text += "**Examples of duplicates found in Dataset 2:**\n\n"
-        num_examples = min(5, num_duplicates)
-        for duplicate_idx in duplicate_indices_in_ds2[:num_examples]:
-            original_idx = duplicate_to_original_mapping[duplicate_idx]
-            original_text = texts1[original_idx]
-            duplicate_text = texts2[duplicate_idx]
-            differences = display_word_differences(original_text, duplicate_text)
-            result_text += f"**Original text (Dataset 1):**\n{original_text}\n\n"
-            result_text += f"**Duplicate text (Dataset 2):**\n{duplicate_text}\n\n"
-            result_text += f"**Differences:**\n{differences}\n"
-            result_text += "-" * 50 + "\n\n"
-        
-        return result_text
+    # Monkey-patch tqdm
+    original_tqdm = tqdm.tqdm
+    tqdm.tqdm = progress.tqdm
+    sys.modules['tqdm'].tqdm = progress.tqdm
+    sys.modules['tqdm.auto'].tqdm = progress.tqdm
+
+    try:
+        # Convert threshold to float
+        threshold = float(threshold)
+        
+        if deduplication_type == "Single dataset":
+            # Check if the dataset is the default one
+            if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split:
+                ds = ds_default1
+            else:
+                ds = load_dataset(dataset1_name, split=dataset1_split)
+            
+            # Extract texts
+            texts = [example[dataset1_text_column] for example in ds]
+            
+            # Compute embeddings
+            embedding_matrix = model.encode(texts, show_progressbar=False)  # Disable internal progress bar
+            
+            # Show progress bar for embedding computation
+            embedding_matrix = progress.tqdm(embedding_matrix, desc="Computing embeddings")
+
+            # Deduplicate
+            deduplicated_indices, duplicate_to_original_mapping = deduplicate(embedding_matrix, threshold, progress=progress)
+            
+            # Prepare the results
+            num_duplicates = len(duplicate_to_original_mapping)
+            num_total = len(texts)
+            num_deduplicated = len(deduplicated_indices)
+            
+            result_text = f"**Total documents:** {num_total}\n"
+            result_text += f"**Number of duplicates found:** {num_duplicates}\n"
+            result_text += f"**Number of unique documents after deduplication:** {num_deduplicated}\n\n"
+            
+            # Show deduplicated examples
+            result_text += "**Examples of duplicates found:**\n\n"
+            num_examples = min(5, num_duplicates)
+            for duplicate_idx, original_idx in list(duplicate_to_original_mapping.items())[:num_examples]:
+                original_text = texts[original_idx]
+                duplicate_text = texts[duplicate_idx]
+                differences = display_word_differences(original_text, duplicate_text)
+                result_text += f"**Original text:**\n{original_text}\n\n"
+                result_text += f"**Duplicate text:**\n{duplicate_text}\n\n"
+                result_text += f"**Differences:**\n{differences}\n"
+                result_text += "-" * 50 + "\n\n"
+            
+            return result_text
+                
+        elif deduplication_type == "Cross-dataset":
+            # Dataset 1
+            if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split:
+                ds1 = ds_default1
+            else:
+                ds1 = load_dataset(dataset1_name, split=dataset1_split)
+            
+            # Dataset 2
+            if dataset2_name == default_dataset2_name and dataset2_split == default_dataset2_split:
+                ds2 = ds_default2
+            else:
+                ds2 = load_dataset(dataset2_name, split=dataset2_split)
+            
+            # Extract texts
+            texts1 = [example[dataset1_text_column] for example in ds1]
+            texts2 = [example[dataset2_text_column] for example in ds2]
+            
+            # Compute embeddings
+            embedding_matrix1 = model.encode(texts1, show_progressbar=False)  # Disable internal progress bar
+            embedding_matrix2 = model.encode(texts2, show_progressbar=False)  # Disable internal progress bar
+            
+            # Show progress bar for embedding computation
+            embedding_matrix1 = progress.tqdm(embedding_matrix1, desc="Computing embeddings for Dataset 1")
+            embedding_matrix2 = progress.tqdm(embedding_matrix2, desc="Computing embeddings for Dataset 2")
+            
+            # Deduplicate across datasets
+            duplicate_indices_in_ds2, duplicate_to_original_mapping = deduplicate_across_datasets(embedding_matrix1, embedding_matrix2, threshold, progress=progress)
+            
+            num_duplicates = len(duplicate_indices_in_ds2)
+            num_total_ds2 = len(texts2)
+            num_unique_ds2 = num_total_ds2 - num_duplicates
+            
+            result_text = f"**Total documents in {dataset2_name}/{dataset2_split}:** {num_total_ds2}\n"
+            result_text += f"**Number of duplicates found in {dataset2_name}/{dataset2_split}:** {num_duplicates}\n"
+            result_text += f"**Number of unique documents in {dataset2_name}/{dataset2_split} after deduplication:** {num_unique_ds2}\n\n"
+            
+            # Show deduplicated examples
+            result_text += "**Examples of duplicates found in Dataset 2:**\n\n"
+            num_examples = min(5, num_duplicates)
+            for duplicate_idx in duplicate_indices_in_ds2[:num_examples]:
+                original_idx = duplicate_to_original_mapping[duplicate_idx]
+                original_text = texts1[original_idx]
+                duplicate_text = texts2[duplicate_idx]
+                differences = display_word_differences(original_text, duplicate_text)
+                result_text += f"**Original text (Dataset 1):**\n{original_text}\n\n"
+                result_text += f"**Duplicate text (Dataset 2):**\n{duplicate_text}\n\n"
+                result_text += f"**Differences:**\n{differences}\n"
+                result_text += "-" * 50 + "\n\n"
+            
+            return result_text
+
+    finally:
+        # Restore original tqdm
+        tqdm.tqdm = original_tqdm
+        sys.modules['tqdm'].tqdm = original_tqdm
+        sys.modules['tqdm.auto'].tqdm = original_tqdm
 
 with gr.Blocks() as demo:
     gr.Markdown("# Semantic Deduplication")
@@ -225,3 +261,232 @@ with gr.Blocks() as demo:
     )
     
 demo.launch()
+
+
+# import gradio as gr
+# from datasets import load_dataset
+# import numpy as np
+# from model2vec import StaticModel
+# from reach import Reach
+# from difflib import ndiff
+
+# def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024, progress=gr.Progress(track_tqdm=True)) -> tuple[np.ndarray, dict[int, int]]:
+#     """
+#     Deduplicate embeddings and return the deduplicated indices and a mapping of removed indices to their corresponding original indices.
+#     """
+#     reach = Reach(vectors=embedding_matrix, items=[str(i) for i in range(len(embedding_matrix))])
+
+#     # Use a set for deduplicated indices and keep track of duplicates
+#     deduplicated_indices = set(range(len(embedding_matrix)))  # Start with all indices as deduplicated
+#     duplicate_to_original_mapping = {}
+
+#     results = reach.nearest_neighbor_threshold(
+#         embedding_matrix, 
+#         threshold=threshold, 
+#         batch_size=batch_size, 
+#         show_progressbar=True
+#     )
+
+#     # Process duplicates
+#     for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates")):
+#         if i not in deduplicated_indices:
+#             continue  # Skip already marked duplicates
+
+#         # Similar items are returned as (index, score), we are only interested in the index
+#         similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i]
+
+#         # Mark similar documents as duplicates and map them to the original
+#         for sim_idx in similar_indices:
+#             if sim_idx in deduplicated_indices:
+#                 deduplicated_indices.remove(sim_idx)
+#                 duplicate_to_original_mapping[sim_idx] = i  # Map duplicate to original
+
+#     return np.array(list(deduplicated_indices)), duplicate_to_original_mapping
+
+# def deduplicate_across_datasets(embedding_matrix_1: np.ndarray, embedding_matrix_2: np.ndarray, threshold: float, batch_size: int = 1024, progress=gr.Progress(track_tqdm=True)) -> tuple[list[int], dict[int, int]]:
+#     """
+#     Deduplicate embeddings across two datasets and return the indices of duplicates between them.
+#     """
+#     reach = Reach(vectors=embedding_matrix_1, items=[str(i) for i in range(len(embedding_matrix_1))])
+
+#     # Keep track of duplicates in the second dataset
+#     duplicate_indices_in_test = []
+#     duplicate_to_original_mapping = {}
+
+#     # Find nearest neighbors from the test set in the train set
+#     results = reach.nearest_neighbor_threshold(
+#         embedding_matrix_2, 
+#         threshold=threshold, 
+#         batch_size=batch_size, 
+#         show_progressbar=True
+#     )
+
+#     # Process duplicates
+#     for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates across datasets")):
+#         # Similar items are returned as (index, score), we are only interested in the index
+#         similar_indices = [int(item[0]) for item in similar_items if item[1] >= threshold]  # Keep those above the threshold
+
+#         # If we find a similar item in the train set, mark it as a duplicate
+#         if similar_indices:
+#             duplicate_indices_in_test.append(i)
+#             duplicate_to_original_mapping[i] = similar_indices[0]  # Map duplicate in test to original in train
+
+#     return duplicate_indices_in_test, duplicate_to_original_mapping
+
+# def display_word_differences(x: str, y: str) -> str:
+#     diff = ndiff(x.split(), y.split())
+#     return " ".join([word for word in diff if word.startswith(('+', '-'))])
+
+# def perform_deduplication(
+#     deduplication_type,
+#     dataset1_name,
+#     dataset1_split,
+#     dataset1_text_column,
+#     dataset2_name="",
+#     dataset2_split="",
+#     dataset2_text_column="",
+#     threshold=0.8,
+#     progress=gr.Progress(track_tqdm=True)
+# ):
+#     # Convert threshold to float
+#     threshold = float(threshold)
+    
+#     if deduplication_type == "Single dataset":
+#         # Load the dataset
+#         ds = load_dataset(dataset1_name, split=dataset1_split)
+        
+#         # Extract texts
+#         texts = [example[dataset1_text_column] for example in ds]
+        
+#         # Compute embeddings
+#         model = StaticModel.from_pretrained("minishlab/M2V_base_output")
+#         embedding_matrix = model.encode(texts, show_progressbar=True)
+        
+#         # Deduplicate
+#         deduplicated_indices, duplicate_to_original_mapping = deduplicate(embedding_matrix, threshold, progress=progress)
+        
+#         # Prepare the results
+#         num_duplicates = len(duplicate_to_original_mapping)
+#         num_total = len(texts)
+#         num_deduplicated = len(deduplicated_indices)
+        
+#         result_text = f"**Total documents:** {num_total}\n"
+#         result_text += f"**Number of duplicates found:** {num_duplicates}\n"
+#         result_text += f"**Number of unique documents after deduplication:** {num_deduplicated}\n\n"
+        
+#         # Show deduplicated examples
+#         result_text += "**Examples of duplicates found:**\n\n"
+#         num_examples = min(5, num_duplicates)
+#         for duplicate_idx, original_idx in list(duplicate_to_original_mapping.items())[:num_examples]:
+#             original_text = texts[original_idx]
+#             duplicate_text = texts[duplicate_idx]
+#             differences = display_word_differences(original_text, duplicate_text)
+#             result_text += f"**Original text:**\n{original_text}\n\n"
+#             result_text += f"**Duplicate text:**\n{duplicate_text}\n\n"
+#             result_text += f"**Differences:**\n{differences}\n"
+#             result_text += "-" * 50 + "\n\n"
+        
+#         return result_text
+        
+#     elif deduplication_type == "Cross-dataset":
+#         # Load datasets
+#         ds1 = load_dataset(dataset1_name, split=dataset1_split)
+#         ds2 = load_dataset(dataset2_name, split=dataset2_split)
+        
+#         # Extract texts
+#         texts1 = [example[dataset1_text_column] for example in ds1]
+#         texts2 = [example[dataset2_text_column] for example in ds2]
+        
+#         # Compute embeddings
+#         model = StaticModel.from_pretrained("minishlab/M2V_base_output")
+#         embedding_matrix1 = model.encode(texts1, show_progressbar=True)
+#         embedding_matrix2 = model.encode(texts2, show_progressbar=True)
+        
+#         # Deduplicate across datasets
+#         duplicate_indices_in_ds2, duplicate_to_original_mapping = deduplicate_across_datasets(embedding_matrix1, embedding_matrix2, threshold, progress=progress)
+        
+#         num_duplicates = len(duplicate_indices_in_ds2)
+#         num_total_ds2 = len(texts2)
+#         num_unique_ds2 = num_total_ds2 - num_duplicates
+        
+#         result_text = f"**Total documents in {dataset2_name}/{dataset2_split}:** {num_total_ds2}\n"
+#         result_text += f"**Number of duplicates found in {dataset2_name}/{dataset2_split}:** {num_duplicates}\n"
+#         result_text += f"**Number of unique documents in {dataset2_name}/{dataset2_split} after deduplication:** {num_unique_ds2}\n\n"
+        
+#         # Show deduplicated examples
+#         result_text += "**Examples of duplicates found in Dataset 2:**\n\n"
+#         num_examples = min(5, num_duplicates)
+#         for duplicate_idx in duplicate_indices_in_ds2[:num_examples]:
+#             original_idx = duplicate_to_original_mapping[duplicate_idx]
+#             original_text = texts1[original_idx]
+#             duplicate_text = texts2[duplicate_idx]
+#             differences = display_word_differences(original_text, duplicate_text)
+#             result_text += f"**Original text (Dataset 1):**\n{original_text}\n\n"
+#             result_text += f"**Duplicate text (Dataset 2):**\n{duplicate_text}\n\n"
+#             result_text += f"**Differences:**\n{differences}\n"
+#             result_text += "-" * 50 + "\n\n"
+        
+#         return result_text
+
+# with gr.Blocks() as demo:
+#     gr.Markdown("# Semantic Deduplication")
+    
+#     deduplication_type = gr.Radio(
+#         choices=["Single dataset", "Cross-dataset"],
+#         label="Deduplication Type",
+#         value="Single dataset"
+#     )
+    
+#     with gr.Row():
+#         dataset1_name = gr.Textbox(value="ag_news", label="Dataset 1 Name")
+#         dataset1_split = gr.Textbox(value="train", label="Dataset 1 Split")
+#         dataset1_text_column = gr.Textbox(value="text", label="Text Column Name")
+    
+#     dataset2_inputs = gr.Column(visible=False)
+#     with dataset2_inputs:
+#         gr.Markdown("### Dataset 2")
+#         with gr.Row():
+#             dataset2_name = gr.Textbox(value="ag_news", label="Dataset 2 Name")
+#             dataset2_split = gr.Textbox(value="test", label="Dataset 2 Split")
+#             dataset2_text_column = gr.Textbox(value="text", label="Text Column Name")
+    
+#     threshold = gr.Slider(
+#         minimum=0.0,
+#         maximum=1.0,
+#         value=0.8,
+#         label="Similarity Threshold"
+#     )
+    
+#     compute_button = gr.Button("Compute")
+    
+#     output = gr.Markdown()
+    
+#     # Function to update the visibility of dataset2_inputs
+#     def update_visibility(deduplication_type_value):
+#         if deduplication_type_value == "Cross-dataset":
+#             return gr.update(visible=True)
+#         else:
+#             return gr.update(visible=False)
+    
+#     deduplication_type.change(
+#         update_visibility,
+#         inputs=deduplication_type,
+#         outputs=dataset2_inputs
+#     )
+    
+#     compute_button.click(
+#         fn=perform_deduplication,
+#         inputs=[
+#             deduplication_type, 
+#             dataset1_name, 
+#             dataset1_split, 
+#             dataset1_text_column,
+#             dataset2_name, 
+#             dataset2_split, 
+#             dataset2_text_column,
+#             threshold
+#         ],
+#         outputs=output
+#     )
+    
+# demo.launch()