Update app.py

app.py

@@ -3,16 +3,25 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
 import spaces  
 
+IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
+IS_SPACE = os.environ.get("SPACE_ID", None) is not None
+
+#device = "cuda" if torch.cuda.is_available() else "cpu"
 #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#dtype = torch.float16
+LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1"
 #print(f"Using device: {device}")
-device="cuda"
+#print(f"Using dtype: {dtype}")
+print(f"low memory: {LOW_MEMORY}")
+
+device = "cuda"
 model_name = "ruslanmv/Medical-Llama3-8B"
-device_map = 'auto'
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+# Move model and tokenizer to the CUDA device
+model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True).to(device)
 tokenizer.pad_token = tokenizer.eos_token
 
-@spaces.GPU  # Decorate the askme function with @spaces.GPU
+@spaces.GPU
 def askme(symptoms, question):
     sys_message = '''\
     You are an AI Medical Assistant trained on a vast dataset of health information. Please be thorough and
@@ -21,7 +30,7 @@ def askme(symptoms, question):
     content = symptoms + " " + question
     messages = [{"role": "system", "content": sys_message}, {"role": "user", "content": content}]
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)  # Ensure inputs are on CUDA device
     outputs = model.generate(**inputs, max_new_tokens=200, use_cache=True)
     response_text = tokenizer.batch_decode(outputs)[0].strip()
     answer = response_text.split('<|im_start|>assistant')[-1].strip()