Updated app with code for deduplication

app.py

@@ -108,24 +108,31 @@ def perform_deduplication(
         # Convert threshold to float
         threshold = float(threshold)
         
+        # Initialize status message
+        status = ""
+
         if deduplication_type == "Single dataset":
             # Load Dataset 1
-            gr.print("Loading Dataset 1...")
+            status = "Loading Dataset 1..."
+            yield status, ""
             if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split:
                 ds = ds_default1
             else:
                 ds = load_dataset(dataset1_name, split=dataset1_split)
             
             # Extract texts
-            gr.print("Extracting texts from Dataset 1...")
+            status = "Extracting texts from Dataset 1..."
+            yield status, ""
             texts = [example[dataset1_text_column] for example in ds]
             
             # Compute embeddings
-            gr.print("Computing embeddings for Dataset 1...")
+            status = "Computing embeddings for Dataset 1..."
+            yield status, ""
             embedding_matrix = model.encode(texts, show_progressbar=True)  # Enable internal progress bar
             
             # Deduplicate
-            gr.print("Deduplicating embeddings...")
+            status = "Deduplicating embeddings..."
+            yield status, ""
             deduplicated_indices, duplicate_to_original_mapping = deduplicate(
                 embedding_matrix, threshold
             )
@@ -154,41 +161,50 @@ def perform_deduplication(
             else:
                 result_text += "No duplicates found."
             
-            return result_text
+            # Final status
+            status = "Deduplication completed."
+            yield status, result_text
                 
         elif deduplication_type == "Cross-dataset":
             # Load Dataset 1
-            gr.print("Loading 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
-            gr.print("Loading 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
-            gr.print("Extracting 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
-            gr.print("Extracting 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
-            gr.print("Computing embeddings for Dataset 1...")
+            status = "Computing embeddings for Dataset 1..."
+            yield status, ""
             embedding_matrix1 = model.encode(texts1, show_progressbar=True)
             
             # Compute embeddings for Dataset 2
-            gr.print("Computing embeddings for Dataset 2...")
+            status = "Computing embeddings for Dataset 2..."
+            yield status, ""
             embedding_matrix2 = model.encode(texts2, show_progressbar=True)
             
             # Deduplicate across datasets
-            gr.print("Deduplicating embeddings 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
             )
@@ -217,7 +233,9 @@ def perform_deduplication(
             else:
                 result_text += "No duplicates found."
             
-            return result_text
+            # Final status
+            status = "Deduplication completed."
+            yield status, result_text
 
     finally:
         # Restore original tqdm
@@ -257,7 +275,8 @@ with gr.Blocks() as demo:
 
     compute_button = gr.Button("Compute")
 
-    output = gr.Markdown()
+    status_output = gr.Markdown()
+    result_output = gr.Markdown()
 
     # Function to update the visibility of dataset2_inputs
     def update_visibility(deduplication_type_value):
@@ -284,9 +303,9 @@ with gr.Blocks() as demo:
             dataset2_text_column,
             threshold
         ],
-        outputs=output
+        outputs=[status_output, result_output]
     )
-    
+        
 demo.launch()