app.py

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/M2V_base_output")

# Default parameters
default_dataset_name = "sst2"
default_dataset_split = "train"
default_text_column = "sentence"
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

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').
    :param text_column: Name of the text column.
    :return: A list of texts from the dataset.
    """
    ds = load_dataset(dataset_name, split=dataset_split)
    return [example[text_column] for example in ds]

def perform_deduplication(
    deduplication_type: str,
    dataset1_name: str,
    dataset1_split: str,
    dataset1_text_column: str,
    dataset2_name: str = "",
    dataset2_split: str = "",
    dataset2_text_column: str = "",
    threshold: float = default_threshold,
    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.
    """
    try:
        threshold = float(threshold)

        # Load and process 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)
            result_text = (
                f"**Total documents:** {len(texts1)}\n\n"
                f"**Duplicates found:** {num_duplicates}\n\n"
                f"**Unique documents after deduplication:** {len(deduplicated_indices)}\n\n"
            )

            if num_duplicates > 0:
                result_text += "**Sample 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"
                    )
            else:
                result_text += "No duplicates found."

            yield "Deduplication completed.", result_text

        else:
            # Load and process Dataset 2
            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
            )

            num_duplicates = len(duplicate_indices)
            result_text = (
                f"**Total documents in {dataset2_name}/{dataset2_split}:** {len(texts2)}\n\n"
                f"**Duplicates found in Dataset 2:** {num_duplicates}\n\n"
                f"**Unique documents after deduplication:** {len(texts2) - num_duplicates}\n\n"
            )

            if num_duplicates > 0:
                result_text += "**Sample 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"
                    )
            else:
                result_text += "No duplicates found."

            yield "Deduplication completed.", result_text

    except Exception as e:
        yield f"An error occurred: {e}", ""
        raise e

# Gradio app with stop button support
with gr.Blocks(css="#status_output { height: 50px; overflow: auto; }") as demo:
    gr.Markdown("# Semantic Deduplication")
    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.
    """)

    deduplication_type = gr.Radio(
        choices=["Single dataset", "Cross-dataset"],
        label="Deduplication Type",
        value="Single dataset",
    )

    with gr.Row():
        dataset1_name = gr.Textbox(value=default_dataset_name, label="Dataset 1 Name")
        dataset1_split = gr.Textbox(value=default_dataset_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_dataset_name, label="Dataset 2 Name")
            dataset2_split = gr.Textbox(value=default_dataset_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")
    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")

    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()



# 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/M2V_base_output")

# # Default parameters
# default_dataset_name = "sst2"
# default_dataset_split = "train"
# default_text_column = "sentence"
# 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."""
#     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

# def display_word_differences(x: str, y: str) -> str:
#     """Display word-level differences between two texts, avoiding Markdown issues."""
#     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."""
#     ds = load_dataset(dataset_name, split=dataset_split)
#     return [example[text_column] for example in ds]

# def perform_deduplication(
#     deduplication_type: str,
#     dataset1_name: str,
#     dataset1_split: str,
#     dataset1_text_column: str,
#     dataset2_name: str = "",
#     dataset2_split: str = "",
#     dataset2_text_column: str = "",
#     threshold: float = default_threshold,
#     progress: gr.Progress = gr.Progress(track_tqdm=True)
# ):
#     """Perform deduplication on one or two datasets."""
#     try:
#         threshold = float(threshold)

#         # Load and process 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)
#             result_text = (
#                 f"**Total documents:** {len(texts1)}\n\n"
#                 f"**Duplicates found:** {num_duplicates}\n\n"
#                 f"**Unique documents after deduplication:** {len(deduplicated_indices)}\n\n"
#             )

#             if num_duplicates > 0:
#                 result_text += "**Sample 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"
#                     )
#             else:
#                 result_text += "No duplicates found."

#             yield "Deduplication completed.", result_text

#         else:
#             # Load and process Dataset 2
#             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
#             )

#             num_duplicates = len(duplicate_indices)
#             result_text = (
#                 f"**Total documents in {dataset2_name}/{dataset2_split}:** {len(texts2)}\n\n"
#                 f"**Duplicates found in Dataset 2:** {num_duplicates}\n\n"
#                 f"**Unique documents after deduplication:** {len(texts2) - num_duplicates}\n\n"
#             )

#             if num_duplicates > 0:
#                 result_text += "**Sample 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"
#                     )
#             else:
#                 result_text += "No duplicates found."

#             yield "Deduplication completed.", result_text

#     except Exception as e:
#         yield f"An error occurred: {e}", ""
#         raise e

# # Gradio app with stop button support
# with gr.Blocks(css="#status_output { height: 50px; overflow: auto; }") as demo:
#     gr.Markdown("# Semantic Deduplication")
#     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.
#     """)

#     deduplication_type = gr.Radio(
#         choices=["Single dataset", "Cross-dataset"],
#         label="Deduplication Type",
#         value="Single dataset",
#     )

#     with gr.Row():
#         dataset1_name = gr.Textbox(value=default_dataset_name, label="Dataset 1 Name")
#         dataset1_split = gr.Textbox(value=default_dataset_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_dataset_name, label="Dataset 2 Name")
#             dataset2_split = gr.Textbox(value=default_dataset_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")
#     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")

#     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()