Update TextGenerationModelTraining.ipynb

Author: jimmytwei

AI-and-Analytics/Features-and-Functionality/IntelTensorFlow_TextGeneration_with_LSTM/TextGenerationModelTraining.ipynb

@@ -47,6 +47,7 @@
    "source": [
     "import string\n",
     "import requests\n",
+    "import os\n",
     "\n",
     "response = requests.get('https://www.gutenberg.org/cache/epub/1497/pg1497.txt')\n",
     "data = response.text.split('\\n')\n",
@@ -253,6 +254,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "num_epochs = 200\n",
+    "# For custom epochs numbers from the environment\n",
+    "if \"ITEX_NUM_EPOCHS\" in os.environ:\n",
+    "    num_epochs = int(os.environ.get('ITEX_NUM_EPOCHS'))\n",
+    "\n",
     "neuron_coef = 4\n",
     "itex_lstm_model = Sequential()\n",
     "itex_lstm_model.add(Embedding(input_dim=vocab_size, output_dim=seq_length, input_length=seq_length))\n",
@@ -262,7 +268,7 @@
     "itex_lstm_model.add(Dense(units=vocab_size, activation='softmax'))\n",
     "itex_lstm_model.summary()\n",
     "itex_lstm_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])\n",
-    "itex_lstm_model.fit(x,y, batch_size=256, epochs=200)"
+    "itex_lstm_model.fit(x,y, batch_size=256, epochs=num_epochs)"
    ]
   },
   {
@@ -296,6 +302,11 @@
     "seq_length = x.shape[1]\n",
     "vocab_size = y.shape[1]\n",
     "\n",
+    "num_epochs = 20\n",
+    "# For custom epochs numbers\n",
+    "if \"KERAS_NUM_EPOCHS\" in os.environ:\n",
+    "    num_epochs = int(os.environ.get('KERAS_NUM_EPOCHS'))\n",
+    "\n",
     "neuron_coef = 1\n",
     "keras_lstm_model = Sequential()\n",
     "keras_lstm_model.add(Embedding(input_dim=vocab_size, output_dim=seq_length, input_length=seq_length))\n",
@@ -305,7 +316,7 @@
     "keras_lstm_model.add(Dense(units=vocab_size, activation='softmax'))\n",
     "keras_lstm_model.summary()\n",
     "keras_lstm_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])\n",
-    "keras_lstm_model.fit(x,y, batch_size=256, epochs=20)"
+    "keras_lstm_model.fit(x,y, batch_size=256, epochs=num_epochs)"
    ]
   },
   {