nlp - 如何保存和加载一个简单的变压器模型？

Question

我编写了一个脚本，在其中使用 simpletransformer 框架训练问答模型。我现在想做预测，但正如我想象的那样，我必须保存训练模型中的权重并将它们加载到我的预测脚本中，但我无法弄清楚如何做到这一点。下面的代码片段是我训练模型的方式。

model = QuestionAnsweringModel("distilbert",
                               "distilbert-base-uncased-distilled-squad",
                               args= {'reprocess_input_data': True,
                                     'overwrite_output_dir': True,
                                     'learning_rate': 5e-5,
                                     'num_train_epochs': 3,
                                     'max_seq_length': 192,
                                     'doc_stride': 64,
                                     'fp16': False,
                                    },
                               use_cuda= cuda_available)


model.train_model(config.QA_TRAIN)