File size: 1,312 Bytes
79e2bbf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
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)
|