wassname/peft
1
0
Fork 0
mirror of https://github.com/wassname/peft.git synced 2026-09-09 11:28:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use on save callback

Browse Source
This commit is contained in:
alvanli
2023-03-08 10:05:53 -05:00
parent 48dc4c624e
commit 5e788b329d
1 changed files with 18 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
+18 -8
View File
@@ -1292,14 +1292,24 @@
}
],
"source": [
"from transformers import Seq2SeqTrainer\n",
"from transformers import Seq2SeqTrainer, TrainerCallback, TrainingArguments, TrainerState, TrainerControl\n",
"from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR\n",
"\n",
"# Save model to local\n",
"class PeftSavingCallback(TrainerCallback):\n",
" def on_train_end(self, args, state, control, **kwargs):\n",
" kwargs[\"model\"].save_pretrained(state.best_model_checkpoint)\n",
" pytorch_model_path = os.path.join(state.best_model_checkpoint, \"pytorch_model.bin\")\n",
" os.remove(pytorch_model_path) if os.path.exists(pytorch_model_path) else None\n",
"class SavePeftModelCallback(TrainerCallback):\n",
" def on_save(\n",
" self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs,\n",
" ):\n",
" checkpoint_folder = os.path.join(\n",
" args.output_dir, f\"{PREFIX_CHECKPOINT_DIR}-{state.global_step}\"\n",
" ) \n",
"\n",
" peft_model_path = os.path.join(checkpoint_folder, \"adapter_model\")\n",
" kwargs[\"model\"].save_pretrained(peft_model_path)\n",
"\n",
" pytorch_model_path = os.path.join(checkpoint_folder, \"pytorch_model.bin\")\n",
" if os.path.exists(pytorch_model_path):\n",
" os.remove(pytorch_model_path)\n",
" return control\n",
"\n",
"trainer = Seq2SeqTrainer(\n",
" args=training_args,\n",
@@ -1309,7 +1319,7 @@
" data_collator=data_collator,\n",
" # compute_metrics=compute_metrics,\n",
" tokenizer=processor.feature_extractor,\n",
" callbacks=[PeftSavingCallback]\n",
" callbacks=[SavePeftModelCallback]\n",
")\n",
"model.config.use_cache = False # silence the warnings. Please re-enable for inference!"
]