Batched fit

app.py

@@ -1,7 +1,304 @@
+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
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+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):
+
+        @gr.render(inputs=[dataset_name, subset_dropdown, split_dropdown])
+        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),
+        }
+
+    @dataset_name.change(
+        inputs=[dataset_name],
+        outputs=[
+            subset_dropdown,
+            split_dropdown,
+            text_column_dropdown,
+            nested_text_column_dropdown,
+        ],
+    )
+    def show_input_from_subset_dropdown(dataset: str) -> dict:
+        return _resolve_dataset_selection(
+            dataset, default_subset="default", default_split="train", text_feature=None
+        )
+
+    @subset_dropdown.change(
+        inputs=[dataset_name, subset_dropdown],
+        outputs=[
+            subset_dropdown,
+            split_dropdown,
+            text_column_dropdown,
+            nested_text_column_dropdown,
+        ],
+    )
+    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
+        )
+
+    @split_dropdown.change(
+        inputs=[dataset_name, subset_dropdown, split_dropdown],
+        outputs=[
+            subset_dropdown,
+            split_dropdown,
+            text_column_dropdown,
+            nested_text_column_dropdown,
+        ],
+    )
+    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
+        )
+
+    @text_column_dropdown.change(
+        inputs=[dataset_name, subset_dropdown, split_dropdown, text_column_dropdown],
+        outputs=[
+            subset_dropdown,
+            split_dropdown,
+            text_column_dropdown,
+            nested_text_column_dropdown,
+        ],
+    )
+    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()

requirements.txt

@@ -0,0 +1,7 @@
+gradio_huggingfacehub_search==0.0.7
+duckdb
+umap-learn
+sentence-transformers
+datamapplot
+bertopic
+pandas