Updates

app.py

@@ -20,6 +20,35 @@ 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)
 
+from tqdm import tqdm as original_tqdm
+# Patch tqdm to use Gradio's progress bar
+def patch_tqdm_for_gradio(progress):
+    class GradioTqdm(original_tqdm):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.progress = progress
+            self.total_batches = kwargs.get('total', len(args[0])) if len(args) > 0 else 1
+        
+        def update(self, n=1):
+            super().update(n)
+            self.progress(self.n / self.total_batches)
+
+    return GradioTqdm
+# Function to patch the original encode function with our Gradio tqdm
+def original_encode_with_tqdm(original_encode_func, patched_tqdm):
+    def new_encode(*args, **kwargs):
+        # Replace tqdm with our patched version
+        original_tqdm_backup = original_tqdm
+        try:
+            # Patch the `tqdm` within encode
+            globals()['tqdm'] = patched_tqdm
+            return original_encode_func(*args, **kwargs)
+        finally:
+            # Restore original tqdm after calling encode
+            globals()['tqdm'] = original_tqdm_backup
+
+    return new_encode
+
 def batch_iterable(iterable, batch_size):
     """Helper function to create batches from an iterable."""
     for i in range(0, len(iterable), batch_size):
@@ -114,15 +143,18 @@ def perform_deduplication(
             yield status, ""
             texts = [example[dataset1_text_column] for example in ds]
 
+            patched_tqdm = patch_tqdm_for_gradio(progress)
+            model.encode = original_encode_with_tqdm(model.encode, patched_tqdm)
             # Compute embeddings
             status = "Computing embeddings for Dataset 1..."
             yield status, ""
-            embedding_matrix = compute_embeddings(
-                texts,
-                batch_size=64,
-                progress=progress,
-                desc="Computing embeddings for Dataset 1",
-            )
+            embedding_matrix = model.encode(texts, show_progressbar=True)
+            # embedding_matrix = compute_embeddings(
+            #     texts,
+            #     batch_size=64,
+            #     progress=progress,
+            #     desc="Computing embeddings for Dataset 1",
+            # )
 
             # Deduplicate
             status = "Deduplicating embeddings..."