Update You_can_still_run_it_directly_on_your_own_machine_if_it's_too_small.ipynb

You_can_still_run_it_directly_on_your_own_machine_if_it's_too_small.ipynb

@@ -193,6 +193,49 @@
       ],
       "id": "PP6KcuhQXeQ1"
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fix1chatFn"
+      },
+      "outputs": [],
+      "source": [
+        "# --- FIX: ham chat() bi thieu, day la ham duoc them lai ---\n",
+        "# Encode/decode qua tokenizers backend, va tu sinh token cho den khi\n",
+        "# gap END_ID (thay vi dung model.generate() goc, vi ham do khong biet\n",
+        "# dung lai o END_TOK va se tra ve ca prompt lan phan sinh ra).\n",
+        "\n",
+        "@torch.no_grad()\n",
+        "def chat(user_input, max_new_tokens=200, temperature=0.8, top_k=40):\n",
+        "    prompt = f\"{USER_TOK}{user_input}{ASSISTANT_TOK}\"\n",
+        "    prompt_ids = tokenizer_backend.encode(prompt).ids\n",
+        "    # cat bot neu prompt dai hon block_size\n",
+        "    prompt_ids = prompt_ids[-model.block_size:]\n",
+        "    idx = torch.tensor([prompt_ids], dtype=torch.long, device=device)\n",
+        "\n",
+        "    generated_ids = []\n",
+        "    for _ in range(max_new_tokens):\n",
+        "        idx_cond = idx[:, -model.block_size:]\n",
+        "        logits, _ = model(idx_cond)\n",
+        "        logits = logits[:, -1, :] / temperature\n",
+        "        if top_k is not None:\n",
+        "            v, _ = torch.topk(logits, top_k)\n",
+        "            logits[logits < v[:, [-1]]] = float(\"-inf\")\n",
+        "        probs = F.softmax(logits, dim=-1)\n",
+        "        next_id = torch.multinomial(probs, num_samples=1)\n",
+        "\n",
+        "        if next_id.item() == END_ID:\n",
+        "            break\n",
+        "\n",
+        "        generated_ids.append(next_id.item())\n",
+        "        idx = torch.cat([idx, next_id], dim=1)\n",
+        "\n",
+        "    reply = tokenizer_backend.decode(generated_ids)\n",
+        "    return reply.strip()\n"
+      ],
+      "id": "fix1chatFn"
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -229,4 +272,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 5
-}
+}