Try to run on zero with custom componetns

app.py

@@ -1,4 +1,4 @@
-# import spaces
+import spaces
 import requests
 import logging
 import duckdb
@@ -73,23 +73,25 @@ prompt = system_prompt + example_prompt + main_prompt
 
 llama2 = TextGeneration(generator, prompt=prompt)
 representation_model = {
-    # "KeyBERT": keybert,
+    "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
-# )
+umap_model = UMAP(
+    n_neighbors=15, n_components=5, min_dist=0.0, metric="cosine", random_state=42
+)
+
+hdbscan_model = HDBSCAN(
+    min_cluster_size=15,
+    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):
@@ -111,19 +113,19 @@ def get_docs_from_parquet(parquet_urls, column, offset, limit):
     return df[column].tolist()
 
 
-# @spaces.GPU
+@spaces.GPU
 def calculate_embeddings(docs):
     return sentence_model.encode(docs, show_progress_bar=True, batch_size=100)
 
 
-# @spaces.GPU
+@spaces.GPU
 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,
+        umap_model=umap_model,
+        hdbscan_model=hdbscan_model,
         representation_model=representation_model,
         # Hyperparameters
         top_n_words=10,
@@ -140,10 +142,7 @@ def fit_model(base_model, docs, embeddings):
     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")
-    # updated_model.set_topic_labels(updated_model.topic_labels_)
-
+    logging.info(f"The following topics are newly found: {new_topics}")
     return updated_model, new_model
 
 
@@ -176,9 +175,7 @@ def generate_topics(dataset, config, split, column, nested_column):
         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)
+        reduced_embeddings = reduce_umap_model.fit_transform(embeddings)
 
         all_docs.extend(docs)
         all_reduced_embeddings = np.vstack((all_reduced_embeddings, reduced_embeddings))