import gradio as gr
from datasets import load_dataset, Dataset
from difflib import ndiff
import pandas as pd
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from semhash import SemHash
from semhash.datamodels import DeduplicationResult
from model2vec import StaticModel
# Default parameters
default_dataset_name = "SetFit/amazon_massive_scenario_en-US"
default_dataset1_split = "train"
default_dataset2_split = "test"
default_text_column = "text"
default_threshold = 0.9
# 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.
"""
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
) -> tuple[list[str], Dataset]:
"""Load texts from a specified dataset split."""
ds = load_dataset(dataset_name, split=dataset_split)
return [example[text_column] for example in ds], 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 create_deduplicated_dataset(
original_dataset: Dataset, deduplicated_texts: list[str], text_column: str
) -> Dataset:
"""Create a new dataset with only the deduplicated texts."""
# Create a mapping from text to original row
text_to_row = {row[text_column]: row for row in original_dataset}
# Build new dataset with deduplicated texts
deduplicated_rows = []
for text in deduplicated_texts:
if text in text_to_row:
deduplicated_rows.append(text_to_row[text])
return Dataset.from_list(deduplicated_rows)
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 using SemHash. This function
streams status updates to Gradio for user feedback.
"""
try:
threshold = float(threshold)
# Load Dataset 1
texts1, dataset1 = load_dataset_texts(
dataset1_name, dataset1_split, dataset1_text_column
)
if deduplication_type == "Single dataset":
# Single-dataset deduplication
result = deduplicate_single_dataset(texts1, threshold=threshold)
# Sort all duplicates by score (ascending for least similar)
for duprec in result.duplicates:
duprec.duplicates.sort(key=lambda x: x[1])
# Create deduplicated dataset
deduplicated_dataset = create_deduplicated_dataset(
dataset1, result.deduplicated, dataset1_text_column
)
# Summarize results
num_duplicates = len(result.duplicates)
deduplicated_count = len(result.deduplicated)
total_docs = len(texts1)
# Create examples table
examples_table = None
if num_duplicates > 0:
# Only show duplicates that actually have near-duplicate records
duplicates_with_data = [
duprec for duprec in result.duplicates if duprec.duplicates
]
# sort duplicates by score (ascending for least similar)
for duprec in result.duplicates:
duprec.duplicates.sort(key=lambda x: x[1])
if duplicates_with_data:
# Create table data for the 5 least similar examples
table_data = []
for duprec in duplicates_with_data[:5]:
dup_text = duprec.record
orig_text, score = duprec.duplicates[0]
table_data.append(
[
orig_text[:200] + "..."
if len(orig_text) > 200
else orig_text,
dup_text[:200] + "..."
if len(dup_text) > 200
else dup_text,
f"{score:.4f}",
]
)
examples_table = pd.DataFrame(
table_data,
columns=["Original Text", "Duplicate Text", "Similarity Score"],
)
# Show success info with stats
gr.Info(
f"Deduplication completed! Found {num_duplicates} duplicates. "
f"Dataset reduced from {total_docs} to {deduplicated_count} unique documents."
)
# Return table with visibility update
if examples_table is not None and not examples_table.empty:
return deduplicated_dataset, gr.update(
visible=True, value=examples_table
)
else:
return deduplicated_dataset, gr.update(visible=False)
else:
# Cross-dataset deduplication
texts2, dataset2 = load_dataset_texts(
dataset2_name, dataset2_split, dataset2_text_column
)
result = deduplicate_two_datasets(texts1, texts2, threshold=threshold)
# Sort duplicates by score (ascending for least similar)
for duprec in result.duplicates:
duprec.duplicates.sort(key=lambda x: x[1])
# Create deduplicated dataset from dataset2
deduplicated_dataset = create_deduplicated_dataset(
dataset2, result.deduplicated, dataset2_text_column
)
num_duplicates = len(result.duplicates)
total_docs2 = len(texts2)
deduplicated_count = len(result.deduplicated)
# Create examples table
examples_table = None
if num_duplicates > 0:
# Again, only show duplicates that have records
duplicates_with_data = [
duprec for duprec in result.duplicates if duprec.duplicates
]
if duplicates_with_data:
# Create table data for the 5 least similar examples
table_data = []
for duprec in duplicates_with_data[:5]:
dup_text = duprec.record
orig_text, score = duprec.duplicates[0]
table_data.append(
[
orig_text[:200] + "..."
if len(orig_text) > 200
else orig_text,
dup_text[:200] + "..."
if len(dup_text) > 200
else dup_text,
f"{score:.4f}",
]
)
examples_table = pd.DataFrame(
table_data,
columns=[
"Original Text (Dataset 1)",
"Duplicate Text (Dataset 2)",
"Similarity Score",
],
)
# Show success info with stats
gr.Info(
f"Deduplication completed! Found {num_duplicates} duplicates in Dataset 2. "
f"Dataset reduced from {total_docs2} to {deduplicated_count} unique documents."
)
# Return table with visibility update
if examples_table is not None and not examples_table.empty:
return deduplicated_dataset, gr.update(
visible=True, value=examples_table
)
else:
return deduplicated_dataset, gr.update(visible=False)
except Exception as e:
gr.Error(f"An error occurred during deduplication: {str(e)}")
return None, gr.update(visible=False)
def push_to_hub(
deduplicated_dataset: Dataset,
output_dataset_name: str,
oauth_profile: gr.OAuthProfile | None,
oauth_token: gr.OAuthToken | None,
progress: gr.Progress = gr.Progress(),
) -> str:
"""Push the deduplicated dataset to Hugging Face Hub."""
if oauth_token is None:
raise gr.Error("Please log in with Hugging Face to push datasets to the Hub.")
if not output_dataset_name.strip():
raise gr.Error("Please provide a dataset name.")
if deduplicated_dataset is None:
raise gr.Error(
"No deduplicated dataset available. Please run deduplication first."
)
try:
progress(0.1, desc="Preparing dataset...")
# Determine the full dataset name (username/dataset_name)
username = oauth_profile.username if oauth_profile else None
if "/" not in output_dataset_name and username:
full_dataset_name = f"{username}/{output_dataset_name}"
else:
full_dataset_name = output_dataset_name
progress(0.3, desc="Pushing to Hub...")
# Push to hub using the OAuth token
deduplicated_dataset.push_to_hub(
full_dataset_name, token=oauth_token.token, private=False
)
progress(1.0, desc="Complete!")
gr.Info(
f"Successfully pushed deduplicated dataset with {len(deduplicated_dataset)} rows to the Hub!"
)
return (
f"✅ **Dataset published:** [{full_dataset_name}]"
f"(https://huggingface.co/datasets/{full_dataset_name})"
)
except Exception as e:
raise gr.Error(f"Failed to push dataset to Hub: {str(e)}")
def get_user_info(oauth_profile: gr.OAuthProfile | None) -> str:
"""Display user login status."""
if oauth_profile is None:
return "Not logged in. Please log in to push datasets to the Hub."
return f"Logged in as: **{oauth_profile.username}**"
def update_push_button_state(oauth_profile: gr.OAuthProfile | None):
"""Update the push button state based on login status."""
is_logged_in = oauth_profile is not None
return gr.update(interactive=is_logged_in)
# --- Gradio App ---
with gr.Blocks(
theme=gr.themes.Ocean(), css="#status_output { height: 50px; overflow: auto; }"
) as demo:
gr.Markdown("# SemDedup-My-Dataset: Semantic Text Deduplication Using SemHash")
gr.Markdown("""
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.
""")
deduplication_type = gr.Radio(
choices=["Cross-dataset", "Single dataset"],
label="Deduplication Type",
value="Cross-dataset", # default
)
with gr.Row():
dataset1_name = HuggingfaceHubSearch(
label="Dataset 1 Name",
placeholder="Search for datasets on HuggingFace Hub",
search_type="dataset",
value=default_dataset_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=True)
with dataset2_inputs:
with gr.Row():
dataset2_name = HuggingfaceHubSearch(
label="Dataset 2 Name",
placeholder="Search for datasets on HuggingFace Hub",
search_type="dataset",
value=default_dataset_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(
0.0, 1.0, value=default_threshold, label="Similarity Threshold"
)
with gr.Row():
compute_button = gr.Button("Deduplicate", variant="primary")
status_output = gr.Markdown(elem_id="status_output")
# Examples table
examples_table = gr.Dataframe(
headers=["Original Text", "Duplicate Text", "Similarity Score"],
datatype=["str", "str", "str"],
)
# Hidden state to store the deduplicated dataset
deduplicated_dataset_state = gr.State()
# Output dataset configuration
gr.Markdown("## Push Deduplicated Dataset to Hub")
with gr.Row():
with gr.Column():
output_dataset_name = gr.Textbox(
label="Output Dataset Name",
placeholder="my-deduplicated-dataset",
info="Will be saved as username/dataset-name",
)
with gr.Column():
push_button = gr.Button(
"Push to Hub", variant="secondary", interactive=False
)
login_button = gr.LoginButton()
# Login section - moved below push to hub
with gr.Row():
user_info = gr.Markdown()
push_output = gr.Markdown()
# HACK: for some reason gradio wants this.
login_button.activate()
def update_visibility(choice: str):
return gr.update(visible=(choice == "Cross-dataset"))
deduplication_type.change(
update_visibility, inputs=deduplication_type, outputs=dataset2_inputs
)
# Update user info and button state when page loads or login status changes
demo.load(get_user_info, inputs=None, outputs=user_info)
demo.load(update_push_button_state, inputs=None, outputs=push_button)
login_button.click(get_user_info, inputs=None, outputs=user_info)
login_button.click(update_push_button_state, inputs=None, outputs=push_button)
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=[deduplicated_dataset_state, examples_table],
)
push_button.click(
fn=push_to_hub,
inputs=[
deduplicated_dataset_state,
output_dataset_name,
],
outputs=push_output,
)
demo.launch(ssr_mode=False)