Update model_utils.py

Add new load_models() model

model_utils.py

@@ -64,47 +64,25 @@ class BiLSTMForTokenClassification(nn.Module):
         return {'loss': loss, 'logits': logits}
 
 def load_models():
-    """
-    Loads the custom BiLSTM model
-
-    Returns:
-        bilstm: The loaded bilstm model.
-    """
-    with open('models/bilstm-model.pkl', 'rb') as f:
-        bilstm_model = pickle.load(f)
-    bilstm_model.eval()
-
-    return bilstm_model
-
-def load_custom_model(model_dir, tokenizer_dir, id2label):
-    """
-    Loads a custom BiLSTM model and tokenizer from local files.
+    model_dir = "./models/bilstm_ner"
+    tokenizer_dir = "./models/tokenizer"
 
-    Args:
-        model_dir (str): Path to the directory containing the model's files.
-        tokenizer_dir (str): Path to the directory containing the tokenizer's files.
-        id2label (dict): Dictionary mapping label IDs to their names.
+    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'
+    }
 
-    Returns:
-        model: Loaded BiLSTMForTokenClassification model.
-        tokenizer: Loaded AutoTokenizer.
-        id2label: Input id2label dictionary.
-    """
     config = AutoConfig.from_pretrained(model_dir, local_files_only=True)
-    config.id2label = id2label
-    config.num_labels = len(id2label)
+    config.id2label = id2label_ner
+    config.num_labels = len(id2label_ner)
 
     model = BiLSTMForTokenClassification(model_name=config._name_or_path, num_labels=config.num_labels)
-    model.config.id2label = id2label
+    model.config.id2label = id2label_ner
     model.load_state_dict(torch.load(os.path.join(model_dir, 'pytorch_model.bin'), map_location=torch.device('cpu')))
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir, local_files_only = True)
-
-    return model, tokenizer, id2label
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir, local_files_only=True)
 
-ner_model_dir = "./models/bilstm_ner"
-tokenizer_dir = "./models/tokenizer"
-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'}
-ner_model, ner_tokenizer, id2label_ner = load_custom_model(ner_model_dir, tokenizer_dir, id2label_ner)
+    return model, tokenizer, id2label_ner
 
 # QA model
 qa_model = pipeline('question-answering', model='deepset/bert-base-cased-squad2')