第9章 LMQL 自定义约束：实现自定义运算符指南
自定义约束
LMQL 的约束语言自带一组可组合的标准运算符。除此之外，你还可以实现自定义运算符来验证更复杂的属性，同时保持与内置运算符的可组合性。本章将通过两个自定义运算符的例子，讲解如何实现自定义的 LMQL 运算符。
实现自定义运算符
要实现自定义运算符，需要定义一个继承自 lmql.ops.Node 的类，并用 @LMQLOp 装饰器修饰，装饰器的参数是运算符名称。之后在 LMQL 查询中就可以直接使用这个名称来引用该自定义运算。
自定义运算符的基本接口由以下方法构成：
python from lmql.ops import LMQLOp, Node, InOpStrInSet
@LMQLOp("custom_operator") class CustomOperator(Node): # 提供运算符的值语义 def forward(self, *args, kwargs): ...
# 提供运算符结果的前瞻（用于下一 token 掩码） def follow(self, *args, kwargs): ... # 提供结果的确定性质（临时或最终） def final(self, *args, kwargs): ...
forward(self, *args, kwargs)：该方法实现运算符的前向语义，即计算 custom_operator(*args) 的结果，判断被验证的属性对给定参数是否成立。此方法会在模型输出的每个 token 上被调用，因此可以在生成过程中实现提前终止。
follow(self, x, kwargs)：该方法实现运算符的 follow 语义，即根据当前参数值对运算符结果进行前瞻。返回值是一个所谓的 follow map，用于区分不同类别的后续 token。此方法在模型输出的每个 token 上至少被调用一次，以确定允许的下一 token 集合。
final(self, x, kwargs)：最后，final 方法实现运算符的 final 语义，即判断运算符结果相对于当前模型输出是否是最终的。这让你能区分临时和最终的验证结果——这一点很重要，可以避免在只是临时违规的情况下过早终止。
可能的返回值为 "var" 和 "fin"，分别表示临时和最终（确定）结果。除了布尔逻辑之外，LMQL 还支持 "inc" 和 "dec"，分别表示单调递增和单调递减的结果。这对字符计数、字符串和列表很有用。例如，输入变量 x 的 final 值可能是 "inc"，表示它是一个正在生成的单调递增字符串；也可能是 "fin"，表示它已达到最终值（查询执行已完成对该值的解码）。
根据你的使用场景，可能只需实现 forward 即可（用于生成过程中的提前终止）。但要获得完整的掩码支持，还需要实现 follow 和 final。
示例：实现一个 foo-bar 约束
为了演示，我们实现一个运算符，强制模型在生成 foo 之后总是生成 bar。明确一下，以下序列分别是合法和非法的：
python "Hello foo!" # 非法 "Hello foo bar!" # 合法 "Hello foo bar foo bar!" # 合法 "Hello foo bar foo!" # 非法
为了实现这一约束，我们定义一个名为 foo_bar 的自定义运算符：
python import re from lmql.ops import LMQLOp, Node, fmap, tset
@LMQLOp("foo_bar") class FooBar(Node): def forward(self, x, kwargs): # 如果 x 尚未生成或其中没有 "foo"，不做限制 if x is None or "foo" not in x: return True
# 确保 foo 后面总是跟着 " bar" bar_segment = x.rsplit("foo",1)[1] return bar_segment.startswith(" bar") or len(bar_segment) == 0
def follow(self, x, kwargs): # (1) 如果 x 尚未生成，不做限制 if x is None: return True # (2) 如果 x 中没有 "foo"，不做限制 if "foo" not in x: return True
# (3) 获取最后一个 "foo" 之后的当前片段 bar_segment = x.rsplit("foo",1)[1] # 3(a) 已经满足 if bar_segment.startswith(" bar"): return True
# 3(b) 强制后续内容与 " bar" 对齐 return fmap( (tset(" bar", regex=True), True), # " bar" -> True ("*", False) # 其他任何内容 -> False )
def final(self, x, result=None, kwargs): # 违规是确定的 if not result: return "fin" # 否则取决于 x 的确定性质 return "fin" if x == "fin" else "var"
forward() 实现：首先实现 forward 方法。为了验证 foo-bar 属性，我们检查紧随潜在 foo 子串之后的字符串片段是否与 " bar" 对齐。这里需要考虑几种情况：x 为 None（尚未生成）、部分匹配（包括空字符串），以及匹配延伸超出 " bar" 的情况。取决于模型的分词方式，forward() 可能会在上述任何一种情况下被调用，因此必须能处理所有情况。
follow() 实现：接下来实现 follow 方法。同样需要考虑多种情况：
如果 x 尚未生成，无需限制下一 token。
如果 x 中没有 foo，无需限制下一 token。
如果 x 末尾正好是 foo，则做如下限制：
(a) 如果该片段已经以 " bar" 开头，无需任何限制。
(b) 否则，把下一 token 限制为 " bar"。为此我们构造一个所谓的 follow map，即从 token 范围到「下一 token 落在该范围内时运算符的未来求值结果」的映射。
在 foo-bar 这个例子里，只需标明：后续为 " bar" 时求值为 True，其他任何后续求值为 False。这由 fmap 函数实现，它根据一组 token 范围及其对应的未来求值结果构造 follow map。构造 token 范围可以用 tset 构造器，它支持按长度、集合、前缀或正则来选取 token，与具体使用的 tokenizer 无关。这里我们使用 regex=True 选项，自动选取完整或部分匹配 " bar" 的所有 token。
final() 实现：最后实现 final 方法。它告诉 LMQL 运行时，自定义运算符的结果相对于当前模型输出是最终的还是临时的。在这个例子中，返回值 False 总可以视为最终的（确定违规，应提前终止）。否则需要考虑 x 当前值的确定性质：如果 x 是最终的，运算符结果也是最终的；否则就是临时的，因为 x 的后续内容仍可能使约束得到满足。
注意：出于演示目的，follow() 实现的第 3b 步简化了一个关于 token 对齐的重要细节。它只考虑了 follow() 恰好在 "foo" 末尾被调用的情况；而实际上，取决于模型行为和分词方式，follow() 也可能在 "foo b" 这样的部分结果上运行，此时正确的 follow map 应指示 "ar" 为合法后续，而不是完整的 " bar"。要处理这种情况，可以直接复用内置 InOpStrInSet（即约束 VAR in [...] 的实现）：把 fmap 调用替换为 InOpStrInSet([]).follow(bar_segment, [" bar"])，它会自动处理所有这类情况。
使用自定义约束运算符
要使用自定义约束运算符，只需在查询上下文中导入它，然后如下使用：
"Say 'foo':[A]" where foo_bar(A)
一般来说，需要确保 @LMQLOp 装饰器在查询被解析之前已在当前进程中执行，例如通过导入包含该运算符实现的模块。
底层原理
有了如上所述的运算符实现，LMQL 运行时就能在生成过程中同时验证模型输出并推导 token 级别的预测掩码。为此，forward、follow 和 final 会在生成过程中被反复调用，并以当前模型输出 x 作为输入。运行时由此既能判断当前模型输出的合法性，也能推导出运算符确定（最终）求值为 False 的下一 token 范围。只要底层实现是正确的，就能可靠地保证模型永远不会选中那些被 follow 掩码排除、且会确定违反约束的 token。
LMQL 约束的表达能力
LMQL 约束在生成过程中通过 token 掩码被即时应用，这意味着模型无法生成任何被约束掩掉的 token。不过，这种方式的表达能力天然受限，因为并非所有关于文本的属性都能逐 token 判定。更具体地说，其表达能力仅限于上下文无关语言的验证。为了安全地使用 token 掩码，LMQL 的 final/follow 语义实现针对 token 掩码提供了可靠性保证（参见 LMQL 论文）。
尽管如此，得益于约束在生成过程中的即时求值，一旦模型输出确定违反约束（即验证结果是最终的），LMQL 约束就会立即触发，避免模型耗费大量计算去生成无效输出。这相比后处理式验证是一大优势——后者往往要等模型已经生成了大量无效输出之后才能发现问题。
