example.py

import os
from nanovllm import LLM, SamplingParams
from transformers import AutoTokenizer


def main():
    """最小离线推理示例：加载模型和 tokenizer，构造 prompt，调用 nano-vLLM 生成文本。"""
    # 模型目录。这里假设你已经按 README 把 Qwen3-0.6B 下载到了 ~/huggingface/Qwen3-0.6B/。
    path = os.path.expanduser("~/huggingface/Qwen3-0.6B/")

    # tokenizer 负责把文本 prompt 转成 token id，也负责应用 Qwen3 的 chat template。
    # nano-vLLM 的 LLM 内部也会加载 tokenizer；这里额外加载一次是为了手动构造聊天格式 prompt。
    tokenizer = AutoTokenizer.from_pretrained(path)

    # 创建推理引擎。
    # enforce_eager=True 表示不启用 CUDA graph，便于理解和调试；tensor_parallel_size=1 表示只用 1 张 GPU。
    ##可以尝试tensor_parallel_size=2，加速
    llm = LLM(path, enforce_eager=True, tensor_parallel_size=1)

    # 采样参数：
    # temperature=0.6 控制随机性，越低越接近确定性输出；
    # max_tokens=256 表示每个 prompt 最多生成 256 个新 token。
    sampling_params = SamplingParams(temperature=0.6, max_tokens=256)

    # 原始用户输入，还不是模型最终看到的完整 chat prompt。
    prompts = [
        "introduce yourself",
        "list all prime numbers within 100",
    ]

    # Qwen3 是聊天模型，需要把用户消息包装成模型训练时使用的对话模板。
    # tokenize=False 表示这里先得到字符串；LLM.generate 之后会在内部 tokenize。
    # add_generation_prompt=True 会追加 assistant 开始生成所需的特殊提示。
    prompts = [
        tokenizer.apply_chat_template(
            [{"role": "user", "content": prompt}],
            tokenize=False,
            add_generation_prompt=True,
        )
        for prompt in prompts
    ]

    # 批量生成。nano-vLLM 会把多个 prompt 放进调度器，动态执行 prefill 和 decode。
    #para上面定义了
    outputs = llm.generate(prompts, sampling_params)

    # outputs 与 prompts 一一对应，每个 output 是 {"text": 生成文本, "token_ids": 生成 token id 列表}。
    for prompt, output in zip(prompts, outputs):
        print("\n")
        print(f"Prompt: {prompt!r}")
        print(f"Completion: {output['text']!r}")


# Python 脚本入口：只有直接运行 example.py 时才调用 main；被其他文件 import 时不会自动执行。
if __name__ == "__main__":
    main()