Update README_zh.md

README_zh.md

@@ -47,40 +47,19 @@ AquilaChat-7B v0.8 在 FlagEval 大模型评测中（ “主观+客观”）相
 ```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)
 ```