{ "cells": [ { "cell_type": "markdown", "id": "6e0dc7da-a5e8-45f5-8f1c-7931368dfd4a", "metadata": {}, "source": [ "# 简介\n", "\n", "> 指导文章：[09. 深入理解 Beam Search：原理, 示例与代码实现](https://github.com/Hoper-J/LLM-Guide-and-Demos-zh_CN/blob/master/09.%20深入理解%20Beam%20Search：原理%2C%20示例与代码实现.md#具体是怎么处理-eos-的)\n", "\n", "在线链接：[Kaggle](https://www.kaggle.com/code/aidemos/07-beam-search) | [Colab](https://colab.research.google.com/drive/1majmMl7z3d7wGBCP4Zy-7g0zfMdZ_G1A?usp=sharing)\n", "\n", "# 示例：过程演示\n", "\n", "![过程演示](../Guide/assets/%E5%9B%BE%E7%89%87%201-6584229.png)\n" ] }, { "cell_type": "code", "execution_count": 1, "id": "85113a9e-ab78-40af-ae24-6397c5d5a471", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "第 1 步:\n", "扩展序列: ，当前得分为 0.0\n", " 候选序列: A，得分为 -0.916290731874155\n", " 候选序列: B，得分为 -1.2039728043259361\n", " 候选序列: C，得分为 -1.6094379124341003\n", " 候选序列: ，得分为 -2.3025850929940455\n", "\n", "选择的 2 个顶束序列:\n", " A，得分为 -0.916290731874155\n", " B，得分为 -1.2039728043259361\n", "\n", "第 2 步:\n", "扩展序列: A，当前得分为 -0.916290731874155\n", " 候选序列: AA，得分为 -2.120263536200091\n", " 候选序列: AB，得分为 -3.2188758248682006\n", " 候选序列: AC，得分为 -1.83258146374831\n", " 候选序列: A，得分为 -2.525728644308255\n", "扩展序列: B，当前得分为 -1.2039728043259361\n", " 候选序列: BA，得分为 -3.506557897319982\n", " 候选序列: BB，得分为 -3.506557897319982\n", " 候选序列: BC，得分为 -2.4079456086518722\n", " 候选序列: B，得分为 -1.8971199848858813\n", "\n", "选择的 2 个顶束序列:\n", " AC，得分为 -1.83258146374831\n", " B，得分为 -1.8971199848858813\n", "\n", "第 3 步:\n", "扩展序列: AC，当前得分为 -1.83258146374831\n", " 候选序列: ACA，得分为 -4.135166556742355\n", " 候选序列: ACB，得分为 -3.4420193761824103\n", " 候选序列: ACC，得分为 -2.525728644308255\n", " 候选序列: AC，得分为 -3.4420193761824103\n", "已完成序列: B，得分为 -1.8971199848858813\n", "\n", "选择的 2 个顶束序列:\n", " ACC，得分为 -2.525728644308255\n", " ACB，得分为 -3.4420193761824103\n", "\n", "第 4 步:\n", "扩展序列: ACC，当前得分为 -2.525728644308255\n", " 候选序列: ACC，得分为 -2.525728644308255\n", "扩展序列: ACB，当前得分为 -3.4420193761824103\n", " 候选序列: ACB，得分为 -3.4420193761824103\n", "\n", "选择的 2 个顶束序列:\n", " ACC，得分为 -2.525728644308255\n", " ACB，得分为 -3.4420193761824103\n", "\n", "第 5 步:\n", "已完成序列: ACC，得分为 -2.525728644308255\n", "已完成序列: ACB，得分为 -3.4420193761824103\n", "\n", "选择的 2 个顶束序列:\n", "\n", "已完成的所有序列:\n", " B，得分为 -1.8971199848858813\n", " ACC，得分为 -2.525728644308255\n", " ACB，得分为 -3.4420193761824103\n", "\n", "最佳序列: B\n" ] } ], "source": [ "import math\n", "\n", "def beam_search(initial_sequence, beam_width, max_length, vocab, get_next_probs):\n", " beam = [(initial_sequence, 0.0)] # (sequence, log_prob)\n", " completed = []\n", "\n", " for step in range(max_length):\n", " print(f\"\\n第 {step + 1} 步:\")\n", " all_candidates = []\n", " for seq, score in beam:\n", " if seq.endswith(''):\n", " completed.append((seq, score))\n", " print(f\"已完成序列: {seq}，得分为 {score}\")\n", " continue\n", " next_probs = get_next_probs(seq)\n", " print(f\"扩展序列: {seq}，当前得分为 {score}\")\n", " for token, prob in next_probs.items():\n", " new_seq = seq + token\n", " new_score = score + math.log(prob)\n", " all_candidates.append((new_seq, new_score))\n", " print(f\" 候选序列: {new_seq}，得分为 {new_score}\")\n", " \n", " # 对所有候选序列按得分降序排列，选择得分最高的 beam_width 个序列\n", " all_candidates.sort(key=lambda x: x[1], reverse=True)\n", " beam = all_candidates[:beam_width]\n", "\n", " # 打印选出的顶束序列\n", " print(f\"\\n选择的 {beam_width} 个顶束序列:\")\n", " for seq, score in beam:\n", " print(f\" {seq}，得分为 {score}\")\n", " \n", " # 如果没有更多序列可以扩展，则退出循环\n", " if not beam:\n", " break\n", "\n", " # 将当前 beam 中剩下的序列加入完成序列中\n", " completed += beam\n", "\n", " # 对完成的序列按得分降序排列，选择得分最高的序列\n", " completed.sort(key=lambda x: x[1], reverse=True)\n", " \n", " print(\"\\n已完成的所有序列:\")\n", " for seq, score in completed:\n", " print(f\" {seq}，得分为 {score}\")\n", " \n", " return completed[0][0]\n", "\n", "# 我们之前示例中设置的概率\n", "def get_next_probs(seq):\n", " probs = {\n", " \"\": {\"A\": 0.4, \"B\": 0.3, \"C\": 0.2, \"\": 0.1},\n", " \"A\": {\"A\": 0.3, \"B\": 0.1, \"C\": 0.4, \"\": 0.2},\n", " \"B\": {\"A\": 0.1, \"B\": 0.1, \"C\": 0.3, \"\": 0.5},\n", " \"AC\": {\"A\": 0.1, \"B\": 0.2, \"C\": 0.5, \"\": 0.2},\n", " }\n", " return probs.get(seq, {\"\": 1.0})\n", "\n", "initial_sequence = \"\"\n", "beam_width = 2\n", "max_length = 5\n", "vocab = {\"A\", \"B\", \"C\", \"\"}\n", "\n", "best_sequence = beam_search(initial_sequence, beam_width, max_length, vocab, get_next_probs)\n", "print(\"\\n最佳序列:\", best_sequence)" ] }, { "cell_type": "markdown", "id": "764213b9-cb71-44a6-9eb2-39774b692153", "metadata": {}, "source": [ "# 示例：使用 Hugging Face Transformers 库" ] }, { "cell_type": "code", "execution_count": 2, "id": "b8900576-5ffb-4577-a33b-993226cbd7b9", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Looking in indexes: http://mirrors.aliyun.com/pypi/simple\n", "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.57.6)\n", "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.25.2)\n", "Requirement already satisfied: huggingface-hub<1.0,>=0.34.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.36.2)\n", "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.26.4)\n", "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (25.0)\n", "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0.3)\n", "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2026.4.4)\n", "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.32.5)\n", "Requirement already satisfied: tokenizers<=0.23.0,>=0.22.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.22.2)\n", "Requirement already satisfied: safetensors>=0.4.3 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.7.0)\n", "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.67.1)\n", "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.34.0->transformers) (2026.2.0)\n", "Requirement already satisfied: hf-xet<2.0.0,>=1.1.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.34.0->transformers) (1.4.3)\n", "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.34.0->transformers) (4.15.0)\n", "Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4.4)\n", "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.11)\n", "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.5.0)\n", "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2025.11.12)\n", "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\u001b[0m\u001b[33m\n", "\u001b[0m\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.3\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m26.1.1\u001b[0m\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpython3 -m pip install --upgrade pip\u001b[0m\n", "Note: you may need to restart the kernel to use updated packages.\n" ] } ], "source": [ "%pip install transformers\n", "#!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118" ] }, { "cell_type": "code", "execution_count": 3, "id": "2563024a-52f0-4469-8835-0f2cc6c66e75", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "生成的文本：\n", "Hello GPT.\n", "\n", "This article was originally published on The Conversation. Read the original article.\n" ] } ], "source": [ "import os\n", "# 设置模型下载镜像（注意，需要在导入 transformers 等模块前进行设置才能起效）\n", "os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'\n", "\n", "import warnings\n", "from transformers import AutoTokenizer, AutoModelForCausalLM\n", "import torch\n", "\n", "# 忽略 FutureWarning 警告\n", "warnings.filterwarnings(\"ignore\", category=FutureWarning)\n", "\n", "# 指定模型名称\n", "model_name = \"distilgpt2\"\n", "\n", "# 加载分词器和模型\n", "tokenizer = AutoTokenizer.from_pretrained(model_name)\n", "model = AutoModelForCausalLM.from_pretrained(model_name)\n", "\n", "# 移动模型到设备\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "model.to(device)\n", "\n", "# 设置模型为评估模式\n", "model.eval()\n", "\n", "# 输入文本\n", "input_text = \"Hello GPT\"\n", "\n", "# 编码输入文本，并生成 attention mask\n", "inputs = tokenizer.encode(input_text, return_tensors=\"pt\").to(device)\n", "attention_mask = torch.ones_like(inputs).to(device)\n", "\n", "# 生成文本，使用 Beam Search\n", "beam_width = 5\n", "with torch.no_grad():\n", " outputs = model.generate(\n", " inputs,\n", " attention_mask=attention_mask,\n", " max_length=50,\n", " num_beams=beam_width, # 你可以看到 beam_width 对应的参数名为 num_beams\n", " no_repeat_ngram_size=2,\n", " early_stopping=True, # 当有 num_beams 个候选序列生成时提前停止（除了 False 外还可以设置为 \"never\"）\n", " pad_token_id=tokenizer.eos_token_id\n", " )\n", "\n", "# 解码生成的文本\n", "generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)\n", "print(\"生成的文本：\")\n", "print(generated_text)\n" ] }, { "cell_type": "markdown", "id": "16486045-9b26-426e-8d3e-a0c3b42dd54b", "metadata": {}, "source": [ "## 对比不同束宽的输出" ] }, { "cell_type": "code", "execution_count": 4, "id": "62cee1a2-d19a-43ea-a9d6-b2dc5e5b58ee", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "The following generation flags are not valid and may be ignored: ['early_stopping']. Set `TRANSFORMERS_VERBOSITY=info` for more details.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "束宽 1 的生成结果：\n", "Hello GPT is a free and open source software project that aims to provide a platform for developers to build and use GPGP-based GPSP based GPCs. GPP is an open-source software development platform that is designed to\n", "--------------------------------------------------\n", "束宽 3 的生成结果：\n", "Hello GPT.\n", "\n", "This article is part of a series of articles on the topic, and will be updated as more information becomes available.\n", "--------------------------------------------------\n", "束宽 5 的生成结果：\n", "Hello GPT.\n", "\n", "This article was originally published on The Conversation. Read the original article.\n", "--------------------------------------------------\n" ] } ], "source": [ "# 输入文本\n", "input_text = \"Hello GPT\"\n", "\n", "# 编码输入文本，并生成 attention mask\n", "inputs = tokenizer.encode(input_text, return_tensors=\"pt\").to(device)\n", "attention_mask = torch.ones_like(inputs).to(device)\n", "\n", "# 设置束宽不同的生成策略\n", "beam_widths = [1, 3, 5] # 使用不同的束宽\n", "\n", "# 生成并打印结果\n", "for beam_width in beam_widths:\n", " with torch.no_grad():\n", " outputs = model.generate(\n", " inputs,\n", " attention_mask=attention_mask,\n", " max_length=50,\n", " num_beams=beam_width,\n", " no_repeat_ngram_size=2,\n", " early_stopping=True,\n", " pad_token_id=tokenizer.eos_token_id\n", " )\n", " generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)\n", " print(f\"束宽 {beam_width} 的生成结果：\")\n", " print(generated_text)\n", " print('-' * 50)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "65252bba-b4d6-4c38-9504-dd588a3659c5", "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.10.12" } }, "nbformat": 4, "nbformat_minor": 5 }