Skip to main content

Module 5 — Streaming responses

Module 4 built a chat that answers in one shot: the user submits, the assistant waits, and after two or ten seconds the whole answer appears. A model that takes three seconds to write two paragraphs is not slow; it feels slow because nothing moves during those three seconds. Streaming rewrites the same interaction at zero cost on the compute side: the first token appears in a few hundred milliseconds, and the rest fills in as it is generated. The change in perceived quality is dramatic, and it is one line of Python.

The generator pattern in Python​

A Python function that uses yield instead of return is a generator. Each yield pauses the function, hands a value to the caller, and resumes when the caller asks for the next value. Gradio consumes generators natively: return one from a chat callback and each yielded value replaces the previous message in the display, in place, until the generator exhausts.

import time

def slow_counter():
    for i in range(1, 6):
        time.sleep(0.5)
        yield f"Counted to {i}"

Calling slow_counter() returns an iterator, not the final value. In a Gradio callback, that is exactly the contract we need: the framework pulls one value at a time from the generator and refreshes the widget between two pulls. From the user's point of view, the message grows on screen.

Streaming with a hosted API​

The OpenAI-compatible chat completions endpoint has a stream=True switch. When set, the endpoint returns not a single response but an iterator over chunks, each carrying a delta. Wrap the loop into a generator that accumulates the assistant's answer, and Gradio does the rest.

import os
import gradio as gr
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

def chat_stream(message: str, history: list[dict], system_prompt: str):
    messages = [{"role": "system", "content": system_prompt}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    stream = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        temperature=0.4,
        stream=True,
    )

    reply = ""
    for chunk in stream:
        delta = chunk.choices[0].delta.content or ""
        reply += delta
        yield reply

gr.ChatInterface(
    fn=chat_stream,
    additional_inputs=[gr.Textbox(value="You are a concise assistant.", label="System prompt")],
    title="Streaming chat",
).launch()

The important detail: each yield carries the cumulative answer, not the incremental delta. Gradio replaces the message in place at every step, so what you yield is what the user sees at that moment. Yielding only the delta would show one token at a time and lose the previous ones — a common first mistake.

Streaming from a local model with transformers​

For a local Hugging Face model, the same idea works with TextIteratorStreamer: it is a helper that turns the model's token-by-token generation loop into a Python iterator. The generation runs in a background thread, and the main thread consumes tokens as they are produced.

import gradio as gr
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import torch

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.3",
    torch_dtype=torch.float16,
    device_map="auto",
)

def chat_stream(message: str, history: list[dict]):
    conversation = history + [{"role": "user", "content": message}]
    inputs = tokenizer.apply_chat_template(
        conversation,
        return_tensors="pt",
        add_generation_prompt=True,
    ).to(model.device)

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    generation_kwargs = dict(
        inputs=inputs,
        streamer=streamer,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.4,
    )
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    reply = ""
    for token in streamer:
        reply += token
        yield reply

gr.ChatInterface(fn=chat_stream, title="Streaming local model").launch()

The pattern is the same: a background loop generates tokens, the callback yields the accumulated text. What changes is that inference now happens on your machine, so latency depends on your hardware, not on a network round-trip. On a laptop CPU the first token can take five to ten seconds; on a GPU it drops to a few hundred milliseconds.

Interruption from the user​

A streaming demo is only pleasant if the user can stop generation when they no longer need the answer. gr.ChatInterface displays a Stop button automatically while a generator is running, and clicking it cancels the callback. In practice this raises a CancelledError inside the generator, propagates back through Gradio, and unblocks the UI.

For the OpenAI client, cancellation stops requesting new chunks, and the server-side generation stops billing shortly after. For a local model, cancellation stops the iteration but the background thread keeps running until the current forward pass ends, so the model keeps producing tokens into the streamer for another turn. That is normal and rarely visible, but if you plan to have hundreds of concurrent users on a modest GPU, wire an explicit stopping_criteria to make cancellation immediate; module 7 will come back to concurrency.

Perceived latency: the numbers that matter​

Two numbers matter more than the total generation time. The time to first token is what feels like the response time: below 500 ms it feels instant, between 500 ms and 2 seconds it feels responsive, above that it feels slow. The tokens per second governs whether the user has time to read as the answer streams: fifteen tokens per second and above is comfortable, five is painful.

A demo that streams at ten tokens per second with a first token in 300 ms will feel fast even if the full paragraph takes eight seconds. The same paragraph delivered in one shot after four seconds feels slower — that is not an intuition, it has been measured in every product that shipped both modes. This is the single most valuable thing streaming buys you, and it is why every serious chat interface streams.

Streaming for classifiers too

Streaming is not only for text. A classifier that returns intermediate confidences can yield partial dictionaries, so the gr.Label bars grow as the model refines its verdict. A speech-to-text pipeline can yield after each processed audio chunk. The pattern generalizes wherever a function produces its output progressively.

In summary​

  • A Python function with yield is a generator; Gradio consumes it natively and refreshes the widget between yields.
  • Each yield should carry the cumulative result, not the delta, because Gradio replaces the message in place.
  • Streaming from an OpenAI-style API uses stream=True; streaming from a local transformers model uses TextIteratorStreamer in a background thread.
  • The metrics that decide perceived speed are time to first token and tokens per second, not the total wall time.

Next module: prefilled examples and caching, so that first-time visitors see the demo working within a second of opening the page.