Update app.py

app.py

@@ -8,11 +8,10 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from dataclasses import dataclass
 
-torch.set_grad_enabled(False)
-model_name = "TheBloke/OpenHermes-2.5-Mistral-7B-GPTQ"
-token = os.environ['hf_token']
-pipe = pipeline("text-generation", model=model_name, device="cuda")
-generate_kwargs = {'max_new_tokens': 20}
+
+chatml_template = """{% for message in messages %}
+    {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
+{% endfor %}"""
 
 
 system_prompt = '''You are given a partial input text for another AI chat interface. 
@@ -26,7 +25,7 @@ Answers should be only the completions themselves. If you have nothing as a comp
 '''
 
 
-extra_prompts = '''
+extra_prompt = '''
 Examples: 
 (1)
 User: "Help me write a sentiment analysis pipeline"
@@ -41,6 +40,14 @@ User: "Help me find a present for my"
 Assistant: "girlfriend;mother;father;friend"
 '''
 
+# setup
+torch.set_grad_enabled(False)
+model_name = "TheBloke/OpenHermes-2.5-Mistral-7B-GPTQ"
+token = os.environ['hf_token']
+pipe = pipeline("text-generation", model=model_name, device="cuda")
+pipe.tokenizer.chat_template = chatml_template # TheBloke says this is the right template for this model
+generate_kwargs = {'max_new_tokens': 20}
+
 # '''
 # You will now get a blank message from the user and then after your answer, the user will give you the text to complete: 
 # Example: