from transformers.pipelines import PIPELINE_REGISTRY,AutoModel,AutoConfig from models.bert import BertCrfModel,BertCrfConfig from pipelines import NERPredictorPipe def register(): PIPELINE_REGISTRY.register_pipeline("ner_predictor", pipeline_class=NERPredictorPipe) AutoConfig.register("bert_crf",BertCrfConfig) AutoModel.register(BertCrfConfig,BertCrfModel)