在上一篇文章《GLM-4-9B-Chat 模型本地推理部署（一）项目Demo部署》中，我们介绍了如何准备环境以支持 GLM-4-9B-Chat 模型的 Demo 本地部署。现在，我们将探讨如何使用 vLLM 来创建一个命令行界面（CLI），并让用户能够与模型进行交互。

什么是 vLLM？

vLLM 是由伯克利大学 LMSYS 组织开源的 LLM 推理和服务框架，其具备快速、易用的特点。该框架运用 PagedAttention 算法，有效提升了 LLM 服务的吞吐量与内存使用效率。同时，vLLM 支持 FastAPI 前端，还对 OpenAI API 接口进行了扩展。凭借内存分页管理与模块化设计，vLLM 在性能方面远超其他同类框架。

vLLM 文档 https://docs.vllm.ai/en/stable/

快速

• 顶尖水准的服务吞吐量
• 通过 PagedAttention 算法实现对注意力KV的高效管理
• 对传入的请求进行连续批处理（Continuous Batching）
• 利用 CUDA/HIP 图实现模型的快速执行
• 量化：GPTQ、AWQ、INT4、INT8 和 FP8
• 优化的 CUDA 内核，包括对 FlashAttention 和 FlashInfer 的集成
• 推测解码（Speculative Decoding）
• 分块预填充（Chunked Prefill）

易用

• 无缝集成 HuggingFace 上的热门模型
• 支持多种解码算法实现高吞吐量服务，如并行采样（Parallel sampling）、集束搜索（Beam Search）
• 支持在分布式推理中使用张量并行（Tensor Parallelism）和流水线并行（Pipeline Parallelism）
• 流式输出
• 与 OpenAI 兼容的 API 服务
• 支持 NVIDIA GPU、AMD CPU/GPU、Intel CPU/GPU、PowerPC CPU、TPU 以及 AWS Neuron 等硬件
• 前缀缓存（Prefix Caching）支持
• 多 LoRA 微调支持

在这篇文章中，我们将使用 Python 脚本来实现一个 CLI 应用程序，该应用程序允许用户通过简单的命令行界面与 GLM-4-9B-Chat 模型进行交互。

首先，请确保已经安装 transformers 和 vllm。如果没有安装，可以通过 pip 安装它们：

pip install transformers vllm

https://tongyi.aliyun.com/?sessionId=f45e199b0da7433994fed3fd4e4d3d9d

创建 CLI 应用

我们的目标是创建一个简单的 CLI 应用，用户可以输入问题，然后从 GLM-4-9B-Chat 模型获取答案。

加载模型和分词器

首先，我们需要加载模型和分词器。我们将使用 vllm 的 AsyncEngineArgs 来配置模型加载选项，并使用 transformers 的 AutoTokenizer 来加载分词器。

import time
import asyncio
from transformers import AutoTokenizer
from vllm import SamplingParams, AsyncEngineArgs, AsyncLLMEngine
from typing import List, Dict

MODEL_PATH = 'THUDM/glm-4-9b-chat'

def load_model_and_tokenizer(model_dir: str):
    # ... [省略部分代码，具体参考官方github]
    engine = AsyncLLMEngine.from_engine_args(engine_args)
    return engine, tokenizer

engine, tokenizer = load_model_and_tokenizer(MODEL_PATH)

定义异步生成函数

接着，定义一个异步生成函数 vllm_gen，它会根据用户的输入和一些采样参数生成文本输出。

async def vllm_gen(messages: List[Dict[str, str]], top_p: float, temperature: float, max_dec_len: int):
    # ... [省略部分代码]
    async for output in engine.generate(inputs=inputs, sampling_params=sampling_params, request_id=f"{time.time()}"):
        yield output.outputs[0].text

实现聊天循环

最后，我们需要一个主循环来处理用户输入和模型输出。这个循环将持续运行，直到用户决定退出。

async def chat():
    # ... [省略部分代码]
    while True:
        # ... [省略部分代码]
        async for output in vllm_gen(messages, top_p, temperature, max_length):
            print(output[current_length:], end="", flush=True)
            current_length = len(output)

if __name__ == "__main__":
    asyncio.run(chat())