def _save_vocab(self, state):
if state.get('vocab', None):
vocab = state['vocab']
if hasattr(vocab, 'save_pretrained'):
vocab.save_pretrained(state['save_dir'])
else:
file_path_name = os.path.join(state['save_dir'], VOCAB_NAME)
if isinstance(vocab, dict):
json_to_text(file_path=file_path_name, data=vocab)
state.pop('vocab')
def save_predict_result(self, file_name, data, file_dir=None):
'''
保存预测信息
'''
if file_dir is None:
file_dir = self.args.output_dir
file_path = os.path.join(file_dir, file_name)
if ".pkl" in file_path:
save_pickle(file_path=file_path, data=data)
elif ".json" in file_path:
json_to_text(file_path=file_path, data=data)
else:
raise ValueError("file type: expected one of (.pkl, .json)")