Spaces:
Sleeping
Sleeping
import requests | |
import logging | |
import duckdb | |
from gradio_huggingfacehub_search import HuggingfaceHubSearch | |
from bertopic import BERTopic | |
import pandas as pd | |
import gradio as gr | |
from bertopic.representation import KeyBERTInspired | |
logging.basicConfig( | |
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" | |
) | |
session = requests.Session() | |
def get_parquet_urls(dataset, config, split): | |
parquet_files = session.get( | |
f"https://datasets-server.huggingface.co/parquet?dataset={dataset}&config={config}&split={split}", | |
timeout=20, | |
).json() | |
if "error" in parquet_files: | |
raise Exception(f"Error fetching parquet files: {parquet_files['error']}") | |
parquet_urls = [file["url"] for file in parquet_files["parquet_files"]] | |
logging.info(f"Parquet files: {parquet_urls}") | |
return ",".join(f"'{url}'" for url in parquet_urls) | |
def get_docs_from_parquet(parquet_urls, column, offset, limit): | |
SQL_QUERY = f"SELECT {column} FROM read_parquet([{parquet_urls}]) LIMIT {limit} OFFSET {offset};" | |
df = duckdb.sql(SQL_QUERY).to_df() | |
logging.debug(f"Dataframe: {df.head(5)}") | |
return df[column].tolist() | |
def generate_topics(dataset, config, split, column, nested_column, progress): | |
logging.info( | |
f"Generating topics for {dataset} with config {config} {split} {column} {nested_column}" | |
) | |
parquet_urls = get_parquet_urls(dataset, config, split) | |
limit = 1_000 | |
chunk_size = 300 | |
offset = 0 | |
representation_model = KeyBERTInspired() | |
docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size) | |
base_model = BERTopic( | |
representation_model=representation_model, min_topic_size=15 | |
).fit(docs) | |
yield base_model.get_topic_info(), base_model.visualize_topics() | |
while True: | |
offset = offset + chunk_size | |
if not docs or offset >= limit: | |
break | |
docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size) | |
logging.info(f"------------> New chunk data {offset=} {chunk_size=}") | |
logging.info(docs[:5]) | |
new_model = BERTopic( | |
"english", representation_model=representation_model, min_topic_size=15 | |
).fit(docs) | |
updated_model = BERTopic.merge_models([base_model, new_model]) | |
nr_new_topics = len(set(updated_model.topics_)) - len(set(base_model.topics_)) | |
new_topics = list(updated_model.topic_labels_.values())[-nr_new_topics:] | |
logging.info("The following topics are newly found:") | |
logging.info(f"{new_topics}\n") | |
# Update the base model | |
base_model = updated_model | |
logging.info(base_model.get_topic_info()) | |
yield base_model.get_topic_info(), base_model.visualize_topics() | |
return base_model.get_topic_info(), base_model.visualize_topics() | |
with gr.Blocks() as demo: | |
gr.Markdown( | |
""" | |
# 💠 Dataset Topic Discovery 🔭 | |
## Select dataset and text column | |
""" | |
) | |
with gr.Row(): | |
with gr.Column(scale=3): | |
dataset_name = HuggingfaceHubSearch( | |
label="Hub Dataset ID", | |
placeholder="Search for dataset id on Huggingface", | |
search_type="dataset", | |
) | |
subset_dropdown = gr.Dropdown(label="Subset", visible=False) | |
split_dropdown = gr.Dropdown(label="Split", visible=False) | |
with gr.Accordion("Dataset preview", open=False): | |
def embed(name, subset, split): | |
html_code = f""" | |
<iframe | |
src="https://huggingface.co/datasets/{name}/embed/viewer/{subset}/{split}" | |
frameborder="0" | |
width="100%" | |
height="600px" | |
></iframe> | |
""" | |
return gr.HTML(value=html_code) | |
with gr.Row(): | |
text_column_dropdown = gr.Dropdown(label="Text column name") | |
nested_text_column_dropdown = gr.Dropdown( | |
label="Nested text column name", visible=False | |
) | |
generate_button = gr.Button("Generate Notebook", variant="primary") | |
gr.Markdown("## Topics info") | |
progress = gr.Progress() | |
topics_df = gr.DataFrame(interactive=False, visible=True) | |
topics_plot = gr.Plot() | |
generate_button.click( | |
generate_topics, | |
inputs=[ | |
dataset_name, | |
subset_dropdown, | |
split_dropdown, | |
text_column_dropdown, | |
nested_text_column_dropdown, | |
progress, | |
], | |
outputs=[topics_df, topics_plot], | |
) | |
# TODO: choose num_rows, random, or offset -> By default limit max to 1176 rows | |
# -> From the article, it could be in GPU 1176/sec | |
def _resolve_dataset_selection( | |
dataset: str, default_subset: str, default_split: str, text_feature | |
): | |
if "/" not in dataset.strip().strip("/"): | |
return { | |
subset_dropdown: gr.Dropdown(visible=False), | |
split_dropdown: gr.Dropdown(visible=False), | |
text_column_dropdown: gr.Dropdown(label="Text column name"), | |
nested_text_column_dropdown: gr.Dropdown(visible=False), | |
} | |
info_resp = session.get( | |
f"https://datasets-server.huggingface.co/info?dataset={dataset}", timeout=20 | |
).json() | |
if "error" in info_resp: | |
return { | |
subset_dropdown: gr.Dropdown(visible=False), | |
split_dropdown: gr.Dropdown(visible=False), | |
text_column_dropdown: gr.Dropdown(label="Text column name"), | |
nested_text_column_dropdown: gr.Dropdown(visible=False), | |
} | |
subsets: list[str] = list(info_resp["dataset_info"]) | |
subset = default_subset if default_subset in subsets else subsets[0] | |
splits: list[str] = list(info_resp["dataset_info"][subset]["splits"]) | |
split = default_split if default_split in splits else splits[0] | |
features = info_resp["dataset_info"][subset]["features"] | |
def _is_string_feature(feature): | |
return isinstance(feature, dict) and feature.get("dtype") == "string" | |
text_features = [ | |
feature_name | |
for feature_name, feature in features.items() | |
if _is_string_feature(feature) | |
] | |
nested_features = [ | |
feature_name | |
for feature_name, feature in features.items() | |
if isinstance(feature, dict) | |
and isinstance(next(iter(feature.values())), dict) | |
] | |
nested_text_features = [ | |
feature_name | |
for feature_name in nested_features | |
if any( | |
_is_string_feature(nested_feature) | |
for nested_feature in features[feature_name].values() | |
) | |
] | |
if not text_feature: | |
return { | |
subset_dropdown: gr.Dropdown( | |
value=subset, choices=subsets, visible=len(subsets) > 1 | |
), | |
split_dropdown: gr.Dropdown( | |
value=split, choices=splits, visible=len(splits) > 1 | |
), | |
text_column_dropdown: gr.Dropdown( | |
choices=text_features + nested_text_features, | |
label="Text column name", | |
), | |
nested_text_column_dropdown: gr.Dropdown(visible=False), | |
} | |
if text_feature in nested_text_features: | |
nested_keys = [ | |
feature_name | |
for feature_name, feature in features[text_feature].items() | |
if _is_string_feature(feature) | |
] | |
return { | |
subset_dropdown: gr.Dropdown( | |
value=subset, choices=subsets, visible=len(subsets) > 1 | |
), | |
split_dropdown: gr.Dropdown( | |
value=split, choices=splits, visible=len(splits) > 1 | |
), | |
text_column_dropdown: gr.Dropdown( | |
choices=text_features + nested_text_features, | |
label="Text column name", | |
), | |
nested_text_column_dropdown: gr.Dropdown( | |
value=nested_keys[0], | |
choices=nested_keys, | |
label="Nested text column name", | |
visible=True, | |
), | |
} | |
return { | |
subset_dropdown: gr.Dropdown( | |
value=subset, choices=subsets, visible=len(subsets) > 1 | |
), | |
split_dropdown: gr.Dropdown( | |
value=split, choices=splits, visible=len(splits) > 1 | |
), | |
text_column_dropdown: gr.Dropdown( | |
choices=text_features + nested_text_features, label="Text column name" | |
), | |
nested_text_column_dropdown: gr.Dropdown(visible=False), | |
} | |
def show_input_from_subset_dropdown(dataset: str) -> dict: | |
return _resolve_dataset_selection( | |
dataset, default_subset="default", default_split="train", text_feature=None | |
) | |
def show_input_from_subset_dropdown(dataset: str, subset: str) -> dict: | |
return _resolve_dataset_selection( | |
dataset, default_subset=subset, default_split="train", text_feature=None | |
) | |
def show_input_from_split_dropdown(dataset: str, subset: str, split: str) -> dict: | |
return _resolve_dataset_selection( | |
dataset, default_subset=subset, default_split=split, text_feature=None | |
) | |
def show_input_from_text_column_dropdown( | |
dataset: str, subset: str, split: str, text_column | |
) -> dict: | |
return _resolve_dataset_selection( | |
dataset, | |
default_subset=subset, | |
default_split=split, | |
text_feature=text_column, | |
) | |
demo.launch() | |