Updated app with code for deduplication

app.py

@@ -26,27 +26,24 @@ def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int
     """
     Deduplicate embeddings and return the deduplicated indices and a mapping of removed indices to their corresponding original indices.
     """
-    # Informative progress bar for building the index
-    progress.tqdm.write("Building search index...")
-    with progress.tqdm(total=1, desc="Building index") as p:
+    # Building the index with a progress bar
+    with progress.tqdm(total=1, desc="Building search index") as p:
         reach = Reach(vectors=embedding_matrix, items=[str(i) for i in range(len(embedding_matrix))])
         p.update(1)
 
     deduplicated_indices = set(range(len(embedding_matrix)))
     duplicate_to_original_mapping = {}
 
-    # Informative progress bar for nearest neighbor search
-    progress.tqdm.write("Finding nearest neighbors...")
+    # Finding nearest neighbors
     results = reach.nearest_neighbor_threshold(
         embedding_matrix,
         threshold=threshold,
         batch_size=batch_size,
-        show_progressbar=False  # Disable internal progress bar
+        show_progressbar=True  # Allow internal progress bar
     )
 
-    total_items = len(embedding_matrix)
     # Processing duplicates with a progress bar
-    progress.tqdm.write("Processing duplicates...")
+    total_items = len(embedding_matrix)
     for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates", total=total_items)):
         if i not in deduplicated_indices:
             continue
@@ -64,27 +61,24 @@ def deduplicate_across_datasets(embedding_matrix_1: np.ndarray, embedding_matrix
     """
     Deduplicate embeddings across two datasets and return the indices of duplicates between them.
     """
-    # Informative progress bar for building the index
-    progress.tqdm.write("Building search index from Dataset 1...")
-    with progress.tqdm(total=1, desc="Building index for Dataset 1") as p:
+    # Building the index from Dataset 1
+    with progress.tqdm(total=1, desc="Building search index from Dataset 1") as p:
         reach = Reach(vectors=embedding_matrix_1, items=[str(i) for i in range(len(embedding_matrix_1))])
         p.update(1)
 
     duplicate_indices_in_test = []
     duplicate_to_original_mapping = {}
 
-    # Informative progress bar for nearest neighbor search
-    progress.tqdm.write("Finding nearest neighbors between datasets...")
+    # 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
+        show_progressbar=True  # Allow internal progress bar
     )
 
     total_items = len(embedding_matrix_2)
     # Processing duplicates with a progress bar
-    progress.tqdm.write("Processing duplicates across datasets...")
     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]
 
@@ -128,14 +122,11 @@ def perform_deduplication(
             else:
                 ds = load_dataset(dataset1_name, split=dataset1_split)
             
-            # Extract texts with progress bar
-            progress.tqdm.write("Extracting texts from Dataset 1...")
-            texts = [example[dataset1_text_column] for example in progress.tqdm(ds, desc="Extracting texts", total=len(ds))]
+            # Extract texts
+            texts = [example[dataset1_text_column] for example in ds]
             
-            # Compute embeddings with progress bar
-            progress.tqdm.write("Computing embeddings for Dataset 1...")
-            embedding_matrix = model.encode(texts, show_progressbar=False)  # Disable internal progress bar
-            embedding_matrix = progress.tqdm(embedding_matrix, desc="Computing embeddings", total=len(texts))
+            # Compute embeddings
+            embedding_matrix = model.encode(texts, show_progressbar=True)  # Enable internal progress bar
             
             # Deduplicate
             result_text = deduplicate_and_prepare_results_single(
@@ -158,22 +149,16 @@ def perform_deduplication(
                 ds2 = load_dataset(dataset2_name, split=dataset2_split)
             
             # Extract texts from Dataset 1
-            progress.tqdm.write("Extracting texts from Dataset 1...")
-            texts1 = [example[dataset1_text_column] for example in progress.tqdm(ds1, desc="Extracting texts from Dataset 1", total=len(ds1))]
+            texts1 = [example[dataset1_text_column] for example in ds1]
             
             # Extract texts from Dataset 2
-            progress.tqdm.write("Extracting texts from Dataset 2...")
-            texts2 = [example[dataset2_text_column] for example in progress.tqdm(ds2, desc="Extracting texts from Dataset 2", total=len(ds2))]
+            texts2 = [example[dataset2_text_column] for example in ds2]
             
             # Compute embeddings for Dataset 1
-            progress.tqdm.write("Computing embeddings for Dataset 1...")
-            embedding_matrix1 = model.encode(texts1, show_progressbar=False)
-            embedding_matrix1 = progress.tqdm(embedding_matrix1, desc="Computing embeddings for Dataset 1", total=len(texts1))
+            embedding_matrix1 = model.encode(texts1, show_progressbar=True)
             
             # Compute embeddings for Dataset 2
-            progress.tqdm.write("Computing embeddings for Dataset 2...")
-            embedding_matrix2 = model.encode(texts2, show_progressbar=False)
-            embedding_matrix2 = progress.tqdm(embedding_matrix2, desc="Computing embeddings for Dataset 2", total=len(texts2))
+            embedding_matrix2 = model.encode(texts2, show_progressbar=True)
             
             # Deduplicate across datasets
             result_text = deduplicate_and_prepare_results_cross(
@@ -322,6 +307,7 @@ with gr.Blocks() as demo:
 demo.launch()
 
 
+
 # import gradio as gr
 # from datasets import load_dataset
 # import numpy as np