Super JSON Mode:加速结构化输出生成的框架
Super JSON Mode 是一个 Python 框架,它通过将目标架构分解为原子组件并并行执行生成,以高效创建来自 LLM 的结构化输出。
它通过 OpenAI 的旧版 completions API 支持最先进的 LLM,同时也支持开源 LLM,例如通过 Hugging Face Transformers 和 vLLM。将来会支持更多的 LLM!
与依赖提示和 HF Transformers 的简单 JSON 生成管道相比,我们发现 Super JSON Mode 生成输出速度最多可快 10 倍。它还更具确定性,并且比简单生成更不容易出现解析问题。
安装非常简单:pip install super-json-mode
它是如何工作的?
结构化输出格式(如 JSON 或 YAML)具有固有的并行或层次结构。
考虑以下非结构化段落(由 GPT-4 生成):
欢迎来到 123 Azure Lane,这是一处令人惊叹的旧金山住宅,拥有极好的现代设计,目前市场价格为 2,500,000 美元。这座物业占地 3,000 平方英尺,结合了精致与舒适,创造出一种真正独特的生活体验。
这是一个小家庭或专业人士的理想家园,我们的独家住所配备了五间宽敞的卧室,每间都充满了温暖和现代优雅。这些卧室经过精心设计,允许充足的自然光和宽敞的存储空间。住宅配有三间优雅设计的全浴室,为住户保证了便利和隐私。
宏伟的入口带您进入宽敞的客厅,为聚会或在火炉旁的安静夜晚提供了绝佳氛围。厨师厨房配备了最先进的电器、定制的橱柜和美丽的花岗岩台面,是任何喜欢烹饪的人的梦想。
如果我们想使用 LLM 提取 地址
、平方英尺
、卧室数量
、浴室数量
和 价格
,我们可以让模型根据描述填写一个架构。
一个潜在的架构(例如由 Pydantic 对象生成的架构)可能是这样的:
{
"address": {
"type": "string"
},
"price": {
"type": "number"
},
"square_feet": {
"type": "integer"
},
"num_beds": {
"type": "integer"
},
"num_baths": {
"type": "integer"
}
}
一个有效的输出可能是这样的:
{
"address": "123 Azure Lane",
"price": 2500000,
"square_feet": 3000,
"num_beds": 5,
"num_baths": 3
}
显而易见的方法是将架构嵌入提示中并要求模型填写。目前大多数团队提取 LLM 结构化输出的方式就是这样。
然而,这种方法效率低下,原因有三。
-
注意到这些键是相互独立的。Super JSON Mode 利用提示并行性,将架构中的每个键值对视为一个独立查询。例如,我们可以在还没有生成
地址
的情况下提取num_baths
! -
请求模型从头生成 JSON 会在可预测的语法(例如括号和键名)上不必要地消耗 tokens(因此也消耗时间),而这些语法已经在输出中预期到。这是生成中的一个强烈先验,我们应该能够利用它来提高延迟。
-
LLM 的并行性非常明显,并且批量运行查询比按顺序运行要快得多。因此,我们可以将架构分成多个查询。LLM 将为每个独立的键并行填写架构,并在一次传递中发出更少的 tokens,从而允许更快的推理时间。
安装
通过 PyPI
运行以下命令:
pip install super-json-mode
手动
- 创建一个 conda 环境
conda create --name superjsonmode python=3.10 -y
conda activate superjsonmode
- 克隆并安装依赖项
git clone https://github.com/varunshenoy/super-json-mode
cd superjsonmode
pip install -r requirements.txt
示例
我们已经尽力使 Super JSON Mode 的使用非常简单。请参见 examples
文件夹,以获取更多使用示例和 vLLM
的使用。
使用 OpenAI 和 gpt-3-instruct-turbo
:
from superjsonmode.integrations.openai import StructuredOpenAIModel
from pydantic import BaseModel
import time
model = StructuredOpenAIModel()
class Character(BaseModel):
name: str
genre: str
age: int
race: str
occupation: str
best_friend: str
home_planet: str
prompt_template = """{prompt}
Please fill in the following information about this character for this key. Keep it succinct. It should be a {type}.
{key}: """
prompt = """Luke Skywalker is a famous character."""
start = time.time()
output = model.generate(
prompt,
extraction_prompt_template=prompt_template,
schema=Character,
batch_size=7,
stop=["\n\n"],
temperature=0,
)
print(f"Total time: {time.time() - start}")
# Total Time: 0.409s
print(output)
# {
# "name": "Luke Skywalker",
# "genre": "Science fiction",
# "age": "23",
# "race": "Human",
# "occupation": "Jedi Knight",
# "best_friend": "Han Solo",
# "home_planet": "Tatooine",
# }
使用 Mistral 7B 和 HuggingFace Transformers:
from transformers import AutoTokenizer, AutoModelForCausalLM
from superjsonmode.integrations.transformers import StructuredOutputForModel
from pydantic import BaseModel
device = "cuda"
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2").to(device)
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
# Create a structured output object
structured_model = StructuredOutputForModel(model, tokenizer)
passage = """..."""
class QuarterlyReport(BaseModel):
company: str
stock_ticker: str
date: str
reported_revenue: str
dividend: str
prompt_template = """[INST]{prompt}
Based on this excerpt, extract the correct value for "{key}". Keep it succinct. It should have a type of `{type}`.[/INST]
{key}: """
output = structured_model.generate(passage,
extraction_prompt_template=prompt_template,
schema=QuarterlyReport,
batch_size=6)
print(json.dumps(output, indent=2))
# {
# "company": "NVIDIA",
# "stock_ticker": "NVDA",
# "date": "2023-10",
# "reported_revenue": "18.12 billion dollars",
# "dividend": "0.04"
# }
路线图
有很多功能可以使 Super JSON Mode 变得更好。以下是一些想法。
-
定性输出分析:我们进行了性能基准测试,但我们应该提出一个更严格的方法来判断 Super JSON Mode 的定性输出。
-
结构化采样:理想情况下,我们应该屏蔽 LLM 的 logits 以强制执行类型约束,类似于 JSONFormer。已经有一些包在做这件事,或者这些包应该集成我们的并行 JSON 生成管道,或者我们应该将其构建到 Super JSON Mode 中。
-
依赖图支持:Super JSON Mode 有一个非常明显的失败案例:当一个键依赖于另一个键时。考虑一个具有两个键
thought
和response
的 JSON 数据块。对于具有大型语言模型的思维链,这种输出是很常见的,并且很明显response
依赖于thought
。我们应该能够传入一个依赖关系图,并按父架构项目完成和传递给子架构项目的方式批量提示。 -
本地模型支持:Super JSON Mode 最适合批量大小通常为 1 的本地情况。您可以利用批处理来减少延迟,类似于推测解码。如果可能的话,我想使用 Ollama 实现这个功能。
-
TRT-LLM 支持:vLLM 很棒且易于使用,但理想情况下,我们应该集成一个更高性能的框架,例如 TRT-LLM。
引用
如果您发现此库对您的工作有用,请引用此 repo,我们将不胜感激:
@misc{ShenoyDerhacobian2024,
author = {Shenoy, Varun 和 Derhacobian, Alex},
title = {Super JSON Mode: A Framework for Accelerated Structured Output Generation},
year = {2024},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/varunshenoy/super-json-mode}}
}
这个项目是为 CS 229:机器学习系统构建的。非常感谢教学团队和助教在整个项目中的指导。