Update README.md

README.md

@@ -45,9 +45,45 @@ from torch.nn.utils.rnn import pad_sequence
 from fengshen.models.deepVAE.deep_vae import Della
 from transformers.models.bert.tokenization_bert import BertTokenizer
 
-tokenizer = BertTokenizer.from_pretrained("IDEA-CCNL/Randeng-DELLA-226M-CVAE-NER-Chinese")
-vae_model = Della.from_pretrained("IDEA-CCNL/Randeng-DELLA-226M-CVAE-NER-Chinese")
-
+tokenizer = BertTokenizer.from_pretrained("IDEA-CCNL/Randeng-DELLA-CVAE-226M-NER-Chinese")
+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
+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 = [('体验中心', '地点/地理位置'), ('昨天', '时间')]
+
+for ent in entity_list:
+    entity_name = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(ent[0]))
+    entity_type = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(ent[1]))
+    decoder_entities.extend(entity_name + [ent_token_type_id] + entity_type + [ent_token_sep_id])
+decoder_entities.extend([bos_token_id]) # for generation 
+decoder_target.append(torch.tensor(decoder_entities, dtype=torch.long))
+inputs = pad_sequence(decoder_target, batch_first=True, padding_value=0)
+
+encoder_outputs = model.encoder(input_ids=inputs.to(device))
+prior_z_list, prior_output_list = model.get_cond_prior_vecs(encoder_outputs.hidden_states[1:], 
+    inputs, sample=True, beta_logvar=0.)
+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)
+
+print(tokenizer.decode(inputs[0]))
+gen_sents = []
+for idx in range(len(outputs)):
+    sent_len= 512 if eos_token_id not in outputs[idx].tolist() else outputs[idx].tolist().index(eos_token_id) + 1
+    start_loc = outputs[idx].tolist().index(bos_token_id) 
+    gen_sent = tokenizer.decode(outputs[idx][start_loc:sent_len]).replace(' ', '')
+    if all([ent[0] in gen_sent for ent in entity_list]):
+        gen_sents.append(gen_sent)
+for s in gen_sents:
+    print(s)
 
 ```