Update method to load yuan2.0-51B

README.md

@@ -30,7 +30,7 @@ Yuan2.0 is a new generation Fundamental Large Language Model developed by IEIT S
 
 
 # <span id="Usage">代码调用/Code Usage</span>
-可以通过如下代码调用 Yuan2-2B 模型来生成文本：
+可以通过如下代码调用 `IEITYuan/Yuan2.0-51B-hf` 模型来生成文本：
 
 You can generate text by invoking the Yuan2-2B model with the following code:
 
@@ -42,11 +42,11 @@ sys.path.append(
 from transformers import AutoModelForCausalLM,AutoTokenizer,LlamaTokenizer
 
 print("Creat tokenizer...")
-tokenizer = LlamaTokenizer.from_pretrained('yuan2-2B')
+tokenizer = LlamaTokenizer.from_pretrained('IEITYuan/Yuan2.0-51B-hf')
 tokenizer.add_tokens(['<sep>', '<pad>', '<mask>', '<predict>', '<FIM_SUFFIX>', '<FIM_PREFIX>', '<FIM_MIDDLE>','<commit_before>','<commit_msg>','<commit_after>','<jupyter_start>','<jupyter_text>','<jupyter_code>','<jupyter_output>','<empty_output>'], special_tokens=True)
 
 print("Creat model...")
-model = AutoModelForCausalLM.from_pretrained('yuan2-2B',device_map='auto',torch_dtype=torch.bfloat16,trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained('IEITYuan/Yuan2.0-51B-hf',device_map='auto',torch_dtype=torch.bfloat16,trust_remote_code=True)
 
 inputs = tokenizer("请问目前最先进的机器学习算法有哪些？", return_tensors="pt")["input_ids"].to("cuda:0")
 outputs = model.generate(inputs,do_sample=False,max_length=100)