From 179a2c2ff66ebf147c562243d7f4b6f37c0cdd23 Mon Sep 17 00:00:00 2001 From: thomwolf Date: Mon, 15 Apr 2019 14:33:23 +0200 Subject: [PATCH] update example to work with new serialization semantic --- examples/run_classifier.py | 25 ++++++++++++++----------- examples/run_openai_gpt.py | 20 +++++++++++++------- examples/run_squad.py | 25 ++++++++++++++----------- examples/run_swag.py | 25 ++++++++++++++----------- 4 files changed, 55 insertions(+), 40 deletions(-) diff --git a/examples/run_classifier.py b/examples/run_classifier.py index 751d581ad9..ba49d18b8d 100644 --- a/examples/run_classifier.py +++ b/examples/run_classifier.py @@ -37,7 +37,7 @@ from sklearn.metrics import matthews_corrcoef, f1_score from pytorch_pretrained_bert.file_utils import PYTORCH_PRETRAINED_BERT_CACHE from pytorch_pretrained_bert.modeling import BertForSequenceClassification, BertConfig, WEIGHTS_NAME, CONFIG_NAME -from pytorch_pretrained_bert.tokenization import BertTokenizer +from pytorch_pretrained_bert.tokenization import BertTokenizer, VOCAB_NAME from pytorch_pretrained_bert.optimization import BertAdam, warmup_linear logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s - %(message)s', @@ -857,18 +857,21 @@ def main(): optimizer.zero_grad() global_step += 1 - # Save a trained model and the associated configuration + # Save a trained model, configuration and tokenizer model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self - output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME) - torch.save(model_to_save.state_dict(), output_model_file) - output_config_file = os.path.join(args.output_dir, CONFIG_NAME) - with open(output_config_file, 'w') as f: - f.write(model_to_save.config.to_json_string()) - # Load a trained model and config that you have fine-tuned - config = BertConfig(output_config_file) - model = BertForSequenceClassification(config, num_labels=num_labels) - model.load_state_dict(torch.load(output_model_file)) + # If we save using the predefined names, we can load using `from_pretrained` + output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME) + output_config_file = os.path.join(args.output_dir, CONFIG_NAME) + output_vocab_file = os.path.join(args.output_dir, VOCAB_NAME) + + torch.save(model_to_save.state_dict(), output_model_file) + model_to_save.config.to_json_file(output_config_file) + tokenizer.save_vocabulary(output_vocab_file) + + # Load a trained model and vocabulary that you have fine-tuned + model = BertForSequenceClassification.from_pretrained(args.output_dir, num_labels=num_labels) + tokenizer = BertTokenizer.from_pretrained(args.output_dir) else: model = BertForSequenceClassification.from_pretrained(args.bert_model, num_labels=num_labels) model.to(device) diff --git a/examples/run_openai_gpt.py b/examples/run_openai_gpt.py index ee30a7a0a4..1686c9eda6 100644 --- a/examples/run_openai_gpt.py +++ b/examples/run_openai_gpt.py @@ -40,6 +40,7 @@ from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler, TensorDataset) from pytorch_pretrained_bert import OpenAIGPTDoubleHeadsModel, OpenAIGPTTokenizer, OpenAIAdam, cached_path +from pytorch_pretrained_bert.modeling_openai import WEIGHTS_NAME, CONFIG_NAME ROCSTORIES_URL = "https://s3.amazonaws.com/datasets.huggingface.co/ROCStories.tar.gz" @@ -218,15 +219,20 @@ def main(): # Save a trained model if args.do_train: + # Save a trained model, configuration and tokenizer model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self - output_model_file = os.path.join(args.output_dir, "pytorch_model.bin") - config = model.config - torch.save(model_to_save.state_dict(), output_model_file) - # Load a trained model that you have fine-tuned - model_state_dict = torch.load(output_model_file) - model = OpenAIGPTDoubleHeadsModel(config) - model.load_state_dict(model_state_dict) + # If we save using the predefined names, we can load using `from_pretrained` + output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME) + output_config_file = os.path.join(args.output_dir, CONFIG_NAME) + + torch.save(model_to_save.state_dict(), output_model_file) + model_to_save.config.to_json_file(output_config_file) + tokenizer.save_vocabulary(args.output_dir) + + # Load a trained model and vocabulary that you have fine-tuned + model = OpenAIGPTDoubleHeadsModel.from_pretrained(args.output_dir) + tokenizer = OpenAIGPTTokenizer.from_pretrained(args.output_dir) model.to(device) if args.do_eval: diff --git a/examples/run_squad.py b/examples/run_squad.py index 043b795326..045c0afe1e 100644 --- a/examples/run_squad.py +++ b/examples/run_squad.py @@ -39,7 +39,7 @@ from pytorch_pretrained_bert.modeling import BertForQuestionAnswering, BertConfi from pytorch_pretrained_bert.optimization import BertAdam, warmup_linear from pytorch_pretrained_bert.tokenization import (BasicTokenizer, BertTokenizer, - whitespace_tokenize) + whitespace_tokenize, VOCAB_NAME) if sys.version_info[0] == 2: import cPickle as pickle @@ -1009,18 +1009,21 @@ def main(): global_step += 1 if args.do_train: - # Save a trained model and the associated configuration + # Save a trained model, configuration and tokenizer model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self - output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME) - torch.save(model_to_save.state_dict(), output_model_file) - output_config_file = os.path.join(args.output_dir, CONFIG_NAME) - with open(output_config_file, 'w') as f: - f.write(model_to_save.config.to_json_string()) - # Load a trained model and config that you have fine-tuned - config = BertConfig(output_config_file) - model = BertForQuestionAnswering(config) - model.load_state_dict(torch.load(output_model_file)) + # If we save using the predefined names, we can load using `from_pretrained` + output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME) + output_config_file = os.path.join(args.output_dir, CONFIG_NAME) + output_vocab_file = os.path.join(args.output_dir, VOCAB_NAME) + + torch.save(model_to_save.state_dict(), output_model_file) + model_to_save.config.to_json_file(output_config_file) + tokenizer.save_vocabulary(output_vocab_file) + + # Load a trained model and vocabulary that you have fine-tuned + model = BertForQuestionAnswering.from_pretrained(args.output_dir) + tokenizer = BertTokenizer.from_pretrained(args.output_dir) else: model = BertForQuestionAnswering.from_pretrained(args.bert_model) diff --git a/examples/run_swag.py b/examples/run_swag.py index f193582640..fa145c29d7 100644 --- a/examples/run_swag.py +++ b/examples/run_swag.py @@ -35,7 +35,7 @@ from tqdm import tqdm, trange from pytorch_pretrained_bert.file_utils import PYTORCH_PRETRAINED_BERT_CACHE from pytorch_pretrained_bert.modeling import (BertForMultipleChoice, BertConfig, WEIGHTS_NAME, CONFIG_NAME) from pytorch_pretrained_bert.optimization import BertAdam, warmup_linear -from pytorch_pretrained_bert.tokenization import BertTokenizer +from pytorch_pretrained_bert.tokenization import BertTokenizer, VOCAB_NAME logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt = '%m/%d/%Y %H:%M:%S', @@ -473,18 +473,21 @@ def main(): if args.do_train: - # Save a trained model and the associated configuration + # Save a trained model, configuration and tokenizer model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self - output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME) - torch.save(model_to_save.state_dict(), output_model_file) - output_config_file = os.path.join(args.output_dir, CONFIG_NAME) - with open(output_config_file, 'w') as f: - f.write(model_to_save.config.to_json_string()) - # Load a trained model and config that you have fine-tuned - config = BertConfig(output_config_file) - model = BertForMultipleChoice(config, num_choices=4) - model.load_state_dict(torch.load(output_model_file)) + # If we save using the predefined names, we can load using `from_pretrained` + output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME) + output_config_file = os.path.join(args.output_dir, CONFIG_NAME) + output_vocab_file = os.path.join(args.output_dir, VOCAB_NAME) + + torch.save(model_to_save.state_dict(), output_model_file) + model_to_save.config.to_json_file(output_config_file) + tokenizer.save_vocabulary(output_vocab_file) + + # Load a trained model and vocabulary that you have fine-tuned + model = BertForMultipleChoice.from_pretrained(args.output_dir, num_choices=4) + tokenizer = BertTokenizer.from_pretrained(args.output_dir) else: model = BertForMultipleChoice.from_pretrained(args.bert_model, num_choices=4) model.to(device)