From b8a034807e97730bc95a5f0bf0ca0763d7c6824c Mon Sep 17 00:00:00 2001 From: hiyouga Date: Fri, 2 Jun 2023 19:02:25 +0800 Subject: [PATCH] tiny fix --- src/export_model.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/export_model.py b/src/export_model.py index 9ba361c8..2d977ae9 100644 --- a/src/export_model.py +++ b/src/export_model.py @@ -4,15 +4,15 @@ from transformers import HfArgumentParser, TrainingArguments -from utils import ModelArguments, load_pretrained +from utils import ModelArguments, FinetuningArguments, load_pretrained def main(): - parser = HfArgumentParser((ModelArguments, TrainingArguments)) - model_args, training_args = parser.parse_args_into_dataclasses() + parser = HfArgumentParser((ModelArguments, TrainingArguments, FinetuningArguments)) + 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") tokenizer.save_pretrained(training_args.output_dir)