register.py

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)