Move spaces.GPU to generate_topics

app.py

@@ -15,8 +15,6 @@ from bertopic import BERTopic
 from bertopic.representation import KeyBERTInspired
 from bertopic.representation import TextGeneration
 
-from cuml.manifold import UMAP
-from cuml.cluster import HDBSCAN
 
 from huggingface_hub import HfApi, SpaceCard
 from sklearn.feature_extraction.text import CountVectorizer
@@ -126,7 +124,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)
 
@@ -137,8 +135,11 @@ 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):
+    from cuml.manifold import UMAP
+    from cuml.cluster import HDBSCAN
+
     umap_model = UMAP(
         n_neighbors=n_neighbors,
         n_components=n_components,
@@ -234,7 +235,10 @@ datasets:
     return repo_id
 
 
+@spaces.GPU(duration=600)
 def generate_topics(dataset, config, split, column, nested_column, plot_type):
+    from cuml.manifold import UMAP
+
     logging.info(
         f"Generating topics for {dataset} with config {config} {split} {column} {nested_column}"
     )