Updates

app.py

@@ -10,7 +10,7 @@ from contextlib import contextmanager
 # Load the model at startup
 model = StaticModel.from_pretrained("minishlab/M2V_base_output")
 
-# Update default dataset to 'sst2' and set default threshold to 0.9
+# Default dataset parameters
 default_dataset1_name = "sst2"
 default_dataset1_split = "train"
 default_dataset2_name = "sst2"
@@ -47,7 +47,6 @@ def deduplicate(
     batch_size: int = 1024,
     progress=None
 ) -> tuple[np.ndarray, dict[int, int]]:
-    # Existing deduplication code remains unchanged
     # Building the index
     progress(0, desc="Building search index...")
     reach = Reach(
@@ -171,18 +170,137 @@ def perform_deduplication(
 
         elif deduplication_type == "Cross-dataset":
             # Similar code for cross-dataset deduplication
-            # Implement similar logic as above for cross-dataset
-            pass
+            # Load Dataset 1
+            status = "Loading Dataset 1..."
+            yield status, ""
+            if (
+                dataset1_name == default_dataset1_name
+                and dataset1_split == default_dataset1_split
+            ):
+                ds1 = ds_default1
+            else:
+                ds1 = load_dataset(dataset1_name, split=dataset1_split)
+
+            # Load Dataset 2
+            status = "Loading Dataset 2..."
+            yield status, ""
+            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 from Dataset 1
+            status = "Extracting texts from Dataset 1..."
+            yield status, ""
+            texts1 = [example[dataset1_text_column] for example in ds1]
+
+            # Extract texts from Dataset 2
+            status = "Extracting texts from Dataset 2..."
+            yield status, ""
+            texts2 = [example[dataset2_text_column] for example in ds2]
+
+            # Compute embeddings for Dataset 1
+            status = "Computing embeddings for Dataset 1..."
+            yield status, ""
+            embedding_matrix1 = compute_embeddings(
+                texts1,
+                batch_size=64,
+                progress=progress,
+                desc="Computing embeddings for Dataset 1",
+            )
+
+            # Compute embeddings for Dataset 2
+            status = "Computing embeddings for Dataset 2..."
+            yield status, ""
+            embedding_matrix2 = compute_embeddings(
+                texts2,
+                batch_size=64,
+                progress=progress,
+                desc="Computing embeddings for Dataset 2",
+            )
+
+            # Deduplicate across datasets
+            status = "Deduplicating embeddings across datasets..."
+            yield status, ""
+            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
+            if num_duplicates > 0:
+                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"
+            else:
+                result_text += "No duplicates found."
+
+            # Final status
+            status = "Deduplication completed."
+            yield status, result_text
 
     except Exception as e:
         yield f"An error occurred: {e}", ""
         raise e
 
-with gr.Blocks() as demo:
-    # Replace 'gr.Markdown' with 'gr.Textbox' for 'status_output'
-    status_output = gr.Textbox().style(height=150)
-    result_output = gr.Markdown()
+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]]:
+    # Building the index from Dataset 1
+    progress(0, desc="Building search index from Dataset 1...")
+    reach = Reach(
+        vectors=embedding_matrix_1, items=[str(i) for i in range(len(embedding_matrix_1))]
+    )
 
+    duplicate_indices_in_test = []
+    duplicate_to_original_mapping = {}
+
+    # Finding nearest neighbors between datasets
+    progress(0, desc="Finding nearest neighbors between datasets...")
+    results = reach.nearest_neighbor_threshold(
+        embedding_matrix_2,
+        threshold=threshold,
+        batch_size=batch_size,
+        show_progressbar=False,  # Disable internal progress bar
+    )
+
+    total_items = len(embedding_matrix_2)
+    # Processing duplicates with a progress bar
+    for i, similar_items in enumerate(
+        progress.tqdm(results, desc="Processing duplicates across datasets", total=total_items)
+    ):
+        similar_indices = [int(item[0]) for item in similar_items if item[1] >= threshold]
+
+        if similar_indices:
+            duplicate_indices_in_test.append(i)
+            duplicate_to_original_mapping[i] = similar_indices[0]
+
+    return duplicate_indices_in_test, duplicate_to_original_mapping
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Semantic Deduplication")
 
     deduplication_type = gr.Radio(
         choices=["Single dataset", "Cross-dataset"],
@@ -209,8 +327,8 @@ with gr.Blocks() as demo:
 
     compute_button = gr.Button("Compute")
 
-    # Adjust the height of the status_output component
-    status_output = gr.Markdown().style(height=150)
+    # Use 'lines' parameter to set the height
+    status_output = gr.Textbox(lines=10, label="Status")
     result_output = gr.Markdown()
 
     # Function to update the visibility of dataset2_inputs