From 48dc4c624e9540d24af41bcde14e66ce744a206b Mon Sep 17 00:00:00 2001
From: alvanli <alvin.ipad2@gmail.com>
Date: Wed, 8 Mar 2023 09:57:13 -0500
Subject: [PATCH] Add callback to save to local

---
 .../peft_bnb_whisper_large_v2_training.ipynb    | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb b/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
index 0f5d5f7..adeeb5e 100644
--- a/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
+++ b/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
@@ -1,7 +1,6 @@
 {
  "cells": [
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "5cefac89",
    "metadata": {},
@@ -10,7 +9,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "090fa3ed",
    "metadata": {},
@@ -22,7 +20,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "625e47a0",
    "metadata": {},
@@ -72,7 +69,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "8a528c1a",
    "metadata": {},
@@ -139,7 +135,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "805b1c56",
    "metadata": {},
@@ -1205,7 +1200,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "3906d436",
    "metadata": {},
@@ -1300,6 +1294,13 @@
    "source": [
     "from transformers import Seq2SeqTrainer\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",
+    "\n",
     "trainer = Seq2SeqTrainer(\n",
     "    args=training_args,\n",
     "    model=model,\n",
@@ -1308,6 +1309,7 @@
     "    data_collator=data_collator,\n",
     "    # compute_metrics=compute_metrics,\n",
     "    tokenizer=processor.feature_extractor,\n",
+    "    callbacks=[PeftSavingCallback]\n",
     ")\n",
     "model.config.use_cache = False  # silence the warnings. Please re-enable for inference!"
    ]
@@ -1585,7 +1587,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "Kzfg2qoXgrhg",
    "metadata": {
@@ -1928,7 +1929,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.10.6"
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {