Update src/merge_lora.py

src/merge_lora.py

@@ -9,17 +9,17 @@ from utils.utils import print_arguments, add_arguments
 
 parser = argparse.ArgumentParser(description=__doc__)
 add_arg = functools.partial(add_arguments, argparser=parser)
-add_arg("lora_model", type=str, default="output/whisper-tiny/checkpoint-best/", help="微调保存的模型路径")
-add_arg('output_dir', type=str, default='models/',    help="合并模型的保存目录")
-add_arg("local_files_only", type=bool, default=False, help="是否只在本地加载模型，不尝试下载")
+add_arg("lora_model", type=str, default="output/whisper-tiny/checkpoint-best/", help="model directory")
+add_arg('output_dir', type=str, default='models/',    help="output directory")
+add_arg("local_files_only", type=bool, default=False, help="Choose local file , if not , find the model on HF")
 args = parser.parse_args()
 print_arguments(args)
 
-assert os.path.exists(args.lora_model), f"模型文件{args.lora_model}不存在"
+assert os.path.exists(args.lora_model), f"model file{args.lora_model}does not exist"
 
 peft_config = PeftConfig.from_pretrained(args.lora_model)
 #
-base_model = WhisperForConditionalGeneration.from_pretrained(peft_config.base_model_name_or_path, device_map={"": "cpu"},
+base_model = WhisperForConditionalGeneration.from_pretrained(peft_config.base_model_name_or_path, device_map={"": "cuda"},
                                                              local_files_only=args.local_files_only)
 
 model = PeftModel.from_pretrained(base_model, args.lora_model, local_files_only=args.local_files_only)
@@ -41,4 +41,4 @@ model.save_pretrained(save_directory)
 feature_extractor.save_pretrained(save_directory)
 tokenizer.save_pretrained(save_directory)
 processor.save_pretrained(save_directory)
-print(f'合并模型保持在：{save_directory}')
+print(f'Merged model save at ：{save_directory}')