Showing tokens as they arrive, and making them stop

Why streaming changes the product

A 300-token answer at 25 tokens per second takes twelve seconds. Delivered at the end, that is a long silence and people reload the page. Delivered as it arrives, the first words appear in under a second and the wait is tolerable. Nothing about the computation changed; the perceived speed did.

This is why time to first token is the latency number worth tracking for anything interactive, and total time is the one worth tracking for anything batch.

The simple streamer

from transformers import TextStreamer

streamer = TextStreamer(tok, skip_prompt=True, skip_special_tokens=True)
model.generate(**enc, max_new_tokens=200, streamer=streamer)

This prints to stdout as tokens decode. It is perfect for a terminal script and useless for a web application, because it writes rather than yields.

The one you use in an app

from transformers import TextIteratorStreamer
from threading import Thread

streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
kwargs = dict(**enc, max_new_tokens=512, streamer=streamer)

Thread(target=model.generate, kwargs=kwargs).start()

for chunk in streamer:
    yield chunk          # into a Gradio generator, or an SSE response

generate blocks, so it runs on a second thread while the main thread consumes the queue. This is the shape behind every streaming chat UI you have used, including the Gradio Spaces in the next module.

Chunks are not tokens and not words. The decoder buffers until it can emit valid text, so you may get " un" then "forgiving". Never assume a chunk boundary is a word boundary, and never split on it for parsing.

Why models do not stop

Generation ends when one of three things happens: the model emits its end-of-sequence token, max_new_tokens is reached, or a stopping criterion fires. When a model rambles past its answer into an invented next turn, the usual cause is that the EOS it was trained to emit is not the EOS configured at generation time: a common mismatch in fine-tunes and conversions, where the template’s end marker and eos_token_id disagree.

Check them against each other:

print(tok.eos_token, tok.eos_token_id)
print(model.generation_config.eos_token_id)

Some models legitimately have several stop tokens, and eos_token_id accepts a list:

model.generate(**enc, eos_token_id=[tok.eos_token_id,
                                    tok.convert_tokens_to_ids("<|eot_id|>")])

max_new_tokens, not max_length

max_length counts prompt plus output, so a 900-token prompt with max_length=1024 leaves room for 124 tokens and the answer gets cut off mid-sentence. max_new_tokens counts only the output and is what you almost always want. Both exist; only one behaves the way people expect.

Always set it. A model with no limit and a broken stop token will generate until it exhausts the context window, which on a 128k model is a long and expensive silence.

Stop strings

When the model reliably emits a marker you want to cut at:

model.generate(**enc, max_new_tokens=256, stop_strings=["\nUser:", "```end"],
               tokenizer=tok)

Stop strings are checked on decoded text, so the tokenizer must be passed. They are a patch over a prompt or template problem rather than a fix, but they are an effective patch and there is no shame in it.

What streaming costs you

Streaming is not free of trade-offs, and two of them bite in production.

You cannot validate output you have already shown. If the answer must be checked (against a schema, a moderation pass, a groundedness test) you either buffer it and lose the benefit, or you display text you may have to retract. Most teams compromise: stream the prose, buffer anything structured.

And a streamed response is harder to cache. A complete answer is one object you can store against a prompt hash. A stream is a sequence of events, so the usual pattern is to assemble the full text as it goes and write that to the cache at the end, serving later identical requests instantly and non-streamed.

Cancellation, which people forget

A user closes the tab and your GPU keeps generating 500 tokens for nobody. In a threaded setup, a custom StoppingCriteria that checks a flag is the mechanism:

from transformers import StoppingCriteria

class Cancelled(StoppingCriteria):
    def __init__(self, flag): self.flag = flag
    def __call__(self, input_ids, scores, **kw): return self.flag.is_set()

On a free Space with one worker, one abandoned long generation blocks everybody in the queue. This is a small piece of code with a large effect on how a shared demo feels.

Do this now

Stream a 300-token answer and time two things: seconds to the first visible character, and seconds to the last. Then run the same prompt without streaming and note only the second number. The gap between those two experiences is the whole argument, and it costs six lines.