INNER CODE UNIT · Python

show_result

mlwithme/BertWithPretrained · Tasks/TaskForChineseNER.py:268

def show_result(sentences, logits, token_ids, entities):
    labels, _ = get_ner_tags(logits, token_ids, entities)
    pretty_print(sentences, labels, entities)


def inference(config, sentences=None):
    model = BertForTokenClassification(config,
                                       config.pretrained_model_dir)
    model_save_path = os.path.join(config.model_save_dir,
                                   config.model_save_name)
    if os.path.exists(model_save_path):
        checkpoint = torch.load(model_save_path)
        loaded_paras = checkpoint['model_state_dict']
        model.load_state_dict(loaded_paras)
        logging.info("## 成功载入已有模型，进行追加训练......")
    else:
        raise ValueError(f" 本地模型{model_save_path}不存在，请先训练模型。")
    model = model.to(config.device)

View source record →