Add files via upload

predict.py

@@ -0,0 +1,156 @@
+import os
+import random
+import argparse
+import numpy as np
+from time import time
+import torch
+from torch.utils.data import TensorDataset, DataLoader
+import bert_codes.tokenization as tokenization
+from train import data2fea, get_entity_dict, compare_entity_dict  # NOTE： it can trigger the utils.check_args() to delete train/dev log files.
+import ipdb
+
+# Configuration
+##############################################################################################
+DICT_LABEL = {"症状": [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]}
+DICT_LABEL_REV = dict()
+for k, v in DICT_LABEL.items():
+    DICT_LABEL_REV.update({v[0]: k})
+    DICT_LABEL_REV.update({v[1]: k})
+
+t_config = time()
+# set args
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu_ids', type=str, default='2, 3')
+# training parameter
+parser.add_argument('--test_batch_size', type=int, default=256)
+parser.add_argument('--float16', type=bool, default=True)  # only sm >= 7.0 (tensorcores)
+parser.add_argument('--seed', type=int, default=42)
+# data and model dir
+parser.add_argument('--test_dir', type=str, default='./datasets/ner_test.json')
+parser.add_argument('--feature_test_dir', type=str, default='./datasets/fea_ner_test.json')
+parser.add_argument('--vocab_file', type=str, default='./pretrained_models/bert_chinese/vocab.txt')
+parser.add_argument('--checkpoint_dir', type=str, default='check_points/base_ner')
+parser.add_argument('--predict_file', type=str, default='predict_log.txt')
+args = parser.parse_args()
+
+# tokenizer initialization
+tokenizer = tokenization.BertTokenizer(vocab_file=args.vocab_file, do_lower_case=True)
+
+# set seed
+random.seed(args.seed)
+np.random.seed(args.seed)
+torch.manual_seed(args.seed)
+
+# set gpu
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
+device = torch.device("cuda")
+n_gpu = torch.cuda.device_count()
+if n_gpu > 0:
+    torch.cuda.manual_seed_all(args.seed)
+print("device %s n_gpu %d" % (device, n_gpu))
+print("device: {} n_gpu: {} 16-bits training: {}".format(device, n_gpu, args.float16))
+
+# load model
+print('***** Loading Model *****')
+model = torch.load(args.checkpoint_dir + "/best_model.pth")
+model.to(device)
+
+print("Configuration Time: {}".format(time() - t_config))
+
+
+# Predict Batch
+##############################################################################################
+def prepare_for_test():
+    global args
+    # get features
+    test_features = data2fea(load_path=args.test_dir, save_path=args.feature_test_dir)
+    # get test dataloader
+    test_input_ids = torch.tensor([f['input_ids'] for f in test_features], dtype=torch.long)
+    test_input_mask = torch.tensor([f['input_mask'] for f in test_features], dtype=torch.long)
+    test_input_segments = torch.tensor([f['input_segments'] for f in test_features], dtype=torch.long)
+    test_input_tags = torch.tensor([f['input_tags'] for f in test_features], dtype=torch.long)
+    test_tensor = TensorDataset(test_input_ids, test_input_mask, test_input_segments, test_input_tags)
+    test_dataloader = DataLoader(test_tensor, batch_size=args.test_batch_size, shuffle=True)
+    print("Test-{}".format(len(test_features)))
+    return test_features, test_dataloader
+
+
+def print_and_save_batch_dict(input_id, d_true, d_pred):
+    global args, tokenizer
+    lst_token = tokenizer.convert_ids_to_tokens(input_id.cpu().numpy())
+    lst_token = [t.replace("##", "") if t.startswith("##") else t for t in lst_token]
+    s_token = "".join(lst_token).replace("[PAD]", "").replace("[CLS]", "").replace("[SEP]", "")
+    with open(args.predict_file, 'a') as aw_dev:
+        aw_dev.write("Sent: {} \n".format(s_token))
+    s_true, s_pred = "", ""
+    for k, lst_v in d_true.items():
+        if len(lst_v) > 0:
+            s_true = s_true + ">" + k + ": "
+            for i_m, m in enumerate(lst_v):
+                lst_m = m.split("_")
+                idx_s, idx_e = int(lst_m[0]), int(lst_m[1])
+                s_true = s_true + "".join(lst_token[idx_s:idx_e]) + "(" + lst_m[0] + ", " + lst_m[1] + "); "
+    for k, lst_v in d_pred.items():
+        if len(lst_v) > 0:
+            s_pred = s_pred + ">" + k + ": "
+            for m in lst_v:
+                lst_m = m.split("_")
+                idx_s, idx_e = int(lst_m[0]), int(lst_m[1])
+                s_pred = s_pred + "".join(lst_token[idx_s:idx_e]) + "(" + lst_m[0] + ", " + lst_m[1] + "); "
+    with open(args.predict_file, 'a') as f:
+        f.write("True: {} \n".format(s_true))
+        f.write("Pred: {} \n".format(s_pred))
+        f.write("\n")
+    return None
+
+
+def predict():
+    global args, model, DICT_LABEL
+    print("***** Preprocessing *****")
+    _, test_dataloader = prepare_for_test()
+
+    print("***** Predict *****")
+    model.eval()
+    d_res = {k: {"precision": [], "recall": [], "f1": []} for k in DICT_LABEL.keys()}
+    with torch.no_grad():
+        for step, batch in enumerate(test_dataloader):
+            batch = tuple(t.to(device) for t in batch)
+            input_ids, input_mask, input_segments, input_tags = batch
+            logits = model(input_ids=input_ids,
+                           token_type_ids=input_segments,
+                           attention_mask=input_mask)
+            logits = logits.detach().cpu().numpy()  # [bs, len, dim]
+            pred_batch = np.argmax(logits, axis=-1)  # get predicted labels: [bs,len]
+            true_batch = input_tags.detach().cpu().numpy()  # get true labels
+            # calculate each sample in the batch
+            batch_size = true_batch.shape[0]
+            for i in range(batch_size):
+                true_batch_now = get_entity_dict(true_batch[i])
+                pred_batch_now = get_entity_dict(pred_batch[i])
+                input_id_now = input_ids[i]
+                res_batch = compare_entity_dict(d_true=true_batch_now, d_pred=pred_batch_now)
+                print_and_save_batch_dict(input_id=input_id_now, d_true=true_batch_now, d_pred=pred_batch_now)
+            for k, v in res_batch.items():
+                d_res[k]["precision"].append(v["precision"])
+                d_res[k]["recall"].append(v["recall"])
+                d_res[k]["f1"].append(v["f1"])
+            with open(args.predict_file, 'a') as f:
+                f.write("Result of batch {} is: \n {} \n".format(step, res_batch))
+                f.write(" ---------------------------------------------------- \n")
+            print("Result of batch {} is: \n {} \n ".format(step, res_batch))
+        # get final scores
+        f1 = np.mean([np.mean(v["f1"]) for k, v in d_res.items()])
+        with open(args.predict_file, 'a') as f:
+            f.write("F1 of all batches is: {:.4f} \n".format(f1))
+            f.write(" ---------------------------------------------------- \n")
+        print("F1 of all batches is: {:.4f} \n".format(f1))
+    return f1
+
+
+# Main
+##############################################################################################
+if __name__ == '__main__':
+    res = predict()
+