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 | |
import sys | |
import tqdm | |
# Load the model at startup | |
model = StaticModel.from_pretrained("minishlab/M2V_base_output") | |
# Update default dataset to 'sst2' and set default threshold to 0.9 | |
default_dataset1_name = "sst2" | |
default_dataset1_split = "train" | |
default_dataset2_name = "sst2" | |
default_dataset2_split = "validation" | |
default_text_column = "sentence" | |
default_threshold = 0.9 | |
# Load the default datasets at startup | |
ds_default1 = load_dataset(default_dataset1_name, split=default_dataset1_split) | |
ds_default2 = load_dataset(default_dataset2_name, split=default_dataset2_split) | |
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. | |
""" | |
# Building the index | |
reach = Reach(vectors=embedding_matrix, items=[str(i) for i in range(len(embedding_matrix))]) | |
deduplicated_indices = set(range(len(embedding_matrix))) | |
duplicate_to_original_mapping = {} | |
# Finding nearest neighbors | |
results = reach.nearest_neighbor_threshold( | |
embedding_matrix, | |
threshold=threshold, | |
batch_size=batch_size, | |
show_progressbar=True # Allow internal progress bar | |
) | |
# Processing duplicates | |
for i, similar_items in enumerate(results): | |
if i not in deduplicated_indices: | |
continue | |
similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i] | |
for sim_idx in similar_indices: | |
if sim_idx in deduplicated_indices: | |
deduplicated_indices.remove(sim_idx) | |
duplicate_to_original_mapping[sim_idx] = i | |
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. | |
""" | |
# Building the index from Dataset 1 | |
reach = Reach(vectors=embedding_matrix_1, items=[str(i) for i in range(len(embedding_matrix_1))]) | |
duplicate_indices_in_test = [] | |
duplicate_to_original_mapping = {} | |
# Finding nearest neighbors between datasets | |
results = reach.nearest_neighbor_threshold( | |
embedding_matrix_2, | |
threshold=threshold, | |
batch_size=batch_size, | |
show_progressbar=True # Allow internal progress bar | |
) | |
# Processing duplicates | |
for i, similar_items in enumerate(results): | |
similar_indices = [int(item[0]) for item in similar_items if item[1] >= threshold] | |
if similar_indices: | |
duplicate_indices_in_test.append(i) | |
duplicate_to_original_mapping[i] = similar_indices[0] | |
return duplicate_indices_in_test, duplicate_to_original_mapping | |
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 perform_deduplication( | |
deduplication_type, | |
dataset1_name, | |
dataset1_split, | |
dataset1_text_column, | |
dataset2_name="", | |
dataset2_split="", | |
dataset2_text_column="", | |
threshold=default_threshold, | |
progress=gr.Progress(track_tqdm=True) | |
): | |
# Deep Monkey-Patching of tqdm | |
original_tqdm = tqdm.tqdm | |
tqdm.tqdm = progress.tqdm | |
for mod_name in list(sys.modules.keys()): | |
if 'tqdm' in mod_name: | |
sys.modules[mod_name].tqdm = progress.tqdm | |
try: | |
# Convert threshold to float | |
threshold = float(threshold) | |
if deduplication_type == "Single dataset": | |
# Load Dataset 1 | |
gr.print("Loading Dataset 1...") | |
if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split: | |
ds = ds_default1 | |
else: | |
ds = load_dataset(dataset1_name, split=dataset1_split) | |
# Extract texts | |
gr.print("Extracting texts from Dataset 1...") | |
texts = [example[dataset1_text_column] for example in ds] | |
# Compute embeddings | |
gr.print("Computing embeddings for Dataset 1...") | |
embedding_matrix = model.encode(texts, show_progressbar=True) # Enable internal progress bar | |
# Deduplicate | |
gr.print("Deduplicating embeddings...") | |
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" | |
# Show deduplicated examples | |
if num_duplicates > 0: | |
result_text += "**Examples of duplicates found:**\n\n" | |
num_examples = min(5, num_duplicates) | |
for duplicate_idx, original_idx in list(duplicate_to_original_mapping.items())[:num_examples]: | |
original_text = texts[original_idx] | |
duplicate_text = texts[duplicate_idx] | |
differences = display_word_differences(original_text, duplicate_text) | |
result_text += f"**Original text:**\n{original_text}\n\n" | |
result_text += f"**Duplicate text:**\n{duplicate_text}\n\n" | |
result_text += f"**Differences:**\n{differences}\n" | |
result_text += "-" * 50 + "\n\n" | |
else: | |
result_text += "No duplicates found." | |
return result_text | |
elif deduplication_type == "Cross-dataset": | |
# Load Dataset 1 | |
gr.print("Loading Dataset 1...") | |
if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split: | |
ds1 = ds_default1 | |
else: | |
ds1 = load_dataset(dataset1_name, split=dataset1_split) | |
# Load Dataset 2 | |
gr.print("Loading Dataset 2...") | |
if dataset2_name == default_dataset2_name and dataset2_split == default_dataset2_split: | |
ds2 = ds_default2 | |
else: | |
ds2 = load_dataset(dataset2_name, split=dataset2_split) | |
# Extract texts from Dataset 1 | |
gr.print("Extracting texts from Dataset 1...") | |
texts1 = [example[dataset1_text_column] for example in ds1] | |
# Extract texts from Dataset 2 | |
gr.print("Extracting texts from Dataset 2...") | |
texts2 = [example[dataset2_text_column] for example in ds2] | |
# Compute embeddings for Dataset 1 | |
gr.print("Computing embeddings for Dataset 1...") | |
embedding_matrix1 = model.encode(texts1, show_progressbar=True) | |
# Compute embeddings for Dataset 2 | |
gr.print("Computing embeddings for Dataset 2...") | |
embedding_matrix2 = model.encode(texts2, show_progressbar=True) | |
# Deduplicate across datasets | |
gr.print("Deduplicating embeddings 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" | |
# Show deduplicated examples | |
if num_duplicates > 0: | |
result_text += "**Examples of duplicates found in Dataset 2:**\n\n" | |
num_examples = min(5, num_duplicates) | |
for duplicate_idx in duplicate_indices_in_ds2[:num_examples]: | |
original_idx = duplicate_to_original_mapping[duplicate_idx] | |
original_text = texts1[original_idx] | |
duplicate_text = texts2[duplicate_idx] | |
differences = display_word_differences(original_text, duplicate_text) | |
result_text += f"**Original text (Dataset 1):**\n{original_text}\n\n" | |
result_text += f"**Duplicate text (Dataset 2):**\n{duplicate_text}\n\n" | |
result_text += f"**Differences:**\n{differences}\n" | |
result_text += "-" * 50 + "\n\n" | |
else: | |
result_text += "No duplicates found." | |
return result_text | |
finally: | |
# Restore original tqdm | |
tqdm.tqdm = original_tqdm | |
for mod_name in list(sys.modules.keys()): | |
if 'tqdm' in mod_name: | |
sys.modules[mod_name].tqdm = original_tqdm | |
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=default_dataset1_name, label="Dataset 1 Name") | |
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=False) | |
with dataset2_inputs: | |
gr.Markdown("### Dataset 2") | |
with gr.Row(): | |
dataset2_name = gr.Textbox(value=default_dataset2_name, label="Dataset 2 Name") | |
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( | |
minimum=0.0, | |
maximum=1.0, | |
value=default_threshold, | |
label="Similarity Threshold" | |
) | |
compute_button = gr.Button("Compute") | |
output = gr.Markdown() | |
# Function to update the visibility of dataset2_inputs | |
def update_visibility(deduplication_type_value): | |
if deduplication_type_value == "Cross-dataset": | |
return gr.update(visible=True) | |
else: | |
return gr.update(visible=False) | |
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=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 | |
# import sys | |
# import tqdm | |
# # Load the model at startup | |
# model = StaticModel.from_pretrained("minishlab/M2V_base_output") | |
# # Update default dataset to 'sst2' and set default threshold to 0.9 | |
# default_dataset1_name = "sst2" | |
# default_dataset1_split = "train" | |
# default_dataset2_name = "sst2" | |
# default_dataset2_split = "validation" | |
# default_text_column = "sentence" | |
# default_threshold = 0.9 | |
# # Load the default datasets at startup | |
# ds_default1 = load_dataset(default_dataset1_name, split=default_dataset1_split) | |
# ds_default2 = load_dataset(default_dataset2_name, split=default_dataset2_split) | |
# def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024, progress=None) -> tuple[np.ndarray, dict[int, int]]: | |
# """ | |
# Deduplicate embeddings and return the deduplicated indices and a mapping of removed indices to their corresponding original indices. | |
# """ | |
# # Update progress to indicate building the index | |
# progress(0, desc="Building search index...") | |
# reach = Reach(vectors=embedding_matrix, items=[str(i) for i in range(len(embedding_matrix))]) | |
# deduplicated_indices = set(range(len(embedding_matrix))) | |
# duplicate_to_original_mapping = {} | |
# # Finding nearest neighbors | |
# progress(0, desc="Finding nearest neighbors...") | |
# results = reach.nearest_neighbor_threshold( | |
# embedding_matrix, | |
# threshold=threshold, | |
# batch_size=batch_size, | |
# show_progressbar=True # Allow internal progress bar | |
# ) | |
# # Processing duplicates with a progress bar | |
# total_items = len(embedding_matrix) | |
# for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates", total=total_items)): | |
# if i not in deduplicated_indices: | |
# continue | |
# similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i] | |
# for sim_idx in similar_indices: | |
# if sim_idx in deduplicated_indices: | |
# deduplicated_indices.remove(sim_idx) | |
# duplicate_to_original_mapping[sim_idx] = i | |
# 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, progress=None) -> tuple[list[int], dict[int, int]]: | |
# """ | |
# Deduplicate embeddings across two datasets and return the indices of duplicates between them. | |
# """ | |
# # Update progress to indicate building the index | |
# progress(0, desc="Building search index from Dataset 1...") | |
# reach = Reach(vectors=embedding_matrix_1, items=[str(i) for i in range(len(embedding_matrix_1))]) | |
# duplicate_indices_in_test = [] | |
# duplicate_to_original_mapping = {} | |
# # Finding nearest neighbors between datasets | |
# progress(0, desc="Finding nearest neighbors between datasets...") | |
# results = reach.nearest_neighbor_threshold( | |
# embedding_matrix_2, | |
# threshold=threshold, | |
# batch_size=batch_size, | |
# show_progressbar=True # Allow internal progress bar | |
# ) | |
# total_items = len(embedding_matrix_2) | |
# # Processing duplicates with a progress bar | |
# for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates across datasets", total=total_items)): | |
# similar_indices = [int(item[0]) for item in similar_items if item[1] >= threshold] | |
# if similar_indices: | |
# duplicate_indices_in_test.append(i) | |
# duplicate_to_original_mapping[i] = similar_indices[0] | |
# return duplicate_indices_in_test, duplicate_to_original_mapping | |
# 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 perform_deduplication( | |
# deduplication_type, | |
# dataset1_name, | |
# dataset1_split, | |
# dataset1_text_column, | |
# dataset2_name="", | |
# dataset2_split="", | |
# dataset2_text_column="", | |
# threshold=default_threshold, | |
# progress=gr.Progress(track_tqdm=True) | |
# ): | |
# # Monkey-patch tqdm | |
# original_tqdm = tqdm.tqdm | |
# original_reach_tqdm = Reach.__dict__['tqdm'] if 'tqdm' in Reach.__dict__ else None | |
# tqdm.tqdm = progress.tqdm | |
# sys.modules['tqdm'].tqdm = progress.tqdm | |
# sys.modules['tqdm.auto'].tqdm = progress.tqdm | |
# Reach.tqdm = progress.tqdm # Monkey-patch reach's tqdm | |
# try: | |
# # Convert threshold to float | |
# threshold = float(threshold) | |
# if deduplication_type == "Single dataset": | |
# # Load Dataset 1 | |
# progress(0, desc="Loading Dataset 1...") | |
# if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split: | |
# ds = ds_default1 | |
# else: | |
# ds = load_dataset(dataset1_name, split=dataset1_split) | |
# # Extract texts | |
# progress(0, desc="Extracting texts from Dataset 1...") | |
# texts = [example[dataset1_text_column] for example in ds] | |
# # Compute embeddings | |
# progress(0, desc="Computing embeddings for Dataset 1...") | |
# embedding_matrix = model.encode(texts, show_progressbar=True) # Enable internal progress bar | |
# # Deduplicate | |
# result_text = deduplicate_and_prepare_results_single( | |
# embedding_matrix, texts, threshold, progress | |
# ) | |
# return result_text | |
# elif deduplication_type == "Cross-dataset": | |
# # Load Dataset 1 | |
# progress(0, desc="Loading Dataset 1...") | |
# if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split: | |
# ds1 = ds_default1 | |
# else: | |
# ds1 = load_dataset(dataset1_name, split=dataset1_split) | |
# # Load Dataset 2 | |
# progress(0, desc="Loading Dataset 2...") | |
# if dataset2_name == default_dataset2_name and dataset2_split == default_dataset2_split: | |
# ds2 = ds_default2 | |
# else: | |
# ds2 = load_dataset(dataset2_name, split=dataset2_split) | |
# # Extract texts from Dataset 1 | |
# progress(0, desc="Extracting texts from Dataset 1...") | |
# texts1 = [example[dataset1_text_column] for example in ds1] | |
# # Extract texts from Dataset 2 | |
# progress(0, desc="Extracting texts from Dataset 2...") | |
# texts2 = [example[dataset2_text_column] for example in ds2] | |
# # Compute embeddings for Dataset 1 | |
# progress(0, desc="Computing embeddings for Dataset 1...") | |
# embedding_matrix1 = model.encode(texts1, show_progressbar=True) | |
# # Compute embeddings for Dataset 2 | |
# progress(0, desc="Computing embeddings for Dataset 2...") | |
# embedding_matrix2 = model.encode(texts2, show_progressbar=True) | |
# # Deduplicate across datasets | |
# result_text = deduplicate_and_prepare_results_cross( | |
# embedding_matrix1, embedding_matrix2, texts1, texts2, threshold, progress, dataset2_name, dataset2_split | |
# ) | |
# return result_text | |
# finally: | |
# # Restore original tqdm | |
# tqdm.tqdm = original_tqdm | |
# sys.modules['tqdm'].tqdm = original_tqdm | |
# sys.modules['tqdm.auto'].tqdm = original_tqdm | |
# # Restore reach's original tqdm | |
# if original_reach_tqdm is not None: | |
# Reach.tqdm = original_reach_tqdm | |
# else: | |
# del Reach.tqdm # If it wasn't originally in Reach's __dict__ | |
# def deduplicate_and_prepare_results_single(embedding_matrix, texts, threshold, progress): | |
# # Deduplicate | |
# deduplicated_indices, duplicate_to_original_mapping = deduplicate( | |
# embedding_matrix, threshold, progress=progress | |
# ) | |
# # 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" | |
# # Show deduplicated examples | |
# if num_duplicates > 0: | |
# result_text += "**Examples of duplicates found:**\n\n" | |
# num_examples = min(5, num_duplicates) | |
# for duplicate_idx, original_idx in list(duplicate_to_original_mapping.items())[:num_examples]: | |
# original_text = texts[original_idx] | |
# duplicate_text = texts[duplicate_idx] | |
# differences = display_word_differences(original_text, duplicate_text) | |
# result_text += f"**Original text:**\n{original_text}\n\n" | |
# result_text += f"**Duplicate text:**\n{duplicate_text}\n\n" | |
# result_text += f"**Differences:**\n{differences}\n" | |
# result_text += "-" * 50 + "\n\n" | |
# else: | |
# result_text += "No duplicates found." | |
# return result_text | |
# def deduplicate_and_prepare_results_cross(embedding_matrix1, embedding_matrix2, texts1, texts2, threshold, progress, dataset2_name, dataset2_split): | |
# # Deduplicate across datasets | |
# duplicate_indices_in_ds2, duplicate_to_original_mapping = deduplicate_across_datasets( | |
# embedding_matrix1, embedding_matrix2, threshold, progress=progress | |
# ) | |
# 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" | |
# # Show deduplicated examples | |
# if num_duplicates > 0: | |
# result_text += "**Examples of duplicates found in Dataset 2:**\n\n" | |
# num_examples = min(5, num_duplicates) | |
# for duplicate_idx in duplicate_indices_in_ds2[:num_examples]: | |
# original_idx = duplicate_to_original_mapping[duplicate_idx] | |
# original_text = texts1[original_idx] | |
# duplicate_text = texts2[duplicate_idx] | |
# differences = display_word_differences(original_text, duplicate_text) | |
# result_text += f"**Original text (Dataset 1):**\n{original_text}\n\n" | |
# result_text += f"**Duplicate text (Dataset 2):**\n{duplicate_text}\n\n" | |
# result_text += f"**Differences:**\n{differences}\n" | |
# result_text += "-" * 50 + "\n\n" | |
# else: | |
# result_text += "No duplicates found." | |
# 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=default_dataset1_name, label="Dataset 1 Name") | |
# 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=False) | |
# with dataset2_inputs: | |
# gr.Markdown("### Dataset 2") | |
# with gr.Row(): | |
# dataset2_name = gr.Textbox(value=default_dataset2_name, label="Dataset 2 Name") | |
# 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( | |
# minimum=0.0, | |
# maximum=1.0, | |
# value=default_threshold, | |
# label="Similarity Threshold" | |
# ) | |
# compute_button = gr.Button("Compute") | |
# output = gr.Markdown() | |
# # Function to update the visibility of dataset2_inputs | |
# def update_visibility(deduplication_type_value): | |
# if deduplication_type_value == "Cross-dataset": | |
# return gr.update(visible=True) | |
# else: | |
# return gr.update(visible=False) | |
# 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=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 | |
# import sys | |
# import tqdm | |
# # Load the model at startup | |
# model = StaticModel.from_pretrained("minishlab/M2V_base_output") | |
# # Load the default datasets at startup | |
# default_dataset1_name = "ag_news" | |
# default_dataset1_split = "train" | |
# default_dataset2_name = "ag_news" | |
# default_dataset2_split = "test" | |
# ds_default1 = load_dataset(default_dataset1_name, split=default_dataset1_split) | |
# ds_default2 = load_dataset(default_dataset2_name, split=default_dataset2_split) | |
# def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024, progress=None) -> 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))]) | |
# deduplicated_indices = set(range(len(embedding_matrix))) | |
# duplicate_to_original_mapping = {} | |
# results = reach.nearest_neighbor_threshold( | |
# embedding_matrix, | |
# threshold=threshold, | |
# batch_size=batch_size, | |
# show_progressbar=True # Allow internal progress bar | |
# ) | |
# # Process duplicates | |
# for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates", total=len(embedding_matrix))): | |
# if i not in deduplicated_indices: | |
# continue | |
# similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i] | |
# for sim_idx in similar_indices: | |
# if sim_idx in deduplicated_indices: | |
# deduplicated_indices.remove(sim_idx) | |
# duplicate_to_original_mapping[sim_idx] = i | |
# 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, progress=None) -> 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))]) | |
# duplicate_indices_in_test = [] | |
# duplicate_to_original_mapping = {} | |
# results = reach.nearest_neighbor_threshold( | |
# embedding_matrix_2, | |
# threshold=threshold, | |
# batch_size=batch_size, | |
# show_progressbar=True # Allow internal progress bar | |
# ) | |
# for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates across datasets", total=len(embedding_matrix_2))): | |
# similar_indices = [int(item[0]) for item in similar_items if item[1] >= threshold] | |
# if similar_indices: | |
# duplicate_indices_in_test.append(i) | |
# duplicate_to_original_mapping[i] = similar_indices[0] | |
# return duplicate_indices_in_test, duplicate_to_original_mapping | |
# 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 perform_deduplication( | |
# deduplication_type, | |
# dataset1_name, | |
# dataset1_split, | |
# dataset1_text_column, | |
# dataset2_name="", | |
# dataset2_split="", | |
# dataset2_text_column="", | |
# threshold=0.8, | |
# progress=gr.Progress(track_tqdm=True) | |
# ): | |
# # Monkey-patch tqdm | |
# original_tqdm = tqdm.tqdm | |
# original_reach_tqdm = Reach.__dict__['tqdm'] if 'tqdm' in Reach.__dict__ else None | |
# tqdm.tqdm = progress.tqdm | |
# sys.modules['tqdm'].tqdm = progress.tqdm | |
# sys.modules['tqdm.auto'].tqdm = progress.tqdm | |
# Reach.tqdm = progress.tqdm # Monkey-patch reach's tqdm | |
# try: | |
# # Convert threshold to float | |
# threshold = float(threshold) | |
# if deduplication_type == "Single dataset": | |
# # Check if the dataset is the default one | |
# if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split: | |
# ds = ds_default1 | |
# else: | |
# ds = load_dataset(dataset1_name, split=dataset1_split) | |
# # Extract texts | |
# texts = [example[dataset1_text_column] for example in ds] | |
# # Compute embeddings | |
# embedding_matrix = model.encode(texts, show_progressbar=True) # Enable internal progress bar | |
# # Deduplicate | |
# deduplicated_indices, duplicate_to_original_mapping = deduplicate(embedding_matrix, threshold, progress=progress) | |
# # 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" | |
# # Show deduplicated examples | |
# result_text += "**Examples of duplicates found:**\n\n" | |
# num_examples = min(5, num_duplicates) | |
# for duplicate_idx, original_idx in list(duplicate_to_original_mapping.items())[:num_examples]: | |
# original_text = texts[original_idx] | |
# duplicate_text = texts[duplicate_idx] | |
# differences = display_word_differences(original_text, duplicate_text) | |
# result_text += f"**Original text:**\n{original_text}\n\n" | |
# result_text += f"**Duplicate text:**\n{duplicate_text}\n\n" | |
# result_text += f"**Differences:**\n{differences}\n" | |
# result_text += "-" * 50 + "\n\n" | |
# return result_text | |
# elif deduplication_type == "Cross-dataset": | |
# # Dataset 1 | |
# if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split: | |
# ds1 = ds_default1 | |
# else: | |
# ds1 = load_dataset(dataset1_name, split=dataset1_split) | |
# # Dataset 2 | |
# if dataset2_name == default_dataset2_name and dataset2_split == default_dataset2_split: | |
# ds2 = ds_default2 | |
# else: | |
# ds2 = load_dataset(dataset2_name, split=dataset2_split) | |
# # Extract texts | |
# texts1 = [example[dataset1_text_column] for example in ds1] | |
# texts2 = [example[dataset2_text_column] for example in ds2] | |
# # Compute embeddings | |
# embedding_matrix1 = model.encode(texts1, show_progressbar=True) # Enable internal progress bar | |
# embedding_matrix2 = model.encode(texts2, show_progressbar=True) # Enable internal progress bar | |
# # Deduplicate across datasets | |
# duplicate_indices_in_ds2, duplicate_to_original_mapping = deduplicate_across_datasets( | |
# embedding_matrix1, embedding_matrix2, threshold, progress=progress) | |
# 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" | |
# # Show deduplicated examples | |
# result_text += "**Examples of duplicates found in Dataset 2:**\n\n" | |
# num_examples = min(5, num_duplicates) | |
# for duplicate_idx in duplicate_indices_in_ds2[:num_examples]: | |
# original_idx = duplicate_to_original_mapping[duplicate_idx] | |
# original_text = texts1[original_idx] | |
# duplicate_text = texts2[duplicate_idx] | |
# differences = display_word_differences(original_text, duplicate_text) | |
# result_text += f"**Original text (Dataset 1):**\n{original_text}\n\n" | |
# result_text += f"**Duplicate text (Dataset 2):**\n{duplicate_text}\n\n" | |
# result_text += f"**Differences:**\n{differences}\n" | |
# result_text += "-" * 50 + "\n\n" | |
# return result_text | |
# finally: | |
# # Restore original tqdm | |
# tqdm.tqdm = original_tqdm | |
# sys.modules['tqdm'].tqdm = original_tqdm | |
# sys.modules['tqdm.auto'].tqdm = original_tqdm | |
# # Restore reach's original tqdm | |
# if original_reach_tqdm is not None: | |
# Reach.tqdm = original_reach_tqdm | |
# else: | |
# del Reach.tqdm # If it wasn't originally in Reach's __dict__ | |
# 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") | |
# dataset1_text_column = gr.Textbox(value="text", 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="ag_news", label="Dataset 2 Name") | |
# dataset2_split = gr.Textbox(value="test", label="Dataset 2 Split") | |
# dataset2_text_column = gr.Textbox(value="text", label="Text Column Name") | |
# 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_inputs | |
# def update_visibility(deduplication_type_value): | |
# if deduplication_type_value == "Cross-dataset": | |
# return gr.update(visible=True) | |
# else: | |
# return gr.update(visible=False) | |
# 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=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 | |
# # import sys | |
# # import tqdm | |
# # # Load the model at startup | |
# # model = StaticModel.from_pretrained("minishlab/M2V_base_output") | |
# # # Load the default datasets at startup | |
# # default_dataset1_name = "ag_news" | |
# # default_dataset1_split = "train" | |
# # default_dataset2_name = "ag_news" | |
# # default_dataset2_split = "test" | |
# # ds_default1 = load_dataset(default_dataset1_name, split=default_dataset1_split) | |
# # ds_default2 = load_dataset(default_dataset2_name, split=default_dataset2_split) | |
# # def deduplicate(embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024, progress=None) -> 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))]) | |
# # deduplicated_indices = set(range(len(embedding_matrix))) | |
# # duplicate_to_original_mapping = {} | |
# # results = reach.nearest_neighbor_threshold( | |
# # embedding_matrix, | |
# # threshold=threshold, | |
# # batch_size=batch_size, | |
# # show_progressbar=True # Allow internal progress bar | |
# # ) | |
# # # Process duplicates | |
# # for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates")): | |
# # if i not in deduplicated_indices: | |
# # continue | |
# # similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i] | |
# # for sim_idx in similar_indices: | |
# # if sim_idx in deduplicated_indices: | |
# # deduplicated_indices.remove(sim_idx) | |
# # duplicate_to_original_mapping[sim_idx] = i | |
# # 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, progress=None) -> 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))]) | |
# # duplicate_indices_in_test = [] | |
# # duplicate_to_original_mapping = {} | |
# # results = reach.nearest_neighbor_threshold( | |
# # embedding_matrix_2, | |
# # threshold=threshold, | |
# # batch_size=batch_size, | |
# # show_progressbar=True # Allow internal progress bar | |
# # ) | |
# # # Process duplicates | |
# # for i, similar_items in enumerate(progress.tqdm(results, desc="Processing duplicates across datasets")): | |
# # similar_indices = [int(item[0]) for item in similar_items if item[1] >= threshold] | |
# # if similar_indices: | |
# # duplicate_indices_in_test.append(i) | |
# # duplicate_to_original_mapping[i] = similar_indices[0] | |
# # return duplicate_indices_in_test, duplicate_to_original_mapping | |
# # 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 perform_deduplication( | |
# # deduplication_type, | |
# # dataset1_name, | |
# # dataset1_split, | |
# # dataset1_text_column, | |
# # dataset2_name="", | |
# # dataset2_split="", | |
# # dataset2_text_column="", | |
# # threshold=0.8, | |
# # progress=gr.Progress(track_tqdm=True) | |
# # ): | |
# # # Monkey-patch tqdm | |
# # original_tqdm = tqdm.tqdm | |
# # tqdm.tqdm = progress.tqdm | |
# # sys.modules['tqdm'].tqdm = progress.tqdm | |
# # sys.modules['tqdm.auto'].tqdm = progress.tqdm | |
# # try: | |
# # # Convert threshold to float | |
# # threshold = float(threshold) | |
# # if deduplication_type == "Single dataset": | |
# # # Check if the dataset is the default one | |
# # if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split: | |
# # ds = ds_default1 | |
# # else: | |
# # ds = load_dataset(dataset1_name, split=dataset1_split) | |
# # # Extract texts | |
# # texts = [example[dataset1_text_column] for example in ds] | |
# # # Compute embeddings | |
# # embedding_matrix = model.encode(texts, show_progressbar=True) # Enable internal progress bar | |
# # # Deduplicate | |
# # deduplicated_indices, duplicate_to_original_mapping = deduplicate(embedding_matrix, threshold, progress=progress) | |
# # # 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" | |
# # # Show deduplicated examples | |
# # result_text += "**Examples of duplicates found:**\n\n" | |
# # num_examples = min(5, num_duplicates) | |
# # for duplicate_idx, original_idx in list(duplicate_to_original_mapping.items())[:num_examples]: | |
# # original_text = texts[original_idx] | |
# # duplicate_text = texts[duplicate_idx] | |
# # differences = display_word_differences(original_text, duplicate_text) | |
# # result_text += f"**Original text:**\n{original_text}\n\n" | |
# # result_text += f"**Duplicate text:**\n{duplicate_text}\n\n" | |
# # result_text += f"**Differences:**\n{differences}\n" | |
# # result_text += "-" * 50 + "\n\n" | |
# # return result_text | |
# # elif deduplication_type == "Cross-dataset": | |
# # # Dataset 1 | |
# # if dataset1_name == default_dataset1_name and dataset1_split == default_dataset1_split: | |
# # ds1 = ds_default1 | |
# # else: | |
# # ds1 = load_dataset(dataset1_name, split=dataset1_split) | |
# # # Dataset 2 | |
# # if dataset2_name == default_dataset2_name and dataset2_split == default_dataset2_split: | |
# # ds2 = ds_default2 | |
# # else: | |
# # ds2 = load_dataset(dataset2_name, split=dataset2_split) | |
# # # Extract texts | |
# # texts1 = [example[dataset1_text_column] for example in ds1] | |
# # texts2 = [example[dataset2_text_column] for example in ds2] | |
# # # Compute embeddings | |
# # embedding_matrix1 = model.encode(texts1, show_progressbar=True) # Enable internal progress bar | |
# # embedding_matrix2 = model.encode(texts2, show_progressbar=True) # Enable internal progress bar | |
# # # Deduplicate across datasets | |
# # duplicate_indices_in_ds2, duplicate_to_original_mapping = deduplicate_across_datasets(embedding_matrix1, embedding_matrix2, threshold, progress=progress) | |
# # 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" | |
# # # Show deduplicated examples | |
# # result_text += "**Examples of duplicates found in Dataset 2:**\n\n" | |
# # num_examples = min(5, num_duplicates) | |
# # for duplicate_idx in duplicate_indices_in_ds2[:num_examples]: | |
# # original_idx = duplicate_to_original_mapping[duplicate_idx] | |
# # original_text = texts1[original_idx] | |
# # duplicate_text = texts2[duplicate_idx] | |
# # differences = display_word_differences(original_text, duplicate_text) | |
# # result_text += f"**Original text (Dataset 1):**\n{original_text}\n\n" | |
# # result_text += f"**Duplicate text (Dataset 2):**\n{duplicate_text}\n\n" | |
# # result_text += f"**Differences:**\n{differences}\n" | |
# # result_text += "-" * 50 + "\n\n" | |
# # return result_text | |
# # finally: | |
# # # Restore original tqdm | |
# # tqdm.tqdm = original_tqdm | |
# # sys.modules['tqdm'].tqdm = original_tqdm | |
# # sys.modules['tqdm.auto'].tqdm = original_tqdm | |
# # 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") | |
# # dataset1_text_column = gr.Textbox(value="text", 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="ag_news", label="Dataset 2 Name") | |
# # dataset2_split = gr.Textbox(value="test", label="Dataset 2 Split") | |
# # dataset2_text_column = gr.Textbox(value="text", label="Text Column Name") | |
# # 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_inputs | |
# # def update_visibility(deduplication_type_value): | |
# # if deduplication_type_value == "Cross-dataset": | |
# # return gr.update(visible=True) | |
# # else: | |
# # return gr.update(visible=False) | |
# # 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=output | |
# # ) | |
# # demo.launch() | |