You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
169 lines
8.4 KiB
169 lines
8.4 KiB
import torch |
|
import numpy as np |
|
from ner.Bert_BiLSTM_CRF import Bert_BiLSTM_CRF |
|
from common import config |
|
import os |
|
import logging.handlers |
|
from concurrent_log_handler import ConcurrentRotatingFileHandler |
|
import time |
|
|
|
try: |
|
log_fmt = '%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s' |
|
formatter = logging.Formatter(log_fmt) |
|
handler = ConcurrentRotatingFileHandler('./logs/log.log', maxBytes=10000000, backupCount=10,encoding='utf-8') |
|
handler.setFormatter(formatter) |
|
logging.basicConfig(level=logging.DEBUG) |
|
log = logging.getLogger(__name__) |
|
log.addHandler(handler) |
|
except Exception as error: |
|
result_dict = {"code": 500,"error_msg":"日志文件打开失败"} |
|
|
|
class panjueshu(): |
|
def __init__(self): |
|
try: |
|
time1=time.time() |
|
VOCAB_list = ['<PAD>', '[CLS]', '[SEP]', 'O'] |
|
self.ner_dict = {} |
|
self.tag_file_path = os.path.join(os.getcwd(),"ner","panjueshu","标签.txt") |
|
self.line_list=[] |
|
with open(self.tag_file_path,"r",encoding="utf-8") as f: |
|
self.line_list=f.readlines() |
|
self.line_len=len(self.line_list) |
|
for line_number in range(self.line_len): |
|
bq_str = self.line_list[line_number].replace("\n","") |
|
self.ner_dict["A" + str(line_number+1)] = bq_str |
|
VOCAB_list.append('B-' + "A" + str(line_number+1)) |
|
VOCAB_list.append('I-' + "A" + str(line_number+1)) |
|
self.VOCAB = tuple(VOCAB_list) |
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
self.model_name="判决书" |
|
tag2idx = {tag: idx for idx, tag in enumerate(self.VOCAB)} |
|
time2 = time.time() |
|
print("ner panjueshu main __init__ time:{}s".format(time2 - time1)) |
|
self.model = Bert_BiLSTM_CRF(tag2idx).cuda() |
|
time3 = time.time() |
|
print("ner panjueshu main __init__ time:{}s".format(time3 - time2)) |
|
self.model.load_state_dict(torch.load(os.path.join(config.model_path, self.model_name,"model.pt"))) |
|
time4 = time.time() |
|
print("ner panjueshu main __init__ time:{}s".format(time4 - time3)) |
|
self.model.eval() |
|
time5 = time.time() |
|
print("ner panjueshu main __init__ time:{}s".format(time5 - time4)) |
|
self.bdjs_str_list = ["。", "；", ".", ";", ":", "：", "，", ","] |
|
except Exception as error: |
|
log.error("ner panjueshu main __init__ error:{}".format(error),exc_info=True) |
|
def deal_input_data(self,input_str_list): |
|
try: |
|
time1=time.time() |
|
output_str_list=[] |
|
vocab_list = self.load_vocab() |
|
all_words_idx_tensor_list = [] |
|
for input_str in input_str_list: |
|
split_input_str_list = [] |
|
while (len(input_str) > 510): |
|
split_start_index = 509 |
|
if (input_str[509] not in self.bdjs_str_list): |
|
split_start_index = 508 |
|
while (split_start_index >= 0): |
|
# print(input_str[split_start_index]) |
|
if (input_str[split_start_index] in self.bdjs_str_list): |
|
break |
|
split_start_index -= 1 |
|
split_input_str = input_str[0:split_start_index + 1] |
|
input_str = input_str[split_start_index + 1:len(input_str)] |
|
split_input_str_list.append(split_input_str) |
|
output_str_list.append(split_input_str) |
|
split_input_str_list.append(input_str) |
|
output_str_list.append(input_str) |
|
for split_input_str in split_input_str_list: |
|
cur_words_idx_list = [] |
|
words_idx_list = [] |
|
cur_words_idx_tensor = [] |
|
words_idx_list.append(101) |
|
for input_word in split_input_str: |
|
input_word = input_word.lower() |
|
if (input_word not in vocab_list): |
|
words_idx_list.append((144)) |
|
else: |
|
word_idx = vocab_list.index(input_word.lower()) |
|
words_idx_list.append((word_idx)) |
|
words_idx_list.append(102) |
|
cur_words_idx_list.append(words_idx_list) |
|
cur_words_idx_array = np.array(cur_words_idx_list) |
|
cur_words_idx_tensor = torch.LongTensor(cur_words_idx_array) |
|
all_words_idx_tensor_list.append(cur_words_idx_tensor) |
|
time2 = time.time() |
|
print("ner panjueshu main deal_input_data time:{}s".format(time2 - time1)) |
|
return all_words_idx_tensor_list,output_str_list |
|
except Exception as error: |
|
log.error("ner panjueshu main deal_input_data error:{}".format(error),exc_info=True) |
|
return [] |
|
def load_vocab(self): |
|
try: |
|
time1=time.time() |
|
vocab_list = [] |
|
with open(os.path.join(config.from_pretrained_path,"vocab.txt"), "r", encoding="utf-8") as reader: |
|
token_list = reader.readlines() |
|
token_list_len = len(token_list) |
|
for token_list_number in range(token_list_len): |
|
token = token_list[token_list_number] |
|
if not token: |
|
break |
|
token = token.strip() |
|
vocab_list.append(token) |
|
time2 = time.time() |
|
print("ner panjueshu main load_vocab time:{}s".format(time2 - time1)) |
|
return vocab_list |
|
except Exception as error: |
|
log.error("ner panjueshu main load_vocab error:{}".format(error),exc_info=True) |
|
return [] |
|
def predict(self,input_str_list): |
|
try: |
|
time1=time.time() |
|
result_ner_json_list=[] |
|
predict_data_list,input_str_list=self.deal_input_data(input_str_list) |
|
with torch.no_grad(): |
|
predict_datas_len = len(predict_data_list) |
|
for predict_datas_number in range(predict_datas_len): |
|
ner_json_list=[] |
|
predict_data = predict_data_list[predict_datas_number] |
|
Y_hat = [] |
|
x = predict_data |
|
x = x.to(self.device) |
|
# print("x".format(x)) |
|
# y = y.to(device) |
|
_, y_hat = self.model(x) # y_hat: (N, T) |
|
# print("_:{} y_hat:{}".format(_,y_hat)) |
|
Y_hat.extend(y_hat.cpu().numpy().tolist()) |
|
Y_hat_len = len(Y_hat) |
|
for Y_hat_number in range(Y_hat_len): |
|
y_hat = Y_hat[Y_hat_number] |
|
input_str = input_str_list[predict_datas_number] |
|
y_hat_len = len(y_hat) |
|
y_hat_number = 0 |
|
#print("原文:{}".format(input_str)) |
|
while (y_hat_number < y_hat_len): |
|
y = y_hat[y_hat_number] |
|
if (y <= 3): |
|
y_hat_number += 1 |
|
continue |
|
else: |
|
ner_eng_name = self.VOCAB[y] |
|
ner_eng_name = ner_eng_name.split("-")[1] |
|
ner_ch_name = self.ner_dict[ner_eng_name] |
|
ner_words = input_str[y_hat_number - 1] |
|
y_hat_number += 1 |
|
while (self.VOCAB[y_hat[y_hat_number]] == "I-" + ner_eng_name): |
|
ner_words += input_str[y_hat_number - 1] |
|
y_hat_number += 1 |
|
ner_json={'tag':ner_ch_name,'term':ner_words} |
|
ner_json_list.append(ner_json) |
|
#print("{}:{}".format(ner_ch_name, ner_words)) |
|
continue |
|
result_ner_json_list.append(ner_json_list) |
|
time2 = time.time() |
|
print("ner panjueshu main predict time:{}s".format(time2 - time1)) |
|
return {'code':200,'data':result_ner_json_list,'msg':""} |
|
except Exception as error: |
|
log.error("ner panjueshu main predict error:{}".format(error),exc_info=True) |
|
return {'code': 500, 'data': [], 'msg': error} |