Updates

app.py

@@ -26,15 +26,19 @@ def batch_iterable(iterable, batch_size):
     for i in range(0, len(iterable), batch_size):
         yield iterable[i:i + batch_size]
 
-def log_time(message, start_time=None):
+def log_time(message, start_time=None, logs=None):
     """Helper function to log the start and end times."""
     current_time = time.time()
     if start_time is not None:
         elapsed = current_time - start_time
-        return f"{message} - Took {elapsed:.2f} seconds"
-    return f"{message} - Started"
-
-def compute_embeddings(texts, batch_size, progress, desc="Computing embeddings"):
+        log_message = f"{message} - Took {elapsed:.2f} seconds"
+    else:
+        log_message = f"{message} - Started"
+    
+    if logs is not None:
+        logs.append(log_message)
+
+def compute_embeddings(texts, batch_size, progress, logs, desc="Computing embeddings"):
     embeddings = []
     total_batches = (len(texts) + batch_size - 1) // batch_size
     for i, batch_texts in enumerate(batch_iterable(texts, batch_size)):
@@ -47,10 +51,11 @@ def deduplicate(
     embedding_matrix: np.ndarray,
     threshold: float,
     batch_size: int = 1024,
-    progress=None
+    progress=None,
+    logs=None
 ) -> tuple[np.ndarray, dict[int, int]]:
     # Building the index
-    progress(0, desc="Building search index...")
+    log_time("Building search index", logs=logs)
     reach = Reach(
         vectors=embedding_matrix, items=[str(i) for i in range(len(embedding_matrix))]
     )
@@ -59,7 +64,7 @@ def deduplicate(
     duplicate_to_original_mapping = {}
 
     # Finding nearest neighbors
-    progress(0, desc="Finding nearest neighbors...")
+    log_time("Finding nearest neighbors", logs=logs)
     results = reach.nearest_neighbor_threshold(
         embedding_matrix,
         threshold=threshold,
@@ -69,6 +74,7 @@ def deduplicate(
 
     # Processing duplicates with a progress bar
     total_items = len(embedding_matrix)
+    log_time("Processing duplicates", logs=logs)
     for i, similar_items in enumerate(
         progress.tqdm(results, desc="Processing duplicates", total=total_items)
     ):
@@ -88,8 +94,9 @@ 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 encode_texts(texts, progress=None):
+def encode_texts(texts, progress=None, logs=None):
     embedding_matrix = model.encode(texts, show_progressbar=False)
+    log_time("Encoding texts completed", logs=logs)
     return embedding_matrix
 
 def perform_deduplication(
@@ -103,18 +110,18 @@ def perform_deduplication(
     threshold=default_threshold,
     progress=gr.Progress(track_tqdm=True),
 ):
+    logs = []  # To store log messages
     try:
         # Convert threshold to float
         threshold = float(threshold)
 
         # Initialize status message
-        status = ""
+        log_time("Deduplication started", logs=logs)
 
         if deduplication_type == "Single dataset":
             # Load Dataset 1
             start_time = time.time()
-            status = log_time("Loading Dataset 1")
-            yield status, ""
+            log_time("Loading Dataset 1", logs=logs)
             if (
                 dataset1_name == default_dataset1_name
                 and dataset1_split == default_dataset1_split
@@ -122,34 +129,27 @@ def perform_deduplication(
                 ds = ds_default1
             else:
                 ds = load_dataset(dataset1_name, split=dataset1_split)
-            status = log_time("Loading Dataset 1 completed", start_time)
-            yield status, ""
+            log_time("Loading Dataset 1 completed", start_time=start_time, logs=logs)
 
             # Extract texts
             start_time = time.time()
-            status = log_time("Extracting texts from Dataset 1")
-            yield status, ""
+            log_time("Extracting texts from Dataset 1", logs=logs)
             texts = [example[dataset1_text_column] for example in ds]
-            status = log_time("Extracting texts from Dataset 1 completed", start_time)
-            yield status, ""
+            log_time("Extracting texts from Dataset 1 completed", start_time=start_time, logs=logs)
 
             # Compute embeddings
             start_time = time.time()
-            status = log_time("Computing embeddings for Dataset 1")
-            yield status, ""
-            embedding_matrix = encode_texts(texts, progress=progress)
-            status = log_time("Computing embeddings for Dataset 1 completed", start_time)
-            yield status, ""
+            log_time("Computing embeddings for Dataset 1", logs=logs)
+            embedding_matrix = encode_texts(texts, progress=progress, logs=logs)
+            log_time("Computing embeddings for Dataset 1 completed", start_time=start_time, logs=logs)
 
             # Deduplicate
             start_time = time.time()
-            status = log_time("Deduplicating embeddings")
-            yield status, ""
+            log_time("Deduplicating embeddings", logs=logs)
             deduplicated_indices, duplicate_to_original_mapping = deduplicate(
-                embedding_matrix, threshold, progress=progress
+                embedding_matrix, threshold, progress=progress, logs=logs
             )
-            status = log_time("Deduplication completed", start_time)
-            yield status, ""
+            log_time("Deduplication completed", start_time=start_time, logs=logs)
 
             # Prepare the results
             num_duplicates = len(duplicate_to_original_mapping)
@@ -177,41 +177,12 @@ def perform_deduplication(
             else:
                 result_text += "No duplicates found."
 
-            # Final status
-            status = log_time("Deduplication process finished")
-            yield status, result_text
-
-        elif deduplication_type == "Cross-dataset":
-            # Similar code for cross-dataset deduplication with time logging
-            start_time = time.time()
-            status = log_time("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)
-            status = log_time("Loading Dataset 1 completed", start_time)
-            yield status, ""
-
-            start_time = time.time()
-            status = log_time("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)
-            status = log_time("Loading Dataset 2 completed", start_time)
-            yield status, ""
-
-            # Similar time logging for embedding computations and deduplication steps
+            log_time("Deduplication process finished", logs=logs)
+            full_log = "\n".join(logs)  # Combine all logs into one output
+            yield full_log, result_text
 
     except Exception as e:
+        full_log = "\n".join(logs)  # Combine all logs into one output in case of an error
         yield f"An error occurred: {e}", ""
         raise e
 
@@ -276,6 +247,7 @@ with gr.Blocks(css="#status_output { height: 150px; overflow: auto; }") as demo:
 
 demo.launch()
 
+
 # import gradio as gr
 # from datasets import load_dataset
 # import numpy as np