joyinning commited on
Commit
24c8832
·
verified ·
1 Parent(s): c9f64e3

Update model_utils.py

Browse files

Add Description of load_models()

Files changed (1) hide show
  1. model_utils.py +9 -0
model_utils.py CHANGED
@@ -64,6 +64,15 @@ class BiLSTMForTokenClassification(nn.Module):
64
  return {'loss': loss, 'logits': logits}
65
 
66
  def load_models():
 
 
 
 
 
 
 
 
 
67
  model_dir = "./models/bilstm_ner"
68
  tokenizer_dir = "./models/tokenizer"
69
 
 
64
  return {'loss': loss, 'logits': logits}
65
 
66
  def load_models():
67
+ """
68
+ Loads the custom BiLSTM NER model and tokenizer from local files.
69
+
70
+ Returns:
71
+ model (BiLSTMForTokenClassification): The loaded BiLSTM NER model.
72
+ tokenizer (AutoTokenizer): The loaded tokenizer associated with the model.
73
+ id2label_ner (dict): A dictionary mapping numerical label indices to NER tags.
74
+ """
75
+
76
  model_dir = "./models/bilstm_ner"
77
  tokenizer_dir = "./models/tokenizer"
78