Fix for small datasets and custom topics

app.py

@@ -152,7 +152,7 @@ def generate_topics(dataset, config, split, column, nested_column):
     base_model = None
     all_docs = []
     reduced_embeddings_list = []
-
+    topics_info, topic_plot = None, None
     while offset < limit:
         docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size)
         if not docs:
@@ -164,11 +164,13 @@ def generate_topics(dataset, config, split, column, nested_column):
 
         embeddings = calculate_embeddings(docs)
         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()
-        ]
-        base_model.set_topic_labels(llama2_labels)
+
+        repr_model_topics = {
+            key: label[0][0].split("\n")[0]
+            for key, label in base_model.get_topics(full=True)["Llama2"].items()
+        }
+
+        base_model.set_topic_labels(repr_model_topics)
 
         reduced_embeddings = reduce_umap_model.fit_transform(embeddings)
         reduced_embeddings_list.append(reduced_embeddings)
@@ -189,7 +191,7 @@ def generate_topics(dataset, config, split, column, nested_column):
         offset += chunk_size
 
     logging.info("Finished processing all data")
-    return base_model.get_topic_info(), base_model.visualize_topics()
+    return topics_info, topic_plot
 
 
 with gr.Blocks() as demo: