Demos/16a. 使用 Transformers 加载量化后的 LLM 大模型（GPTQ & AWQ）.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4fdab2a6-be55-492e-b563-5f62ccc534af",
   "metadata": {},
   "source": [
    "# a. 使用 Transformers 加载量化后的 LLM 大模型（GPTQ & AWQ） \n",
    "\n",
    "> 引导文章：[19a. 从加载到对话：使用 Transformers 本地运行量化 LLM 大模型（GPTQ & AWQ）](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/19a.%20从加载到对话：使用%20Transformers%20本地运行量化%20LLM%20大模型（GPTQ%20%26%20AWQ）.md)。\n",
    "\n",
    "代码文件没有显卡要求，在个人计算机上均可进行对话。\n",
    "\n",
    "**模型文件约为 4 GB**。\n",
    "\n",
    "这里还有一个简单的 [🎡 AI Chat 脚本](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/CodePlayground/chat.py)供你尝试，详见：[CodePlayground](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/CodePlayground/README.md#当前的玩具)，点击 `►` 或对应的文本展开。\n",
    "\n",
    "Llama-cpp-python 关于 GGUF 文件加载的相关链接：[文章 19b](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/19b.%20从加载到对话：使用%20Llama-cpp-python%20本地运行量化%20LLM%20大模型（GGUF）.md) | [代码文件 16b](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Demos/16b.%20使用%20Llama-cpp-python%20加载量化后的%20LLM%20大模型（GGUF）.ipynb)。\n",
    "\n",
    "在线链接：[Kaggle - a](https://www.kaggle.com/code/aidemos/16a-transformers-llm-gptq) | [Colab - a](https://colab.research.google.com/drive/1cmIDjHriW8aQ5mIsV6ZeTqdnqYe6PoOv?usp=sharing)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "468c340d-1694-4a80-bc61-3e97e00e9e42",
   "metadata": {},
   "source": [
    "## 手动下载模型（推荐）\n",
    "\n",
    "来试试多线程指定文件下载，对于 Linux，这里给出配置命令，其余系统可以参照[《a. 使用 HFD 加快 Hugging Face 模型和数据集的下载》](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/a.%20使用%20HFD%20加快%20Hugging%20Face%20模型和数据集的下载.md)先进行环境配置。你也可以跳过这部分，后面会介绍自动下载。\n",
    "\n",
    "```bash\n",
    "sudo apt-get update\n",
    "sudo apt-get install git git-lfs wget aria2\n",
    "git lfs install\n",
    "```\n",
    "\n",
    "下载并配置 HFD 脚本：\n",
    "\n",
    "```bash\n",
    "wget https://huggingface.co/hfd/hfd.sh\n",
    "chmod a+x hfd.sh\n",
    "export HF_ENDPOINT=https://hf-mirror.com\n",
    "```\n",
    "\n",
    "使用多线程下载指定模型。\n",
    "\n",
    "### GPTQ\n",
    "\n",
    "命令遵循 `./hfd.sh <model_path> --tool aria2c -x <线程数>`的格式：\n",
    "\n",
    "```bash\n",
    "./hfd.sh neuralmagic/Mistral-7B-Instruct-v0.3-GPTQ-4bit --tool aria2c -x 16\n",
    "```\n",
    "\n",
    "### AWQ\n",
    "\n",
    "命令遵循 `./hfd.sh <model_path> --tool aria2c -x <线程数>`的格式：\n",
    "\n",
    "```python\n",
    "./hfd.sh solidrust/Mistral-7B-Instruct-v0.3-AWQ --tool aria2c -x 16\n",
    "```\n",
    "\n",
    "### GGUF\n",
    "\n",
    "使用多线程下载指定模型，命令遵循 `./hfd.sh <model_path> --include <file_name> --tool aria2c -x <线程数>`的格式：\n",
    "\n",
    "```python\n",
    "./hfd.sh bartowski/Mistral-7B-Instruct-v0.3-GGUF --include Mistral-7B-Instruct-v0.3-Q4_K_M.gguf --tool aria2c -x 16\n",
    "```\n",
    "\n",
    "下载完成你应该可以看到类似的输出：\n",
    "\n",
    "```\n",
    "Download Results:\n",
    "gid   |stat|avg speed  |path/URI\n",
    "======+====+===========+=======================================================\n",
    "145eba|OK  |   6.8MiB/s|./Mistral-7B-Instruct-v0.3-Q4_K_M.gguf\n",
    "\n",
    "Status Legend:\n",
    "(OK):download completed.\n",
    "Downloaded https://huggingface.co/bartowski/Mistral-7B-Instruct-v0.3-GGUF/resolve/main/Mistral-7B-Instruct-v0.3-Q4_K_M.gguf successfully.\n",
    "Download completed successfully.\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9dea4a5-9f06-40c1-b443-8d51b2be20d0",
   "metadata": {},
   "source": [
    "## Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be3f9a99-176d-429c-af1d-63a2cb46aee9",
   "metadata": {},
   "source": [
    "### 环境配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a085c448-afe2-451b-897c-68ed7ad131d9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install numpy==1.24.4\n",
    "!pip install --upgrade transformers\n",
    "!pip install optimum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0e3eb61-65f5-4043-9a37-41f1a7e9da1f",
   "metadata": {},
   "source": [
    "你需要注意的是，如果安装不正确，GPTQ 将无法正确使用 GPU 进行推理，也就是说无法进行加速，即便 print(model.device) 显示为 \"cuda\"。类似的问题见 [Is This Inference Speed Slow?  #130](https://github.com/AutoGPTQ/AutoGPTQ/issues/130) 和 [CUDA extension not installed #694](https://github.com/AutoGPTQ/AutoGPTQ/issues/694)。\n",
    "\n",
    "这个问题是普遍存在的，当你直接使用 `pip install auto-gptq` 进行安装时，可能就会出现。\n",
    "\n",
    "你可以通过以下命令检查已安装的版本：\n",
    "\n",
    "```bash\n",
    "pip list | grep auto-gptq\n",
    "```\n",
    "\n",
    "如果发现之前安装的版本不带 cuda 标识，卸载它，从源码重新进行安装（推理速度将提升为原来的 15 倍以上）。\n",
    "\n",
    "```bash\n",
    "pip uninstall auto-gptq\n",
    "git clone https://github.com/PanQiWei/AutoGPTQ.git && cd AutoGPTQ\n",
    "```\n",
    "\n",
    "```bash\n",
    "# 以下两种方式任选一种进行安装即可，经测试均有效\n",
    "pip install -vvv --no-build-isolation -e .\n",
    "# >> Successfully installed auto-gptq-0.8.0.dev0+cu121\n",
    "\n",
    "python setup.py install\n",
    "# >> Finished processing dependencies for auto-gptq==0.8.0.dev0+cu121\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16e82aff-a3a9-4ea5-a510-f935628ac045",
   "metadata": {},
   "source": [
    "### GPTQ\n",
    "\n",
    "#### 导入库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4edf71d3-8ea9-4cb7-b5f9-25bd1c2bc8b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c74c5e0-8d78-4800-af07-290302b054fd",
   "metadata": {},
   "source": [
    "下面介绍两种导入模型的方法，实际执行时本地/自动导入二选一。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d0d16b3-1d23-486f-bd32-a60d245d7d87",
   "metadata": {},
   "source": [
    "#### 设置模型路径\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45027d9a-0f4a-4bf0-8416-9f20fb010825",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 如果你已经配置过了，可以直接在 Notebook 中执行下面的命令下载。\n",
    "!export HF_ENDPOINT=https://hf-mirror.com\n",
    "!./hfd.sh neuralmagic/Mistral-7B-Instruct-v0.3-GPTQ-4bit --tool aria2c -x 16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e3ae2a7-fd89-4b8a-b7d9-b47a44c613be",
   "metadata": {},
   "source": [
    "如果已经在本地下载了模型，可以通过指定模型路径来加载模型。以下示例假设模型位于当前目录的 `Mistral-7B-Instruct-v0.3-GPTQ-4bit` 文件夹下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61f7636b-e20d-4a38-880d-8915dd0b0216",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 指定本地模型的路径\n",
    "model_path = \"./Mistral-7B-Instruct-v0.3-GPTQ-4bit\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f97d93f-4c5a-4f6d-b90c-9640e74a3a17",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "如果没有本地模型，设置远程路径（`id` + `/` + `model_name`），导入的时候会自动从 Hugging Face 下载模型："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13d0d8bc-3948-4f99-be76-41c8858a7eec",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 指定远程模型的路径\n",
    "model_path = \"neuralmagic/Mistral-7B-Instruct-v0.3-GPTQ-4bit\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f5510a5-4d7a-40a5-8859-54bbf2557a4e",
   "metadata": {},
   "source": [
    "#### 加载模型\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f05f48-b9a7-48f5-996c-153227b93822",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载分词器\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
    "\n",
    "# 下载并加载模型\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_path,\n",
    "    torch_dtype=\"auto\",  # 自动选择模型的权重数据类型\n",
    "    device_map=\"auto\"    # 自动选择可用的设备（CPU/GPU）\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44f31a6f-94e4-48bb-a8fd-78eac7b2f3af",
   "metadata": {},
   "source": [
    "#### 推理测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b25eaaa6-59e6-4a6f-a0ed-aa9db261cb7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 输入文本\n",
    "input_text = \"Hello, World!\"\n",
    "\n",
    "# 将输入文本编码为模型可接受的格式\n",
    "input_ids = tokenizer.encode(input_text, return_tensors=\"pt\").to(model.device)\n",
    "\n",
    "# 生成输出\n",
    "with torch.no_grad():\n",
    "    output_ids = model.generate(\n",
    "        input_ids=input_ids,\n",
    "        max_length=50,\n",
    "        pad_token_id=tokenizer.eos_token_id\n",
    "    )\n",
    "\n",
    "# 解码生成的输出\n",
    "output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)\n",
    "\n",
    "# 打印生成的文本\n",
    "print(output_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c093ae9a-0049-4eb7-995d-78958aaa4004",
   "metadata": {},
   "source": [
    "### AWQ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58ec1b61-708a-440e-9eee-91d360676bea",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install autoawq autoawq-kernels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58a036ce-ce60-4efe-aa9b-86ef1c1d9d61",
   "metadata": {},
   "source": [
    "#### 导入库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da97c440-7083-4a09-9203-4d4a4bec7102",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from awq import AutoAWQForCausalLM\n",
    "from transformers import AutoTokenizer, TextStreamer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02c534d8-3792-41a5-97ee-e7859e2aa094",
   "metadata": {},
   "source": [
    "下面介绍两种导入模型的方法，实际执行时本地/自动导入二选一。\n",
    "\n",
    "#### 设置模型路径\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7e43cbe-4282-4c05-bf6c-881d14e572a6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 如果你已经配置过了，可以直接在 Notebook 中执行下面的命令下载。\n",
    "!export HF_ENDPOINT=https://hf-mirror.com\n",
    "!./hfd.sh solidrust/Mistral-7B-Instruct-v0.3-AWQ --tool aria2c -x 16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "275cbbd2-03a1-4246-8b89-ae4b5e7d5363",
   "metadata": {},
   "source": [
    "如果已经在本地下载了模型，可以通过指定模型路径来加载模型。以下示例假设模型位于当前目录的 `Mistral-7B-Instruct-v0.3-AWQ` 文件夹下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6362c6e6-4a72-45a8-a51f-539550798f62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 指定本地模型的路径\n",
    "model_path = \"./Mistral-7B-Instruct-v0.3-AWQ\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f136eca1-3f7f-4731-bf79-0ac659b5566d",
   "metadata": {},
   "source": [
    "如果没有本地模型，设置远程路径（`id` + `/` + `model_name`），导入的时候会自动从 Hugging Face 下载模型：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d5c9108-fa94-4cf1-a26e-3b9e24ce29f0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 指定远程模型的路径\n",
    "model_path = \"solidrust/Mistral-7B-Instruct-v0.3-AWQ\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a752722-eb3f-4e15-aafa-b89ec8cb27d7",
   "metadata": {},
   "source": [
    "#### 加载模型\n",
    "\n",
    "一些权重不会被加载，对于当前任务来说这是预期的行为。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "177e052c-d8bb-41d2-bd0a-8113ec1597b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载分词器\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_path,\n",
    "    trust_remote_code=True\n",
    ")\n",
    "\n",
    "# 下载并加载模型\n",
    "model = AutoAWQForCausalLM.from_quantized(\n",
    "    model_path,\n",
    "    fuse_layers=True  # 融合部分模型层以提高推理速度\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a909e888-42a1-4a3a-954a-c090bf1afdcf",
   "metadata": {},
   "source": [
    "#### 推理测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72bf0b23-4328-439f-8265-c48f82ae2a2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置设备\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# 输入文本\n",
    "input_text = \"Hello, World!\"\n",
    "\n",
    "# 将输入文本编码为模型可接受的格式\n",
    "input_ids = tokenizer.encode(input_text, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# 生成输出\n",
    "with torch.no_grad():\n",
    "    output_ids = model.generate(\n",
    "        input_ids=input_ids,\n",
    "        max_length=50,\n",
    "        pad_token_id=tokenizer.eos_token_id\n",
    "    )\n",
    "\n",
    "# 解码生成的输出\n",
    "output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)\n",
    "\n",
    "# 打印生成的文本\n",
    "print(output_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "261f1f2d-6ed3-4caa-b3e4-04e8ac1cd452",
   "metadata": {},
   "source": [
    "### 统一方式加载\n",
    "\n",
    "值得一提的是，你可能已经注意到了 GPTQ 并不是使用 Auto-GPTQ 库，而是直接使用 Transformers（虽然背后仍然依赖于 Auto-GPTQ）。\n",
    "\n",
    "“要是所有的模型都能统一就好了，这样就不用查阅其他库的文档了。”\n",
    "\n",
    "事实上，AWQ 也可以通过 Transformers 来加载（如果你对参数设置没有更细致的要求）："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23672cc5-e408-49c6-b0d9-12be2e06fd86",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "# 指定模型，切换为对应的 GPTQ 或 AWQ 模型路径，可以是远程路径，会自动下载\n",
    "model_path = \"neuralmagic/Mistral-7B-Instruct-v0.3-GPTQ-4bit\"\n",
    "\n",
    "# 加载分词器\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
    "\n",
    "# 加载模型\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_path,\n",
    "    torch_dtype=\"auto\",  # 自动选择模型的权重数据类型\n",
    "    device_map=\"auto\"    # 自动选择可用的设备（CPU/GPU）\n",
    ")\n",
    "\n",
    "# 设置设备\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# 输入文本\n",
    "input_text = \"Hello, World!\"\n",
    "\n",
    "# 将输入文本编码为模型可接受的格式\n",
    "input_ids = tokenizer.encode(input_text, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# 生成输出\n",
    "with torch.no_grad():\n",
    "    output_ids = model.generate(\n",
    "        input_ids=input_ids,\n",
    "        max_length=50,\n",
    "        pad_token_id=tokenizer.eos_token_id\n",
    "    )\n",
    "\n",
    "# 解码生成的输出\n",
    "output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)\n",
    "\n",
    "# 打印生成的文本\n",
    "print(output_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "532ea31a-a5d8-47c3-9e7c-5fe6ee5c290f",
   "metadata": {},
   "source": [
    "### 了解提示词模版（prompt_template）\n",
    "\n",
    "其实非常简单，就是曾经提到的占位符（下图对于 `{{question}}` 的应用）。\n",
    "\n",
    "![占位符](../Guide/assets/%E5%8D%A0%E4%BD%8D%E7%AC%A6-6055722.png)\n",
    "\n",
    "举个直观的例子："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "212a83d1-e986-4dae-a7b7-962e0444a466",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义 Prompt Template\n",
    "prompt_template = \"问：{question}\\n答：\"\n",
    "\n",
    "# 定义问题\n",
    "question = \"人工智能的未来发展方向是什么？\"\n",
    "\n",
    "# 使用 Prompt Template 生成完整的提示\n",
    "prompt = prompt_template.format(question=question)\n",
    "print(prompt)\n",
    "# print(\"\\n\")\n",
    "# print(f\"问：{question}\\n答：\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05446cce-b0fa-4c49-a602-bcae318fd314",
   "metadata": {},
   "source": [
    "#### tokenizer.chat_template\n",
    "\n",
    "查看模型的 `chat_template`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc212d22-5316-4bf9-bc2d-58212e533b5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 打印 chat_template 信息（如果存在的话）\n",
    "if hasattr(tokenizer, 'chat_template'):\n",
    "    print(tokenizer.chat_template)\n",
    "else:\n",
    "    print(\"Tokenizer 没有 'chat_template' 属性。\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbe9a3a9-cac7-47ea-890d-430b1ba20fd4",
   "metadata": {},
   "source": [
    "### 流式输出\n",
    "\n",
    "在项目初期认识 API 的时候，文章[《01. 初识 LLM API：环境配置与多轮对话演示》](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/01.%20初识%20LLM%20API：环境配置与多轮对话演示.md#流式输出)有提到过流式输出，这是我们一直以来见到的大模型输出方式：逐字（token）打印而非等全部生成完打印。\n",
    "\n",
    "执行下面的代码试试（无论之前导入的是哪种模型，都可以继续）："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "545b6478-3937-46e7-9235-7f91c02bfedd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TextStreamer\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# 创建 TextStreamer 实例\n",
    "streamer = TextStreamer(\n",
    "    tokenizer, \n",
    "    skip_prompt=True,         # 在输出时跳过输入的提示部分，仅显示生成的文本\n",
    "    skip_special_tokens=True  # 忽略生成过程中的特殊标记（比如 <pad> / <eos> ...）\n",
    ")\n",
    "\n",
    "# 将提示编码为模型输入\n",
    "input_ids = tokenizer.encode(prompt, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# 设置生成参数\n",
    "generation_kwargs = {\n",
    "    \"input_ids\": input_ids,  # 模型的输入 ID，注意，这不是 Embedding\n",
    "    \"max_length\": 200,       # 生成的最大 token 数\n",
    "    \"streamer\": streamer,    # 使用 TextStreamer 实现生成过程中逐步输出文本\n",
    "    \"pad_token_id\": tokenizer.eos_token_id  # 默认行为，消除 open-end 警告\n",
    "}\n",
    "\n",
    "# 开始生成文本\n",
    "with torch.no_grad():\n",
    "    # ** 是 Python 中的解包操作符，它将字典中的键值对解包为函数的关键字参数。\n",
    "    # 在这里，**generation_kwargs 将字典中的参数逐一传递给 model.generate() 方法，\n",
    "    # 等效于直接写出所有参数：\n",
    "    # model.generate(input_ids=input_ids, max_length=200, do_sample=True, ...)\n",
    "    # 你需要注意到，这和之前采用了不同的传参方式，但本质是一样的。\n",
    "    # 在后续的教程中，会较少地使用这种方式进行传参。\n",
    "    # 因为这很好的分离了参数，所以也增加了乍一看之下的抽象度，为了初见的直观，将减少使用。\n",
    "    model.generate(**generation_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b235a01-1780-4573-b49e-3d889b0d4b76",
   "metadata": {},
   "source": [
    "### 单轮对话\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "056aae13-8ba2-4654-8825-db9e97dd17cc",
   "metadata": {},
   "source": [
    "（如果重新启动内核的话，遵循 `导入库`-> `导入模型` -> `当前代码块` 的顺序执行。）\n",
    "\n",
    "让我们直接设计 `messages`，并应用 `chat_template` 进行对话：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16da6253-3280-4761-aa07-1b6b39336a6e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from transformers import TextStreamer\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# 定义输入\n",
    "prompt = input(\"User: \")\n",
    "\n",
    "# 定义消息列表\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": prompt}\n",
    "]\n",
    "\n",
    "# 使用 tokenizer.apply_chat_template() 生成模型输入\n",
    "input_ids = tokenizer.apply_chat_template(messages, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# 创建 TextStreamer 实例\n",
    "streamer = TextStreamer(\n",
    "    tokenizer, \n",
    "    skip_prompt=True,         # 在输出时跳过输入的提示部分，仅显示生成的文本\n",
    "    skip_special_tokens=True  # 忽略生成过程中的特殊标记（比如 <pad> / <eos> ...）\n",
    ")\n",
    "\n",
    "# 设置生成参数\n",
    "generation_kwargs = {\n",
    "    \"input_ids\": input_ids,  # 模型的输入 ID，注意，这不是 Embedding\n",
    "    \"max_length\": 500,      # 生成的最大 token 数\n",
    "    \"streamer\": streamer,    # 使用 TextStreamer 实现生成过程中逐步输出文本\n",
    "    \"pad_token_id\": tokenizer.eos_token_id  # 默认行为，消除 open-end 警告\n",
    "}\n",
    "\n",
    "# 开始生成文本\n",
    "with torch.no_grad():\n",
    "    model.generate(**generation_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b0bbef9-d9c7-4840-b6f0-8b042a811a2c",
   "metadata": {},
   "source": [
    "### 多轮对话\n",
    "\n",
    "如果重新启动内核的话，遵循 `导入库`-> `导入模型` -> `当前代码块` 的顺序执行。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dea07ea5-e61d-43a9-abee-6aab5322cb15",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TextStreamer\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# 初始化对话历史\n",
    "messages = []\n",
    "\n",
    "# 开始多轮对话\n",
    "while True:\n",
    "    # 获取输入\n",
    "    prompt = input(\"User: \")\n",
    "    \n",
    "    # 退出对话条件（当然，你也可以直接终止代码块）\n",
    "    if prompt.lower() in [\"exit\", \"quit\", \"bye\"]:\n",
    "        print(\"Goodbye!\")\n",
    "        break\n",
    "    \n",
    "    # 将输入添加到对话历史\n",
    "    messages.append({\"role\": \"user\", \"content\": prompt})\n",
    "    \n",
    "    # 使用 tokenizer.apply_chat_template() 生成模型输入\n",
    "    input_ids = tokenizer.apply_chat_template(messages, return_tensors=\"pt\").to(device)\n",
    "    \n",
    "    # 创建 TextStreamer 实例\n",
    "    streamer = TextStreamer(\n",
    "        tokenizer, \n",
    "        skip_prompt=True,         # 在输出时跳过输入的提示部分，仅显示生成的文本\n",
    "        skip_special_tokens=True  # 忽略生成过程中的特殊标记（比如 <pad> / <eos> ...）\n",
    "    )\n",
    "    \n",
    "    # 设置生成参数\n",
    "    generation_kwargs = {\n",
    "        \"input_ids\": input_ids,                  # 模型的输入 ID\n",
    "        \"max_length\": input_ids.shape[1] + 500,  # 生成的最大 token 数，input_ids.shape[1] 即输入对应的 tokens 数量\n",
    "        \"streamer\": streamer,                    # 使用 TextStreamer 实现生成过程中逐步输出文本\n",
    "        \"pad_token_id\": tokenizer.eos_token_id   # 默认行为，消除警告\n",
    "    }\n",
    "    \n",
    "    # 开始生成回复\n",
    "    with torch.no_grad():\n",
    "        output_ids = model.generate(**generation_kwargs)\n",
    "    \n",
    "    # 获取生成的回复文本\n",
    "    assistant_reply = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)\n",
    "    \n",
    "    # 将模型的回复添加到对话历史\n",
    "    messages.append({\"role\": \"assistant\", \"content\": assistant_reply})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7a66a9f-776d-4d6e-81d5-77c6a60525c0",
   "metadata": {},
   "source": [
    "注意，这里有一个小坑，你不能简单使用 `output_ids[0]` 来保存回复，因为`output_ids` 中实际上包含了 `input_ids`，打印它们：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c31f7cc7-1c07-4715-bbed-386d06ad77a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenizer.decode(input_ids[0], skip_special_tokens=True))\n",
    "print(tokenizer.decode(output_ids[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d32c70b9-bf14-46ad-acc8-8bd493c31a89",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}