add eos_id option for ch07

ch05/01_main-chapter-code/ch05.ipynb

@@ -1852,7 +1852,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def generate(model, idx, max_new_tokens, context_size, temperature, top_k=None):\n",
+    "def generate(model, idx, max_new_tokens, context_size, temperature, top_k=None, eos_id=None):\n",
     "\n",
     "    # For-loop is the same as before: Get logits, and only focus on last time step\n",
     "    for _ in range(max_new_tokens):\n",
@@ -1882,6 +1882,9 @@
     "        else:\n",
     "            idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch_size, 1)\n",
     "\n",
+    "        if idx_next == eos_id:  # Stop generating early if end-of-sequence token is encountered and eos_id is specified\n",
+    "            break\n",
+    "\n",
     "        # Same as before: append sampled index to the running sequence\n",
     "        idx = torch.cat((idx, idx_next), dim=1)  # (batch_size, num_tokens+1)\n",
     "\n",
@@ -2372,7 +2375,7 @@
     "\n",
     "token_ids = generate(\n",
     "    model=gpt,\n",
-    "    idx=text_to_token_ids(\"Every effort moves you\", tokenizer),\n",
+    "    idx=text_to_token_ids(\"Every effort moves you\", tokenizer).to(device),\n",
     "    max_new_tokens=25,\n",
     "    context_size=NEW_CONFIG[\"context_length\"],\n",
     "    top_k=50,\n",
@@ -2439,7 +2442,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.11.4"
   }
  },
  "nbformat": 4,