执行摘要
PR#50200 修复了 Rust 前端 matches_stop_string 函数的一个行为差异:原实现按 stop 列表顺序选择第一个匹配项,而 Python 前端选择在文本中最早完成的 stop 项。修复后将 find_map 替换为 filter_map + min_by_key,按完成偏移量择优,并对齐行为。仅修改一个文件,新增三个测试,风险低。
功能与动机
当多个 stop string 同时出现在新生成的文本窗口中时(例如多 token 解码或推测解码),Python 前端(已在 #49391 修复)会选择最早完成的那一个,而 Rust 前端仍按列表顺序取第一个。这导致相同请求下两个前端输出不同:输出文本可能仍包含 stop string,且 stop_reason 错误。该修复是 Rust 前端功能对等(#44280)的一部分。
实现拆解
- 核心逻辑变更:
matches_stop_string 函数从使用 find_map(返回第一个匹配)改为 filter_map 收集所有匹配,再通过 min_by_key(|&(_, _, end)| end) 选择完成偏移量最小的 match。min_by_key 保留第一个最小值,因此平局时自动按列表顺序 break tie。
- 接口不变:函数签名
Option<(usize, usize)> 不变,调用方无需修改。
- 测试覆盖:新增三个单元测试:
stop_string_earliest_completing_wins_regardless_of_list_order:两个 stop 在不同位置完成时验证最早完成的胜出(无论列表顺序)。
stop_string_ties_broken_by_list_order:完成偏移量相同时验证列表顺序生效。
stop_string_completion_position_not_start_position:验证选择基于完成位置而非起始位置。
另外更新了现有测试注释以准确反映场景。
rust/src/text/src/output/decoded.rs
核心修复文件:修改了 matches_stop_string 函数的实现逻辑并新增三个测试用例。
fn matches_stop_string(stops: &[String], output: &str, new_bytes: usize) -> Option<(usize, usize)> {
// We compare byte subslices to avoid utf8 boundary problem
let output = output.as_bytes();
let next_off = (output.len() + 1) - new_bytes;
stops
.iter()
.map(|ss| (ss.as_bytes(), ss.len(), next_off.saturating_sub(ss.len())))
.enumerate()
// filter_map 收集所有匹配项,不再短路
.filter_map(|(ss_idx, (ss, len, start_off))| {
output[start_off..]
.windows(len)
.position(|w| w == ss)
.map(|pos| (ss_idx, start_off + pos, start_off + pos + len))
})
// min_by_key 按完成偏移量选择最早完成的 stop string
// 平局时保留第一个最小值,即按列表顺序
.min_by_key(|&(_, _, end)| end)
.map(|(ss_idx, start, _)| (ss_idx, start))
}
评论区精华
无 discussion。审核者 njhill 直接批准(LGTM)。
风险与影响
- 风险:低。变更局限于单个私有函数,接口不变,已有测试全部通过。新行为与 Python 对齐,无回归。
- 影响:使用 Rust 前端时,多 stop string 场景下的输出和 stop_reason 与 Python 前端一致,提升了用户体验。对性能无显著影响(匹配项数通常很小)。
关联脉络
-
49391:Python 端相同 bug 的修复,此 PR 是其 Rust 对等实现。
-
44280:Rust 前端功能对等 roadmap,此 PR 弥合了一个已知差距。
参与讨论