Update README.md

README.md

@@ -51,12 +51,13 @@ vae_model = Della.from_pretrained("IDEA-CCNL/Randeng-DELLA-CVAE-226M-NER-Chinese
 special_tokens_dict = {'bos_token': '<BOS>', 'eos_token': '<EOS>', 'additional_special_tokens': ['<ENT>', '<ENS>']}
 tokenizer.add_special_tokens(special_tokens_dict)
 
-model = vae_model.model
+device = 0
+model = vae_model.model.to(device)
 ent_token_type_id = tokenizer.additional_special_tokens_ids[0]
 ent_token_sep_id = tokenizer.additional_special_tokens_ids[1]
 bos_token_id, eos_token_id = tokenizer.bos_token_id, tokenizer.eos_token_id
 decoder_target, decoder_entities = [], []
-entity_list = [('体验中心', '地点/地理位置'), ('昨天', '时间')]
+entity_list = [('深圳', '地点/地理位置'), ('昨天', '时间')]
 
 for ent in entity_list:
     entity_name = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(ent[0]))
@@ -72,7 +73,7 @@ prior_z_list, prior_output_list = model.get_cond_prior_vecs(encoder_outputs.hidd
 outputs = model.decoder.generate(input_ids=inputs.to(device), layer_latent_vecs=prior_z_list, labels=None,
     label_ignore=model.pad_token_id, num_return_sequences=32, max_new_tokens=256,
     eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id,
-    no_repeat_ngram_size=-1, do_sample=True, top_p=0.5)
+    no_repeat_ngram_size=-1, do_sample=True, top_p=0.8)
 
 print(tokenizer.decode(inputs[0]))
 gen_sents = []
@@ -84,7 +85,6 @@ for idx in range(len(outputs)):
         gen_sents.append(gen_sent)
 for s in gen_sents:
     print(s)
-
 ```
 
 ## 引用 Citation