| 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('.') | |