joyinning commited on
Commit
099ddf3
1 Parent(s): 07a5491

Fix model_utils.py

Browse files
Files changed (1) hide show
  1. model_utils.py +1 -1
model_utils.py CHANGED
@@ -95,7 +95,7 @@ def load_custom_model(model_dir, tokenizer_dir, id2label):
95
  ner_model_dir = "joyinning/chatbot-info-extraction/models/bilstm_ner"
96
  tokenizer_dir = "joyinning/chatbot-info-extraction/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 = load_custom_model(ner_model_dir, tokenizer_dir, id2label_ner)
99
 
100
  # QA model
101
  qa_model = pipeline('question-answering', model='deepset/bert-base-cased-squad2')
 
95
  ner_model_dir = "joyinning/chatbot-info-extraction/models/bilstm_ner"
96
  tokenizer_dir = "joyinning/chatbot-info-extraction/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
 
100
  # QA model
101
  qa_model = pipeline('question-answering', model='deepset/bert-base-cased-squad2')