使用LM Format Enforcer确保语言模型输出的格式正确

最新推荐文章于 2024-10-06 17:48:10 发布

qq_37836323

最新推荐文章于 2024-10-06 17:48:10 发布

阅读量304

点赞数 5

文章标签：语言模型深度学习人工智能 python

本文链接：https://blog.csdn.net/qq_29929123/article/details/142373324

版权

引言

在使用语言模型生成文本时，格式正确性往往是一个挑战。为了确保输出符合特定格式，我们可以使用 LM Format Enforcer 库。这篇文章将深入探讨如何利用该库通过过滤token来强制语言模型的输出格式。

主要内容

LM Format Enforcer简介

LM Format Enforcer 通过结合字符级解析器与tokenizer前缀树，只允许那些包含潜在有效格式字符序列的token。它支持批量生成并仍处于实验阶段。

设置模型

我们将设置一个LLama2模型并初始化所需的输出格式：

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():
    config = AutoConfig.from_pretrained(model_id)
    config.pretraining_tp = 1
    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

基线输出

为了建立一个质性基线，我们先查看模型在没有结构化解码时的输出。

DEFAULT_SYSTEM_PROMPT = """..."""  # 系统提示语
prompt = """..."""  # JSON格式的提示语

from langchain_huggingface import HuggingFacePipeline
from transformers import pipeline

hf_model = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=200)
original_model = HuggingFacePipeline(pipeline=hf_model)

generated = original_model.predict(get_prompt("Michael Jordan"))
print(generated)

代码示例

通过LM Format Enforcer强制输出格式

使用 LMFormatEnforcer 可以确保输出精确符合指定的JSON模式。

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)

批处理

LMFormatEnforcer 也支持批处理模式：

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)