import spaces 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, MaximalMarginalRelevance, TextGeneration, ) from umap import UMAP import numpy as np from torch import cuda from torch import bfloat16 from transformers import ( # BitsAndBytesConfig, AutoTokenizer, AutoModelForCausalLM, pipeline, ) from prompts import system_prompt, example_prompt, main_prompt from umap import UMAP from hdbscan import HDBSCAN # from cuml.cluster import HDBSCAN # from cuml.manifold import UMAP from sentence_transformers import SentenceTransformer logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) session = requests.Session() sentence_model = SentenceTransformer("all-MiniLM-L6-v2") keybert = KeyBERTInspired() mmr = MaximalMarginalRelevance(diversity=0.3) model_id = "meta-llama/Llama-2-7b-chat-hf" device = f"cuda:{cuda.current_device()}" if cuda.is_available() else "cpu" logging.info(device) # bnb_config = BitsAndBytesConfig( # load_in_4bit=True, # 4-bit quantization # bnb_4bit_quant_type="nf4", # Normalized float 4 # bnb_4bit_use_double_quant=True, # Second quantization after the first # bnb_4bit_compute_dtype=bfloat16, # Computation type # ) tokenizer = AutoTokenizer.from_pretrained(model_id) # Llama 2 Model model = AutoModelForCausalLM.from_pretrained( model_id, trust_remote_code=True, # quantization_config=bnb_config, device_map="auto", ) generator = pipeline( model=model, tokenizer=tokenizer, task="text-generation", temperature=0.1, max_new_tokens=500, repetition_penalty=1.1, ) prompt = system_prompt + example_prompt + main_prompt llama2 = TextGeneration(generator, prompt=prompt) representation_model = { "KeyBERT": keybert, "Llama2": llama2, # "MMR": mmr, } umap_model = UMAP( n_neighbors=15, n_components=5, min_dist=0.0, metric="cosine", random_state=42 ) hdbscan_model = HDBSCAN( min_cluster_size=15, metric="euclidean", cluster_selection_method="eom", prediction_data=True, ) reduce_umap_model = UMAP( n_neighbors=15, n_components=2, min_dist=0.0, metric="cosine", random_state=42 ) 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.debug(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() @spaces.GPU def calculate_embeddings(docs): return sentence_model.encode(docs, show_progress_bar=True, batch_size=100) @spaces.GPU def fit_model(base_model, docs, embeddings): new_model = BERTopic( "english", # Sub-models embedding_model=sentence_model, umap_model=umap_model, hdbscan_model=hdbscan_model, representation_model=representation_model, # Hyperparameters top_n_words=10, verbose=True, min_topic_size=15, ) logging.info("Fitting new model") new_model.fit(docs, embeddings) logging.info("End fitting new model") if base_model is None: return new_model, new_model 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(f"The following topics are newly found: {new_topics}") return updated_model, new_model def generate_topics(dataset, config, split, column, nested_column): 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 base_model = None all_docs = [] all_reduced_embeddings = np.empty((0, 2)) while True: docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size) logging.info( f"------------> New chunk data {offset=} {chunk_size=} with {len(docs)} docs" ) embeddings = calculate_embeddings(docs) offset = offset + chunk_size if not docs or offset >= limit: break base_model, _ = fit_model(base_model, docs, embeddings) llama2_labels = [ label[0][0].split("\n")[0] for label in base_model.get_topics(full=True)["Llama2"].values() ] logging.info(f"Topics: {llama2_labels}") base_model.set_topic_labels(llama2_labels) reduced_embeddings = reduce_umap_model.fit_transform(embeddings) all_docs.extend(docs) all_reduced_embeddings = np.vstack((all_reduced_embeddings, reduced_embeddings)) topics_info = base_model.get_topic_info() topic_plot = base_model.visualize_documents( all_docs, reduced_embeddings=all_reduced_embeddings, custom_labels=True ) logging.info(f"Topics for merged model: {base_model.topic_labels_}") yield topics_info, topic_plot logging.info("Finished processing all data") return base_model.get_topic_info(), base_model.visualize_topics() with gr.Blocks() as demo: gr.Markdown("# 💠 Dataset Topic Discovery 🔭") gr.Markdown("## Select dataset and text column") with gr.Accordion("Data details", open=True): 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""" """ 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("## Datamap") topics_plot = gr.Plot() with gr.Accordion("Topics Info", open=False): topics_df = gr.DataFrame(interactive=False, visible=True) generate_button.click( generate_topics, inputs=[ dataset_name, subset_dropdown, split_dropdown, text_column_dropdown, nested_text_column_dropdown, ], 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()