mem

app/app.py

@@ -4,6 +4,7 @@ import utils as ut
 import elements as el
 import os
 import torch
+import psutil
 
 persist_directory = load_model.persist_directory
 st.title('myRetrievalGPT')
@@ -19,6 +20,7 @@ else:
     import torch
     torch.cuda.empty_cache()
 
+    st.write(str( torch.cuda.is_available()) + str(psutil.virtual_memory()))
     model_type = st.selectbox(
         'Select the Documents to be used to answer your question',
         ('OpenAI', 'decapoda-research/llama-7b-hf (gpu+cpu)', 'llama-7b 4bit (cpu only)',) ) 
@@ -34,7 +36,10 @@ else:
         llm= load_model.load_openai_model()
     elif model_type=='decapoda-research/llama-7b-hf (gpu+cpu)':
         # Add more models here
-        llm = load_model.load_gpu_model("decapoda-research/llama-7b-hf")
+        if not torch.cuda.is_available() and psutil.virtual_memory().available< 17254768640:
+            st.write('You do not have enough memory to use this model:' + str(psutil.virtual_memory().available))
+        else:
+            llm = load_model.load_gpu_model("decapoda-research/llama-7b-hf")
     else:
         llm = load_model.load_cpu_model()