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

add local_files_only

Browse files
Files changed (1) hide show
  1. model_utils.py +2 -2
model_utils.py CHANGED
@@ -81,14 +81,14 @@ def load_models():
81
  return bert_model, bilstm_model
82
 
83
  def load_custom_model(model_dir, tokenizer_dir, id2label):
84
- config = AutoConfig.from_pretrained(model_dir)
85
  config.id2label = id2label
86
  config.num_labels = len(id2label)
87
 
88
  model = BiLSTMForTokenClassification(model_name=config._name_or_path, num_labels=config.num_labels)
89
  model.config.id2label = id2label
90
  model.load_state_dict(torch.load(os.path.join(model_dir, 'pytorch_model.bin'), map_location=torch.device('cpu')))
91
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
92
 
93
  return model, tokenizer
94
 
 
81
  return bert_model, bilstm_model
82
 
83
  def load_custom_model(model_dir, tokenizer_dir, id2label):
84
+ config = AutoConfig.from_pretrained(model_dir, local_files_only=True)
85
  config.id2label = id2label
86
  config.num_labels = len(id2label)
87
 
88
  model = BiLSTMForTokenClassification(model_name=config._name_or_path, num_labels=config.num_labels)
89
  model.config.id2label = id2label
90
  model.load_state_dict(torch.load(os.path.join(model_dir, 'pytorch_model.bin'), map_location=torch.device('cpu')))
91
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir, local_files_only = True)
92
 
93
  return model, tokenizer
94