Update app.py

app.py

@@ -20,12 +20,15 @@ llm = Llama(
 
 # Build inference prompt according to your dataset format
 def format_prompt(user_message):
-    return f"""<|system|>
-You are a helpful assistant.
-<|user|>
-{user_message}
-<|assistant|>
-"""
+    return (
+        "<|begin_of_text|>"
+        "<|start_header_id|>system<|end_header_id|>\n"
+        "You are a helpful assistant.\n"
+        "<|start_header_id|>user<|end_header_id|>\n"
+        f"{user_message}\n"
+        "<|start_header_id|>assistant<|end_header_id|>\n"
+    )
+
 
 def respond(user_input):
     prompt = format_prompt(user_input)