Add chat template

README.md

@@ -120,27 +120,25 @@ tokenizer = transformers.AutoTokenizer.from_pretrained(
         use_fast=False,
     )
 
-system_message = "You are Orca, an AI language model created by Microsoft. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior."
-user_message = "How can you determine if a restaurant is popular among locals or mainly attracts tourists, and why might this information be useful?"
+messages = [
+    {"role": "system", "content": "You are Orca, an AI language model created by Microsoft. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior."},
+    {"role": "user", "content": "How can you determine if a restaurant is popular among locals or mainly attracts tourists, and why might this information be useful?"}
+]
+inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
 
-prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant"
-
-inputs = tokenizer(prompt, return_tensors='pt')
-output_ids = model.generate(inputs["input_ids"],)
-answer = tokenizer.batch_decode(output_ids)[0]
+output_ids = model.generate(inputs)
+answer = tokenizer.decode(output_ids[0])
 
 print(answer)
 
 # This example continues showing how to add a second turn message by the user to the conversation
-second_turn_user_message = "Give me a list of the key points of your first answer."
-
-# we set add_special_tokens=False because we dont want to automatically add a bos_token between messages
-second_turn_message_in_markup = f"\n<|im_start|>user\n{second_turn_user_message}<|im_end|>\n<|im_start|>assistant"
-second_turn_tokens = tokenizer(second_turn_message_in_markup, return_tensors='pt', add_special_tokens=False)
-second_turn_input = torch.cat([output_ids, second_turn_tokens['input_ids']], dim=1)
+messages.append(
+    {"role": "user", "content": "Give me a list of the key points of your first answer."}
+)
+second_turn_input = tokenizer.apply_chat_template(messages, return_tensors='pt')
 
-output_ids_2 = model.generate(second_turn_input,)
-second_turn_answer = tokenizer.batch_decode(output_ids_2)[0]
+output_ids_2 = model.generate(second_turn_input)
+second_turn_answer = tokenizer.decode(output_ids_2[0])
 
 print(second_turn_answer)
 ```
@@ -209,16 +207,16 @@ tokenizer = transformers.AutoTokenizer.from_pretrained(
     add_special_tokens=False,
 )
 
-system_message = "You are Orca, an AI language model created by Microsoft. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior."
-user_message = "\" \n :You can't just say, \"\"that's crap\"\" and remove it without gaining a consensus. You already know this, based on your block history. —/ \" \nIs the comment obscene? \nOptions : Yes, No."
-
-prompt =  f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant"
+messages = [
+    {"role": "system", "content":"You are Orca, an AI language model created by Microsoft. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior."},
+    {"role": "user", "content": "\" \n :You can't just say, \"\"that's crap\"\" and remove it without gaining a consensus. You already know this, based on your block history. —/ \" \nIs the comment obscene? \nOptions : Yes, No."}
+]
 
-inputs = tokenizer(prompt, return_tensors='pt')
+inputs = tokenizer.apply_chat_template(messages, return_tensors='pt')
 inputs = inputs.to(device)
 
-output_ids = model.generate(inputs["input_ids"], max_length=4096, do_sample=False, temperature=0.0, use_cache=True)
-sequence_length = inputs["input_ids"].shape[1]
+output_ids = model.generate(inputs, max_length=4096, do_sample=False, temperature=0.0, use_cache=True)
+sequence_length = inputs.shape[1]
 new_output_ids = output_ids[:, sequence_length:]
 answers = tokenizer.batch_decode(new_output_ids, skip_special_tokens=True)
 final_output = answers[0] if not should_filter_out(answers[0]) else "[Content Filtered]"

tokenizer_config.json

@@ -9,6 +9,7 @@
     "rstrip": false,
     "single_word": false
   },
+  "chat_template": "{{ bos_token }} {% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
   "clean_up_tokenization_spaces": false,
   "eos_token": {
     "__type": "AddedToken",