| 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 a semantic deduplication process where we identify duplicate texts within a single dataset or across two datasets. | |
| The deduplication is based on cosine similarity between the embeddings of the texts. | |
| You can adjust the similarity threshold to control the strictness of the deduplication. | |
| """) | |
| 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("Compute") | |
| stop_button = gr.Button("Stop") | |
| 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], | |
| ) | |
| # Stop button functionality | |
| stop_button.click(lambda: demo.stop(), None, None) | |
| 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. | |
| # :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()) | |
| # # Wrap differences in a code block to prevent interpretation as Markdown | |
| # formatted_diff = "\n".join(word for word in diff if word.startswith(("+", "-"))) | |
| # return f"```\n{formatted_diff}\n```" | |
| # # def display_word_differences(x: str, y: str) -> str: | |
| # # """ | |
| # # Display the word-level differences between two texts. | |
| # # :param x: First text. | |
| # # :param y: Second text. | |
| # # :return: A string showing word-level differences. | |
| # # """ | |
| # # diff = ndiff(x.split(), y.split()) | |
| # # return " ".join(word for word in diff if word.startswith(("+", "-"))) | |
| # 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 | |
| # 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("Compute") | |
| # 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() | |