Update Model/NER/VLSP2016/Load_model.py

Model/NER/VLSP2016/Load_model.py

@@ -1,34 +1,34 @@
-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('E:/demo_datn/pythonProject1/Model/NER/VLSP2016/best_model.pt', map_location='cpu')
-else:
-        checkpoint_data = torch.load('E:/demo_datn/pythonProject1/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)
-
-
+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)
+
+