()
| 14 | return parser.parse_args() |
| 15 | |
| 16 | def main(): |
| 17 | args = get_args() |
| 18 | |
| 19 | base_model = AutoModelForCausalLM.from_pretrained( |
| 20 | args.base_model_name_or_path, |
| 21 | return_dict=True, |
| 22 | torch_dtype=torch.float16 |
| 23 | ) |
| 24 | |
| 25 | model = PeftModel.from_pretrained(base_model, args.peft_model_path) |
| 26 | model = model.merge_and_unload() |
| 27 | |
| 28 | tokenizer = AutoTokenizer.from_pretrained(args.base_model_name_or_path) |
| 29 | |
| 30 | if args.push_to_hub: |
| 31 | print(f"Saving to hub ...") |
| 32 | model.push_to_hub(f"{args.base_model_name_or_path}-merged", use_temp_dir=False, private=True) |
| 33 | tokenizer.push_to_hub(f"{args.base_model_name_or_path}-merged", use_temp_dir=False, private=True) |
| 34 | else: |
| 35 | model.save_pretrained(f"{args.base_model_name_or_path}-merged") |
| 36 | tokenizer.save_pretrained(f"{args.base_model_name_or_path}-merged") |
| 37 | print(f"Model saved to {args.base_model_name_or_path}-merged") |
| 38 | |
| 39 | if __name__ == "__main__" : |
| 40 | main() |
no test coverage detected