Update app.py

app.py

@@ -17,6 +17,28 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 @spaces.GPU
 def askme(symptoms, question):
+    custom_template = [
+        {"role": "system", "content": "You are an AI Medical Assistant trained on a vast dataset of health information. Please be thorough and provide an informative answer. If you don't know the answer to a specific medical inquiry, advise seeking professional help."},
+        {"role": "user", "content": "Symptoms: {symptoms}\nQuestion: {question}\n"},
+        {"role": "assistant", "content": "{assistant_response}\n"}
+    ]
+
+    sys_message = custom_template[0]["content"]
+    content = custom_template[1]["content"].format(symptoms=symptoms, question=question)
+    messages = [{"role": "system", "content": sys_message}, {"role": "user", "content": content}]
+
+    prompt = tokenizer.apply_chat_template(messages, template=custom_template, tokenize=False, add_generation_prompt=True)
+    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, skip_special_tokens=True)[0].strip()
+
+    # Extract only the assistant's response
+    assistant_response = response_text.split("<|im_start|>assistant")[-1].strip()
+    return assistant_response
+
+
+
+def askmeold(symptoms, question):
     sys_message = '''\
     You are an AI Medical Assistant trained on a vast dataset of health information. Please be thorough and
     provide an informative answer. If you don't know the answer to a specific medical inquiry, advise seeking professional help.