使用LM Format Enforcer优化语言模型输出格式

最新推荐文章于 2024-09-26 11:55:47 发布

llzwxh888

最新推荐文章于 2024-09-26 11:55:47 发布

阅读量728

点赞数 17

文章标签：语言模型 python 深度学习

本文链接：https://blog.csdn.net/ppoojjj/article/details/141731670

版权

使用LM Format Enforcer优化语言模型输出格式

引言

在使用大型语言模型(LLM)时,我们经常需要模型输出特定格式的内容,如JSON或符合特定模式的文本。然而,LLM的输出并不总是能够精确地遵循我们期望的格式。这就是LM Format Enforcer库发挥作用的地方。本文将介绍如何使用LM Format Enforcer来强制LLM输出符合预定义格式的内容,从而提高输出质量和可用性。

LM Format Enforcer简介

LM Format Enforcer是一个用于强制语言模型输出格式的库。它通过结合字符级解析器和分词器前缀树,只允许包含可能导致有效格式的字符序列的令牌。这个库支持批量生成,但需要注意的是,它仍处于实验阶段。

安装和设置

首先,让我们安装必要的库:

# 使用API代理服务提高访问稳定性
!pip install --upgrade --quiet lm-format-enforcer langchain-huggingface > /dev/null

接下来,我们需要设置模型和定义输出格式:

import logging
from langchain_experimental.pydantic_v1 import BaseModel
import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

logging.basicConfig(level=logging.ERROR)

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

model_id = "meta-llama/Llama-2-7b-chat-hf"
device = "cuda" if torch.cuda.is_available() else "cpu"

# 加载模型和分词器
config = AutoConfig.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, config=config, torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_id)

使用LM Format Enforcer

现在,让我们看看如何使用LM Format Enforcer来强制模型输出符合我们定义的JSON格式:

from langchain_experimental.llms import LMFormatEnforcer
from transformers import pipeline

# 创建HuggingFace pipeline
hf_model = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=200)

# 创建LMFormatEnforcer实例
lm_format_enforcer = LMFormatEnforcer(json_schema=PlayerInformation.schema(), pipeline=hf_model)

# 定义提示词
prompt = """Please give me information about Michael Jordan. You must respond using JSON format, according to the following schema:

{schema}
"""

# 获取结果
results = lm_format_enforcer.predict(prompt.format(schema=PlayerInformation.schema_json()))
print(results)

输出将是一个符合PlayerInformation模式的JSON对象,包含Michael Jordan的信息。

批量处理

LM Format Enforcer还支持批量处理,这在需要处理多个提示词时非常有用:

prompts = [
    prompt.format(schema=PlayerInformation.schema_json()) 
    for _ in ["Michael Jordan", "Kareem Abdul Jabbar", "Tim Duncan"]
]
results = lm_format_enforcer.generate(prompts)
for generation in results.generations:
    print(generation[0].text)

使用正则表达式

除了JSON格式,LM Format Enforcer还支持使用正则表达式来过滤输出:

question_prompt = "When was Michael Jordan Born? Please answer in mm/dd/yyyy format."
date_regex = r"(0?[1-9]|1[0-2])\/(0?[1-9]|1\d|2\d|3[01])\/(19|20)\d{2}"
answer_regex = " In mm/dd/yyyy format, Michael Jordan was born in " + date_regex

lm_format_enforcer = LMFormatEnforcer(regex=answer_regex, pipeline=hf_model)
print(lm_format_enforcer.predict(question_prompt))