outlines 如何让 LLM 输出符合正则表达式的结果
date
Jan 26, 2024
slug
用上 “算法了”
status
Published
tags
LLM
summary
用上 “算法了”
type
Post
最近在调研如何使 LLM 能够输出带约束的结果，例如只在固定的候选答案中选择一个回答、只输出数字、输出符合 regex 的结果、输出符合 json schema 的结果等，候选的项目有：
按照几个项目 readme 里面的 demo 跑了下，感觉 outlines 的体验最好，细看了下 outlines 中是如何实现让 LLM 的输出符合 regex 约束，大致流程如下：
- 根据正则表达式字符串创建 FSM（有限状态机），并转 DFSM（确定性有限状态机）
- 根据 FSM 的 states_to_token_maps 获得当前 state 允许（符合正则表达式）模型输出的 allowed_token_ids
- LLM 生成 logits 后，根据 allowed_token_ids 对 logits 加 mask（不允许的输出的 token 的 logit 置为 -inf），只允许输出 allowed_token_ids 中的 token
- 根据模型生成的 new_token_id 更新 FSM 当前的状态，更新方式为，看 new_token_id 命中了 states_to_token_maps[state] 中的哪个新状态，如果下一个状态是 final state 则停止生成
- 重复 2~4 的步骤直到停止状态或者达到最大输出长度
如何构建 FSM 的 states_to_token_maps？
- FSM 自己有一个 alphabet_symbol_mapping，key 为正则允许的字符，value 为对应的状态
- 核心的构建逻辑在
state_scan_tokens和_walk_fsm，输入初始 state 和 alphabet_symbol_mapping，遍历 tokenizer 的 vocabs，获得当前 state 下可能生成的token_ids_end_states
下面的 demo 代码中我去掉了 outlines 中的一些抽象和各种 edge case，并模拟了模型输出，可以更直观地了解整个过程
输出