Update src/backend/chatbot.py

src/backend/chatbot.py

@@ -18,7 +18,7 @@ from llama_index.llms.llama_utils import (
     messages_to_prompt,
     completion_to_prompt,
 )
-import subprocess
+from huggingface_hub import HfApi, SpaceHardware
 
 # set version
 # st.session_state.demo_lite = False
@@ -33,11 +33,12 @@ print("BP 4 ")
 def init_llm(model, demo_lite):
     # st.write("BP 4.1: model: ", model)
     if demo_lite == False:
-    
-        # Set environment variables
-        subprocess.run(['powershell', '$env:FORCE_CMAKE=1'], shell=True)
-        subprocess.run(['powershell', '$env:CMAKE_ARGS="-DLLAMA_CUBLAS=on"'], shell=True)
-
+        api = HfApi(token=HF_TOKEN)
+        runtime = api.get_space_runtime(repo_id=TRAINING_SPACE_ID)
+        # Check if Space is loaded with a GPU.
+        if runtime.hardware == SpaceHardware.T4_SMALL:
+            print("using nvidia T4 small")
+                
         print("BP 5 : running full demo")
         if model == "Llama2-7b_CPP":
             model_path = "src/models/llama-2-7b-chat.Q4_K_M.gguf"