Update README.md

README.md

@@ -40,25 +40,19 @@ The additional details of the Aquila model will be presented in the official tec
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-
+device = torch.device("cuda")
 model_info = "BAAI/Aquila2-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")
-
-text = "汽车EDR是什么"
-
+model.to(device)
+text = "请给出10个要到北京旅游的理由。"
 tokens = tokenizer.encode_plus(text)['input_ids'][:-1]
-
-tokens = torch.tensor(tokens)[None,].to("cuda:0")
-
-
+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())
-
     print(out)
 ```