说明

我之前写过一个博客《使用vllm部署大语言模型》介绍使用vllm将大模型服务化。实际上vllm也支持离线推理。

我写了一个离线推理的例子，并且可以让用户多次输入。

代码

from vllm import LLM, SamplingParams
import readline  # 导入readline模块

# 加载模型
llm = LLM(
    model="~/my_models/llama-7b",
    tensor_parallel_size=1,
    gpu_memory_utilization=0.9, trust_remote_code=True
)

# 定义生成参数
sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.9,
    max_tokens=256
)
chat_template="{%- if messages[0]['role'] == 'system' -%}    {%- set system_message = messages[0]['content'] -%}    {%- set messages = messages[1:] -%}{%- else -%}    {% set system_message = '' -%}{%- endif -%}{{ bos_token + system_message }}{%- for message in messages -%}    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}        {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}    {%- endif -%}    {%- if message['role'] == 'user' -%}        {{ 'USER: ' + message['content'] + '\n' }}    {%- elif message['role'] == 'assistant' -%}        {{ 'ASSISTANT: ' + message['content'] + eos_token + '\n' }}    {%- endif -%}{%- endfor -%}{%- if add_generation_prompt -%}    {{ 'ASSISTANT:' }} {% endif %}"

# 在循环开始前初始化readline
readline.parse_and_bind("tab: complete")  # 启用基本编辑功能
readline.set_completer(None)  # 禁用自动补全

while True:
    try:
        prompt = input("请输入您的问题（输入Q退出）：")
        if prompt.upper() == "Q":
            print("程序退出")
            break
            
        # 执行推理generate
        prompts = [prompt]
        outputs = llm.generate(prompts, sampling_params)

        #或者用chat接口
        #prompts = [{"role": "user", "content":prompt}]
        #outputs = llm.chat(prompts, sampling_params, chat_template=chat_template)

        # 输出结果
        for output in outputs:
            print(output.outputs[0].text)
        print("\n" + "="*50 + "\n")  # 添加分隔线，使每次对话更清晰
    except KeyboardInterrupt:
        print("\n输入中断，请重新输入")
        continue

可以使用llm.generate()或者llm.chat()进行推理。使用后者的时候必须传入chat_template参数。

参考资料

使用vllm部署大语言模型