Changing sentence transformer

app.py

@@ -6,7 +6,11 @@ from bertopic import BERTopic
 import pandas as pd
 import gradio as gr
 from bertopic.representation import KeyBERTInspired
-import spaces
+from umap import UMAP
+
+# 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"
@@ -24,7 +28,7 @@ def get_parquet_urls(dataset, config, split):
     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.info(f"Parquet files: {parquet_urls}")
+    logging.debug(f"Parquet files: {parquet_urls}")
     return ",".join(f"'{url}'" for url in parquet_urls)
 
 
@@ -34,7 +38,7 @@ def get_docs_from_parquet(parquet_urls, column, offset, limit):
     logging.debug(f"Dataframe: {df.head(5)}")
     return df[column].tolist()
 
-@spaces.GPU
+
 def generate_topics(dataset, config, split, column, nested_column):
     logging.info(
         f"Generating topics for {dataset} with config {config} {split} {column} {nested_column}"
@@ -45,39 +49,60 @@ def generate_topics(dataset, config, split, column, nested_column):
     chunk_size = 300
     offset = 0
     representation_model = KeyBERTInspired()
-
-    docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size)
-
-    base_model = BERTopic(
-        representation_model=representation_model, min_topic_size=15
-    ).fit(docs)
-
-    yield base_model.get_topic_info(), base_model.visualize_topics()
-
+    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)
+    sentence_model = SentenceTransformer("all-MiniLM-L6-v2", device="cuda")
     while True:
+        docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size)
+        logging.info(f"------------> New chunk data {offset=} {chunk_size=}")
+        embeddings = sentence_model.encode(docs, show_progress_bar=True, batch_size=100)
+        logging.info(f"Embeddings shape: {embeddings.shape}")
         offset = offset + chunk_size
         if not docs or offset >= limit:
             break
 
-        docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size)
-        logging.info(f"------------> New chunk data {offset=} {chunk_size=}")
-        logging.info(docs[:5])
-
         new_model = BERTopic(
-            "english", representation_model=representation_model, min_topic_size=15
-        ).fit(docs)
-        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")
-
-        # Update the base model
-        base_model = updated_model
-
+            "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())
-        yield base_model.get_topic_info(), base_model.visualize_topics()
-
+        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  
+        ) 
+    logging.info("Finished processing all data")
     return base_model.get_topic_info(), base_model.visualize_topics()
 
 

requirements.txt

@@ -4,4 +4,6 @@ umap-learn
 sentence-transformers
 datamapplot
 bertopic
-pandas
+pandas
+torch
+cuml-cu11