Update app.py

app.py

@@ -17,6 +17,30 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 @spaces.GPU
 def askme(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"},
+        {"role": "assistant", "content": "{assistant_response}\n"}
+    ]
+
+    prompt = tokenizer.apply_chat_template(template, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(**inputs, max_new_tokens=300, use_cache=True)
+    response_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0].strip()
+
+    #start_idx = response_text.find("<|im_start|>assistant")
+    #end_idx = response_text.find("<|im_end|>", start_idx)
+    #assistant_response = response_text[start_idx + len("<|im_start|>assistant"):end_idx]
+
+    #return assistant_response.split(". ")[0] + "
+    return response_text
+
+
+
+
+
+
+def askmeold(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"},