Update README.md

README.md

@@ -85,11 +85,12 @@ res_prefix.append(EOS_TOKEN_ID)
 l_rp=len(res_prefix)
 
 tokenized=tokenizer.encode(plain_text,add_special_tokens=False,truncation=True,max_length=self.max_seq_length-2-l_rp)
-
 tokenized+=res_prefix
+batch=[tokenized]*2
+input_ids=torch.tensor(np.array(batch),dtype=torch.long)
 
 # Generate answer
-pred_ids = model.generate(input_ids=tokenized,max_new_token=self.max_target_length,do_sample=True,top_p=0.9)
+pred_ids = model.generate(input_ids=input_ids,max_new_token=self.max_target_length,do_sample=True,top_p=0.9)
 pred_tokens=tokenizer.batch_decode(pred_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
 res=pred_tokens.replace('<extra_id_0>','').replace('有答案:','')
 ```