Views
No views yet
1
2from transformers import AutoTokenizer, AutoModelForTokenClassification
3import torch
4
5class NER:
6 """
7 实体命名实体识别
8 """
9 def __init__(self,model_path) -> None:
10 """
11 Args:
12 model_path:模型地址
13 """
14
15 self.model_path = model_path
16 self.tokenizer = AutoTokenizer.from_pretrained(model_path)
17 self.model = AutoModelForTokenClassification.from_pretrained(model_path)
18
19 def ner(self,sentence:str) -> list:
20 """
21 命名实体识别
22 Args:
23 sentence:要识别的句子
24 Return:
25 实体列表:[{'type':'LOC','tokens':[...]},...]
26 """
27 ans = []
28 for i in range(0,len(sentence),500):
29 ans = ans + self._ner(sentence[i:i+500])
30 return ans
31
32 def _ner(self,sentence:str) -> list:
33 if len(sentence) == 0: return []
34 inputs = self.tokenizer(
35 sentence, add_special_tokens=True, return_tensors="pt"
36 )
37
38 if torch.cuda.is_available():
39 self.model = self.model.to(torch.device('cuda:0'))
40 for key in inputs:
41 inputs[key] = inputs[key].to(torch.device('cuda:0'))
42
43 with torch.no_grad():
44 logits = self.model(**inputs).logits
45 predicted_token_class_ids = logits.argmax(-1)
46 predicted_tokens_classes = [self.model.config.id2label[t.item()] for t in predicted_token_class_ids[0]]
47 entities = []
48 entity = {}
49 for idx, token in enumerate(self.tokenizer.tokenize(sentence,add_special_tokens=True)):
50 if 'B-' in predicted_tokens_classes[idx] or 'S-' in predicted_tokens_classes[idx]:
51 if len(entity) != 0:
52 entities.append(entity)
53 entity = {}
54 entity['type'] = predicted_tokens_classes[idx].replace('B-','').replace('S-','')
55 entity['tokens'] = [token]
56 elif 'I-' in predicted_tokens_classes[idx] or 'E-' in predicted_tokens_classes[idx] or 'M-' in predicted_tokens_classes[idx]:
57 if len(entity) == 0:
58 entity['type'] = predicted_tokens_classes[idx].replace('I-','').replace('E-','').replace('M-','')
59 entity['tokens'] = []
60 entity['tokens'].append(token)
61 else:
62 if len(entity) != 0:
63 entities.append(entity)
64 entity = {}
65 if len(entity) > 0:
66 entities.append(entity)
67 return entities
68
69ner_model = NER('lixin12345/chinese-medical-ner')
70text = """
71患者既往慢阻肺多年;冠心病史6年，平素规律服用心可舒、保心丸等控制可;双下肢静脉血栓3年，保守治疗效果可;左侧腹股沟斜疝无张力修补术后2年。否认"高血压、糖尿病"等慢性病病史，否认"肝炎、结核"等传染病病史及其密切接触史，否认其他手术、重大外伤、输血史，否认"食物、药物、其他"等过敏史，预防接种史随社会。
72"""
73ans = ner_model.ner(text)
74# ans
75
76# DiseaseNameOrComprehensiveCertificate
77# 慢阻肺
78
79# DiseaseNameOrComprehensiveCertificate
80# 冠心病
81
82# Drug
83# 心可舒
84
85# Drug
86# 保心丸
87
88# DiseaseNameOrComprehensiveCertificate
89# 双下肢静脉血栓
90
91# DiseaseNameOrComprehensiveCertificate
92# 左侧腹股沟斜疝
93
94# TreatmentOrPreventionProcedures
95# 无张力修补术
96
97# DiseaseNameOrComprehensiveCertificate
98# 高血压
99
100# DiseaseNameOrComprehensiveCertificate
101# 糖尿病
102
103# DiseaseNameOrComprehensiveCertificate
104# 肝炎
105
106# DiseaseNameOrComprehensiveCertificate
107# 结核