-
Notifications
You must be signed in to change notification settings - Fork 0
/
qa_model_inference.py
67 lines (52 loc) · 2.74 KB
/
qa_model_inference.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
# -*- coding: utf-8 -*-
import sys
from collections import OrderedDict
import torch
from transformers import AutoConfig, AutoTokenizer
from PraticeOfTransformers.CustomModelForNSPQABILSTM import CustomModelForNSPQABILSTM
from PraticeOfTransformers.CustomModelForNer import BertForNerAppendBiLstmAndCrf
model_name = 'bert-base-chinese'
tokenizer = AutoTokenizer.from_pretrained(model_name)
path = "./model_path/qa_path/ultimate_dict_nsp_qa_lstm_union_epoch_1000"
'''
问答部分
'''
if path != "":
config = AutoConfig.from_pretrained(pretrained_model_name_or_path=model_name, num_labels=2)
model = CustomModelForNSPQABILSTM(config)
# 因为后面的参数没有初始化,所以采用非强制性约束
state_dict = torch.load(path)
new_state_dict = OrderedDict()
for k, v in state_dict.items(): # k为module.xxx.weight, v为权重
if k.startswith('module.'):
name = k[7:] # 截取`module.`后面的xxx.weight
new_state_dict[name] = v
else:
new_state_dict[k] = v
# 因为后面的参数没有初始化,所以采用非强制性约束,多GPu的加载到单GPU上需要, map_location='cuda:0'
model.load_state_dict(new_state_dict, strict=True)
else:
model = CustomModelForNSPQABILSTM.from_pretrained(
pretrained_model_name_or_path=model_name) # num_labels 测试用一下,看看参数是否传递
# 看是否用cpu或者gpu训练
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
def get_qa_result(question_passage_list):
with torch.no_grad():
encoded_dict = tokenizer.batch_encode_plus(
batch_text_or_text_pairs=question_passage_list,
# 输入文本对 # 输入文本,采用list[tuple(text,question)]的方式进行输入
add_special_tokens=True, # 添加 '[CLS]' 和 '[SEP]'
max_length=512, # 填充 & 截断长度
truncation=True,
padding='longest',
return_attention_mask=True, # 返回 attn. masks.
)
model_output = model(input_ids=torch.tensor(encoded_dict['input_ids']).to(device),
attention_mask=torch.tensor(encoded_dict['attention_mask']).to(device),
token_type_ids=torch.tensor(encoded_dict['token_type_ids']).to(device))
qa_nsp_logits=model_output.nsp_relationship_scores.to("cpu")
qa_start_logits = model_output.qa_start_logits.to("cpu")
qa_end_logits = model_output.qa_end_logits.to("cpu")
returniputs=[tokenizer.convert_ids_to_tokens(input_id) for input_id in encoded_dict['input_ids']]
return (qa_nsp_logits,qa_start_logits,qa_end_logits,returniputs)