Update app.py

app.py

@@ -26,11 +26,13 @@ def askme(symptoms, question):
     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() #skip_special_tokens=True
-    # Remove system messages and content
-   #  Extract only the assistant's response
-    assistant_response =response_text.split("assistant")[1].strip().split("user")[0].strip()   
-    return assistant_response
+    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
 
 # Example usage
 symptoms = '''\
@@ -87,8 +89,8 @@ css = """
 }
 """
 
-welcome_message = """# AI Medical Chatbot
-Ask any medical question and get answers from our AI Medical Chatbot
+welcome_message = """# AI Medical Llama 3 Chatbot
+Ask any medical question giving first your symptoms and get answers from our AI Medical Llama3 Chatbot
 Developed by Ruslan Magana. Visit [https://ruslanmv.com/](https://ruslanmv.com/) for more information."""
 
 
@@ -106,9 +108,6 @@ with gr.Blocks(css=css) as interface:
 
 interface.launch()
 
-
-
-
 '''
 iface = gr.Interface(
     fn=askme,