引导式生成（Guided Generation）：基于有限状态机（FSM）强制模型输出符合JSON Schema

引导式生成：基于有限状态机（FSM）强制模型输出符合JSON Schema

大家好，今天我们来聊聊一个非常实用且具有挑战性的主题：引导式生成，特别是如何利用有限状态机（FSM）来强制模型输出符合预定义的JSON Schema。在自然语言处理和生成式AI领域，确保输出结果的结构化和有效性至关重要。JSON Schema作为一种标准的结构化数据描述语言，为我们提供了定义数据结构的强大工具。而FSM则为我们提供了一种控制生成流程的机制，确保输出始终符合Schema的约束。

1. 问题背景：结构化输出的重要性

在许多应用场景中，我们不仅仅需要模型生成流畅的文本，更需要模型生成结构化的数据。例如：

• API调用: 模型需要生成包含特定参数的JSON请求，以便调用外部API。
• 数据提取: 模型需要从文本中提取信息，并以JSON格式组织这些信息。
• 配置生成: 模型需要生成配置文件，这些文件必须符合特定的格式和约束。

如果模型生成的JSON不符合Schema，会导致程序出错，数据丢失，甚至安全问题。传统的生成方法，例如基于Transformer的模型，虽然能够生成高质量的文本，但很难保证输出的结构化和有效性。因此，我们需要一种方法来引导模型的生成过程，使其始终符合预定义的JSON Schema。

2. 解决方案：有限状态机（FSM）的引入

有限状态机（FSM）是一种计算模型，它包含一组状态、一组输入和一组状态转移函数。在每个状态下，FSM根据当前的输入，转移到下一个状态。我们可以利用FSM来表示JSON Schema的结构，并控制模型的生成过程，确保输出始终符合Schema的约束。

具体来说，我们可以将JSON Schema的每个元素（例如对象、数组、字符串、数字等）映射到一个状态。状态之间的转移则由Schema的结构关系决定。例如，如果一个对象包含一个名为"name"的字符串属性，那么FSM就会从"对象开始"状态转移到"name属性"状态，然后转移到"字符串值"状态。

3. FSM的构建：从JSON Schema到状态转移图

构建FSM的第一步是将JSON Schema转换为状态转移图。这个过程需要仔细分析Schema的结构，并确定每个元素对应的状态和转移规则。

下面是一个简单的JSON Schema示例：

{
  "type": "object",
  "properties": {
    "name": {
      "type": "string"
    },
    "age": {
      "type": "integer",
      "minimum": 0,
      "maximum": 120
    },
    "city": {
      "type": "string",
      "enum": ["Beijing", "Shanghai", "Guangzhou"]
    }
  },
  "required": ["name", "age"]
}

根据这个Schema，我们可以构建如下的状态转移图：

状态	描述	输入	下一个状态
ROOT	根状态，表示Schema的开始	{	OBJECT_START
OBJECT_START	对象开始状态	"name":	NAME_KEY
NAME_KEY	"name"属性的键	字符串值	NAME_VALUE
NAME_VALUE	"name"属性的值	,	OBJECT_CONTINUE
OBJECT_CONTINUE	对象继续状态	"age":	AGE_KEY
AGE_KEY	"age"属性的键	数字值	AGE_VALUE
AGE_VALUE	"age"属性的值	,	OBJECT_CONTINUE2
OBJECT_CONTINUE2	对象继续状态	"city":	CITY_KEY
CITY_KEY	"city"属性的键	字符串值	CITY_VALUE
CITY_VALUE	"city"属性的值	}	OBJECT_END
OBJECT_END	对象结束状态		END
END	结束状态，表示Schema的结束

4. 生成过程的控制：FSM驱动的解码器

有了状态转移图，我们就可以构建一个FSM驱动的解码器，用于控制模型的生成过程。解码器在每个步骤都会检查当前状态，并根据状态转移图，限制模型可以生成的Token。

例如，如果当前状态是NAME_VALUE，那么解码器就会限制模型只能生成字符串值。如果当前状态是AGE_VALUE，那么解码器就会限制模型只能生成整数值，并且必须在指定的范围内（minimum和maximum）。

下面是一个简化的Python代码示例，展示了如何使用FSM来控制模型的生成过程：

import json
from typing import List, Dict, Any

class FSM:
    def __init__(self, schema: Dict[str, Any]):
        self.schema = schema
        self.state = "ROOT"
        self.path = [] # 用来跟踪当前在JSON中的路径
        self.transitions = self.build_transitions(schema) # 根据schema构建状态转移表

    def build_transitions(self, schema: Dict[str, Any]) -> Dict[str, Dict[str, str]]:
        """
        根据JSON Schema构建状态转移表。
        这只是一个简化的示例，需要根据具体的Schema进行扩展。
        """
        transitions = {
            "ROOT": {"{": "OBJECT_START"},
            "OBJECT_START": {}, # 动态填充
            "OBJECT_CONTINUE": {}, # 动态填充
            "NAME_KEY": {"string": "NAME_VALUE"},
            "NAME_VALUE": {",": "OBJECT_CONTINUE", "}": "OBJECT_END"},
            "AGE_KEY": {"integer": "AGE_VALUE"},
            "AGE_VALUE": {",": "OBJECT_CONTINUE", "}": "OBJECT_END"},
            "CITY_KEY": {"string": "CITY_VALUE"},
            "CITY_VALUE": {"}": "OBJECT_END"},
            "OBJECT_END": {},
            "END": {}
        }

        # 根据Schema的properties动态填充OBJECT_START和OBJECT_CONTINUE的状态转移
        if schema.get("type") == "object" and "properties" in schema:
            for key, value in schema["properties"].items():
                transitions["OBJECT_START"][f'"{key}":'] = f"{key.upper()}_KEY"
                transitions["OBJECT_CONTINUE"][f'"{key}":'] = f"{key.upper()}_KEY"

        return transitions

    def get_valid_tokens(self) -> List[str]:
        """
        返回当前状态下有效的Token。
        """
        return list(self.transitions.get(self.state, {}).keys())

    def next_state(self, token: str) -> None:
        """
        根据Token更新状态。
        """
        if token in self.transitions.get(self.state, {}):
            self.state = self.transitions[self.state][token]
        else:
            raise ValueError(f"Invalid token '{token}' in state '{self.state}'")

# 示例Schema
schema = {
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "age": {"type": "integer"}
    },
    "required": ["name", "age"]
}

# 初始化FSM
fsm = FSM(schema)

# 模拟生成过程
tokens = []
try:
    tokens.append("{")
    fsm.next_state("{")

    tokens.append('"name":')
    fsm.next_state('"name":')

    tokens.append('"John Doe"')
    fsm.next_state("string") #  模拟生成了一个字符串

    tokens.append(",")
    fsm.next_state(",")

    tokens.append('"age":')
    fsm.next_state('"age":')

    tokens.append("30")
    fsm.next_state("integer") # 模拟生成了一个整数

    tokens.append("}")
    fsm.next_state("}")

    print("Generated JSON:", "".join(tokens))

except ValueError as e:
    print("Error:", e)
    print("Current state:", fsm.state)
    print("Valid tokens:", fsm.get_valid_tokens())

在这个示例中，FSM类负责维护状态和状态转移规则。get_valid_tokens方法返回当前状态下有效的Token，next_state方法根据Token更新状态。在生成过程中，我们可以使用get_valid_tokens方法来限制模型可以生成的Token，确保输出始终符合Schema的约束。

5. 模型的集成：在解码过程中应用FSM

要将FSM集成到现有的生成模型中，我们需要修改模型的解码过程。传统的解码过程通常是贪婪搜索或Beam Search，在每个步骤中选择概率最高的Token。我们需要修改这个过程，使其在选择Token时，首先检查FSM的约束，然后只选择满足约束的Token。

具体来说，我们可以使用以下步骤：

1. 获取当前状态: 在每个解码步骤开始时，获取FSM的当前状态。
2. 获取有效Token: 使用get_valid_tokens方法获取当前状态下有效的Token列表。
3. 过滤候选Token: 使用有效Token列表过滤模型的候选Token。只保留在有效Token列表中的Token。
4. 选择Token: 在过滤后的候选Token中，选择概率最高的Token。
5. 更新状态: 使用next_state方法更新FSM的状态。
6. 生成下一个Token: 重复步骤1-5，直到生成完整的JSON。

6. 高级技巧：处理复杂的JSON Schema

上面的示例只是一个简化的演示，实际的JSON Schema可能会更加复杂，例如包含嵌套对象、数组、oneOf、anyOf等。要处理这些复杂的Schema，我们需要对FSM进行扩展。

• 嵌套对象和数组: 可以使用递归的方式来处理嵌套对象和数组。每当遇到一个嵌套对象或数组时，就创建一个新的FSM，并将其与父FSM连接起来。
• oneOf和anyOf: 可以为每个oneOf或anyOf创建一个新的状态分支。在生成过程中，根据模型的输出，选择一个分支进行后续生成。
• enum: 可以将enum中的每个值都作为一个单独的状态转移。

7. 局限性与挑战

尽管FSM可以有效地引导模型生成符合JSON Schema的输出，但它也存在一些局限性和挑战：

• 复杂性: 构建和维护FSM可能非常复杂，特别是对于复杂的JSON Schema。
• 性能: 在解码过程中应用FSM会增加计算开销，降低生成速度。
• 灵活性: FSM的约束过于严格，可能会限制模型的创造性。

8. 替代方案

除了FSM之外，还有一些其他的方案可以用于引导模型生成结构化的输出：

• Prompt工程: 通过精心设计的Prompt，引导模型生成符合Schema的输出。
• 微调: 使用符合Schema的数据集对模型进行微调，使其更好地理解Schema的约束。
• 混合方法: 将FSM与其他方法结合使用，例如使用Prompt工程来初始化生成过程，然后使用FSM来保证输出的结构化。

9. 代码示例：更完整的FSM实现

以下是一个更完整的FSM实现，可以处理一些更复杂的JSON Schema特征，例如嵌套对象和数组：

import json
from typing import List, Dict, Any, Optional

class FSM:
    def __init__(self, schema: Dict[str, Any]):
        self.schema = schema
        self.state = "ROOT"
        self.path = [] # 用来跟踪当前在JSON中的路径
        self.transitions = self.build_transitions(schema) # 根据schema构建状态转移表
        self.stack = [] # 用于处理嵌套对象和数组

    def build_transitions(self, schema: Dict[str, Any], prefix: str = "") -> Dict[str, Dict[str, str]]:
        """
        递归地根据JSON Schema构建状态转移表。
        """
        transitions = {}

        if schema.get("type") == "object":
            transitions[f"{prefix}OBJECT_START"] = {"{": f"{prefix}OBJECT_BEGIN"}
            transitions[f"{prefix}OBJECT_BEGIN"] = {}  # 动态填充
            transitions[f"{prefix}OBJECT_CONTINUE"] = {} # 动态填充
            transitions[f"{prefix}OBJECT_END"] = {"}": self.get_return_state(prefix)}  # 返回上一级状态

            if "properties" in schema:
                for key, value in schema["properties"].items():
                    transitions[f"{prefix}OBJECT_BEGIN"][f'"{key}":'] = f"{prefix}{key.upper()}_KEY"
                    transitions[f"{prefix}OBJECT_CONTINUE"][f'"{key}":'] = f"{prefix}{key.upper()}_KEY"

                    key_type = value.get("type")
                    if key_type == "string":
                        transitions[f"{prefix}{key.upper()}_KEY"] = {"string": f"{prefix}{key.upper()}_VALUE"}
                        transitions[f"{prefix}{key.upper()}_VALUE"] = {",": f"{prefix}OBJECT_CONTINUE", "}": f"{prefix}OBJECT_END"}
                    elif key_type == "integer":
                        transitions[f"{prefix}{key.upper()}_KEY"] = {"integer": f"{prefix}{key.upper()}_VALUE"}
                        transitions[f"{prefix}{key.upper()}_VALUE"] = {",": f"{prefix}OBJECT_CONTINUE", "}": f"{prefix}OBJECT_END"}
                    elif key_type == "object":
                        # 递归处理嵌套对象
                        nested_prefix = f"{prefix}{key.upper()}_"
                        nested_transitions = self.build_transitions(value, nested_prefix)
                        transitions.update(nested_transitions)
                        transitions[f"{prefix}{key.upper()}_KEY"] = {"{": f"{nested_prefix}OBJECT_START"}
                        transitions[f"{nested_prefix}OBJECT_END"] = {",": f"{prefix}OBJECT_CONTINUE", "}": f"{prefix}OBJECT_END"} # 嵌套对象结束，返回上一级
                    elif key_type == "array":
                         # 递归处理数组
                        nested_prefix = f"{prefix}{key.upper()}_ARRAY_"
                        transitions[f"{prefix}{key.upper()}_KEY"] = {"[": f"{nested_prefix}ARRAY_START"}
                        transitions[f"{nested_prefix}ARRAY_START"] = {}

                        items_schema = value.get("items", {})
                        items_type = items_schema.get("type")

                        if items_type == "string":
                            transitions[f"{nested_prefix}ARRAY_START"] = {"string": f"{nested_prefix}STRING_VALUE", "]": f"{nested_prefix}ARRAY_END"} # 数组为空的情况
                            transitions[f"{nested_prefix}STRING_VALUE"] = {",": f"{nested_prefix}ARRAY_CONTINUE", "]": f"{nested_prefix}ARRAY_END"}
                            transitions[f"{nested_prefix}ARRAY_CONTINUE"] = {"string": f"{nested_prefix}STRING_VALUE"}
                        elif items_type == "integer":
                            transitions[f"{nested_prefix}ARRAY_START"] = {"integer": f"{nested_prefix}INTEGER_VALUE", "]": f"{nested_prefix}ARRAY_END"}  # 数组为空的情况
                            transitions[f"{nested_prefix}INTEGER_VALUE"] = {",": f"{nested_prefix}ARRAY_CONTINUE", "]": f"{nested_prefix}ARRAY_END"}
                            transitions[f"{nested_prefix}ARRAY_CONTINUE"] = {"integer": f"{nested_prefix}INTEGER_VALUE"}

                        transitions[f"{nested_prefix}ARRAY_END"] = {",": f"{prefix}OBJECT_CONTINUE", "}": f"{prefix}OBJECT_END"}  # 数组结束，返回上一级
        elif schema.get("type") == "array":
            #  可以处理根级别的数组，但示例中主要关注对象
            pass

        # 添加ROOT状态
        if prefix == "":
            transitions["ROOT"] = {"{": "OBJECT_START"}  # 假设根节点是对象
            transitions["END"] = {}

        return transitions

    def get_return_state(self, prefix: str) -> str:
        """
        根据前缀计算返回状态。
        """
        if not prefix:
            return "END" # 根对象结束，返回END
        parts = prefix[:-1].split("_") # 去掉最后一个_
        parent_prefix = "_".join(parts[:-1]) + "_" if len(parts) > 1 else "" # 再次去掉最后一个KEY或ARRAY，得到父对象的前缀

        if "ARRAY" in prefix:
            return f"{parent_prefix}OBJECT_CONTINUE" if parent_prefix else "OBJECT_END"
        else:
             return f"{parent_prefix}OBJECT_CONTINUE" if parent_prefix else "OBJECT_END" # 父对象继续

    def get_valid_tokens(self) -> List[str]:
        """
        返回当前状态下有效的Token。
        """
        return list(self.transitions.get(self.state, {}).keys())

    def next_state(self, token: str) -> None:
        """
        根据Token更新状态。
        """
        if token in self.transitions.get(self.state, {}):
            self.state = self.transitions[self.state][token]
        else:
            raise ValueError(f"Invalid token '{token}' in state '{self.state}'")

# 示例Schema
schema = {
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "age": {"type": "integer"},
        "address": {
            "type": "object",
            "properties": {
                "city": {"type": "string"},
                "zipcode": {"type": "integer"}
            }
        },
        "hobbies": {
            "type": "array",
            "items": {"type": "string"}
        }
    },
    "required": ["name", "age"]
}

# 初始化FSM
fsm = FSM(schema)

# 模拟生成过程
tokens = []
try:
    tokens.append("{")
    fsm.next_state("{")

    tokens.append('"name":')
    fsm.next_state('"name":')

    tokens.append('"John Doe"')
    fsm.next_state("string") #  模拟生成了一个字符串

    tokens.append(",")
    fsm.next_state(",")

    tokens.append('"age":')
    fsm.next_state('"age":')

    tokens.append("30")
    fsm.next_state("integer") # 模拟生成了一个整数

    tokens.append(",")
    fsm.next_state(",")

    tokens.append('"address":')
    fsm.next_state('"address":')

    tokens.append("{")
    fsm.next_state("{")

    tokens.append('"city":')
    fsm.next_state('"city":')

    tokens.append('"New York"')
    fsm.next_state("string")

    tokens.append(",")
    fsm.next_state(",")

    tokens.append('"zipcode":')
    fsm.next_state('"zipcode":')

    tokens.append("10001")
    fsm.next_state("integer")

    tokens.append("}")
    fsm.next_state("}")

    tokens.append(",")
    fsm.next_state(",")

    tokens.append('"hobbies":')
    fsm.next_state('"hobbies":')

    tokens.append("[")
    fsm.next_state("[")

    tokens.append('"reading"')
    fsm.next_state("string")

    tokens.append(",")
    fsm.next_state(",")

    tokens.append('"hiking"')
    fsm.next_state("string")

    tokens.append("]")
    fsm.next_state("]")

    tokens.append("}")
    fsm.next_state("}")

    print("Generated JSON:", "".join(tokens))

except ValueError as e:
    print("Error:", e)
    print("Current state:", fsm.state)
    print("Valid tokens:", fsm.get_valid_tokens())

10. 总结：FSM引导生成，结构化输出成为可能

通过引入有限状态机（FSM），我们可以有效地引导模型生成符合JSON Schema的结构化输出。FSM通过状态转移图表示Schema的结构，并限制模型可以生成的Token，从而保证输出的结构化和有效性。尽管FSM存在一些局限性和挑战，但它仍然是解决结构化输出问题的一种有效方法。通过不断地改进和扩展FSM，我们可以更好地满足各种应用场景的需求，让结构化输出成为可能。

11. 展望：未来的发展方向

未来，我们可以探索以下几个方向来改进基于FSM的引导式生成：

• 自动化FSM构建: 开发自动化工具，根据JSON Schema自动构建FSM。
• 混合方法: 将FSM与其他方法结合使用，例如Prompt工程和微调，以提高生成质量和灵活性。
• 自适应FSM: 根据模型的输出动态调整FSM的状态转移规则，使其更加灵活和智能。

希望今天的分享能够帮助大家更好地理解和应用基于FSM的引导式生成技术。谢谢大家！