INNER CODE UNIT · Python

inference

mlwithme/BertWithPretrained · Tasks/TaskForChineseNER.py:273

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)
    data_loader = LoadChineseNERDataset(
        entities=config.entities,
        num_labels=config.num_labels,
        ignore_idx=config.ignore_idx,
        vocab_path=config.vocab_path,

View source record →