Update app.py

app.py

@@ -33,7 +33,15 @@ if checkpoint in mistral_models:
     chat_template = chat_template.replace('    ', '').replace('\n', '')
     tokenizer.chat_template = chat_template
     
-
+def format_prompt_mistral(message, chatbot, system_prompt=SYSTEM_PROMPT+SYSTEM_PROMPT_NOUS):
+    messages = [{"role": "system", "content": system_prompt}]
+    for user_message, bot_message in chatbot:
+        messages.append({"role": "user", "content": user_message})
+        messages.append({"role": "assistant", "content": bot_message})
+    messages.append({"role": "user", "content": message})
+    newPrompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True, return_tensors="pt")
+    return newPrompt
 
 def format_prompt_cohere(message, chatbot, system_prompt=SYSTEM_PROMPT):
     messages = [{"role": "system", "content": system_prompt}]
@@ -70,7 +78,7 @@ match checkpoint:
     case "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO":
         format_prompt=format_prompt_nous
     case "mistralai/Mixtral-8x7B-Instruct-v0.1":
-        format_prompt=format_prompt_cohere
+        format_prompt=format_prompt_mistral
 
 def inference(message, history, temperature=0.9, maxTokens=512, topP=0.9, repPenalty=1.1):
     # Updating the settings for the generation