Update app.py

app.py

@@ -17,6 +17,25 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 @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
+    provide an informative answer. If you don't know the answer to a specific medical inquiry, advise seeking professional help.
+    '''
+    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)  # 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() 
+    # Remove system messages and content    
+    # Extract and return the generated text, removing the prompt
+    # Extract only the assistant's response
+    answer = response_text.split('<|im_start|>assistant')[-1].strip()
+    #answer =response_text.split("assistant")[1].strip().split("user")[0].strip()   
+    return answer
+
+
+def askmeold(symptoms, question):
     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": f"Symptoms: {symptoms}\nQuestion: {question}\n"},
@@ -30,14 +49,6 @@ def askme(symptoms, question):
     return response_text
 
 
-examples = [
-    ["headache", "What are the possible causes of a headache?"],
-    ["fever", "How can I treat a fever at home?"],
-    ["cough", "What are the symptoms of a cough?"],
-    ["chest pain", "What are the possible causes of chest pain?"],
-]
-   
-
 # Example usage
 symptoms = '''\
 I'm a 35-year-old male and for the past few months, I've been experiencing fatigue,