from transformers import AutoModelForTokenClassification, AutoTokenizer from helpers import helper MODEL_NAME = 'boda/ANER' # Load model and tokenizer model = AutoModelForTokenClassification.from_pretrained('.') tokenizer = AutoTokenizer.from_pretrained('.') # change in the model labels # model.config.id2label = {i: ' '+v+' ' for i, v in model.config.id2label.items() if i != 0 } # model.config.id2label[0] = 'O' # model.config.label2id = {label: i for i, label in model.config.id2label.items()} # save model after finish # model.save_pretrained('.')