model)
eval(model, ner_dev_data, dev_data, dev_bl, graph, entity_linking,
args)
if args.do_predict:
utils.load_checkpoint(os.path.join(args.model_dir, 'best.pth.tar'),
model)
model.eval()
logging.info('Loading graph and entity linking...')
graph = pickle.load(open('graph/graph.pkl', 'rb'))
entity_linking = pickle.load(open('graph/entity_linking.pkl', 'rb'))
while True:
try:
logging.info('请输入问题：')
question = input()
ner_data = bl.build_ner_data(question)
ner_bl = bl.batch_loader(ner_data,
None,
args.ner_max_len,
args.re_max_len,
args.batch_size,
is_train=False)
for batch_data in ner_bl:
batch_data = tuple(
tmp.to(args.device) for tmp in batch_data)
head_logits, tail_logits = (
tmp.cpu()
for tmp in model(batch_data, 0, False)) # (bz, seqlen)
head = head_logits[0].argmax().item()
tail = tail_logits[0][head:].argmax().item() + head
tokens = ner_data[0]['tokens']
