# import gradio as gr # def greet(name): # return "Hello " + name + "!!" # demo = gr.Interface(fn=greet, inputs="text", outputs="text") # demo.launch() import gradio as gr from datasets import load_dataset import numpy as np from model2vec import StaticModel from reach import Reach from tqdm import tqdm def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024) -> tuple[np.ndarray, dict[int, int]]: """ Deduplicate embeddings and return the deduplicated indices and a mapping of removed indices to their corresponding original indices. """ reach = Reach(vectors=embedding_matrix, items=[str(i) for i in range(len(embedding_matrix))]) # Use a set for deduplicated indices and keep track of duplicates deduplicated_indices = set(range(len(embedding_matrix))) # Start with all indices as deduplicated duplicate_to_original_mapping = {} results = reach.nearest_neighbor_threshold( embedding_matrix, threshold=threshold, batch_size=batch_size, show_progressbar=True ) # Process duplicates for i, similar_items in enumerate(tqdm(results)): if i not in deduplicated_indices: continue # Skip already marked duplicates # Similar items are returned as (index, score), we are only interested in the index similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i] # Mark similar documents as duplicates and map them to the original for sim_idx in similar_indices: if sim_idx in deduplicated_indices: deduplicated_indices.remove(sim_idx) duplicate_to_original_mapping[sim_idx] = i # Map duplicate to original return np.array(list(deduplicated_indices)), duplicate_to_original_mapping def deduplicate_across_datasets(embedding_matrix_1: np.ndarray, embedding_matrix_2: np.ndarray, threshold: float, batch_size: int = 1024) -> tuple[list[int], dict[int, int]]: """ Deduplicate embeddings across two datasets and return the indices of duplicates between them. """ reach = Reach(vectors=embedding_matrix_1, items=[str(i) for i in range(len(embedding_matrix_1))]) # Keep track of duplicates in the second dataset duplicate_indices_in_test = [] duplicate_to_original_mapping = {} # Find nearest neighbors from the test set in the train set results = reach.nearest_neighbor_threshold( embedding_matrix_2, threshold=threshold, batch_size=batch_size, show_progressbar=True ) # Process duplicates for i, similar_items in enumerate(tqdm(results)): # Similar items are returned as (index, score), we are only interested in the index similar_indices = [int(item[0]) for item in similar_items if item[1] >= threshold] # Keep those above the threshold # If we find a similar item in the train set, mark it as a duplicate if similar_indices: duplicate_indices_in_test.append(i) duplicate_to_original_mapping[i] = similar_indices[0] # Map duplicate in test to original in train return duplicate_indices_in_test, duplicate_to_original_mapping def perform_deduplication( deduplication_type, dataset1_name, dataset1_split, dataset2_name, dataset2_split, threshold ): # Convert threshold to float threshold = float(threshold) if deduplication_type == "Single dataset": # Load the dataset ds = load_dataset(dataset1_name, split=dataset1_split) # Extract texts texts = [example['text'] for example in ds] # Compute embeddings model = StaticModel.from_pretrained("minishlab/M2V_base_output") embedding_matrix = model.encode(texts, show_progressbar=True) # Deduplicate deduplicated_indices, duplicate_to_original_mapping = deduplicate(embedding_matrix, threshold) # Prepare the results num_duplicates = len(duplicate_to_original_mapping) num_total = len(texts) num_deduplicated = len(deduplicated_indices) result_text = f"**Total documents:** {num_total}\n" result_text += f"**Number of duplicates found:** {num_duplicates}\n" result_text += f"**Number of unique documents after deduplication:** {num_deduplicated}\n\n" result_text += f"**Deduplicated indices:** {deduplicated_indices.tolist()}\n\n" result_text += f"**Duplicate to original mapping:** {duplicate_to_original_mapping}\n" return result_text elif deduplication_type == "Cross-dataset": # Load datasets ds1 = load_dataset(dataset1_name, split=dataset1_split) ds2 = load_dataset(dataset2_name, split=dataset2_split) # Extract texts texts1 = [example['text'] for example in ds1] texts2 = [example['text'] for example in ds2] # Compute embeddings model = StaticModel.from_pretrained("minishlab/M2V_base_output") embedding_matrix1 = model.encode(texts1, show_progressbar=True) embedding_matrix2 = model.encode(texts2, show_progressbar=True) # Deduplicate across datasets duplicate_indices_in_ds2, duplicate_to_original_mapping = deduplicate_across_datasets(embedding_matrix1, embedding_matrix2, threshold) num_duplicates = len(duplicate_indices_in_ds2) num_total_ds2 = len(texts2) num_unique_ds2 = num_total_ds2 - num_duplicates result_text = f"**Total documents in {dataset2_name}/{dataset2_split}:** {num_total_ds2}\n" result_text += f"**Number of duplicates found in {dataset2_name}/{dataset2_split}:** {num_duplicates}\n" result_text += f"**Number of unique documents in {dataset2_name}/{dataset2_split} after deduplication:** {num_unique_ds2}\n\n" result_text += f"**Duplicate indices in {dataset2_name}/{dataset2_split}:** {duplicate_indices_in_ds2}\n\n" result_text += f"**Duplicate to original mapping:** {duplicate_to_original_mapping}\n" return result_text with gr.Blocks() as demo: gr.Markdown("# Semantic Deduplication") deduplication_type = gr.Radio(choices=["Single dataset", "Cross-dataset"], label="Deduplication Type", value="Single dataset") with gr.Row(): dataset1_name = gr.Textbox(value="ag_news", label="Dataset 1 Name") dataset1_split = gr.Textbox(value="train", label="Dataset 1 Split") dataset2_row = gr.Row(visible=False) with dataset2_row: dataset2_name = gr.Textbox(value="ag_news", label="Dataset 2 Name") dataset2_split = gr.Textbox(value="test", label="Dataset 2 Split") threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.8, label="Similarity Threshold") compute_button = gr.Button("Compute") output = gr.Markdown() # Function to update the visibility of dataset2_row def update_visibility(deduplication_type): if deduplication_type == "Cross-dataset": return {dataset2_row: gr.update(visible=True)} else: return {dataset2_row: gr.update(visible=False)} deduplication_type.change(update_visibility, inputs=deduplication_type, outputs=[dataset2_row]) compute_button.click( fn=perform_deduplication, inputs=[deduplication_type, dataset1_name, dataset1_split, dataset2_name, dataset2_split, threshold], outputs=output ) demo.launch()