Update notebooks/QLoRA.ipynb

notebooks/QLoRA.ipynb

@@ -109,7 +109,9 @@
     "dataset_name = \"kunishou/databricks-dolly-15k-ja\"\n",
     "dataset = datasets.load_dataset(dataset_name)\n",
     "dataset = dataset.map(encode)\n",
+    "dataset = dataset[\"train\"].train_test_split(0.1)\n",
     "train_dataset = dataset[\"train\"]\n",
+    "eval_dataset = dataset[\"test\"]\n",
     "\n",
     "# load model\n",
     "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=\"auto\", quantization_config=bnb_config, torch_dtype=torch.float16)\n",
@@ -147,16 +149,23 @@
     "    learning_rate=2e-4,\n",
     "    per_device_train_batch_size=1,\n",
     "    gradient_accumulation_steps=16,\n",
+    "    per_device_eval_batch_size=1,\n",
     "    num_train_epochs=0.1,\n",
     "    logging_strategy='steps',\n",
     "    logging_steps=10,\n",
-    "    save_strategy='epoch'\n",
+    "    save_strategy='epoch',\n",
+    "    evaluation_strategy='epoch',\n",
+    "    load_best_model_at_end=True,\n",
+    "    metric_for_best_model=\"eval_loss\",\n",
+    "    greater_is_better=False,\n",
+    "    save_total_limit=2\n",
     ")\n",
     "\n",
     "trainer = Trainer(\n",
     "    model=model,\n",
     "    args=training_args,\n",
     "    train_dataset=train_dataset,\n",
+    "    eval_dataset=eval_dataset,\n",
     "    data_collator=get_collator(tokenizer, 256)\n",
     ")\n",
     "\n",