Spaces:
Sleeping
Sleeping
joyinning
commited on
Commit
•
b78f5bc
1
Parent(s):
099ddf3
add local_files_only
Browse files- 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 |
|