Updates

app.py

@@ -4,7 +4,7 @@ import numpy as np
 from model2vec import StaticModel
 from reach import Reach
 from difflib import ndiff
-import tqdm
+import asyncio
 
 # Load the model at startup
 model = StaticModel.from_pretrained("minishlab/M2V_base_output")
@@ -30,7 +30,54 @@ def display_word_differences(x: str, y: str) -> str:
     diff = ndiff(x.split(), y.split())
     return " ".join([word for word in diff if word.startswith(('+', '-'))])
 
-def perform_deduplication(
+async def compute_embeddings_async(texts, batch_size, progress, desc):
+    embeddings = []
+    total_batches = (len(texts) + batch_size - 1) // batch_size
+    for i, batch_texts in enumerate(batch_iterable(texts, batch_size)):
+        batch_embeddings = await asyncio.to_thread(model.encode, batch_texts, show_progressbar=False)
+        embeddings.append(batch_embeddings)
+        progress((i + 1) / total_batches, desc=desc)
+        await asyncio.sleep(0)
+    embedding_matrix = np.concatenate(embeddings, axis=0)
+    return embedding_matrix
+
+async def deduplicate_async(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024, progress=None) -> tuple[np.ndarray, dict[int, int]]:
+    """
+    Deduplicate embeddings asynchronously.
+    """
+    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 = {}
+
+    progress(0, desc="Finding nearest neighbors...")
+    results = await asyncio.to_thread(reach.nearest_neighbor_threshold,
+                                      embedding_matrix,
+                                      threshold=threshold,
+                                      batch_size=batch_size,
+                                      show_progressbar=False)
+
+    total_items = len(embedding_matrix)
+    for i, similar_items in enumerate(results):
+        if i not in deduplicated_indices:
+            continue
+
+        similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i]
+
+        for sim_idx in similar_indices:
+            if sim_idx in deduplicated_indices:
+                deduplicated_indices.remove(sim_idx)
+                duplicate_to_original_mapping[sim_idx] = i
+
+        if i % 100 == 0:
+            progress(i / total_items, desc="Processing duplicates")
+            await asyncio.sleep(0)
+
+    progress(1, desc="Processing duplicates")
+    return np.array(list(deduplicated_indices)), duplicate_to_original_mapping
+
+async def perform_deduplication(
     deduplication_type,
     dataset1_name,
     dataset1_split,
@@ -65,19 +112,12 @@ def perform_deduplication(
             # Compute embeddings
             status = "Computing embeddings for Dataset 1..."
             yield status, ""
-            embeddings = []
-            batch_size = 64
-            total_batches = (len(texts) + batch_size - 1) // batch_size
-            # Use progress.tqdm without yielding inside the loop
-            for batch_texts in progress.tqdm(batch_iterable(texts, batch_size), desc="Computing embeddings for Dataset 1", total=total_batches):
-                batch_embeddings = model.encode(batch_texts, show_progressbar=False)
-                embeddings.append(batch_embeddings)
-            embedding_matrix = np.concatenate(embeddings, axis=0)
+            embedding_matrix = await compute_embeddings_async(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(
+            deduplicated_indices, duplicate_to_original_mapping = await deduplicate_async(
                 embedding_matrix, threshold, progress=progress
             )
 
@@ -110,6 +150,7 @@ def perform_deduplication(
             yield status, result_text
 
         elif deduplication_type == "Cross-dataset":
+            # Similar code for cross-dataset deduplication, using async functions
             # Load Dataset 1
             status = "Loading Dataset 1..."
             yield status, ""
@@ -139,28 +180,17 @@ def perform_deduplication(
             # Compute embeddings for Dataset 1
             status = "Computing embeddings for Dataset 1..."
             yield status, ""
-            embeddings1 = []
-            batch_size = 64
-            total_batches1 = (len(texts1) + batch_size - 1) // batch_size
-            for batch_texts in progress.tqdm(batch_iterable(texts1, batch_size), desc="Computing embeddings for Dataset 1", total=total_batches1):
-                batch_embeddings = model.encode(batch_texts, show_progressbar=False)
-                embeddings1.append(batch_embeddings)
-            embedding_matrix1 = np.concatenate(embeddings1, axis=0)
+            embedding_matrix1 = await compute_embeddings_async(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, ""
-            embeddings2 = []
-            total_batches2 = (len(texts2) + batch_size - 1) // batch_size
-            for batch_texts in progress.tqdm(batch_iterable(texts2, batch_size), desc="Computing embeddings for Dataset 2", total=total_batches2):
-                batch_embeddings = model.encode(batch_texts, show_progressbar=False)
-                embeddings2.append(batch_embeddings)
-            embedding_matrix2 = np.concatenate(embeddings2, axis=0)
+            embedding_matrix2 = await compute_embeddings_async(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(
+            duplicate_indices_in_ds2, duplicate_to_original_mapping = await deduplicate_across_datasets_async(
                 embedding_matrix1, embedding_matrix2, threshold, progress=progress
             )
 
@@ -196,132 +226,30 @@ def perform_deduplication(
         yield f"An error occurred: {e}", ""
         raise e
 
-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=False  # Disable internal progress bar
-    )
-
-    # 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
-
-        similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i]
-
-        for sim_idx in similar_indices:
-            if sim_idx in deduplicated_indices:
-                deduplicated_indices.remove(sim_idx)
-                duplicate_to_original_mapping[sim_idx] = i
-
-    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, progress=None) -> tuple[list[int], dict[int, int]]:
+async def deduplicate_across_datasets_async(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.
+    Deduplicate embeddings across two datasets asynchronously.
     """
-    # 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
-    )
+    results = await asyncio.to_thread(reach.nearest_neighbor_threshold,
+                                      embedding_matrix_2,
+                                      threshold=threshold,
+                                      batch_size=batch_size,
+                                      show_progressbar=False)
 
     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)):
+    for i, similar_items in enumerate(results):
         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"],
-        label="Deduplication Type",
-        value="Single dataset"
-    )
-
-    with gr.Row():
-        dataset1_name = gr.Textbox(value=default_dataset1_name, label="Dataset 1 Name")
-        dataset1_split = gr.Textbox(value=default_dataset1_split, label="Dataset 1 Split")
-        dataset1_text_column = gr.Textbox(value=default_text_column, label="Text Column Name")
-
-    dataset2_inputs = gr.Column(visible=False)
-    with dataset2_inputs:
-        gr.Markdown("### Dataset 2")
-        with gr.Row():
-            dataset2_name = gr.Textbox(value=default_dataset2_name, label="Dataset 2 Name")
-            dataset2_split = gr.Textbox(value=default_dataset2_split, label="Dataset 2 Split")
-            dataset2_text_column = gr.Textbox(value=default_text_column, label="Text Column Name")
-
-    threshold = gr.Slider(
-        minimum=0.0,
-        maximum=1.0,
-        value=default_threshold,
-        label="Similarity Threshold"
-    )
-
-    compute_button = gr.Button("Compute")
-
-    status_output = gr.Markdown()
-    result_output = gr.Markdown()
-
-    # Function to update the visibility of dataset2_inputs
-    def update_visibility(deduplication_type_value):
-        if deduplication_type_value == "Cross-dataset":
-            return gr.update(visible=True)
-        else:
-            return gr.update(visible=False)
-
-    deduplication_type.change(
-        update_visibility,
-        inputs=deduplication_type,
-        outputs=dataset2_inputs
-    )
-
-    compute_button.click(
-        fn=perform_deduplication,
-        inputs=[
-            deduplication_type,
-            dataset1_name,
-            dataset1_split,
-            dataset1_text_column,
-            dataset2_name,
-            dataset2_split,
-            dataset2_text_column,
-            threshold
-        ],
-        outputs=[status_output, result_output]
-    )
-
-demo.launch()
+        if i % 100 == 0:
+            progress(i / total_items, desc="Processing duplicates across datasets")