Updated app with code for deduplication

app.py

@@ -1,18 +1,14 @@
-# import gradio as gr
-
-# def greet(name):
-#     return "Hello " + name + "!!"
-
-# demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-# demo.launch()
-
-
 import gradio as gr
 from datasets import load_dataset
 import numpy as np
 from model2vec import StaticModel
 from reach import Reach
 from tqdm import tqdm
+import difflib
+
+def display_word_differences(x: str, y: str) -> str:
+    diff = difflib.ndiff(x.split(), y.split())
+    return " ".join([word for word in diff if word.startswith(('+', '-'))])
 
 def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024) -> tuple[np.ndarray, dict[int, int]]:
     """
@@ -28,11 +24,11 @@ def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int
         embedding_matrix, 
         threshold=threshold, 
         batch_size=batch_size, 
-        show_progressbar=True
+        show_progressbar=False  # Disable internal progress bar
     )
 
     # Process duplicates
-    for i, similar_items in enumerate(tqdm(results)):
+    for i, similar_items in enumerate(tqdm(results, desc="Processing duplicates")):
         if i not in deduplicated_indices:
             continue  # Skip already marked duplicates
 
@@ -62,11 +58,11 @@ def deduplicate_across_datasets(embedding_matrix_1: np.ndarray, embedding_matrix
         embedding_matrix_2, 
         threshold=threshold, 
         batch_size=batch_size, 
-        show_progressbar=True
+        show_progressbar=False  # Disable internal progress bar
     )
 
     # Process duplicates
-    for i, similar_items in enumerate(tqdm(results)):
+    for i, similar_items in enumerate(tqdm(results, desc="Processing 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 item[1] >= threshold]  # Keep those above the threshold
 
@@ -83,100 +79,144 @@ def perform_deduplication(
     dataset1_split,
     dataset2_name,
     dataset2_split,
+    text_column_name,
     threshold
 ):
     # 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['text'] 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)
-        
-        # 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"
-        result_text += f"**Deduplicated indices:** {deduplicated_indices.tolist()}\n\n"
-        result_text += f"**Duplicate to original mapping:** {duplicate_to_original_mapping}\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['text'] for example in ds1]
-        texts2 = [example['text'] 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)
-        
-        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"
-        result_text += f"**Duplicate indices in {dataset2_name}/{dataset2_split}:** {duplicate_indices_in_ds2}\n\n"
-        result_text += f"**Duplicate to original mapping:** {duplicate_to_original_mapping}\n"
-        
-        return result_text
+
+    with gr.Progress(track_tqdm=True) as progress:
+        if deduplication_type == "Single dataset":
+            # Load the dataset
+            ds = load_dataset(dataset1_name, split=dataset1_split)
+
+            # Extract texts
+            try:
+                texts = [example[text_column_name] for example in ds]
+            except KeyError:
+                return f"Error: Text column '{text_column_name}' not found in dataset."
+
+            # Compute embeddings
+            progress(0.1, desc="Loading model and computing embeddings...")
+            model = StaticModel.from_pretrained("minishlab/M2V_base_output")
+            embedding_matrix = model.encode(texts, show_progressbar=False)
+
+            # Deduplicate
+            progress(0.5, desc="Performing deduplication...")
+            deduplicated_indices, duplicate_to_original_mapping = deduplicate(embedding_matrix, threshold)
+
+            # 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 sample duplicates
+            result_text += "### Sample Duplicate Pairs with Differences:\n\n"
+            num_examples = min(5, num_duplicates)
+            if num_examples > 0:
+                sample_duplicates = list(duplicate_to_original_mapping.items())[:num_examples]
+                for duplicate_idx, original_idx in sample_duplicates:
+                    original_text = texts[original_idx]
+                    duplicate_text = texts[duplicate_idx]
+                    differences = display_word_differences(original_text, duplicate_text)
+                    result_text += f"**Original Text (Index {original_idx}):**\n{original_text}\n\n"
+                    result_text += f"**Duplicate Text (Index {duplicate_idx}):**\n{duplicate_text}\n\n"
+                    result_text += f"**Differences:**\n{differences}\n\n"
+                    result_text += "---\n\n"
+            else:
+                result_text += "No duplicates found.\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
+            try:
+                texts1 = [example[text_column_name] for example in ds1]
+                texts2 = [example[text_column_name] for example in ds2]
+            except KeyError:
+                return f"Error: Text column '{text_column_name}' not found in one of the datasets."
+
+            # Compute embeddings
+            progress(0.1, desc="Computing embeddings for Dataset 1...")
+            model = StaticModel.from_pretrained("minishlab/M2V_base_output")
+            embedding_matrix1 = model.encode(texts1, show_progressbar=False)
+
+            progress(0.5, desc="Computing embeddings for Dataset 2...")
+            embedding_matrix2 = model.encode(texts2, show_progressbar=False)
+
+            # Deduplicate across datasets
+            progress(0.7, desc="Performing cross-dataset deduplication...")
+            duplicate_indices_in_ds2, duplicate_to_original_mapping = deduplicate_across_datasets(embedding_matrix1, embedding_matrix2, threshold)
+
+            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 sample duplicates
+            result_text += "### Sample Duplicate Pairs with Differences:\n\n"
+            num_examples = min(5, num_duplicates)
+            if num_examples > 0:
+                sample_duplicates = list(duplicate_to_original_mapping.items())[:num_examples]
+                for duplicate_idx, original_idx in sample_duplicates:
+                    original_text = texts1[original_idx]
+                    duplicate_text = texts2[duplicate_idx]
+                    differences = display_word_differences(original_text, duplicate_text)
+                    result_text += f"**Original Text in {dataset1_name}/{dataset1_split} (Index {original_idx}):**\n{original_text}\n\n"
+                    result_text += f"**Duplicate Text in {dataset2_name}/{dataset2_split} (Index {duplicate_idx}):**\n{duplicate_text}\n\n"
+                    result_text += f"**Differences:**\n{differences}\n\n"
+                    result_text += "---\n\n"
+            else:
+                result_text += "No duplicates found.\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")
-    
-    dataset2_row = gr.Row(visible=False)
+
+    dataset2_row = gr.Column(visible=False)
     with dataset2_row:
         dataset2_name = gr.Textbox(value="ag_news", label="Dataset 2 Name")
         dataset2_split = gr.Textbox(value="test", label="Dataset 2 Split")
-    
+
+    text_column_name = 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_row
-    def update_visibility(deduplication_type):
-        if deduplication_type == "Cross-dataset":
+    def update_visibility(choice):
+        if choice == "Cross-dataset":
             return {dataset2_row: gr.update(visible=True)}
         else:
             return {dataset2_row: gr.update(visible=False)}
     
     deduplication_type.change(update_visibility, inputs=deduplication_type, outputs=[dataset2_row])
-    
+
     compute_button.click(
         fn=perform_deduplication,
-        inputs=[deduplication_type, dataset1_name, dataset1_split, dataset2_name, dataset2_split, threshold],
+        inputs=[deduplication_type, dataset1_name, dataset1_split, dataset2_name, dataset2_split, text_column_name, threshold],
         outputs=output
     )
-    
+
 demo.launch()