Enable Zero GPU

app.py

@@ -26,7 +26,7 @@ from transformers import (
 from torch import bfloat16
 
 # These imports at the end because of torch/datamapplot issue in Zero GPU
-# import spaces
+import spaces
 import gradio as gr
 
 from prompts import REPRESENTATION_PROMPT
@@ -128,7 +128,7 @@ 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=32)
 
@@ -139,7 +139,7 @@ def calculate_n_neighbors_and_components(n_rows):
     return n_neighbors, n_components
 
 
-# @spaces.GPU
+@spaces.GPU
 def fit_model(docs, embeddings, n_neighbors, n_components):
     global global_topic_model
 
@@ -204,6 +204,7 @@ def _push_to_hub(
 
 
 def generate_topics(dataset, config, split, column, nested_column, plot_type):
+    global global_topic_model
     logging.info(
         f"Generating topics for {dataset} with config {config} {split} {column} {nested_column}"
     )