LLaMA-Factory/src/export_model.py
hiyouga d21a64160d tiny fix
Former-commit-id: 994f2a79831a1dea8425e3eff62f0bc8238b78d6
2023-06-27 23:54:24 +08:00

21 lines
682 B
Python

# coding=utf-8
# Exports the fine-tuned model.
# Usage: python export_model.py --checkpoint_dir path_to_checkpoint --output_dir path_to_save_model
from utils import load_pretrained, prepare_args
def main():
model_args, _, training_args, finetuning_args = prepare_args(stage="sft")
model, tokenizer = load_pretrained(model_args, finetuning_args)
model.save_pretrained(training_args.output_dir, max_shard_size="10GB")
tokenizer.save_pretrained(training_args.output_dir)
print("model and tokenizer have been saved at:", training_args.output_dir)
print("Remember to copy the *.py files from the original directory.")
if __name__ == "__main__":
main()