MAI / AI · Coding Interview

Microsoft MAI 流式 Stop Token:跨块识别与尽早安全输出

输出流中 stop token 之前的文本,token 可横跨 chunk;用 KMP 前缀状态仅暂存可能组成 token 的后缀。

题目重点:输出流中 stop token 之前的文本,token 可横跨 chunk;用 KMP 前缀状态仅暂存可能组成 token 的后缀。

题意与假设

这道练习题保留了面经中值得训练的核心约束,并将需要澄清的条件单独列出。现场拿到变体时,先和面试官确认这些条件。

  • token 非空,按字符精确匹配,大小写敏感;只认第一个完整出现的 token。
  • 跨块的部分 token 前缀需要暂存;仅在完整 token 到达时停止,EOF 残缺前缀在 finish() 时输出。若面试官将 EOF 残缺 token 也视为停止,规则需另改。
  • feed(chunk) 返回当前可确定安全的输出字符串;被 stop 命中的 chunk 后半段和其后 chunks 都不输出。

解题路线

  1. 先构建 token 的 KMP 前缀函数。处理流字符时维护 matched 个字符,pending 恰好是当前输入末尾与 token 前缀相同、尚不能安全输出的字符。
  2. 新字符失配时按前缀函数回退,把已经不可能属于任何未来 token 的 pending 前端字符立即输出;不回退扫描全部历史。
  3. 新字符匹配则放入 pending;matched 达 token 长度时已找到首个完整 token,丢弃 pending 并永久停止。
  4. 若流结束且未见完整 token,finish() 输出仍暂存的 pending。该缓冲最多 token 长度−1,不需把所有 chunks 合并。
  5. 输出接口若是 callback/async iterator,可把每次 feed 的安全片段立即发出,避免等待整条流;背压与下游错误是另一个工程问题。

手动走一遍

token=<END>,先 feed('hello') 可立即输出 hello;feed(' how<E') 输出 ' how',暂存 '<E';feed('ND>HAPPY') 补全 token,输出空并停止,HAPPY 丢弃。若最后直接 EOF 而没有 ND>,finish 会输出 '<E',因为没有完整 token。

Python 参考实现

下面的代码按上文假设实现,适合本地运行并追加自己的测试用例。

代码区可横向滚动;使用键盘时,先聚焦代码区再按方向键查看。

from collections import deque


class StopTokenEmitter:
    """Emit safe prefix immediately while retaining only possible token prefix."""

    def __init__(self, token: str):
        if not token:
            raise ValueError('empty token')
        self.token = token
        self.prefix = [0] * len(token)
        matched = 0
        for i in range(1, len(token)):
            while matched and token[i] != token[matched]:
                matched = self.prefix[matched - 1]
            if token[i] == token[matched]:
                matched += 1
            self.prefix[i] = matched
        self.pending: deque[str] = deque()
        self.matched = 0
        self.stopped = False

    def feed(self, chunk: str) -> str:
        if self.stopped:
            return ''
        output = []
        for char in chunk:
            while self.matched and char != self.token[self.matched]:
                fallback = self.prefix[self.matched - 1]
                for _ in range(self.matched - fallback):
                    output.append(self.pending.popleft())
                self.matched = fallback
            if char == self.token[self.matched]:
                self.pending.append(char)
                self.matched += 1
                if self.matched == len(self.token):
                    self.stopped = True
                    self.pending.clear()
                    break
            else:
                output.append(char)
        return ''.join(output)

    def finish(self) -> str:
        if self.stopped:
            return ''
        remaining = ''.join(self.pending)
        self.pending.clear()
        self.matched = 0
        return remaining


if __name__ == '__main__':
    emitter = StopTokenEmitter('<END>')
    assert ''.join([emitter.feed('hello'), emitter.feed(' how<E'),
                    emitter.feed('ND>HAPPY'), emitter.finish()]) == 'hello how'

复杂度

预处理 token 长度 m 为 O(m);每个输入字符至多进入/退出 pending 一次,KMP 回退总摊 O(L),L 是已处理字符数;额外状态 O(m),不含每次返回的输出。单次 feed 的输出长度与该 chunk 及先前最多 m−1 个暂存字符有关。

易错点与追问

容易失分的边界

  • 每 chunk 单独 find(token),漏掉跨边界 <E/ND>。
  • 等所有 chunks 到齐才 join,违反尽早输出要求且内存随总流长增长。
  • 把 '<E' 立即输出,一旦下块补 'ND>' 就无法撤回。
  • 把题目的 truncate 示例误写成 EOF 部分 token 也算完整 stop。
  • 遇重叠前缀如 token=abab 时错误清空所有匹配状态,漏检后续 token。

继续深挖

  • 若同时有多个 stop token,如何用 trie/Aho–Corasick 管理候选和优先级?
  • token 匹配按 UTF-8 字节还是 Unicode 字符?
  • 若输出到不可撤回网络流,下游背压与取消如何传播?
✓已复制微信号!