Update README.md

README.md

@@ -40,8 +40,15 @@ messages = [{
 inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True)
 inputs = inputs[:, 1:] # remove bos token
 
-outputs = model.generate(input_ids=inputs.cuda(),
-                         max_new_tokens=20)
+terminators = [
+    tokenizer.eos_token_id,
+    tokenizer.convert_tokens_to_ids("<|eot_id|>")
+]
+
+outputs = model.generate(input_ids=inputs.cuda()[:, 1:],
+                         max_new_tokens=20,
+                         eos_token_id=terminators)
+
 response = tokenizer.decode(outputs[0])
 
 outputs = model.generate(inputs=input_ids.cuda(), attention_mask=attention_mask.cuda(), max_new_tokens=10)