Enable spaces

app.py

@@ -1,4 +1,4 @@
-# import spaces
+import spaces
 import requests
 import logging
 import duckdb
@@ -57,7 +57,7 @@ model = AutoModelForCausalLM.from_pretrained(
     offload_folder="offload",  # Offloading part of the model to CPU to save GPU memory
 )
 
-# Enable gradient checkpointing for memory efficiency during backprop
+# Enable gradient checkpointing for memory efficiency during backprop?
 model.gradient_checkpointing_enable()
 
 generator = pipeline(
@@ -122,13 +122,13 @@ def get_docs_from_parquet(parquet_urls, column, offset, limit):
     return df[column].tolist()
 
 
-# @spaces.GPU
+@spaces.GPU
 # TODO: Modify batch size to reduce memory consumption during embedding calculation, which value is better?
 def calculate_embeddings(docs):
     return sentence_model.encode(docs, show_progress_bar=True, batch_size=32)
 
 
-# @spaces.GPU
+@spaces.GPU
 def fit_model(base_model, docs, embeddings):
     new_model = BERTopic(
         "english",
@@ -195,12 +195,11 @@ def generate_topics(dataset, config, split, column, nested_column):
         all_docs.extend(docs)
 
         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,
-        # )
-        topic_plot = base_model.visualize_barchart()
+        topic_plot = base_model.visualize_documents(
+            all_docs,
+            reduced_embeddings=np.vstack(reduced_embeddings_list),
+            custom_labels=True,
+        )
 
         logging.info(f"Topics: {repr_model_topics}")