Transformers streamer 怎样处理生成停止信号

Transformers 里的 streamer 有两个容易混淆的“停止”：一个是生成已经结束，streamer 向消费者放入结束标记；另一个是用户点击停止后，让 model.generate() 本身尽快退出。前者由 TextIteratorStreamer 的队列协议完成，后者要交给 stop_strings 或 StoppingCriteria。只在消费循环里 break，并不会自动取消后台生成线程。

官方地址：https://huggingface.co/docs/transformers/

把 streamer 的结束 sentinel 当成“消费者可以收尾”，把 StoppingCriteria 当成“生成循环应该停止”。两者可以同时使用，但职责不能互换。

先看清 streamer 的结束信号

generate() 生成 token 时会调用 streamer 的 put()，生成结束时调用 end()。对 TextIteratorStreamer 来说，结束处理会先把尚未刷出的文本放进队列，再放入内部的 stop_signal；消费者取到它后，迭代器通过 StopIteration 正常结束。这个标记是队列内部协议，不是一个让模型中途停下的取消按钮。

因此，正常结束时应该让消费者把队列读完；如果生成线程抛出异常，消费者还需要通过线程状态、共享异常容器或上层任务状态拿到错误，不能把“没有更多文本”一律当成成功。

同步消费时，timeout 只负责发现异常

最常见的用法是把 generate() 放到线程里，主线程遍历 streamer。设置 timeout 后，队列长时间没有新文本会暴露为超时异常，避免界面永远卡在等待状态；它不会改变模型的停止条件。

from queue import Empty
from threading import Thread
from transformers import TextIteratorStreamer

# timeout 只限制消费者等待队列的时间，不会取消 generate。
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, timeout=10.0)
generation_kwargs = {
    **inputs,
    "streamer": streamer,
    "max_new_tokens": 256,
}

# 后台线程负责生成，当前线程只负责读取已经整理好的文本片段。
worker = Thread(target=model.generate, kwargs=generation_kwargs)
worker.start()

try:
    for piece in streamer:
        # 这里可以把片段推给 WebSocket、终端或其他消费者。
        yield piece
except Empty:
    # 队列超时只说明暂时没有新片段，应用应记录状态并决定是否取消任务。
    raise RuntimeError("生成线程长时间没有产生文本")
finally:
    # 只有生成线程真正结束后才收尾，避免留下仍在占用模型的后台任务。
    worker.join()

这里的关键判断是：for piece in streamer 结束，代表队列协议已经发出结束信号；timeout 触发，则代表读取等待超时。两种情况的业务含义不同，日志和接口状态也应分开记录。

异步消费不要把同步 sentinel 改成字符串

需要 asyncio 时使用 AsyncTextIteratorStreamer。它内部使用 asyncio 队列，消费者用 async for 读取；生成完成后会抛出 StopAsyncIteration，等待超过 timeout 则是 TimeoutError。不要自己向队列塞一个普通字符串来模拟结束，因为正文片段本身也可能包含同样的字符串。

import asyncio
from threading import Thread
from transformers import AsyncTextIteratorStreamer

async def stream_response(model, tokenizer, inputs):
    # 必须在正在运行的事件循环中创建异步 streamer。
    streamer = AsyncTextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        timeout=10.0,
    )
    kwargs = {**inputs, "streamer": streamer, "max_new_tokens": 256}
    worker = Thread(target=model.generate, kwargs=kwargs)
    worker.start()

    try:
        async for piece in streamer:
            # 片段可以直接交给异步响应层或 WebSocket。
            yield piece
    except TimeoutError:
        # 超时是消费侧异常，不等于模型已经收到取消信号。
        raise RuntimeError("异步生成读取超时")
    finally:
        # 生产代码还应把 worker 的异常传回任务状态。
        await asyncio.to_thread(worker.join)

真正停止生成要使用停止条件

如果需求是“遇到某个字符串就停”，可以把 stop_strings 和 tokenizer 一起传给 generate()。如果需求是“用户点了停止按钮就停”，则应实现一个读取共享事件的 StoppingCriteria。它在生成循环中返回每个 batch 行是否停止，和 streamer 的消费者循环是两条不同的控制路径。

import threading
import torch
from transformers import StoppingCriteria, StoppingCriteriaList

class StopByEvent(StoppingCriteria):
    def __init__(self, stop_event: threading.Event):
        self.stop_event = stop_event

    def __call__(self, input_ids, scores, **kwargs):
        # 返回每个 batch 行的布尔结果；事件设置后让生成循环停止。
        should_stop = self.stop_event.is_set()
        return torch.full(
            (input_ids.shape[0],),
            should_stop,
            dtype=torch.bool,
            device=input_ids.device,
        )

stop_event = threading.Event()
criteria = StoppingCriteriaList([StopByEvent(stop_event)])

# UI 的停止按钮或上层取消逻辑应调用 stop_event.set()。
generation_kwargs = {
    **inputs,
    "streamer": streamer,
    "stopping_criteria": criteria,
    "max_new_tokens": 256,
}
worker = Thread(target=model.generate, kwargs=generation_kwargs)
worker.start()

# 业务层停止时设置事件；生成线程随后按检查点退出并结束 streamer。
stop_event.set()

固定终止词也可以直接写成 model.generate(**inputs, stop_strings=[""], tokenizer=tokenizer)。注意它是“生成内容匹配后停止”，不是让消费者忽略剩余队列。对外部取消来说，设置事件后仍要等待线程收尾，并处理模型正在执行的当前一步可能带来的短暂延迟。

常见误区与速查

• 只 break 消费循环：只能停止当前读取，不能保证后台生成停止；应同时设计停止条件和线程收尾。
• 把 timeout 当取消：timeout 是队列等待边界，发生后要有取消策略，不能假设模型已退出。
• 手写结束字符串：使用类自身的 sentinel 和 StopIteration/StopAsyncIteration，避免和真实文本冲突。
• 忽略生成异常：消费者看到结束前，要有任务状态或异常传递通道，区分自然结束、超时和失败。

小结

Transformers streamer 的“停止信号”首先是队列协议：生成侧结束时放入 sentinel，消费者据此结束迭代。它不等于取消生成。同步场景用 TextIteratorStreamer，异步场景用 AsyncTextIteratorStreamer；用户主动停止、固定终止词或预算限制，则通过 StoppingCriteria、stop_strings 等生成侧条件完成。把两层职责拆开，才能同时得到及时输出、可控取消和清晰的异常状态。

相关问题

TextIteratorStreamer 为什么一直等不到结束？

先确认生成线程确实调用了 generate() 并能正常返回，再检查是否在消费者侧错误地吞掉了线程异常或设置了过短的 timeout。正常完成时，streamer 会在最后一段文本之后放入结束 sentinel。

停止按钮按下后为什么还会多出一小段文字？

停止条件是在生成循环检查点生效，模型已经计算中的当前步可能仍会进入队列。前端应在收到取消状态后停止展示后续片段，同时等待生成线程完成清理；不要把这段短暂尾部输出误判为 streamer 失效。