# %% import argparse import os.path import pickle import unicodedata import torch from tqdm import tqdm import NER_medNLP as ner import utils from EntityNormalizer import EntityNormalizer, EntityDictionary, DefaultDiseaseDict, DefaultDrugDict device = torch.device("mps" if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu') # %% global変数として使う dict_key = {} # %% def to_xml(data, id_to_tags): with open("key_attr.pkl", "rb") as tf: key_attr = pickle.load(tf) text = data['text'] count = 0 for i, entities in enumerate(data['entities_predicted']): if entities == "": return span = entities['span'] try: type_id = id_to_tags[entities['type_id']].split('_') except: print("out of rage type_id", entities) continue tag = type_id[0] if not type_id[1] == "": attr = ' ' + value_to_key(type_id[1], key_attr) + '=' + '"' + type_id[1] + '"' else: attr = "" if 'norm' in entities: attr = attr + ' norm="' + str(entities['norm']) + '"' add_tag = "<" + str(tag) + str(attr) + ">" text = text[:span[0] + count] + add_tag + text[span[0] + count:] count += len(add_tag) add_tag = "" text = text[:span[1] + count] + add_tag + text[span[1] + count:] count += len(add_tag) return text def predict_entities(model, tokenizer, sentences_list): # entities_list = [] # 正解の固有表現を追加していく entities_predicted_list = [] # 抽出された固有表現を追加していく text_entities_set = [] for dataset in sentences_list: text_entities = [] for sample in tqdm(dataset, desc='Predict', leave=False): text = sample encoding, spans = tokenizer.encode_plus_untagged( text, return_tensors='pt' ) encoding = {k: v.to(device) for k, v in encoding.items()} with torch.no_grad(): output = model(**encoding) scores = output.logits scores = scores[0].cpu().numpy().tolist() # 分類スコアを固有表現に変換する entities_predicted = tokenizer.convert_bert_output_to_entities( text, scores, spans ) # entities_list.append(sample['entities']) entities_predicted_list.append(entities_predicted) text_entities.append({'text': text, 'entities_predicted': entities_predicted}) text_entities_set.append(text_entities) return text_entities_set def combine_sentences(text_entities_set, id_to_tags, insert: str): documents = [] for text_entities in text_entities_set: document = [] for t in text_entities: document.append(to_xml(t, id_to_tags)) documents.append('\n'.join(document)) return documents def value_to_key(value, key_attr): # attributeから属性名を取得 global dict_key if dict_key.get(value) != None: return dict_key[value] for k in key_attr.keys(): for v in key_attr[k]: if value == v: dict_key[v] = k return k # %% def normalize_entities(text_entities_set, id_to_tags, disease_dict=None, disease_candidate_col=None, disease_normalization_col=None, disease_matching_threshold=None, drug_dict=None, drug_candidate_col=None, drug_normalization_col=None, drug_matching_threshold=None): if disease_dict: disease_dict = EntityDictionary(disease_dict, disease_candidate_col, disease_normalization_col) else: disease_dict = DefaultDiseaseDict() disease_normalizer = EntityNormalizer(disease_dict, matching_threshold=disease_matching_threshold) if drug_dict: drug_dict = EntityDictionary(drug_dict, drug_candidate_col, drug_normalization_col) else: drug_dict = DefaultDrugDict() drug_normalizer = EntityNormalizer(drug_dict, matching_threshold=drug_matching_threshold) for entry in tqdm(text_entities_set, desc='Normalization', leave=False): for text_entities in entry: entities = text_entities['entities_predicted'] for entity in entities: tag = id_to_tags[entity['type_id']].split('_')[0] normalizer = drug_normalizer if tag == 'm-key' \ else disease_normalizer if tag == 'd' \ else None if normalizer is None: continue normalization, score = normalizer.normalize(entity['name']) entity['norm'] = str(normalization) def run(model, input, output=None, normalize=False, **kwargs): with open("id_to_tags.pkl", "rb") as tf: id_to_tags = pickle.load(tf) len_num_entity_type = len(id_to_tags) # Load the model and tokenizer classification_model = ner.BertForTokenClassification_pl.from_pretrained_bin(model_path=model, num_labels=2 * len_num_entity_type + 1) bert_tc = classification_model.bert_tc.to(device) tokenizer = ner.NER_tokenizer_BIO.from_pretrained( 'tohoku-nlp/bert-base-japanese-whole-word-masking', num_entity_type=len_num_entity_type # Entityの数を変え忘れないように！ ) # Load input files if (os.path.isdir(input)): files = [os.path.join(input, f) for f in os.listdir(input) if os.path.isfile(os.path.join(input, f))] else: files = [input] for file in tqdm(files, desc="Input file"): try: with open(file) as f: articles_raw = f.read() article_norm = unicodedata.normalize('NFKC', articles_raw) sentences_raw = utils.split_sentences(articles_raw) sentences_norm = utils.split_sentences(article_norm) text_entities_set = predict_entities(bert_tc, tokenizer, [sentences_norm]) for i, texts_ent in enumerate(text_entities_set[0]): texts_ent['text'] = sentences_raw[i] if normalize: normalize_entities(text_entities_set, id_to_tags, **kwargs) documents = combine_sentences(text_entities_set, id_to_tags, '\n') tqdm.write(f"File: {file}") tqdm.write(documents[0]) tqdm.write("") if output: with open(file.replace(input, output), 'w') as f: f.write(documents[0]) except Exception as e: tqdm.write("Error while processing file: {}".format(file)) tqdm.write(str(e)) tqdm.write("") if __name__ == '__main__': parser = argparse.ArgumentParser(description='Predict entities from text') parser.add_argument('-m', '--model', type=str, default='pytorch_model.bin', help='Path to model checkpoint') parser.add_argument('-i', '--input', type=str, default='text.txt', help='Path to text file or directory') parser.add_argument('-o', '--output', type=str, default=None, help='Path to output file or directory') parser.add_argument('-n', '--normalize', action=argparse.BooleanOptionalAction, help='Enable entity normalization', default=False) # Dictionary override arguments parser.add_argument("--drug-dict", help="File path for overriding the default drug dictionary") parser.add_argument("--drug-candidate-col", type=int, help="Column name for drug candidates in the CSV file (required if --drug-dict is specified)") parser.add_argument("--drug-normalization-col", type=int, help="Column name for drug normalization in the CSV file (required if --drug-dict is specified") parser.add_argument('--disease-matching-threshold', type=int, default=50, help='Matching threshold for disease dictionary') parser.add_argument("--disease-dict", help="File path for overriding the default disease dictionary") parser.add_argument("--disease-candidate-col", type=int, help="Column name for disease candidates in the CSV file (required if --disease-dict is specified)") parser.add_argument("--disease-normalization-col", type=int, help="Column name for disease normalization in the CSV file (required if --disease-dict is specified)") parser.add_argument('--drug-matching-threshold', type=int, default=50, help='Matching threshold for drug dictionary') args = parser.parse_args() argument_dict = vars(args) run(**argument_dict)