Update README.md

README.md

@@ -35,12 +35,10 @@ tokenizer = AutoTokenizer.from_pretrained("hiyouga/Qwen-14B-Chat-LLaMAfied")
 model = AutoModelForCausalLM.from_pretrained("hiyouga/Qwen-14B-Chat-LLaMAfied", torch_dtype="auto", device_map="auto")
 streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-query = (
-    "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-    "<|im_start|>user\nWho are you?<|im_end|>\n"
-    "<|im_start|>assistant\n"
-)
-inputs = tokenizer([query], return_tensors="pt")
+messages = [
+    {"role": "user", "content": "Who are you?"}
+]
+inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
 inputs = inputs.to("cuda")
 generate_ids = model.generate(**inputs, eos_token_id=[151643, 151645], max_new_tokens=256, streamer=streamer)
 ```