Trying to load in cuda only within spaces environment to enable zero GPU space to run successfully

app.py

@@ -1,10 +1,11 @@
+import spaces
 import os
 import socket
 import gradio as gr
 import pandas as pd
 import numpy as np
 from datetime import datetime
-import spaces
+
 
 from funcs.topic_core_funcs import pre_clean, optimise_zero_shot, extract_topics, reduce_outliers, represent_topics, visualise_topics, save_as_pytorch_model, change_default_vis_col
 from funcs.helper_functions import initial_file_load, custom_regex_load, ensure_output_folder_exists, output_folder, get_connection_params, get_or_create_env_var

funcs/embeddings.py

@@ -2,29 +2,18 @@ import time
 import numpy as np
 import os
 import spaces
-from torch import cuda, backends, version
+
 from sentence_transformers import SentenceTransformer
 from sklearn.pipeline import make_pipeline
 from sklearn.decomposition import TruncatedSVD
 from sklearn.feature_extraction.text import TfidfVectorizer
 
-# Check for torch cuda
+
 # If you want to disable cuda for testing purposes
 #os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
 
-print("Is CUDA enabled? ", cuda.is_available())
-print("Is a CUDA device available on this computer?", backends.cudnn.enabled)
-if cuda.is_available():
-    torch_device = "gpu"
-    print("Cuda version installed is: ", version.cuda)
-    high_quality_mode = "Yes"
-    os.system("nvidia-smi")
-else: 
-    torch_device =  "cpu"
-    high_quality_mode = "No"
 
-
-@spaces.GPU
+@spaces.GPU(duration=120)
 def make_or_load_embeddings(docs: list, file_list: list, embeddings_out: np.ndarray, embeddings_super_compress: str, high_quality_mode_opt: str, embeddings_name:str="mixedbread-ai/mxbai-embed-xsmall-v1", random_seed:int=42) -> np.ndarray:
     """
     Create or load embeddings for the given documents.
@@ -41,6 +30,20 @@ def make_or_load_embeddings(docs: list, file_list: list, embeddings_out: np.ndar
         np.ndarray: The generated or loaded embeddings.
     """
 
+    # Check for torch cuda
+    from torch import cuda, backends, version
+
+    print("Is CUDA enabled? ", cuda.is_available())
+    print("Is a CUDA device available on this computer?", backends.cudnn.enabled)
+    if cuda.is_available():
+        torch_device = "gpu"
+        print("Cuda version installed is: ", version.cuda)
+        high_quality_mode = "Yes"
+        os.system("nvidia-smi")
+    else: 
+        torch_device =  "cpu"
+        high_quality_mode = "No"
+
     if high_quality_mode_opt == "Yes":
     # Define a list of possible local locations to search for the model
         local_embeddings_locations = [

funcs/representation_model.py

@@ -2,7 +2,7 @@ import os
 from bertopic.representation import LlamaCPP
 
 from pydantic import BaseModel
-import torch.cuda
+
 from huggingface_hub import hf_hub_download
 from gradio import Warning
 
@@ -19,6 +19,19 @@ random_seed = 42
 RUNNING_ON_AWS = get_or_create_env_var('RUNNING_ON_AWS', '0')
 print(f'The value of RUNNING_ON_AWS is {RUNNING_ON_AWS}')
 
+from torch import cuda, backends, version, get_num_threads
+
+print("Is CUDA enabled? ", cuda.is_available())
+print("Is a CUDA device available on this computer?", backends.cudnn.enabled)
+if cuda.is_available():
+    torch_device = "gpu"
+    print("Cuda version installed is: ", version.cuda)
+    high_quality_mode = "Yes"
+    os.system("nvidia-smi")
+else: 
+    torch_device =  "cpu"
+    high_quality_mode = "No"
+
 # Currently set n_gpu_layers to 0 even with cuda due to persistent bugs in implementation with cuda
 print("torch device for representation functions:", torch_device)
 if torch_device == "gpu":
@@ -29,7 +42,7 @@ else: #     torch_device =  "cpu"
     n_gpu_layers = 0
 
 #print("Running on device:", torch_device)
-n_threads = torch.get_num_threads()
+n_threads = get_num_threads()
 print("CPU n_threads:", n_threads)
 
 # Default Model parameters

funcs/topic_core_funcs.py

@@ -14,9 +14,9 @@ PandasDataFrame = Type[pd.DataFrame]
 from funcs.clean_funcs import initial_clean, regex_clean
 from funcs.anonymiser import expand_sentences_spacy
 from funcs.helper_functions import read_file, zip_folder, delete_files_in_folder, save_topic_outputs, output_folder, get_or_create_env_var, custom_regex_load
-from funcs.embeddings import make_or_load_embeddings, torch_device
+from funcs.embeddings import make_or_load_embeddings
 from funcs.bertopic_vis_documents import visualize_documents_custom, visualize_hierarchical_documents_custom, hierarchical_topics_custom, visualize_hierarchy_custom
-from funcs.representation_model import create_representation_model, llm_config, chosen_start_tag, random_seed, RUNNING_ON_AWS
+
 from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
 from sklearn.decomposition import TruncatedSVD
 import funcs.anonymiser as anon
@@ -31,14 +31,14 @@ today = datetime.now().strftime("%d%m%Y")
 today_rev = datetime.now().strftime("%Y%m%d")
 
 # Load embeddings
-if RUNNING_ON_AWS=="0":
-    embeddings_name = "mixedbread-ai/mxbai-embed-large-v1" #"mixedbread-ai/mxbai-embed-xsmall-v1" #"mixedbread-ai/mxbai-embed-large-v1"
-else:
-    embeddings_name = "mixedbread-ai/mxbai-embed-large-v1" #"mixedbread-ai/mxbai-embed-xsmall-v1"
+# if RUNNING_ON_AWS=="0":
+#     embeddings_name = "mixedbread-ai/mxbai-embed-large-v1" #"mixedbread-ai/mxbai-embed-xsmall-v1" #"mixedbread-ai/mxbai-embed-large-v1"
+# else:
+embeddings_name = "mixedbread-ai/mxbai-embed-large-v1" #"mixedbread-ai/mxbai-embed-xsmall-v1"
 
 # LLM model used for representing topics
-hf_model_name = "bartowski/Llama-3.2-3B-Instruct-GGUF" #"bartowski/Phi-3.1-mini-128k-instruct-GGUF"
-hf_model_file = "Llama-3.2-3B-Instruct-Q5_K_M.gguf" #"Phi-3.1-mini-128k-instruct-Q4_K_M.gguf"
+hf_model_name = "unsloth/gemma-2-it-GGUF"#"bartowski/Llama-3.2-3B-Instruct-GGUF" #"bartowski/Phi-3.1-mini-128k-instruct-GGUF"
+hf_model_file = "gemma-2-2b-it.q8_0.gguf" #"Llama-3.2-3B-Instruct-Q5_K_M.gguf" #"Phi-3.1-mini-128k-instruct-Q4_K_M.gguf"
 
 # When topic modelling column is chosen, change the default visualisation column to the same
 def change_default_vis_col(in_colnames:List[str]):
@@ -573,6 +573,8 @@ def represent_topics(topic_model: BERTopic, docs: List[str], data_file_name_no_e
 
     progress(0.1, desc= "Loading model and creating new topic representation")
 
+    from funcs.representation_model import create_representation_model, llm_config, chosen_start_tag
+
     representation_model = create_representation_model(representation_type, llm_config, hf_model_name, hf_model_file, chosen_start_tag, high_quality_mode)  
 
     progress(0.3, desc= "Updating existing topics")