Update README.md

README.md

@@ -58,40 +58,19 @@ For detailed evaluation results, please refer to the website http://flageval.baa
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-
-device = torch.device("cuda:1")
-
+device = torch.device("cuda")
 model_info = "BAAI/AquilaChat-7B"
 tokenizer = AutoTokenizer.from_pretrained(model_info, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True)
 model.eval()
 model.to(device)
-
 text = "请给出10个要到北京旅游的理由。"
-
 tokens = tokenizer.encode_plus(text)['input_ids'][:-1]
-
 tokens = torch.tensor(tokens)[None,].to(device)
-
-
+stop_tokens = ["###", "[UNK]", "</s>"]
 with torch.no_grad():
-    out = model.generate(tokens, do_sample=True, max_length=512, eos_token_id=100007)[0]
-
+    out = model.generate(tokens, do_sample=True, max_length=512, eos_token_id=100007, bad_words_ids=[[tokenizer.encode(token)[0] for token in stop_tokens]])[0]
     out = tokenizer.decode(out.cpu().numpy().tolist())
-    if "###" in out:
-        special_index = out.index("###")
-        out = out[: special_index]
-
-    if "[UNK]" in out:
-        special_index = out.index("[UNK]")
-        out = out[:special_index]
-
-    if "</s>" in out:
-        special_index = out.index("</s>")
-        out = out[: special_index]
-
-    if len(out) > 0 and out[0] == " ":
-        out = out[1:]
     print(out)
 ```