Update app.py

app.py

@@ -169,6 +169,18 @@ model = MllamaForConditionalGeneration.from_pretrained(
 processor = AutoProcessor.from_pretrained(model_id)
 
 
+def extract_assistant_reply(input_string):
+    # Define the tag that indicates the start of the assistant's reply
+    start_tag = "<|start_header_id|>assistant<|end_header_id|>"
+    # Find the position where the assistant's reply starts
+    start_index = input_string.find(start_tag)
+    if start_index == -1:
+        return "Assistant's reply not found."
+    start_index += len(start_tag)
+    # Extract everything after the start tag
+    assistant_reply = input_string[start_index:].strip()
+    return assistant_reply
+
 def extract_json_from_markdown(markdown_text):
     """Extract JSON or code block from markdown text."""
     try:
@@ -252,6 +264,8 @@ def diagnose_router(image):
     #     </ul>
     # </div>
     # """
+    
+    markdown_text=extract_assistant_reply(markdown_text)
     html_output = markdown.markdown(markdown_text)
     return html_output