Update main.py

main.py

@@ -17,12 +17,13 @@ class Item(BaseModel):
     top_p: float = 0.15
     repetition_penalty: float = 1.0
 
+# <s> [INST] Instruction [/INST] Model answer</s> [INST] Follow-up instruction [/INST]
 def format_prompt(message, history):
     prompt = "<s>"
     for user_prompt, bot_response in history:
         prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
+        prompt += f" {bot_response} "
+    prompt += f"</s>[INST] {message} [/INST]"
     return prompt
 
 def generate(item: Item):