update notebook

train.ipynb

@@ -65,7 +65,7 @@
     "    import os\n",
     "    import wandb\n",
     "\n",
-    "    num_epochs = int(os.getenv(\"NUM_EPOCHS\", 3))\n",
+    "    num_epochs = int(os.getenv(\"NUM_EPOCHS\", 1))\n",
     "    output_dir = os.getenv(\"OUTPUT_DIR\", \"./model\")\n",
     "    seed = int(os.getenv(\"RANDOM_SEED\", 42))\n",
     "    model_id = os.getenv(\"MODEL_ID\")\n",
@@ -73,8 +73,8 @@
     "    save_steps = float(os.getenv(\"SAVE_STEPS\", 100))\n",
     "    if save_steps.is_integer():\n",
     "        save_steps = int(save_steps)\n",
-    "    train_size = float(os.getenv(\"TRAIN_SIZE\", 4_000_000))\n",
-    "    test_size = float(os.getenv(\"TEST_SIZE\", 400_000))\n",
+    "    train_size = float(os.getenv(\"TRAIN_SIZE\", 12_000_000))\n",
+    "    test_size = float(os.getenv(\"TEST_SIZE\", 1_200_000))\n",
     "    if train_size.is_integer():\n",
     "        train_size = int(train_size)\n",
     "    if test_size.is_integer():\n",