Update README.md

README.md

@@ -41,7 +41,21 @@ gen_tokens = model.generate(input_ids, do_sample=True, max_new_tokens=128)
 print("-"*20 + "Output for model"  + 20 * '-')
 print(tokenizer.batch_decode(gen_tokens)[0])
 ```
+Alternatively, you can construct the prompt by applying the chat template of tokenizer on input conversation:
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("LLM360/K2-Chat")
+model = AutoModelForCausalLM.from_pretrained("LLM360/K2-Chat")
+
+messages = [{"role": "user", "content": "what is the highest mountain on earth?"}]
 
+input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+gen_tokens = model.generate(input_ids, do_sample=True, max_new_tokens=128)
+
+print("-"*20 + "Output for model"  + 20 * '-')
+print(tokenizer.batch_decode(gen_tokens)[0])
+```
 ## LLM360 Developer Suite
 We provide step-by-step finetuning tutorials for tech enthusiasts, AI practitioners and academic or industry researchers [here](https://www.llm360.ai/developer.html).