使用LM Format Enforcer确保AI输出格式正确

最新推荐文章于 2024-10-04 22:36:54 发布

afTFODguAKBF

最新推荐文章于 2024-10-04 22:36:54 发布

阅读量126

点赞数 3

文章标签：人工智能深度学习 pytorch python

本文链接：https://blog.csdn.net/afTFODguAKBF/article/details/142683277

版权

## 引言

在使用语言模型（如LLM）生成结构化数据时，确保输出格式符合预期是至关重要的。然而，模型往往会产生不符合规定格式的输出。本文将介绍一个强大的工具——`LM Format Enforcer`，它能帮助开发者确保语言模型生成的输出符合指定格式。我们将探讨其工作原理、具体使用方法以及潜在的挑战和解决方案。

## 主要内容

### LM Format Enforcer的工作原理

`LM Format Enforcer`结合了字符级解析器与分词器前缀树，只允许那些包含潜在有效格式序列的token输出。这种方法能有效减少生成不合规格式的风险。同时，它支持批量生成，使得处理多条输入变得更高效。

### 设置语言模型

首先，我们需要设置一个LLama2模型，并初始化所需的输出格式。请注意，访问Llama2模型需要获得许可。

```python
import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-2-7b-chat-hf"
device = "cuda"

if torch.cuda.is_available():
    config = AutoConfig.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, config=config, torch_dtype=torch.float16, load_in_8bit=True, device_map="auto")
else:
    raise Exception("GPU not available")

tokenizer = AutoTokenizer.from_pretrained(model_id)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

使用LM Format Enforcer

接下来，我们使用LM Format Enforcer来确保输出符合指定的JSON Schema格式。

from langchain_experimental.llms import LMFormatEnforcer

lm_format_enforcer = LMFormatEnforcer(json_schema=PlayerInformation.schema(), pipeline=hf_model)
results = lm_format_enforcer.predict(get_prompt("Michael Jordan"))
print(results)

批量处理

LM Format Enforcer还支持批量处理，这对于处理大量请求非常有用。

prompts = [get_prompt(name) for name in ["Michael Jordan", "Kareem Abdul Jabbar", "Tim Duncan"]]
results = lm_format_enforcer.generate(prompts)
for generation in results.generations:
    print(generation[0].text)

代码示例

以下是一个完整的使用示例，包括模型初始化和格式强制执行：

from langchain_experimental.pydantic_v1 import BaseModel

class PlayerInformation(BaseModel):
    first_name: str
    last_name: str
    num_seasons_in_nba: int
    year_of_birth: int

# 使用API代理服务提高访问稳定性
model_id = "meta-llama/Llama-2-7b-chat-hf"