Round of fixes

Author: fchollet

chapter11_part04_sequence-to-sequence-learning.ipynb

@@ -405,6 +405,8 @@
     "            padding_mask = tf.cast(\n",
     "                mask[:, tf.newaxis, :], dtype=\"int32\")\n",
     "            padding_mask = tf.minimum(padding_mask, causal_mask)\n",
+    "        else:\n",
+    "            padding_mask = mask\n",
     "        attention_output_1 = self.attention_1(\n",
     "            query=inputs,\n",
     "            value=inputs,\n",

chapter12_part01_text-generation.ipynb

@@ -293,6 +293,8 @@
     "            padding_mask = tf.cast(\n",
     "                mask[:, tf.newaxis, :], dtype=\"int32\")\n",
     "            padding_mask = tf.minimum(padding_mask, causal_mask)\n",
+    "        else:\n",
+    "            padding_mask = mask\n",
     "        attention_output_1 = self.attention_1(\n",
     "            query=inputs,\n",
     "            value=inputs,\n",
@@ -391,6 +393,8 @@
     "        self.model_input_length = model_input_length\n",
     "        self.temperatures = temperatures\n",
     "        self.print_freq = print_freq\n",
+    "        vectorized_prompt = text_vectorization([prompt])[0].numpy()\n",
+    "        self.prompt_length = np.nonzero(vectorized_prompt == 0)[0][0]\n",
     "\n",
     "    def on_epoch_end(self, epoch, logs=None):\n",
     "        if (epoch + 1) % self.print_freq != 0:\n",
@@ -401,7 +405,9 @@
     "            for i in range(self.generate_length):\n",
     "                tokenized_sentence = text_vectorization([sentence])\n",
     "                predictions = self.model(tokenized_sentence)\n",
-    "                next_token = sample_next(predictions[0, i, :])\n",
+    "                next_token = sample_next(\n",
+    "                    predictions[0, self.prompt_length - 1 + i, :]\n",
+    "                )\n",
     "                sampled_token = tokens_index[next_token]\n",
     "                sentence += \" \" + sampled_token\n",
     "            print(sentence)\n",