fix finetuning bug by adding model.train() after evaluation

urchade · Mar 19, 2024 · e15c22a · e15c22a
1 parent abb3881
commit e15c22a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/examples/finetune.ipynb b/examples/finetune.ipynb
@@ -139,6 +139,9 @@
     "        pbar.set_description(description)\n",
     "\n",
     "        if (step + 1) % config.eval_every == 0:\n",
+    "\n",
+    "            model.eval()\n",
+    "            \n",
     "            if eval_data is not None:\n",
     "                results, f1 = model.evaluate(eval_data[\"samples\"], flat_ner=True, threshold=0.5, batch_size=12,\n",
     "                                     entity_types=eval_data[\"entity_types\"])\n",
@@ -148,7 +151,9 @@
     "            if not os.path.exists(config.save_directory):\n",
     "                os.makedirs(config.save_directory)\n",
     "                \n",
-    "            model.save_pretrained(f\"{config.save_directory}/finetuned_{step}\")"
+    "            model.save_pretrained(f\"{config.save_directory}/finetuned_{step}\")\n",
+    "\n",
+    "            model.train()"
    ]
   },
   {