mirror of
https://github.com/wassname/peft.git
synced 2026-09-09 11:28:32 +08:00
minor changes
This commit is contained in:
@@ -1295,13 +1295,16 @@
|
||||
"from transformers import Seq2SeqTrainer, TrainerCallback, TrainingArguments, TrainerState, TrainerControl\n",
|
||||
"from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"class SavePeftModelCallback(TrainerCallback):\n",
|
||||
" def on_save(\n",
|
||||
" self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs,\n",
|
||||
" self,\n",
|
||||
" args: TrainingArguments,\n",
|
||||
" state: TrainerState,\n",
|
||||
" control: TrainerControl,\n",
|
||||
" **kwargs,\n",
|
||||
" ):\n",
|
||||
" checkpoint_folder = os.path.join(\n",
|
||||
" args.output_dir, f\"{PREFIX_CHECKPOINT_DIR}-{state.global_step}\"\n",
|
||||
" ) \n",
|
||||
" checkpoint_folder = os.path.join(args.output_dir, f\"{PREFIX_CHECKPOINT_DIR}-{state.global_step}\")\n",
|
||||
"\n",
|
||||
" peft_model_path = os.path.join(checkpoint_folder, \"adapter_model\")\n",
|
||||
" kwargs[\"model\"].save_pretrained(peft_model_path)\n",
|
||||
@@ -1311,6 +1314,7 @@
|
||||
" os.remove(pytorch_model_path)\n",
|
||||
" return control\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"trainer = Seq2SeqTrainer(\n",
|
||||
" args=training_args,\n",
|
||||
" model=model,\n",
|
||||
@@ -1319,7 +1323,7 @@
|
||||
" data_collator=data_collator,\n",
|
||||
" # compute_metrics=compute_metrics,\n",
|
||||
" tokenizer=processor.feature_extractor,\n",
|
||||
" callbacks=[SavePeftModelCallback]\n",
|
||||
" callbacks=[SavePeftModelCallback],\n",
|
||||
")\n",
|
||||
"model.config.use_cache = False # silence the warnings. Please re-enable for inference!"
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user