prompt format

app.py

@@ -27,11 +27,25 @@ pipe = Llama(
     model_path=model_path
 )
 
+def format_prompt(message: str, history: List[List[str]]):
+    prompt = "<|begin_of_text|>"  # Start with the begin of text token
+    prompt += "<|im_start|>system\n<|im_end|>\n" # Assuming there's no system prompt here, just adding system role tags
+
+    for user_msg, assistant_msg in history:
+         prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
+         if assistant_msg:
+             prompt += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
+    
+    prompt += f"<|im_start|>user\n{message}<|im_end|>\n"
+    prompt += "<|im_start|>assistant\n" # Start of the Assistant's part 
+
+    return prompt
+
 def predict(message: str, history: List[List[str]], max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS, progress=gr.Progress()):
     if not message:
         return "", history, ""
 
-    prompt = message
+    prompt = format_prompt(message, history)
     history.append([message, ""])
 
     # Initialize reply for this round