Adding plot type selection

app.py

@@ -1,29 +1,39 @@
-import requests
 import logging
+import os
+
+import datamapplot
 import duckdb
 import numpy as np
+import requests
+
 from torch import cuda
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from bertopic import BERTopic
 from bertopic.representation import KeyBERTInspired
-
 from cuml.manifold import UMAP
 from cuml.cluster import HDBSCAN
 
 from sklearn.feature_extraction.text import CountVectorizer
-
 from sentence_transformers import SentenceTransformer
 
 from dotenv import load_dotenv
-import os
 
+# These imports at the end because of torch/datamapplot issue in Zero GPU
 # import spaces
 import gradio as gr
 
 
 """
 TODOs:
-- Try for small dataset <1000 rows
+- Hide params panel when generating plot
+
+- Improve DataMapPlot plot arguments
+- Add export button for final plot
+- Export and serve an interactive HTML plot?
+- Try with more rows
+
+- Add TextGenerationLayer
+- Make it run on Zero GPU
 """
 
 load_dotenv()
@@ -137,7 +147,7 @@ def fit_model(docs, embeddings, n_neighbors, n_components):
     logging.info("Global model updated")
 
 
-def generate_topics(dataset, config, split, column, nested_column):
+def generate_topics(dataset, config, split, column, nested_column, plot_type):
     logging.info(
         f"Generating topics for {dataset} with config {config} {split} {column} {nested_column}"
     )
@@ -202,12 +212,65 @@ def generate_topics(dataset, config, split, column, nested_column):
         reduced_embeddings_list.append(reduced_embeddings)
 
         all_docs.extend(docs)
+        reduced_embeddings_array = np.vstack(reduced_embeddings_list)
 
         topics_info = base_model.get_topic_info()
-        topic_plot = base_model.visualize_documents(
-            all_docs,
-            reduced_embeddings=np.vstack(reduced_embeddings_list),
-            custom_labels=True,
+        all_topics, _ = base_model.transform(all_docs)
+        all_topics = np.array(all_topics)
+        # topic_plot, _ = datamapplot.create_plot(
+        #     data_map_coords=reduced_embeddings_array,
+        #     labels=all_topics.astype(str),
+        #     use_medoids=True,
+        #     figsize=(12, 12),
+        #     dpi=100,
+        #     title="PubMed - Literature review",
+        #     sub_title="A data map of papers representing artificial intelligence and machine learning in ophthalmology",
+        #     title_keywords={"fontsize": 36, "fontfamily": "Roboto Black"},
+        #     sub_title_keywords={
+        #         "fontsize": 18,
+        #     },
+        #     highlight_label_keywords={
+        #         "fontsize": 12,
+        #         "fontweight": "bold",
+        #         "bbox": {"boxstyle": "round"},
+        #     },
+        #     label_font_size=8,
+        #     label_wrap_width=16,
+        #     label_linespacing=1.25,
+        #     label_direction_bias=1.3,
+        #     label_margin_factor=2.0,
+        #     label_base_radius=15.0,
+        #     point_size=4,
+        #     marker_type="o",
+        #     arrowprops={
+        #         "arrowstyle": "wedge,tail_width=0.5",
+        #         "connectionstyle": "arc3,rad=0.05",
+        #         "linewidth": 0,
+        #         "fc": "#33333377",
+        #     },
+        #     add_glow=True,
+        #     glow_keywords={
+        #         "kernel_bandwidth": 0.75,  # controls how wide the glow spreads.
+        #         "kernel": "cosine",  # controls the kernel type. Default is "gaussian". See https://scikit-learn.org/stable/modules/density.html#kernel-density.
+        #         "n_levels": 32,  # controls how many "levels" there are in the contour plot.
+        #         "max_alpha": 0.9,  # controls the translucency of the glow.
+        #     },
+        #     darkmode=False,
+        # )
+
+        topic_plot = (
+            base_model.visualize_document_datamap(
+                docs=all_docs,
+                reduced_embeddings=reduced_embeddings_array,
+                title=f"<b>{dataset}</b>",
+            )
+            if plot_type == "DataMapPlot"
+            else base_model.visualize_documents(
+                docs=all_docs,
+                reduced_embeddings=reduced_embeddings_array,
+                custom_labels=True,
+                title=f"<b>{dataset}</b>",
+            )
         )
 
         rows_processed += len(docs)
@@ -228,6 +291,7 @@ def generate_topics(dataset, config, split, column, nested_column):
         offset += CHUNK_SIZE
 
     logging.info("Finished processing all data")
+
     yield (
         topics_info,
         topic_plot,
@@ -271,7 +335,12 @@ with gr.Blocks() as demo:
             nested_text_column_dropdown = gr.Dropdown(
                 label="Nested text column name", visible=False
             )
-
+            plot_type_radio = gr.Radio(
+                ["DataMapPlot", "Plotly"],
+                value="Plotly",
+                label="Choose the plot type",
+                interactive=True,
+            )
         generate_button = gr.Button("Generate Topics", variant="primary")
 
     gr.Markdown("## Data map")
@@ -287,8 +356,13 @@ with gr.Blocks() as demo:
             split_dropdown,
             text_column_dropdown,
             nested_text_column_dropdown,
+            plot_type_radio,
+        ],
+        outputs=[
+            topics_df,
+            topics_plot,
+            full_topics_generation_label,
         ],
-        outputs=[topics_df, topics_plot, full_topics_generation_label],
     )
 
     def _resolve_dataset_selection(