add trained_folder note

README.md

@@ -82,7 +82,7 @@ ai = aitextgen(tokenizer_file=tokenizer_file, config=config)
 # which automatically processes the dataset with the appropriate size.
 data = TokenDataset(file_name, tokenizer_file=tokenizer_file, block_size=64)
 
-# Train the model! It will save pytorch_model.bin periodically and after completion.
+# Train the model! It will save pytorch_model.bin periodically and after completion to the `trained_model` folder.
 # On a 2020 8-core iMac, this took ~25 minutes to run.
 ai.train(data, batch_size=8, num_steps=50000, generate_every=5000, save_every=5000)
 

notebooks/training_hello_world.ipynb

@@ -131,7 +131,7 @@
   },
   {
    "source": [
-    "Train the model! It will save pytorch_model.bin periodically and after completion. On a 2020 8-core iMac, this took ~25 minutes to run.\n",
+    "Train the model! It will save pytorch_model.bin periodically and after completion to the `trained_model` folder. On a 2020 8-core iMac, this took ~25 minutes to run.\n",
     "\n",
     "The configuration below processes 400,000 subsets of tokens (8 * 50000), which is about just one pass through all the data (1 epoch). Ideally you'll want multiple passes through the data and a training loss less than `2.0` for coherent output; when training a model from scratch, that's more difficult, but with long enough training you can get there!"
    ],
@@ -317,9 +317,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ai2 = aitextgen(model=\"trained_model/pytorch_model.bin\",\n",
-    "                tokenizer_file=\"aitextgen.tokenizer.json\",\n",
-    "                config=\"trained_model/config.json\")"
+    "ai2 = aitextgen(model_folder=\"trained_model\",\n",
+    "                tokenizer_file=\"aitextgen.tokenizer.json\")"
    ]
   },
   {