Update README.md

README.md

@@ -70,6 +70,12 @@ We evaluate the model on [Passkey Retrieval](https://arxiv.org/abs/2309.12307) t
 
 <img src="data/passkey.png"></img>
 
+# Environment
+```bash
+torch>=2.1.1
+transformers==4.39.3
+```
+
 # Usage
 ```python
 import json
@@ -85,9 +91,9 @@ model = model.cuda().eval()
 
 with torch.no_grad():
   # short context
-  text = "Tell me about yourself."
-  inputs = tokenizer(text, return_tensors="pt").to("cuda")
-  outputs = model.generate(**inputs, max_new_tokens=20)
+  messages = [{"role": "user", "content": "Tell me about yourself."}]
+  inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to("cuda")
+  outputs = model.generate(**inputs, max_new_tokens=50)
   print(f"Input Length: {inputs['input_ids'].shape[1]}")
   print(f"Output:       {tokenizer.decode(outputs[0], skip_special_tokens=True)}")
 
@@ -97,7 +103,8 @@ with torch.no_grad():
   # long context
   with open("data/infbench.json", encoding="utf-8") as f:
     example = json.load(f)
-  inputs = tokenizer(example["context"], return_tensors="pt").to("cuda")
+  messages = [{"role": "user", "content": example["context"]}]
+  inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to("cuda")
   outputs = model.generate(**inputs, do_sample=False, top_p=1, temperature=1, max_new_tokens=20)[:, inputs["input_ids"].shape[1]:]
   print("*"*20)
   print(f"Input Length: {inputs['input_ids'].shape[1]}")