Demos/16b. 使用 Llama-cpp-python 加载量化后的 LLM 大模型（GGUF）.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "186f8d20-3165-4ea1-96d3-6dfa9b8859b4",
   "metadata": {},
   "source": [
    "# b. 使用 Llama-cpp-python 加载量化后的 LLM 大模型（GGUF）\n",
    "\n",
    "> 引导文章：[19b. 从加载到对话：使用 Llama-cpp-python 本地运行量化 LLM 大模型（GGUF）](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/19b.%20从加载到对话：使用%20Llama-cpp-python%20本地运行量化%20LLM%20大模型（GGUF）.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",
    "Transformers 关于 GPTQ & AWQ 加载的相关链接：[文章 19a](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) | [代码文件 16a](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Demos/16a.%20使用%20Transformers%20加载量化后的%20LLM%20大模型（GPTQ%20%26%20AWQ）.ipynb)\n",
    "\n",
    "在线链接：[Kaggle - b](https://www.kaggle.com/code/aidemos/16b-llama-cpp-python-llm-gguf) | [Colab - b](https://colab.research.google.com/drive/1AhgC0qDaqWBXAI9eSbwTStGgvgFfLOpf?usp=sharing)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a3d1707-d7e2-4d0b-8fa6-807892c52bc2",
   "metadata": {},
   "source": [
    "## Llama-cpp-python "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b34d8b2a",
   "metadata": {},
   "source": [
    "### 环境配置\n",
    "\n",
    "为了确保之后的 \"offload\" 正常工作，需要进行一些额外的工作。\n",
    "\n",
    "首先，找到 CUDA 的安装路径："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d83c9b96",
   "metadata": {},
   "outputs": [],
   "source": [
    "!find /usr/local -name \"cuda\" -exec readlink -f {} \\;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48546e57",
   "metadata": {},
   "source": [
    "复制对应（最短）的路径，修改 `cuda_home` 环境变量："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55d687df",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import subprocess\n",
    "\n",
    "# 设置环境变量\n",
    "cuda_home = \"/usr/local/cuda-12.1\"  # 请替换为你的路径\n",
    "\n",
    "command = f\"\"\"\n",
    "CMAKE_ARGS=\"-DGGML_CUDA=on \\\n",
    "            -DCUDA_PATH={cuda_home} \\\n",
    "            -DCUDAToolkit_ROOT={cuda_home} \\\n",
    "            -DCUDAToolkit_INCLUDE_DIR={cuda_home}/include \\\n",
    "            -DCUDAToolkit_LIBRARY_DIR={cuda_home}/lib64 \\\n",
    "            -DCMAKE_CUDA_COMPILER={cuda_home}/bin/nvcc\" \\\n",
    "FORCE_CMAKE=1 \\\n",
    "pip install --upgrade --force-reinstall llama-cpp-python --no-cache-dir --verbose\n",
    "\"\"\"\n",
    "\n",
    "subprocess.run(command, shell=True, check=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d352a092",
   "metadata": {},
   "source": [
    "### GGUF"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5047589",
   "metadata": {},
   "source": [
    "#### 安装库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd3af47c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install gguf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "368d751d",
   "metadata": {},
   "source": [
    "#### 导入库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "460326a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_cpp import Llama"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5bd1c6e",
   "metadata": {},
   "source": [
    "下面介绍两种导入模型的方法，实际执行时本地/自动导入二选一。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52283fcc",
   "metadata": {},
   "source": [
    "#### 本地导入模型\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fe8627e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 如果你已经配置过了，可以直接在 Notebook 中执行下面的命令下载。\n",
    "!export HF_ENDPOINT=https://hf-mirror.com\n",
    "!./hfd.sh bartowski/Mistral-7B-Instruct-v0.3-GGUF --include Mistral-7B-Instruct-v0.3-Q4_K_M.gguf --tool aria2c -x 16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05857047",
   "metadata": {},
   "source": [
    "根据模型路径导入模型，注意，文件位于 `<repo_id>` 文件夹下，以当前下载的文件为例："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0b244a9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 指定本地模型的路径\n",
    "model_path = \"./Mistral-7B-Instruct-v0.3-GGUF/Mistral-7B-Instruct-v0.3-Q4_K_M.gguf\"\n",
    "\n",
    "# 加载模型\n",
    "llm = Llama(\n",
    "    model_path=model_path,\n",
    "    # n_gpu_layers=-1,  # 取消注释使用 GPU 加速\n",
    "    # verbose=False,  # 取消注释减少模型方面的信息打印\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94dfa7fc",
   "metadata": {},
   "source": [
    "#### 自动下载并导入\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4d1c736",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 指定仓库的名称和文件名\n",
    "repo_id = \"bartowski/Mistral-7B-Instruct-v0.3-GGUF\"\n",
    "filename = \"Mistral-7B-Instruct-v0.3-Q4_K_M.gguf\"\n",
    "#filename = \"*Q4_K_M.gguf\"  # 使用通配符也是可以的\n",
    "\n",
    "\n",
    "# 下载并加载模型\n",
    "llm = Llama.from_pretrained(\n",
    "    repo_id=repo_id,\n",
    "    filename=filename,\n",
    "    #n_gpu_layers=-1,  # 取消注释使用 GPU 加速\n",
    "    #verbose=False,  # 取消注释减少模型方面的信息打印\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d7f7cf7",
   "metadata": {},
   "source": [
    "#### 推理测试\n",
    "\n",
    "简单使用以下命令进行推理测试："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4e190a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 输入文本\n",
    "input_text = \"Hello, World!\"\n",
    "\n",
    "# 生成输出\n",
    "output = llm(input_text, max_tokens=50)\n",
    "\n",
    "# 打印生成的文本\n",
    "print(output['choices'][0]['text'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "314951dd-5b1e-4dc1-ac87-d920a4ab9e48",
   "metadata": {},
   "source": [
    "### 流式输出\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5355e797-8025-4019-890f-4d5b10bf49d3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "prompt = \"人工智能的未来发展方向是什么？\"\n",
    "\n",
    "output = llm.create_chat_completion(\n",
    "    messages=[{\n",
    "        \"role\": \"user\",\n",
    "        \"content\": prompt\n",
    "    }],\n",
    "    max_tokens=200,\n",
    "    stream=True\n",
    ")\n",
    "\n",
    "for chunk in output:\n",
    "    delta = chunk['choices'][0]['delta']\n",
    "    if 'role' in delta:\n",
    "        print(delta['role'], end=': ', flush=True)\n",
    "    elif 'content' in delta:\n",
    "        print(delta['content'], end='', flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec7f36bc-82f8-47da-84b4-785ef804ad61",
   "metadata": {},
   "source": [
    "查看 output 的构造："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81f847f8-73dd-4ba8-b524-7b4a85e96665",
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import islice\n",
    "\n",
    "prompt = \"人工智能的未来发展方向是什么？\"\n",
    "\n",
    "output = llm.create_chat_completion(\n",
    "    messages=[{\n",
    "        \"role\": \"user\",\n",
    "        \"content\": prompt\n",
    "    }],\n",
    "    max_tokens=200,\n",
    "    stream=True\n",
    ")\n",
    "\n",
    "print(type(output))\n",
    "\n",
    "# 将生成器转换为列表\n",
    "output_list = list(islice(output, 3))\n",
    "\n",
    "# 获取前 3 个条目\n",
    "output_list[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d1806f8-9509-4dbf-9662-c1bf0934fef1",
   "metadata": {},
   "source": [
    "将刚刚对于流式输出的处理抽象为函数便于后续调用："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9145b07-1cd3-43df-8eba-21794ac91fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def handle_stream_output(output):\n",
    "    \"\"\"\n",
    "    处理流式输出，将生成的内容逐步打印出来。\n",
    "    \n",
    "    参数：\n",
    "        output: 生成器对象，来自 create_chat_completion 的流式输出\n",
    "    \"\"\"\n",
    "    for chunk in output:\n",
    "        delta = chunk['choices'][0]['delta']\n",
    "        if 'role' in delta:\n",
    "            print(f\"{delta['role']}: \", end='', flush=True)\n",
    "        elif 'content' in delta:\n",
    "            print(delta['content'], end='', flush=True)\n",
    "            \n",
    "# 使用示例\n",
    "prompt = \"人工智能的未来发展方向是什么？\"\n",
    "\n",
    "output = llm.create_chat_completion(\n",
    "    messages=[{\n",
    "        \"role\": \"user\",\n",
    "        \"content\": prompt\n",
    "    }],\n",
    "    max_tokens=200,\n",
    "    stream=True\n",
    ")\n",
    "\n",
    "handle_stream_output(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "554686f8-2b6f-422a-8701-e8af93f1c2b5",
   "metadata": {},
   "source": [
    "### 多轮对话\n",
    "\n",
    "让我们自定义一个交互的对话类（需要注意到 handle_stream_output() 有所修改）。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5c8d8f6-2656-43af-a53e-b49cca1ec166",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_cpp import Llama\n",
    "\n",
    "def handle_stream_output(output):\n",
    "    \"\"\"\n",
    "    处理流式输出，将生成的内容逐步打印出来，并收集完整的回复。\n",
    "\n",
    "    参数：\n",
    "        output: 生成器对象，来自 create_chat_completion 的流式输出\n",
    "\n",
    "    返回：\n",
    "        response: 完整的回复文本\n",
    "    \"\"\"\n",
    "    response = \"\"\n",
    "    for chunk in output:\n",
    "        delta = chunk['choices'][0]['delta']\n",
    "        if 'role' in delta:\n",
    "            print(f\"{delta['role']}: \", end='', flush=True)\n",
    "        elif 'content' in delta:\n",
    "            content = delta['content']\n",
    "            print(content, end='', flush=True)\n",
    "            response += content\n",
    "    return response\n",
    "\n",
    "class ChatSession:\n",
    "    def __init__(self, llm):\n",
    "        self.llm = llm\n",
    "        self.messages = []\n",
    "\n",
    "    def add_message(self, role, content):\n",
    "        \"\"\"\n",
    "        添加一条消息到会话中。\n",
    "\n",
    "        参数：\n",
    "            role: 消息角色，通常为 'user' 或 'assistant'\n",
    "            content: 消息内容\n",
    "        \"\"\"\n",
    "        self.messages.append({\"role\": role, \"content\": content})\n",
    "\n",
    "    def get_response_stream(self, user_input):\n",
    "        \"\"\"\n",
    "        获取模型对用户输入的响应（流式输出）。\n",
    "\n",
    "        参数：\n",
    "            user_input: 用户输入的文本\n",
    "\n",
    "        返回：\n",
    "            response: 完整的回复文本\n",
    "        \"\"\"\n",
    "        self.add_message(\"user\", user_input)\n",
    "        \n",
    "        try:\n",
    "            output = self.llm.create_chat_completion(\n",
    "                messages=self.messages,\n",
    "                stream=True  # 开启流式输出\n",
    "            )\n",
    "            \n",
    "            response = handle_stream_output(output)  # 同时打印和收集回复\n",
    "            \n",
    "            self.add_message(\"assistant\", response.strip())\n",
    "            # print(len(response),len(response.strip()))\n",
    "            return response.strip()\n",
    "        except Exception as e:\n",
    "            print(f\"\\n发生错误: {e}\")\n",
    "            return \"\"\n",
    "\n",
    "# 初始化模型（假设使用本地路径）\n",
    "model_path = \"./Mistral-7B-Instruct-v0.3-GGUF/Mistral-7B-Instruct-v0.3-Q4_K_M.gguf\"\n",
    "llm = Llama(\n",
    "    model_path=model_path,\n",
    "    n_gpu_layers=-1,  # 根据需要卸载到 GPU\n",
    "    n_ctx=4096,       # 设置上下文窗口大小，尝试注释这行进行多轮对话，看看会发生什么\n",
    "    verbose=False,    # 禁用详细日志输出\n",
    ")\n",
    "\n",
    "# 创建会话实例\n",
    "chat = ChatSession(llm)\n",
    "        \n",
    "# 开始对话\n",
    "while True:\n",
    "    prompt = input(\"User: \")\n",
    "    # 退出对话条件（当然，你也可以直接终止代码块）\n",
    "    if prompt.lower() in [\"exit\", \"quit\", \"bye\"]:\n",
    "        print(\"Goodbye!\")\n",
    "        break\n",
    "    chat.get_response_stream(prompt)\n",
    "    print()  # 换行以便下一次输入，这是因为之前的 print 都设置了 end=''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62026caf-d39c-4e86-80f9-526d49765c47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用 bye 退出上面的对话后打印查看\n",
    "print(chat.messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c516583-336c-411b-afc3-0e53f3b20ddd",
   "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
}