Adding Llama2 topics

app.py

@@ -6,9 +6,24 @@ from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from bertopic import BERTopic
 import pandas as pd
 import gradio as gr
-from bertopic.representation import KeyBERTInspired
+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
@@ -21,6 +36,60 @@ logging.basicConfig(
 
 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=150,
+#     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):
@@ -44,22 +113,27 @@ def get_docs_from_parquet(parquet_urls, column, offset, limit):
 
 # @spaces.GPU
 def calculate_embeddings(docs):
-    embeddings = sentence_model.encode(docs, show_progress_bar=True, batch_size=100)
-    logging.info(f"Embeddings shape: {embeddings.shape}")
-    return embeddings
+    return sentence_model.encode(docs, show_progress_bar=True, batch_size=100)
 
 
 # @spaces.GPU
-def fit_model(base_model, sentence_model, representation_model, docs, embeddings):
+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,
-        min_topic_size=15,  # umap_model=umap_model, hdbscan_model=hdbscan_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
 
@@ -68,6 +142,8 @@ def fit_model(base_model, sentence_model, representation_model, docs, embeddings
     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")
+    # updated_model.set_topic_labels(updated_model.topic_labels_)
+
     return updated_model, new_model
 
 
@@ -80,7 +156,6 @@ def generate_topics(dataset, config, split, column, nested_column):
     limit = 1_000
     chunk_size = 300
     offset = 0
-    representation_model = KeyBERTInspired()
     base_model = None
     all_docs = []
     all_reduced_embeddings = np.empty((0, 2))
@@ -93,22 +168,25 @@ def generate_topics(dataset, config, split, column, nested_column):
         offset = offset + chunk_size
         if not docs or offset >= limit:
             break
-        base_model, _ = fit_model(
-            base_model, sentence_model, representation_model, docs, embeddings
-        )
+        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 = UMAP(
             n_neighbors=10, n_components=2, min_dist=0.0, metric="cosine"
         ).fit_transform(embeddings)
-        logging.info(f"Reduced embeddings shape: {reduced_embeddings.shape}")
 
         all_docs.extend(docs)
         all_reduced_embeddings = np.vstack((all_reduced_embeddings, reduced_embeddings))
-        logging.info(f"Stacked embeddings shape: {all_reduced_embeddings.shape}")
         topics_info = base_model.get_topic_info()
         topic_plot = base_model.visualize_documents(
-            all_docs, reduced_embeddings=all_reduced_embeddings
+            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")
@@ -116,47 +194,45 @@ def generate_topics(dataset, config, split, column, nested_column):
 
 
 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",
+    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"""
+                <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
             )
-        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")
+        generate_button = gr.Button("Generate Notebook", variant="primary")
 
-    gr.Markdown("## Topics info")
-    topics_df = gr.DataFrame(interactive=False, visible=True)
+    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=[

prompts.py

@@ -0,0 +1,29 @@
+system_prompt = """
+<s>[INST] <<SYS>>
+You are a helpful, respectful and honest assistant for labeling topics.
+<</SYS>>
+"""
+
+example_prompt = """
+I have a topic that contains the following documents:
+- Traditional diets in most cultures were primarily plant-based with a little meat on top, but with the rise of industrial style meat production and factory farming, meat has become a staple food.
+- Meat, but especially beef, is the word food in terms of emissions.
+- Eating meat doesn't make you a bad person, not eating meat doesn't make you a good one.
+
+The topic is described by the following keywords: 'meat, beef, eat, eating, emissions, steak, food, health, processed, chicken'.
+
+Based on the information about the topic above, please create a short label of this topic. Make sure you to only return the label and nothing more.
+
+[/INST] Environmental impacts of eating meat
+"""
+
+main_prompt = """
+[INST]
+I have a topic that contains the following documents:
+[DOCUMENTS]
+
+The topic is described by the following keywords: '[KEYWORDS]'.
+
+Based on the information about the topic above, please create a short label of this topic. Make sure you to only return the label and nothing more.
+[/INST]
+"""