Updated app with code for deduplication

app.py

@@ -4,11 +4,7 @@ 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(('+', '-'))])
+from difflib import ndiff
 
 def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024) -> tuple[np.ndarray, dict[int, int]]:
     """
@@ -24,11 +20,11 @@ def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int
         embedding_matrix, 
         threshold=threshold, 
         batch_size=batch_size, 
-        show_progressbar=False  # Disable internal progress bar
+        show_progressbar=True
     )
 
     # Process duplicates
-    for i, similar_items in enumerate(tqdm(results, desc="Processing duplicates")):
+    for i, similar_items in enumerate(tqdm(results)):
         if i not in deduplicated_indices:
             continue  # Skip already marked duplicates
 
@@ -58,11 +54,11 @@ def deduplicate_across_datasets(embedding_matrix_1: np.ndarray, embedding_matrix
         embedding_matrix_2, 
         threshold=threshold, 
         batch_size=batch_size, 
-        show_progressbar=False  # Disable internal progress bar
+        show_progressbar=True
     )
 
     # Process duplicates
-    for i, similar_items in enumerate(tqdm(results, desc="Processing duplicates")):
+    for i, similar_items in enumerate(tqdm(results)):
         # 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
 
@@ -73,150 +69,154 @@ def deduplicate_across_datasets(embedding_matrix_1: np.ndarray, embedding_matrix
 
     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([f"{word}" for word in diff if word.startswith(('+', '-'))])
+
 def perform_deduplication(
     deduplication_type,
     dataset1_name,
     dataset1_split,
+    dataset1_text_column,
     dataset2_name,
     dataset2_split,
-    text_column_name,
+    dataset2_text_column,
     threshold
 ):
     # Convert threshold to float
     threshold = float(threshold)
-
-    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...")
+    
+    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
+        with gr.Progress(track_tqdm=True):
             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...")
+        
+        # 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)
+        examples_shown = 0
+        for duplicate_idx, original_idx in duplicate_to_original_mapping.items():
+            if examples_shown >= num_examples:
+                break
+            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"
+            examples_shown += 1
+        
+        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
+        with gr.Progress(track_tqdm=True):
             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
+        
+        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)
+        examples_shown = 0
+        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"
+            examples_shown += 1
+        
+        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.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")
-
+    
+    with gr.Tab("Dataset 1"):
+        with gr.Row():
+            dataset1_name = gr.Textbox(value="ag_news", label="Dataset Name")
+            dataset1_split = gr.Textbox(value="train", label="Split")
+            dataset1_text_column = gr.Textbox(value="text", label="Text Column Name")
+    
+    dataset2_tab = gr.Tab("Dataset 2", visible=False)
+    with dataset2_tab:
+        with gr.Row():
+            dataset2_name = gr.Textbox(value="ag_news", label="Dataset Name")
+            dataset2_split = gr.Textbox(value="test", label="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_row
-    def update_visibility(choice):
-        if choice == "Cross-dataset":
-            return {dataset2_row: gr.update(visible=True)}
+    
+    # Function to update the visibility of dataset2_tab
+    def update_visibility(deduplication_type):
+        if deduplication_type == "Cross-dataset":
+            return {dataset2_tab: gr.update(visible=True)}
         else:
-            return {dataset2_row: gr.update(visible=False)}
+            return {dataset2_tab: gr.update(visible=False)}
+    
+    deduplication_type.change(update_visibility, inputs=deduplication_type, outputs=[dataset2_tab])
     
-    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, text_column_name, threshold],
+        inputs=[
+            deduplication_type, 
+            dataset1_name, 
+            dataset1_split, 
+            dataset1_text_column,
+            dataset2_name, 
+            dataset2_split, 
+            dataset2_text_column,
+            threshold
+        ],
         outputs=output
     )
-
+    
 demo.launch()