Spaces:
Sleeping
Sleeping
joyinning
commited on
Commit
·
9a1456d
1
Parent(s):
24d2e66
revise file path
Browse files- model_utils.py +3 -3
model_utils.py
CHANGED
@@ -71,7 +71,7 @@ def load_models():
|
|
71 |
Returns:
|
72 |
bert_model: The loaded BERT model.
|
73 |
"""
|
74 |
-
bert_model = AutoModelForTokenClassification.from_pretrained("
|
75 |
bert_model.eval()
|
76 |
|
77 |
with open('models/bilstm-model.pkl', 'rb') as f:
|
@@ -92,8 +92,8 @@ def load_custom_model(model_dir, tokenizer_dir, id2label):
|
|
92 |
|
93 |
return model, tokenizer
|
94 |
|
95 |
-
ner_model_dir = "
|
96 |
-
tokenizer_dir = "
|
97 |
id2label_ner = {0: 'O', 1: 'I-art', 2: 'B-org', 3: 'B-geo', 4: 'I-per', 5: 'B-eve', 6: 'I-geo', 7: 'B-per', 8: 'I-nat', 9: 'B-art', 10: 'B-tim', 11: 'I-gpe', 12: 'I-tim', 13: 'B-nat', 14: 'B-gpe', 15: 'I-org', 16: 'I-eve'}
|
98 |
ner_model, ner_tokenizer, id2label_ner = load_custom_model(ner_model_dir, tokenizer_dir, id2label_ner)
|
99 |
|
|
|
71 |
Returns:
|
72 |
bert_model: The loaded BERT model.
|
73 |
"""
|
74 |
+
bert_model = AutoModelForTokenClassification.from_pretrained(".models/bert-model")
|
75 |
bert_model.eval()
|
76 |
|
77 |
with open('models/bilstm-model.pkl', 'rb') as f:
|
|
|
92 |
|
93 |
return model, tokenizer
|
94 |
|
95 |
+
ner_model_dir = "./models/bilstm_ner"
|
96 |
+
tokenizer_dir = "./models/tokenizer"
|
97 |
id2label_ner = {0: 'O', 1: 'I-art', 2: 'B-org', 3: 'B-geo', 4: 'I-per', 5: 'B-eve', 6: 'I-geo', 7: 'B-per', 8: 'I-nat', 9: 'B-art', 10: 'B-tim', 11: 'I-gpe', 12: 'I-tim', 13: 'B-nat', 14: 'B-gpe', 15: 'I-org', 16: 'I-eve'}
|
98 |
ner_model, ner_tokenizer, id2label_ner = load_custom_model(ner_model_dir, tokenizer_dir, id2label_ner)
|
99 |
|