Daobook xgrammar

Daobook xgrammar#

XGrammar [DRC+24] 开源解决方案，旨在实现灵活、便携且快速的结构化生成。该项目的使命是将灵活的无开销结构化生成带到每个角落。

XGrammar 是由陈天奇团队推出的开源软件库，能为⼤型语⾔模型（LLM）提供⾼效、灵活且可移植的结构化数据⽣成能⼒。基于上下⽂⽆关语法（CFG）定义结构，⽀持递归组合以表⽰复杂结构，适合⽣成JSON、SQL等格式数据。XGrammar ⽤字节级下推⾃动机优化解释 CFG，减少每 token 延迟，实现百倍加速，⼏乎⽆额外开销。XGrammar 集成多种系统优化，如⾃适应 token 掩码缓存、上下⽂扩展等，提⾼掩码⽣成速度并减少预处理时间。XGrammar 的 C++ 后端设计易于集成，并⽀持在 LLM 推理中实现零开销的结构化⽣成。

XGrammar 的主要功能

  • ⾼效结构化⽣成：⽀持上下⽂⽆关语法（CFG），⽀持定义和⽣成遵循特定格式（如JSON、SQL）的结构化数据。

  • 灵活性：基于CFG的递归规则，能灵活地表⽰复杂的结构，适应多样的结构化数据需求。

  • 零开销集成：XGrammar与LLM推理引擎共同设计，能在LLM推理中实现零开销的结构化⽣成。

  • 快速执⾏：基于系统优化，显著提⾼结构化⽣成的执⾏速度，相⽐于SOTA⽅法，每token延迟减少多达100倍。

  • 跨平台部署：具有最⼩且可移植的C++后端，能轻松集成到多个环境和框架中。

  • ⾃适应token掩码缓存：在预处理阶段⽣成，加快运⾏时的掩码⽣成。

XGrammar 的技术原理

  • 字节级下推⾃动机（PDA）：⽤字节级PDA解释CFG，⽀持每个字符边缘包含⼀个或多个字节，处理不规则的token边界，⽀持包含sub-UTF8字符的token。

  • 预处理和运⾏时优化：在预处理阶段，⽣成⾃适应token掩码缓存，基于预先计算与上下⽂⽆关的token加快运⾏时的掩码⽣成。

  • 上下⽂⽆关与相关token的区分：区分上下⽂⽆关token和上下⽂相关token，预先计算PDA中每个位置的上下⽂⽆关token的有效性，并将它们存储在⾃适应token掩码缓存中。

  • 语法编译：基于语法编译过程，预先计算掩码中相当⼀部分token，加快掩码⽣成速度。

  • 算法和系统优化：包括上下⽂扩展、持续性执⾏堆栈、下推⾃动机结构优化等，进⼀步提⾼掩码⽣成速度并减少预处理时间。

  • 掩码⽣成与LLM推理重叠：将CPU上的掩码⽣成过程与GPU上的LLM推理过程并⾏化，消除约束解码的开销

XGrammar 的应⽤场景

  • 编程语⾔辅助：⽤于辅助编写和调试代码，⾃动⽣成符合特定编程语⾔规范的代码⽚段，提⾼开发效率。

  • 数据库操作：⽣成符合SQL语法的查询语句，帮助开发者或应⽤程序⾃动构建数据库查询，减少⼿动编写SQL语句的⼯作量。

  • ⾃然语⾔处理（NLP）：⽣成结构化的训练数据，⽤于训练和优化NLP模型，提⾼模型对结构化信息的处理能⼒。

  • Web开发：⾃动⽣成前端代码和API⽂档，确保⽂档与代码的⼀致性，提⾼开发效率和维护性。

  • 配置⽂件和模板：⽣成和填充配置⽂件及模板，如⾃动化⽣成系统配置、填充邮件模板等，提⾼⾃动化⽔平。

from set_env import temp_dir

安装#

pip 安装

pip install xgrammar

当使用 NVIDIA GPU 时，请同时安装这些额外的依赖项以启用 CUDA 支持来应用位掩码：

pip install cuda-python nvidia-cuda-nvrtc-cu12

测试#

下载模型：

git clone https://www.modelscope.cn/LLM-Research/Llama-3.2-1B-Instruct.git {temp_dir}/LLM-Research/Llama-3.2-1B-Instruct

实例化模型、分词器和输入：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

device = "cuda"  # Or "cpu", etc.
model_name = f"{temp_dir}/LLM-Research/Llama-3.2-1B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.float32, device_map=device
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
config = AutoConfig.from_pretrained(model_name)

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Introduce yourself in JSON briefly."},
]
texts = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
model_inputs = tokenizer(texts, return_tensors="pt").to(model.device)

2025-01-06 14:55:49.127058: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1736146549.148811 3636295 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1736146549.155424 3636295 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-01-06 14:55:49.178656: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.

编译语法#

构建 GrammarCompiler 并编译语法。

语法可以是内置的 JSON 语法、JSON 模式字符串或EBNF字符串。EBNF 提供了更高的自定义灵活性。有关规范，请参阅 GBNF 文档。

import xgrammar as xgr
tokenizer_info = xgr.TokenizerInfo.from_huggingface(tokenizer, vocab_size=config.vocab_size)
grammar_compiler = xgr.GrammarCompiler(tokenizer_info)
compiled_grammar = grammar_compiler.compile_builtin_json_grammar()
# Other ways: provide a json schema string
# compiled_grammar = grammar_compiler.compile_json_schema(json_schema_string)
# Or provide an EBNF string
# compiled_grammar = grammar_compiler.compile_grammar(ebnf_string)

使用语法生成#

使用 LogitsProcessor 结合语法进行生成。

xgr_logits_processor = xgr.contrib.hf.LogitsProcessor(compiled_grammar)
generated_ids = model.generate(
    **model_inputs, max_new_tokens=512, logits_processor=[xgr_logits_processor]
)
generated_ids = generated_ids[0][len(model_inputs.input_ids[0]) :]
print(tokenizer.decode(generated_ids, skip_special_tokens=True))

Setting `pad_token_id` to `eos_token_id`:None for open-end generation.

{
  "name": "Assistant",
  "id": "AI-Assistant",
  "description": "A helpful assistant designed to assist with a wide range of tasks and questions",
  "image": "https://example.com/assistant-icon.png",
  "skills": ["Conversational AI", "Text-based interaction", "Language understanding"]
}