| 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) |
|
|
|
|
|
|
