Update README.md

README.md

@@ -25,20 +25,18 @@ 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")
+device = torch.device("cuda:0")
 model_info = "BAAI/AquilaChat2-34B"
 tokenizer = AutoTokenizer.from_pretrained(model_info, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.bfloat16)
 model.eval()
 model.to(device)
 text = "请给出10个要到北京旅游的理由。"
-tokens = tokenizer.encode_plus(text)['input_ids']
-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, bad_words_ids=[[tokenizer.encode(token)[0] for token in stop_tokens]])[0]
-    out = tokenizer.decode(out.cpu().numpy().tolist())
-    print(out)
+from predict import predict
+out = predict(model, text, tokenizer=tokenizer, max_gen_len=200, top_p=0.95,
+              seed=1234, topk=100, temperature=0.9, sft=True, device=device,
+              model_name="AquilaChat2-34B")
+print(out)
 ```