Update README.md

README.md

@@ -1,3 +1,117 @@
----
-license: apache-2.0
----
+
+### Usage
+
+```python
+
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+import torch
+
+class NER:
+    """
+    实体命名实体识别
+    """
+    def __init__(self,model_path) -> None:
+        """
+        Args:
+            model_path:模型地址
+        """
+
+        self.model_path = model_path
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForTokenClassification.from_pretrained(model_path)
+
+    def ner(self,sentence:str) -> list:
+        """
+        命名实体识别
+        Args:
+            sentence:要识别的句子
+        Return:
+            实体列表:[{'type':'LOC','tokens':[...]},...]
+        """
+        ans = []
+        for i in range(0,len(sentence),500):
+            ans = ans + self._ner(sentence[i:i+500])
+        return ans
+    
+    def _ner(self,sentence:str) -> list:
+        if len(sentence) == 0: return []
+        inputs = self.tokenizer(
+            sentence, add_special_tokens=True, return_tensors="pt"
+        )
+        
+        if torch.cuda.is_available():
+            self.model = self.model.to(torch.device('cuda:0'))
+            for key in inputs:
+                inputs[key] = inputs[key].to(torch.device('cuda:0'))
+            
+        with torch.no_grad():
+            logits = self.model(**inputs).logits
+        predicted_token_class_ids = logits.argmax(-1)
+        predicted_tokens_classes = [self.model.config.id2label[t.item()] for t in predicted_token_class_ids[0]]
+        entities = []
+        entity = {}
+        for idx, token in enumerate(self.tokenizer.tokenize(sentence,add_special_tokens=True)):
+            if 'B-' in predicted_tokens_classes[idx] or 'S-' in predicted_tokens_classes[idx]:
+                if len(entity) != 0:
+                    entities.append(entity)
+                entity = {}
+                entity['type'] = predicted_tokens_classes[idx].replace('B-','').replace('S-','')
+                entity['tokens'] = [token]
+            elif 'I-' in predicted_tokens_classes[idx] or 'E-' in predicted_tokens_classes[idx] or 'M-' in predicted_tokens_classes[idx]:
+                if len(entity) == 0:
+                    entity['type'] = predicted_tokens_classes[idx].replace('I-','').replace('E-','').replace('M-','')
+                    entity['tokens'] = []
+                entity['tokens'].append(token)
+            else:
+                if len(entity) != 0:
+                    entities.append(entity)
+                    entity = {}
+        if len(entity) > 0:
+            entities.append(entity)
+        return entities
+
+ner_model = NER('lixin12345/chinese-medical-ner')
+text = """
+患者既往慢阻肺多年;冠心病史6年，平素规律服用心可舒、保心丸等控制可;双下肢静脉血栓3年，保守治疗效果可;左侧腹股沟斜疝无张力修补术后2年。否认"高血压、糖尿病"等慢性病病史，否认"肝炎、结核"等传染病病史及其密切接触史，否认其他手术、重大外伤、输血史，否认"食物、药物、其他"等过敏史，预防接种史随社会。
+"""
+ans = ner_model.ner(text)
+# ans
+
+# DiseaseNameOrComprehensiveCertificate
+# 慢阻肺
+
+# DiseaseNameOrComprehensiveCertificate
+# 冠心病
+
+# Drug
+# 心可舒
+
+# Drug
+# 保心丸
+
+# DiseaseNameOrComprehensiveCertificate
+# 双下肢静脉血栓
+
+# DiseaseNameOrComprehensiveCertificate
+# 左侧腹股沟斜疝无张力修补术
+
+# DiseaseNameOrComprehensiveCertificate
+# 高血压
+
+# DiseaseNameOrComprehensiveCertificate
+# 糖尿病
+
+# DiseaseNameOrComprehensiveCertificate
+# 肝炎
+
+# DiseaseNameOrComprehensiveCertificate
+# 结核
+```
+
+### Source
+
+From hit [wi](https://wi.hit.edu.cn/gk.htm) 
+
+---
+license: apache-2.0
+---