From 382afc3822c3170262b2af5a892965b1da56e499 Mon Sep 17 00:00:00 2001 From: hiyouga Date: Fri, 2 Jun 2023 19:02:25 +0800 Subject: [PATCH] tiny fix Former-commit-id: d479aaf07a0997737da90c00558c42b32109928f --- 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)