使用JSONFormer进行结构化JSON解码的实践

在AI技术的不断发展中，生成式模型的多样化应用也越来越广泛。JSONFormer是一个可以让本地Hugging Face管道模型与JSON Schema进行结构化解码的库。这篇文章将带你深入了解如何使用JSONFormer来改进模型的输出格式。

技术背景介绍

通常，我们使用生成式模型来完成各种自然语言处理任务。然而，生成输出的无序性和格式不对齐是常见的问题，尤其是当我们需要模型输出符合特定结构时。JSONFormer通过填充结构化的token并从模型中采样内容token，有效解决了这一难题。

核心原理解析

JSONFormer与Hugging Face的pipeline模型协作，将模型生成的输出与预定义的JSON Schema结构结合，从而确保输出在结构上的正确性。这种方式不仅减少了后处理的复杂性，还提升了可靠性。

代码实现演示

首先，让我们尝试不使用JSONFormer进行生成任务，看看传统模型的输出格式：

import logging
import json
import os
import requests
from langchain_core.tools import tool
from langchain_huggingface import HuggingFacePipeline
from transformers import pipeline

# 配置日志
logging.basicConfig(level=logging.ERROR)

# 环境变量中获取Hugging Face API Key
HF_TOKEN = os.environ.get("HUGGINGFACE_API_KEY")

@tool
def ask_star_coder(query: str, temperature: float = 1.0, max_new_tokens: float = 250):
    """调用BigCode StarCoder模型进行编程问题查询。"""
    url = "https://api-inference.huggingface.co/models/bigcode/starcoder"
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "content-type": "application/json",
    }
    payload = {
        "inputs": f"{query}\n\nAnswer:",
        "temperature": temperature,
        "max_new_tokens": int(max_new_tokens),
    }
    response = requests.post(url, headers=headers, data=json.dumps(payload))
    response.raise_for_status()
    return json.loads(response.content.decode("utf-8"))

# 创建文本生成的pipeline
hf_model = pipeline(
    "text-generation", model="cerebras/Cerebras-GPT-590M", max_new_tokens=200
)

# 原始模型的预测
original_model = HuggingFacePipeline(pipeline=hf_model)

prompt = """You must respond using JSON format, with a single action and single...
AI Assistant:"""

generated = original_model.predict(prompt, stop=["Observation:", "Human:"])
print(generated)

输出结果未能按照预期的JSON格式返回。

接下来，我们使用JSONFormer来保证输出格式的正确性：

from langchain_experimental.llms import JsonFormer

# 定义JSON schema
decoder_schema = {
    "title": "Decoding Schema",
    "type": "object",
    "properties": {
        "action": {"type": "string", "default": ask_star_coder.name},
        "action_input": {
            "type": "object",
            "properties": ask_star_coder.args,
        },
    },
}

# 使用JSONFormer进行预测
json_former = JsonFormer(json_schema=decoder_schema, pipeline=hf_model)

results = json_former.predict(prompt, stop=["Observation:", "Human:"])
print(results)

JSONFormer成功保障了输出的结构性！

应用场景分析

JSONFormer特别适用于需要生成结构化数据的场景，例如数据库填充、自动化报告生成、数据解析等。这种方式不仅提高了开发效率，还减少了数据格式错误。

实践建议

1. 在设置JSON Schema时，务必清晰定义所有需要的字段及其类型。
2. 针对不同任务，调节生成模型的参数如max_new_tokens和temperature，以获取最优结果。
3. 结合实际需求灵活应用JSONFormer，提升模型生成结构化数据的能力。

如果遇到问题欢迎在评论区交流。

—END—