Adding viz for merged model

app.py

@@ -1,4 +1,4 @@
-import spaces
+# import spaces
 import requests
 import logging
 import duckdb
@@ -8,6 +8,7 @@ import pandas as pd
 import gradio as gr
 from bertopic.representation import KeyBERTInspired
 from umap import UMAP
+import numpy as np
 
 # from cuml.cluster import HDBSCAN
 # from cuml.manifold import UMAP
@@ -41,14 +42,14 @@ def get_docs_from_parquet(parquet_urls, column, offset, limit):
     return df[column].tolist()
 
 
-@spaces.GPU
+# @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
 
 
-@spaces.GPU
+# @spaces.GPU
 def fit_model(base_model, sentence_model, representation_model, docs, embeddings):
     new_model = BERTopic(
         "english",
@@ -81,59 +82,35 @@ def generate_topics(dataset, config, split, column, nested_column):
     offset = 0
     representation_model = KeyBERTInspired()
     base_model = None
-    # docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size)
-
-    # base_model = BERTopic(
-    #     "english", representation_model=representation_model, min_topic_size=15
-    # )
-    # base_model.fit_transform(docs)
-
-    # yield base_model.get_topic_info(), base_model.visualize_topics()
-    # Create instances of GPU-accelerated UMAP and HDBSCAN
-    # umap_model = UMAP(n_components=5, n_neighbors=15, min_dist=0.0)
-    # hdbscan_model = HDBSCAN(min_samples=10, gen_min_span_tree=True)
+    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=}")
+        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
-
-        # new_model = BERTopic(
-        #     "english",
-        #     embedding_model=sentence_model,
-        #     representation_model=representation_model,
-        #     min_topic_size=15,  # umap_model=umap_model, hdbscan_model=hdbscan_model
-        # )
-        # logging.info("Fitting new model")
-        # new_model.fit(docs, embeddings)
-        # logging.info("End fitting new model")
-        # if base_model is not None:
-        #     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")
-        #     base_model = updated_model
-        # else:
-        #     base_model = new_model
-        # logging.info(base_model.get_topic_info())
-        base_model, new_model = fit_model(
+        base_model, _ = fit_model(
             base_model, sentence_model, representation_model, docs, embeddings
         )
-        # 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}")
-        yield (
-            base_model.get_topic_info(),
-            new_model.visualize_documents(
-                docs, embeddings=embeddings
-            ),  # TODO: Visualize the merged models
+        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
         )
+
+        yield topics_info, topic_plot
+
     logging.info("Finished processing all data")
     return base_model.get_topic_info(), base_model.visualize_topics()
 

requirements.txt

@@ -7,4 +7,5 @@ sentence-transformers
 datamapplot
 bertopic
 pandas
-torch
+torch
+numpy