From dfca8661ea61ea27e925049b7cf166ee0f239c27 Mon Sep 17 00:00:00 2001 From: Guoqiang QI Date: Tue, 1 Aug 2023 20:42:41 +0800 Subject: [PATCH] Remove unnecessary duplicated model loading. --- ptuning/web_demo.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/ptuning/web_demo.py b/ptuning/web_demo.py index 43d0c82..39b3d8a 100644 --- a/ptuning/web_demo.py +++ b/ptuning/web_demo.py @@ -136,17 +136,16 @@ def main(): config.pre_seq_len = model_args.pre_seq_len config.prefix_projection = model_args.prefix_projection + model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) + if model_args.ptuning_checkpoint is not None: print(f"Loading prefix_encoder weight from {model_args.ptuning_checkpoint}") - model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin")) new_prefix_state_dict = {} for k, v in prefix_state_dict.items(): if k.startswith("transformer.prefix_encoder."): new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict) - else: - model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) if model_args.quantization_bit is not None: print(f"Quantized to {model_args.quantization_bit} bit") @@ -163,4 +162,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main()