Update README.md

README.md

@@ -46,27 +46,26 @@ We will continue to release improved versions of Aquila model as open source. Fo
 ```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]
@@ -74,17 +73,13 @@ with torch.no_grad():
     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:]
-
-        convert_tokens = convert_tokens[1:]
-        probs = probs[1:]
-
     print(out)
 ```