This commit is contained in:
hiyouga 2023-06-02 19:02:25 +08:00
parent e3aaef7d4a
commit b8a034807e
1 changed files with 4 additions and 4 deletions

View File

@ -4,15 +4,15 @@
from transformers import HfArgumentParser, TrainingArguments from transformers import HfArgumentParser, TrainingArguments
from utils import ModelArguments, load_pretrained from utils import ModelArguments, FinetuningArguments, load_pretrained
def main(): def main():
parser = HfArgumentParser((ModelArguments, TrainingArguments)) parser = HfArgumentParser((ModelArguments, TrainingArguments, FinetuningArguments))
model_args, training_args = parser.parse_args_into_dataclasses() model_args, training_args, finetuning_args = parser.parse_args_into_dataclasses()
model, tokenizer = load_pretrained(model_args) model, tokenizer = load_pretrained(model_args, finetuning_args)
model.save_pretrained(training_args.output_dir, max_shard_size="1GB") model.save_pretrained(training_args.output_dir, max_shard_size="1GB")
tokenizer.save_pretrained(training_args.output_dir) tokenizer.save_pretrained(training_args.output_dir)