Update README_zh.md

README_zh.md

@@ -39,27 +39,40 @@ license: other
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from cyg_conversation import covert_prompt_to_input_ids_with_history
 
-tokenizer = AutoTokenizer.from_pretrained("BAAI/AquilaChat-7B")
-model = AutoModelForCausalLM.from_pretrained("BAAI/AquilaChat-7B")
+device = torch.device("cuda:1")
+
+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("cuda:0")
-vocab = tokenizer.vocab
-print(len(vocab))
+model.to(device)
 
 text = "请给出10个要到北京旅游的理由。"
 
-tokens = covert_prompt_to_input_ids_with_history(text, history=[], tokenizer=tokenizer, max_token=512)
+tokens = tokenizer.encode_plus(text)['input_ids'][:-1]
 
-tokens = torch.tensor(tokens)[None,].to("cuda:0")
+tokens = torch.tensor(tokens)[None,].to(device)
 
 
 with torch.no_grad():
     out = model.generate(tokens, do_sample=True, max_length=512, eos_token_id=100007)[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)
 ```