Updated code with SemHash

app.py

@@ -1,88 +1,55 @@
 import gradio as gr
 from datasets import load_dataset
-import numpy as np
-from model2vec import StaticModel
-from reach import Reach
 from difflib import ndiff
 
-# Load the model
-model = StaticModel.from_pretrained("minishlab/potion-base-8M")
+from semhash import SemHash
+from semhash.datamodels import DeduplicationResult
+
+from model2vec import StaticModel
 
 # Default parameters
 default_dataset_name = "ag_news"
-default_dataset1_split = "train"  # Default for the first dataset is "train"
-default_dataset2_split = "test"   # Default for the second dataset is "test"
+default_dataset1_split = "train"
+default_dataset2_split = "test"
 default_text_column = "text"
 default_threshold = 0.9
 
-def deduplicate_embeddings(
-    embeddings_a: np.ndarray,
-    embeddings_b: np.ndarray = None,
-    threshold: float = 0.9,
-    batch_size: int = 1024,
-    progress=None
-) -> tuple[np.ndarray, dict[int, int]]:
-    """
-    Deduplicate embeddings within one dataset or across two datasets.
-
-    :param embeddings_a: Embeddings of Dataset 1.
-    :param embeddings_b: Optional, embeddings of Dataset 2.
-    :param threshold: Similarity threshold for deduplication.
-    :param batch_size: Batch size for similarity computation.
-    :param progress: Gradio progress tracker for feedback.
-    :return: Deduplicated indices and a mapping of removed indices to their original counterparts.
-    """
-    if embeddings_b is None:
-        reach = Reach(vectors=embeddings_a, items=[str(i) for i in range(len(embeddings_a))])
-        duplicate_to_original = {}
-        results = reach.nearest_neighbor_threshold(
-            embeddings_a, threshold=threshold, batch_size=batch_size, show_progressbar=False
-        )
-        for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates", total=len(embeddings_a))):
-            for sim_idx, _ in similar_items:
-                sim_idx = int(sim_idx)
-                if sim_idx != i and sim_idx not in duplicate_to_original:
-                    duplicate_to_original[sim_idx] = i
-        deduplicated_indices = set(range(len(embeddings_a))) - set(duplicate_to_original.keys())
-        return deduplicated_indices, duplicate_to_original
-    else:
-        reach = Reach(vectors=embeddings_a, items=[str(i) for i in range(len(embeddings_a))])
-        duplicate_indices_in_b = []
-        duplicate_to_original = {}
-        results = reach.nearest_neighbor_threshold(
-            embeddings_b, threshold=threshold, batch_size=batch_size, show_progressbar=False
-        )
-        for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates", total=len(embeddings_b))):
-            if similar_items:
-                duplicate_indices_in_b.append(i)
-                duplicate_to_original[i] = int(similar_items[0][0])
-        return duplicate_indices_in_b, duplicate_to_original
+# Load the model to use
+model = StaticModel.from_pretrained("minishlab/potion-base-8M")
+
 
 def display_word_differences(x: str, y: str) -> str:
     """
     Display the word-level differences between two texts, formatted to avoid
     misinterpretation of Markdown syntax.
-
-    :param x: First text.
-    :param y: Second text.
-    :return: A string showing word-level differences, wrapped in a code block.
     """
     diff = ndiff(x.split(), y.split())
     formatted_diff = "\n".join(word for word in diff if word.startswith(("+", "-")))
     return f"```\n{formatted_diff}\n```"
 
-def load_dataset_texts(dataset_name: str, dataset_split: str, text_column: str) -> list[str]:
-    """
-    Load texts from a specified dataset and split.
 
-    :param dataset_name: Name of the dataset.
-    :param dataset_split: Split of the dataset (e.g., 'train', 'validation', 'test').
-    :param text_column: Name of the text column.
-    :return: A list of texts from the dataset.
-    """
+def load_dataset_texts(dataset_name: str, dataset_split: str, text_column: str) -> list[str]:
+    """Load texts from a specified dataset split."""
     ds = load_dataset(dataset_name, split=dataset_split)
     return [example[text_column] for example in ds]
 
+
+def deduplicate_single_dataset(texts: list[str], threshold: float) -> DeduplicationResult:
+    """Deduplicate within a single dataset using SemHash, treating each text as a raw string record."""
+    # Build a SemHash index from the raw texts
+    semhash = SemHash.from_records(records=texts, model=model)
+    # Deduplicate the entire dataset
+    return semhash.self_deduplicate(threshold=threshold)
+
+
+def deduplicate_two_datasets(texts1: list[str], texts2: list[str], threshold: float) -> DeduplicationResult:
+    """Deduplicate dataset2 against dataset1, both as raw strings, using SemHash."""
+    # Build SemHash index on dataset1
+    semhash = SemHash.from_records(records=texts1, model=model)
+    # Deduplicate texts2 against dataset1
+    return semhash.deduplicate(records=texts2, threshold=threshold)
+
+
 def perform_deduplication(
     deduplication_type: str,
     dataset1_name: str,
@@ -95,93 +62,107 @@ def perform_deduplication(
     progress: gr.Progress = gr.Progress(track_tqdm=True)
 ):
     """
-    Perform deduplication on one or two datasets based on the deduplication type.
-
-    :param deduplication_type: 'Single dataset' or 'Cross-dataset'.
-    :param dataset1_name: Name of the first dataset.
-    :param dataset1_split: Split of the first dataset.
-    :param dataset1_text_column: Text column of the first dataset.
-    :param dataset2_name: Optional, name of the second dataset (for cross-dataset deduplication).
-    :param dataset2_split: Optional, split of the second dataset.
-    :param dataset2_text_column: Optional, text column of the second dataset.
-    :param threshold: Similarity threshold for deduplication.
-    :param progress: Gradio progress tracker.
-    :return: Status updates and result text for the Gradio interface.
+    Perform deduplication on one or two datasets using SemHash. This function
+    streams status updates to Gradio for user feedback.
     """
     try:
         threshold = float(threshold)
 
-        # Load and process Dataset 1
+        # Load Dataset 1
         yield "Loading Dataset 1...", ""
         texts1 = load_dataset_texts(dataset1_name, dataset1_split, dataset1_text_column)
-        yield "Computing embeddings for Dataset 1...", ""
-        embeddings1 = model.encode(texts1, show_progressbar=True)
 
         if deduplication_type == "Single dataset":
-            # Deduplicate within Dataset 1
-            yield "Deduplicating within Dataset 1...", ""
-            deduplicated_indices, duplicate_mapping = deduplicate_embeddings(
-                embeddings1, threshold=threshold, progress=progress
-            )
-
-            num_duplicates = len(duplicate_mapping)
+            # Single-dataset deduplication
+            yield "Deduplicating within Dataset 1 (SemHash)...", ""
+            result = deduplicate_single_dataset(texts1, threshold=threshold)
+
+            # Sort all duplicates in descending order of their highest score
+            for duprec in result.duplicates:
+                duprec.duplicates.sort(key=lambda x: x[1], reverse=True)
+
+            # Summarize results
+            num_duplicates = len(result.duplicates)
+            deduplicated_count = len(result.deduplicated)
+            total_docs = len(texts1)
+            
             result_text = (
-                f"**Total documents:** {len(texts1)}\n\n"
+                f"**Total documents (Dataset 1):** {total_docs}\n\n"
                 f"**Duplicates found:** {num_duplicates}\n\n"
-                f"**Unique documents after deduplication:** {len(deduplicated_indices)}\n\n"
+                f"**Unique documents after deduplication:** {deduplicated_count}\n\n"
                 + "-" * 50 + "\n\n"
             )
 
+            # Show example duplicates
             if num_duplicates > 0:
                 result_text += "**Example duplicates:**\n\n"
-                for dup_idx, orig_idx in list(duplicate_mapping.items())[:5]:
-                    orig_text = texts1[orig_idx]
-                    dup_text = texts1[dup_idx]
-                    differences = display_word_differences(orig_text, dup_text)
-                    result_text += (
-                        f"**Original:**\n{orig_text}\n\n"
-                        f"**Duplicate:**\n{dup_text}\n\n"
-                        f"**Differences:**\n{differences}\n"
-                        + "-" * 50 + "\n\n"
-                    )
+                for duprec in result.duplicates[:5]:
+                    dup_text = duprec.record 
+                    if duprec.duplicates:
+                        orig_text, score = duprec.duplicates[0]
+                        differences = display_word_differences(orig_text, dup_text)
+                        result_text += (
+                            f"**Original:**\n{orig_text}\n\n"
+                            f"**Duplicate:**\n{dup_text}\n\n"
+                            f"**Similarity Score:** {score:.4f}\n"
+                            f"**Differences:**\n{differences}\n"
+                            + "-" * 50 + "\n\n"
+                        )
+                    else:
+                        # Possibly an exact duplicate cluster
+                        result_text += (
+                            f"**Duplicate:**\n{dup_text}\n\n"
+                            "No near-duplicate details available.\n"
+                            + "-" * 50 + "\n\n"
+                        )
             else:
                 result_text += "No duplicates found."
 
             yield "Deduplication completed.", result_text
 
         else:
-            # Load and process Dataset 2
+            # Cross-dataset deduplication
             yield "Loading Dataset 2...", ""
             texts2 = load_dataset_texts(dataset2_name, dataset2_split, dataset2_text_column)
-            yield "Computing embeddings for Dataset 2...", ""
-            embeddings2 = model.encode(texts2, show_progressbar=True)
 
-            # Deduplicate Dataset 2 against Dataset 1
-            yield "Deduplicating Dataset 2 against Dataset 1...", ""
-            duplicate_indices, duplicate_mapping = deduplicate_embeddings(
-                embeddings1, embeddings_b=embeddings2, threshold=threshold, progress=progress
-            )
+            yield "Deduplicating Dataset 2 against Dataset 1 (SemHash)...", ""
+            result = deduplicate_two_datasets(texts1, texts2, threshold=threshold)
+
+            # Sort duplicates in descending order of their highest score
+            for duprec in result.duplicates:
+                duprec.duplicates.sort(key=lambda x: x[1], reverse=True)
+
+            num_duplicates = len(result.duplicates)
+            total_docs2 = len(texts2)
+            deduplicated_count = len(result.deduplicated)
 
-            num_duplicates = len(duplicate_indices)
             result_text = (
-                f"**Total documents in {dataset2_name}/{dataset2_split}:** {len(texts2)}\n\n"
+                f"**Total documents in {dataset2_name}/{dataset2_split}:** {total_docs2}\n\n"
                 f"**Duplicates found in Dataset 2:** {num_duplicates}\n\n"
-                f"**Unique documents after deduplication:** {len(texts2) - num_duplicates}\n\n"
+                f"**Unique documents after deduplication:** {deduplicated_count}\n\n"
                 + "-" * 50 + "\n\n"
             )
 
             if num_duplicates > 0:
                 result_text += "**Example duplicates from Dataset 2:**\n\n"
-                for idx in duplicate_indices[:5]:
-                    orig_text = texts1[duplicate_mapping[idx]]
-                    dup_text = texts2[idx]
-                    differences = display_word_differences(orig_text, dup_text)
-                    result_text += (
-                        f"**Original (Dataset 1):**\n{orig_text}\n\n"
-                        f"**Duplicate (Dataset 2):**\n{dup_text}\n\n"
-                        f"**Differences:**\n{differences}\n"
-                        + "-" * 50 + "\n\n"
-                    )
+                for duprec in result.duplicates[:5]:
+                    dup_text = duprec.record  # The "duplicate" text from dataset2
+                    if duprec.duplicates:
+                        orig_text, score = duprec.duplicates[0]
+                        differences = display_word_differences(orig_text, dup_text)
+                        result_text += (
+                            f"**Original (Dataset 1):**\n{orig_text}\n\n"
+                            f"**Duplicate (Dataset 2):**\n{dup_text}\n\n"
+                            f"**Similarity Score:** {score:.4f}\n"
+                            f"**Differences:**\n{differences}\n"
+                            + "-" * 50 + "\n\n"
+                        )
+                    else:
+                        result_text += (
+                            f"**Potential Duplicate (Dataset 2):**\n{dup_text}\n\n"
+                            "No near-duplicate details available.\n"
+                            + "-" * 50 + "\n\n"
+                        )
             else:
                 result_text += "No duplicates found."
 
@@ -191,44 +172,47 @@ def perform_deduplication(
         yield f"An error occurred: {e}", ""
         raise e
 
-# Gradio app with stop button support
+
+# --- Gradio App ---
 with gr.Blocks(theme=gr.themes.Ocean(), css="#status_output { height: 50px; overflow: auto; }") as demo:
-    gr.Markdown("# Semantic Deduplication")
+    gr.Markdown("# Semantic Text Deduplication Using SemHash")
     gr.Markdown("""
-    This demo showcases semantic deduplication using Model2Vec for HuggingFace datasets.
-    It can be used to identify duplicate texts within a single dataset or across two datasets.
-    You can adjust the similarity threshold to control the strictness of the deduplication.\n
-    NOTE: this demo runs on a free CPU backend, so it may be slow for large datasets. For faster results, please run the code locally.
+    This demo showcases **semantic deduplication** using [SemHash](https://github.com/MinishLab/semhash) for HuggingFace datasets, using a [Model2Vec](https://github.com/MinishLab/model2vec) encoder.
+    It can be used to identify duplicate texts within a **single dataset** or across **two datasets**.
+    You can adjust the similarity threshold to control the strictness of the deduplication.
+
+    **NOTE**: This demo runs on a free CPU backend, so it may be slow for large datasets.
+    For faster results, please run the code locally.
     """)
 
     deduplication_type = gr.Radio(
-        choices=["Cross-dataset", "Single dataset"],  # Swapped "Cross-dataset" to the left
+        choices=["Cross-dataset", "Single dataset"],
         label="Deduplication Type",
-        value="Cross-dataset",  # Set "Cross-dataset" as the default value
+        value="Cross-dataset",  # default
     )
 
     with gr.Row():
         dataset1_name = gr.Textbox(value=default_dataset_name, label="Dataset 1 Name")
-        dataset1_split = gr.Textbox(value=default_dataset1_split, label="Dataset 1 Split")  # Default split is "train"
+        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=True)  # Make dataset2_inputs visible by default
+    dataset2_inputs = gr.Column(visible=True)
     with dataset2_inputs:
         with gr.Row():
             dataset2_name = gr.Textbox(value=default_dataset_name, label="Dataset 2 Name")
-            dataset2_split = gr.Textbox(value=default_dataset2_split, label="Dataset 2 Split")  # Default split is "test"
+            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(0.0, 1.0, value=default_threshold, label="Similarity Threshold")
 
-    with gr.Row():  # Placing the button in the same row for better alignment
+    with gr.Row():
         compute_button = gr.Button("Deduplicate")
 
     status_output = gr.Markdown(elem_id="status_output")
     result_output = gr.Markdown()
 
     def update_visibility(choice: str):
-        return gr.update(visible=choice == "Cross-dataset")
+        return gr.update(visible=(choice == "Cross-dataset"))
 
     deduplication_type.change(update_visibility, inputs=deduplication_type, outputs=dataset2_inputs)
 
@@ -247,5 +231,5 @@ with gr.Blocks(theme=gr.themes.Ocean(), css="#status_output { height: 50px; over
         outputs=[status_output, result_output],
     )
 
-
 demo.launch()
+

requirements.txt

@@ -1,5 +1,4 @@
-reach < 5
-model2vec
+semhash>=0.2.0
 numpy
 datasets
 tqdm