| from transformers import RobertaConfig, AutoConfig | |
| from transformers import AutoTokenizer, AutoModelForTokenClassification | |
| from Model.NER.VLSP2021.Ner_CRF import PhoBertCrf,PhoBertSoftmax,PhoBertLstmCrf | |
| from Model.NER.VLSP2021.Predict_Ner import ViTagger | |
| import torch | |
| from spacy import displacy | |
| import re | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| MODEL_MAPPING = { | |
| 'vinai/phobert-base': { | |
| 'softmax': PhoBertSoftmax, | |
| 'crf': PhoBertCrf, | |
| 'lstm_crf': PhoBertLstmCrf | |
| }, | |
| } | |
| if device == 'cpu': | |
| checkpoint_data = torch.load('/Model/NER/VLSP2016/best_model.pt', map_location='cpu') | |
| else: | |
| checkpoint_data = torch.load('/Model/NER/VLSP2016/best_model.pt') | |
| configs = checkpoint_data['args'] | |
| print(configs.model_name_or_path) | |
| tokenizer = AutoTokenizer.from_pretrained(configs.model_name_or_path) | |
| model_clss = MODEL_MAPPING[configs.model_name_or_path][configs.model_arch] | |
| config = AutoConfig.from_pretrained(configs.model_name_or_path, | |
| num_labels=len(checkpoint_data['classes']), | |
| finetuning_task=configs.task) | |
| model = model_clss(config=config) | |
| model.resize_token_embeddings(len(tokenizer)) | |
| model.to(device) | |
| model.load_state_dict(checkpoint_data['model'],strict=False) | |
| print(model) | |