Update app.py

app.py

@@ -22,17 +22,23 @@ DESCRIPTION = os.environ.get("DESCRIPTION")
 @spaces.GPU()
 def predict(message, history, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p):
     # Format history with a given chat template
-    if CHAT_TEMPLATE == "ChatML":
+    if CHAT_TEMPLATE == "Auto":
+        stop_tokens = [tokenizer.eos_token_id]
+        instruction = []
+        for user, assistant in history:
+            instruction.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
+        instruction.append({"role": "user", "content": message})
+    elif CHAT_TEMPLATE == "ChatML":
         stop_tokens = ["<|endoftext|>", "<|im_end|>"]
         instruction = '<|im_start|>system\n' + system_prompt + '\n<|im_end|>\n'
-        for human, assistant in history:
-            instruction += '<|im_start|>user\n' + human + '\n<|im_end|>\n<|im_start|>assistant\n' + assistant
+        for user, assistant in history:
+            instruction += '<|im_start|>user\n' + user + '\n<|im_end|>\n<|im_start|>assistant\n' + assistant
         instruction += '\n<|im_start|>user\n' + message + '\n<|im_end|>\n<|im_start|>assistant\n'
     elif CHAT_TEMPLATE == "Mistral Instruct":
         stop_tokens = ["</s>", "[INST]", "[INST] ", "<s>", "[/INST]", "[/INST] "]
         instruction = '<s>[INST] ' + system_prompt
-        for human, assistant in history:
-            instruction += human + ' [/INST] ' + assistant + '</s>[INST]'
+        for user, assistant in history:
+            instruction += user + ' [/INST] ' + assistant + '</s>[INST]'
         instruction += ' ' + message + ' [/INST]'
     else:
         raise Exception("Incorrect chat template, select 'ChatML' or 'Mistral Instruct'")