Update training code

training.ipynb

@@ -19,7 +19,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {
     "id": "t-jXeSJKE1WM"
    },
@@ -32,15 +32,19 @@
     "import torch\n",
     "from transformers import (\n",
     "    PreTrainedTokenizerFast,\n",
+    "    AutoTokenizer,\n",
     "    DataCollatorForSeq2Seq,\n",
     "    Seq2SeqTrainingArguments,\n",
-    "    BertJapaneseTokenizer,\n",
     "    Trainer\n",
     ")\n",
     "from transformers.models.encoder_decoder.modeling_encoder_decoder import EncoderDecoderModel\n",
     "\n",
     "from datasets import load_dataset\n",
     "\n",
+    "import os\n",
+    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
+    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n",
+    "\n",
     "# encoder_model_name = \"xlm-roberta-base\"\n",
     "encoder_model_name = \"cl-tohoku/bert-base-japanese-v2\"\n",
     "decoder_model_name = \"skt/kogpt2-base-v2\""
@@ -48,31 +52,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "id": "nEW5trBtbykK"
    },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "(device(type='cpu'), 0)"
-      ]
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
-    "# device = torch.device(\"cpu\")\n",
-    "device, torch.cuda.device_count()"
+    "# device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "# # device = torch.device(\"cpu\")\n",
+    "# torch.cuda.set_device(device)\n",
+    "# device, torch.cuda.device_count()"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {
     "id": "5ic7pUUBFU_v"
    },
@@ -82,9 +76,9 @@
     "    def build_inputs_with_special_tokens(self, token_ids: List[int]) -> List[int]:\n",
     "        return token_ids + [self.eos_token_id]        \n",
     "\n",
-    "src_tokenizer = BertJapaneseTokenizer.from_pretrained(encoder_model_name)\n",
-    "trg_tokenizer = GPT2Tokenizer.from_pretrained(decoder_model_name, bos_token='</s>', eos_token='</s>', unk_token='<unk>',\n",
-    "  pad_token='<pad>', mask_token='<mask>')"
+    "src_tokenizer = AutoTokenizer.from_pretrained(encoder_model_name, use_fast=False)\n",
+    "trg_tokenizer = GPT2Tokenizer.from_pretrained(decoder_model_name, use_fast=False,\n",
+    "  bos_token='</s>', eos_token='</s>', unk_token='<unk>', pad_token='<pad>', mask_token='<mask>')"
    ]
   },
   {
@@ -98,25 +92,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {
-    "collapsed": false
-   },
-   "outputs": [],
-   "source": [
-    "dataset = load_dataset(\"sappho192/Tatoeba-Challenge-jpn-kor\")\n",
-    "# dataset = load_dataset(\"D:\\\\REPO\\\\Tatoeba-Challenge-jpn-kor\")\n",
-    "\n",
-    "train_dataset = dataset['train']\n",
-    "test_dataset = dataset['test']\n",
-    "\n",
-    "train_first_row = train_dataset[0]\n",
-    "test_first_row = test_dataset[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {
     "id": "65L4O1c5FLKt"
    },
@@ -124,7 +100,7 @@
    "source": [
     "class PairedDataset:\n",
     "    def __init__(self, \n",
-    "        source_tokenizer: PreTrainedTokenizerFast, target_tokenizer: PreTrainedTokenizerFast,\n",
+    "        source_tokenizer: AutoTokenizer, target_tokenizer: GPT2Tokenizer,\n",
     "        file_path: str = None,\n",
     "        dataset_raw: datasets.Dataset = None\n",
     "    ):\n",
@@ -132,7 +108,7 @@
     "        self.trg_tokenizer = target_tokenizer\n",
     "        \n",
     "        if file_path is not None:\n",
-    "            with open(file_path, 'r') as fd:\n",
+    "            with open(file_path, 'r', encoding=\"utf-8\") as fd:\n",
     "                reader = csv.reader(fd)\n",
     "                next(reader)\n",
     "                self.data = [row for row in reader]\n",
@@ -159,52 +135,66 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {
     "collapsed": false
    },
    "outputs": [],
    "source": [
-    "DATA_ROOT = './output'\n",
-    "FILE_FFAC_FULL = 'ffac_full.csv'\n",
-    "FILE_FFAC_TEST = 'ffac_test.csv'\n",
-    "FILE_JA_KO_TRAIN = 'ja_ko_train.csv'\n",
-    "FILE_JA_KO_TEST = 'ja_ko_test.csv'\n",
+    "# DATASET_TARGET = \"TATOEBA_2023\"\n",
+    "# DATASET_TARGET = \"FFAC\"\n",
+    "DATASET_TARGET = \"AIHUB\"\n",
+    "\n",
+    "if (DATASET_TARGET == \"TATOEBA_2023\"):\n",
+    "    # dataset = load_dataset(\"sappho192/Tatoeba-Challenge-jpn-kor\")\n",
+    "    dataset = load_dataset(\"/home/akalive/dataset/Tatoeba-Challenge-jpn-kor\")\n",
+    "\n",
+    "    train_dataset = dataset['train']\n",
+    "    test_dataset = dataset['test']\n",
+    "\n",
+    "    train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, dataset_raw=train_dataset)\n",
+    "    eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, dataset_raw=test_dataset)\n",
+    "elif (DATASET_TARGET == \"FFAC\"):\n",
+    "    DATA_ROOT = '/home/akalive/dataset/ffac/output'\n",
+    "    FILE_FFAC_FULL = 'ffac_full.csv'\n",
+    "    FILE_FFAC_TEST = 'ffac_test.csv'\n",
+    "    FILE_JA_KO_TRAIN = 'tteb_train.csv'\n",
+    "    FILE_JA_KO_TEST = 'tteb_test.csv'\n",
+    "\n",
+    "    # train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_FFAC_FULL}')\n",
+    "    # eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_FFAC_TEST}') \n",
+    "\n",
+    "    train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_JA_KO_TRAIN}')\n",
+    "    eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_JA_KO_TEST}')\n",
+    "elif (DATASET_TARGET == \"AIHUB\"):\n",
+    "    # AIHUB dataset spent 25~33GB of VRAM with batch_size=30 while training.\n",
+    "    DATA_ROOT = '/home/akalive/dataset/jkpair/data'\n",
+    "    FILE_TRAIN = 'train.csv'\n",
+    "    FILE_VAL = 'validation.csv'\n",
+    "\n",
+    "    train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_TRAIN}')\n",
+    "    eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_VAL}')\n",
     "\n",
-    "# train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_FFAC_FULL}')\n",
-    "# eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_FFAC_TEST}') \n",
+    "train_first_row = train_dataset[0]\n",
+    "eval_first_row = eval_dataset[0]\n",
     "\n",
-    "# train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_JA_KO_TRAIN}')\n",
-    "# eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_JA_KO_TEST}')"
+    "print(train_first_row)\n",
+    "print(eval_first_row)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {
-    "collapsed": false
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'input_ids': [2, 33, 2181, 1402, 893, 15200, 893, 13507, 881, 933, 882, 829, 3], 'labels': [9085, 10936, 10993, 23363, 9134, 18368, 8006, 389, 1]}"
-      ]
-     },
-     "execution_count": 7,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
-    "train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, dataset_raw=train_dataset)\n",
-    "eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, dataset_raw=test_dataset)\n",
-    "eval_dataset[0]"
+    "print(train_dataset)\n",
+    "train_dataset[0]"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -226,20 +216,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {
     "id": "I7uFbFYJFje8"
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Some weights of GPT2LMHeadModel were not initialized from the model checkpoint at skt/kogpt2-base-v2 and are newly initialized: ['transformer.h.0.crossattention.c_attn.bias', 'transformer.h.0.crossattention.c_attn.weight', 'transformer.h.0.crossattention.c_proj.bias', 'transformer.h.0.crossattention.c_proj.weight', 'transformer.h.0.crossattention.q_attn.bias', 'transformer.h.0.crossattention.q_attn.weight', 'transformer.h.0.ln_cross_attn.bias', 'transformer.h.0.ln_cross_attn.weight', 'transformer.h.1.crossattention.c_attn.bias', 'transformer.h.1.crossattention.c_attn.weight', 'transformer.h.1.crossattention.c_proj.bias', 'transformer.h.1.crossattention.c_proj.weight', 'transformer.h.1.crossattention.q_attn.bias', 'transformer.h.1.crossattention.q_attn.weight', 'transformer.h.1.ln_cross_attn.bias', 'transformer.h.1.ln_cross_attn.weight', 'transformer.h.10.crossattention.c_attn.bias', 'transformer.h.10.crossattention.c_attn.weight', 'transformer.h.10.crossattention.c_proj.bias', 'transformer.h.10.crossattention.c_proj.weight', 'transformer.h.10.crossattention.q_attn.bias', 'transformer.h.10.crossattention.q_attn.weight', 'transformer.h.10.ln_cross_attn.bias', 'transformer.h.10.ln_cross_attn.weight', 'transformer.h.11.crossattention.c_attn.bias', 'transformer.h.11.crossattention.c_attn.weight', 'transformer.h.11.crossattention.c_proj.bias', 'transformer.h.11.crossattention.c_proj.weight', 'transformer.h.11.crossattention.q_attn.bias', 'transformer.h.11.crossattention.q_attn.weight', 'transformer.h.11.ln_cross_attn.bias', 'transformer.h.11.ln_cross_attn.weight', 'transformer.h.2.crossattention.c_attn.bias', 'transformer.h.2.crossattention.c_attn.weight', 'transformer.h.2.crossattention.c_proj.bias', 'transformer.h.2.crossattention.c_proj.weight', 'transformer.h.2.crossattention.q_attn.bias', 'transformer.h.2.crossattention.q_attn.weight', 'transformer.h.2.ln_cross_attn.bias', 'transformer.h.2.ln_cross_attn.weight', 'transformer.h.3.crossattention.c_attn.bias', 'transformer.h.3.crossattention.c_attn.weight', 'transformer.h.3.crossattention.c_proj.bias', 'transformer.h.3.crossattention.c_proj.weight', 'transformer.h.3.crossattention.q_attn.bias', 'transformer.h.3.crossattention.q_attn.weight', 'transformer.h.3.ln_cross_attn.bias', 'transformer.h.3.ln_cross_attn.weight', 'transformer.h.4.crossattention.c_attn.bias', 'transformer.h.4.crossattention.c_attn.weight', 'transformer.h.4.crossattention.c_proj.bias', 'transformer.h.4.crossattention.c_proj.weight', 'transformer.h.4.crossattention.q_attn.bias', 'transformer.h.4.crossattention.q_attn.weight', 'transformer.h.4.ln_cross_attn.bias', 'transformer.h.4.ln_cross_attn.weight', 'transformer.h.5.crossattention.c_attn.bias', 'transformer.h.5.crossattention.c_attn.weight', 'transformer.h.5.crossattention.c_proj.bias', 'transformer.h.5.crossattention.c_proj.weight', 'transformer.h.5.crossattention.q_attn.bias', 'transformer.h.5.crossattention.q_attn.weight', 'transformer.h.5.ln_cross_attn.bias', 'transformer.h.5.ln_cross_attn.weight', 'transformer.h.6.crossattention.c_attn.bias', 'transformer.h.6.crossattention.c_attn.weight', 'transformer.h.6.crossattention.c_proj.bias', 'transformer.h.6.crossattention.c_proj.weight', 'transformer.h.6.crossattention.q_attn.bias', 'transformer.h.6.crossattention.q_attn.weight', 'transformer.h.6.ln_cross_attn.bias', 'transformer.h.6.ln_cross_attn.weight', 'transformer.h.7.crossattention.c_attn.bias', 'transformer.h.7.crossattention.c_attn.weight', 'transformer.h.7.crossattention.c_proj.bias', 'transformer.h.7.crossattention.c_proj.weight', 'transformer.h.7.crossattention.q_attn.bias', 'transformer.h.7.crossattention.q_attn.weight', 'transformer.h.7.ln_cross_attn.bias', 'transformer.h.7.ln_cross_attn.weight', 'transformer.h.8.crossattention.c_attn.bias', 'transformer.h.8.crossattention.c_attn.weight', 'transformer.h.8.crossattention.c_proj.bias', 'transformer.h.8.crossattention.c_proj.weight', 'transformer.h.8.crossattention.q_attn.bias', 'transformer.h.8.crossattention.q_attn.weight', 'transformer.h.8.ln_cross_attn.bias', 'transformer.h.8.ln_cross_attn.weight', 'transformer.h.9.crossattention.c_attn.bias', 'transformer.h.9.crossattention.c_attn.weight', 'transformer.h.9.crossattention.c_proj.bias', 'transformer.h.9.crossattention.c_proj.weight', 'transformer.h.9.crossattention.q_attn.bias', 'transformer.h.9.crossattention.q_attn.weight', 'transformer.h.9.ln_cross_attn.bias', 'transformer.h.9.ln_cross_attn.weight']\n",
-      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "model = EncoderDecoderModel.from_encoder_decoder_pretrained(\n",
     "    encoder_model_name,\n",
@@ -251,174 +232,69 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CustomTrainingArguments(Seq2SeqTrainingArguments):\n",
+    "    def __init__(self,*args, **kwargs):\n",
+    "        super(CustomTrainingArguments, self).__init__(*args, **kwargs)\n",
+    "\n",
+    "    @property\n",
+    "    def device(self) -> \"torch.device\":\n",
+    "        \"\"\"\n",
+    "        The device used by this process.\n",
+    "        Name the device the number you use.\n",
+    "        \"\"\"\n",
+    "        return torch.device(\"cuda:0\")\n",
+    "\n",
+    "    @property\n",
+    "    def n_gpu(self):\n",
+    "        \"\"\"\n",
+    "        The number of GPUs used by this process.\n",
+    "        Note:\n",
+    "            This will only be greater than one when you have multiple GPUs available but are not using distributed\n",
+    "            training. For distributed training, it will always be 1.\n",
+    "        \"\"\"\n",
+    "        # Make sure `self._n_gpu` is properly setup.\n",
+    "        # _ = self._setup_devices\n",
+    "        # I set to one manullay\n",
+    "        self._n_gpu = 1\n",
+    "        return self._n_gpu\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
     "id": "YFq2GyOAUV0W"
    },
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "Finishing last run (ID:1vwqqxps) before initializing another..."
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a82aa19a250b43f28d7ecc72eeebc88d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "VBox(children=(Label(value='0.001 MB of 0.010 MB uploaded\\r'), FloatProgress(value=0.10972568578553615, max=1.…"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       " View run <strong style=\"color:#cdcd00\">jbert+kogpt2</strong> at: <a href='https://wandb.ai/sappho192/fftr-poc1/runs/1vwqqxps' target=\"_blank\">https://wandb.ai/sappho192/fftr-poc1/runs/1vwqqxps</a><br/>Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Find logs at: <code>.\\wandb\\run-20240131_135356-1vwqqxps\\logs</code>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Successfully finished last run (ID:1vwqqxps). Initializing new run:<br/>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c2cd7f6fb5b1428b98b80a3cc82ec303",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011288888888884685, max=1.0…"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Tracking run with wandb version 0.16.2"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Run data is saved locally in <code>d:\\REPO\\ffxiv-ja-ko-translator\\wandb\\run-20240131_135421-etxsdxw2</code>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Syncing run <strong><a href='https://wandb.ai/sappho192/fftr-poc1/runs/etxsdxw2' target=\"_blank\">jbert+kogpt2</a></strong> to <a href='https://wandb.ai/sappho192/fftr-poc1' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       " View project at <a href='https://wandb.ai/sappho192/fftr-poc1' target=\"_blank\">https://wandb.ai/sappho192/fftr-poc1</a>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       " View run at <a href='https://wandb.ai/sappho192/fftr-poc1/runs/etxsdxw2' target=\"_blank\">https://wandb.ai/sappho192/fftr-poc1/runs/etxsdxw2</a>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
     "# for Trainer\n",
     "import wandb\n",
     "\n",
     "collate_fn = DataCollatorForSeq2Seq(src_tokenizer, model)\n",
-    "wandb.init(project=\"fftr-poc1\", name='jbert+kogpt2')\n",
+    "wandb.init(project=\"aihub-gt-2023\", name='jbert+kogpt2')\n",
     "\n",
     "arguments = Seq2SeqTrainingArguments(\n",
+    "# arguments = CustomTrainingArguments(\n",
     "    output_dir='dump',\n",
     "    do_train=True,\n",
     "    do_eval=True,\n",
     "    evaluation_strategy=\"epoch\",\n",
     "    save_strategy=\"epoch\",\n",
-    "    num_train_epochs=3,\n",
+    "    num_train_epochs=5, # for 40GB\n",
     "    # num_train_epochs=25,\n",
-    "    per_device_train_batch_size=1,\n",
-    "    # per_device_train_batch_size=30, # takes 40GB\n",
-    "    # per_device_train_batch_size=64,\n",
-    "    per_device_eval_batch_size=1,\n",
-    "    # per_device_eval_batch_size=30,\n",
-    "    # per_device_eval_batch_size=64,\n",
+    "    # per_device_train_batch_size=15,\n",
+    "    per_device_train_batch_size=30, # takes 40GB\n",
+    "    # per_device_eval_batch_size=10,\n",
+    "    per_device_eval_batch_size=10,\n",
     "    warmup_ratio=0.1,\n",
     "    gradient_accumulation_steps=4,\n",
     "    save_total_limit=5,\n",
     "    dataloader_num_workers=1,\n",
-    "    # fp16=True, # ENABLE if CUDA is enabled\n",
+    "    fp16=True, # ENABLE if CUDA is enabled\n",
     "    load_best_model_at_end=True,\n",
     "    report_to='wandb'\n",
     ")\n",
@@ -454,26 +330,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {
     "id": "7vTqAgW6Ve3J"
    },
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0afe460e9f614d9a90379cf99fcf8af3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/9671328 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
     "trainer.train()\n",
     "\n",
@@ -484,12 +345,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# import wandb\n",
-    "# wandb.finish()"
+    "wandb.finish()"
    ]
   }
  ],