Updated app with code for deduplication

app.py

@@ -4,9 +4,7 @@ import numpy as np
 from model2vec import StaticModel
 from reach import Reach
 from difflib import ndiff
-import sys
 import tqdm
-from tqdm.utils import format_interval, format_num, format_sizeof
 
 # Load the model at startup
 model = StaticModel.from_pretrained("minishlab/M2V_base_output")
@@ -23,26 +21,41 @@ default_threshold = 0.9
 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) -> tuple[np.ndarray, dict[int, int]]:
+def batch_iterable(iterable, batch_size):
+    """Helper function to create batches from an iterable."""
+    for i in range(0, len(iterable), batch_size):
+        yield iterable[i:i + batch_size]
+
+def compute_embeddings(texts, batch_size, progress, desc="Computing embeddings"):
+    embeddings = []
+    for batch in progress.tqdm(batch_iterable(texts, batch_size), total=(len(texts) + batch_size - 1) // batch_size, desc=desc):
+        batch_embeddings = model.encode(batch, show_progressbar=False)
+        embeddings.append(batch_embeddings)
+    return np.concatenate(embeddings, axis=0)
+
+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.
     """
     # Building the index
+    progress(0, desc="Building search index...")
     reach = Reach(vectors=embedding_matrix, items=[str(i) for i in range(len(embedding_matrix))])
 
     deduplicated_indices = set(range(len(embedding_matrix)))
     duplicate_to_original_mapping = {}
 
     # Finding nearest neighbors
+    progress(0, desc="Finding nearest neighbors...")
     results = reach.nearest_neighbor_threshold(
         embedding_matrix,
         threshold=threshold,
         batch_size=batch_size,
-        show_progressbar=True  # Allow internal progress bar
+        show_progressbar=False  # Disable internal progress bar
     )
 
-    # Processing duplicates
-    for i, similar_items in enumerate(results):
+    # Processing duplicates with a progress bar
+    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
 
@@ -55,26 +68,29 @@ def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int
 
     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) -> 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.
     """
     # 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=True  # Allow internal progress bar
+        show_progressbar=False  # Disable internal progress bar
     )
 
-    # Processing duplicates
-    for i, similar_items in enumerate(results):
+    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:
@@ -98,31 +114,6 @@ def perform_deduplication(
     threshold=default_threshold,
     progress=gr.Progress(track_tqdm=True)
 ):
-    # Custom tqdm class that wraps progress.tqdm and includes module-level attributes
-    class TqdmWrapper(tqdm.std.tqdm):
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs)
-
-    # Copy module-level attributes from original tqdm module
-    TqdmWrapper.format_interval = staticmethod(format_interval)
-    TqdmWrapper.format_num = staticmethod(format_num)
-    TqdmWrapper.format_sizeof = staticmethod(format_sizeof)
-
-    # Monkey-patch tqdm.tqdm with our wrapper
-    original_tqdm_tqdm = tqdm.tqdm
-    tqdm.tqdm = progress.tqdm
-
-    # Monkey-patch model2vec's tqdm reference if needed
-    import model2vec.model
-    if hasattr(model2vec.model, 'tqdm'):
-        original_model2vec_tqdm = model2vec.model.tqdm
-        model2vec.model.tqdm = TqdmWrapper
-
-    # Monkey-patch reach's tqdm reference if needed
-    if hasattr(Reach, 'tqdm'):
-        original_reach_tqdm = Reach.tqdm
-        Reach.tqdm = TqdmWrapper
-
     try:
         # Convert threshold to float
         threshold = float(threshold)
@@ -147,13 +138,13 @@ def perform_deduplication(
             # Compute embeddings
             status = "Computing embeddings for Dataset 1..."
             yield status, ""
-            embedding_matrix = model.encode(texts, show_progressbar=True)  # Enable internal progress bar
+            embedding_matrix = compute_embeddings(texts, batch_size=64, progress=progress, desc="Computing embeddings for Dataset 1")
 
             # Deduplicate
             status = "Deduplicating embeddings..."
             yield status, ""
             deduplicated_indices, duplicate_to_original_mapping = deduplicate(
-                embedding_matrix, threshold
+                embedding_matrix, threshold, progress=progress
             )
 
             # Prepare the results
@@ -214,18 +205,18 @@ def perform_deduplication(
             # Compute embeddings for Dataset 1
             status = "Computing embeddings for Dataset 1..."
             yield status, ""
-            embedding_matrix1 = model.encode(texts1, show_progressbar=True)
+            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 = model.encode(texts2, show_progressbar=True)
+            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
+                embedding_matrix1, embedding_matrix2, threshold, progress=progress
             )
 
             num_duplicates = len(duplicate_indices_in_ds2)
@@ -256,13 +247,9 @@ def perform_deduplication(
             status = "Deduplication completed."
             yield status, result_text
 
-    finally:
-        # Restore original tqdm functions
-        tqdm.tqdm = original_tqdm_tqdm
-        if hasattr(model2vec.model, 'tqdm'):
-            model2vec.model.tqdm = original_model2vec_tqdm
-        if hasattr(Reach, 'tqdm'):
-            Reach.tqdm = original_reach_tqdm
+    except Exception as e:
+        yield f"An error occurred: {e}", ""
+        raise e
 
 with gr.Blocks() as demo:
     gr.Markdown("# Semantic Deduplication")
@@ -330,6 +317,7 @@ demo.launch()
 
 
 
+
 # import gradio as gr
 # from datasets import load_dataset
 # import numpy as np