View the streaming example on GitHub

In this example, we go through a Truss that serves Qwen2.5 7B Instruct and streams the output to the client.

Why streaming?

LLMs generate tokens in sequence, so you can return useful output to users before the full response is ready. Truss supports streaming output to do this.

Set up the imports

In this example, we use the HuggingFace transformers library to build a text generation model.

model/model.py

from threading import Thread
from typing import Dict

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    GenerationConfig,
    TextIteratorStreamer,
)

Define the load function

In the load function of the Truss, we implement logic involved in downloading the instruct version of the Qwen2.5 7B model and loading it into memory.

model/model.py

class Model:
    def __init__(self, **kwargs):
        self.model = None
        self.tokenizer = None

    def load(self):
        self.tokenizer = AutoTokenizer.from_pretrained(
            "Qwen/Qwen2.5-7B-Instruct"
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            "Qwen/Qwen2.5-7B-Instruct", dtype=torch.bfloat16, device_map="auto"
        ).eval()

Define the preprocess function

In the preprocess function of the Truss, we set up a generate_args dictionary with some generation arguments from the inference request to be used in the predict function.

model/model.py

    def preprocess(self, request: dict) -> dict:
        generate_args = {
            "max_new_tokens": request.get("max_new_tokens", 512),
            "temperature": request.get("temperature", 0.5),
            "top_p": request.get("top_p", 0.95),
            "top_k": request.get("top_k", 40),
            "repetition_penalty": 1.05,
            "no_repeat_ngram_size": 0,
            "use_cache": True,
            "do_sample": True,
            "eos_token_id": self.tokenizer.eos_token_id,
            "pad_token_id": self.tokenizer.pad_token_id,
        }
        request["generate_args"] = generate_args
        return request

Define the predict function

In the predict function of the Truss, we implement the actual inference logic. The two main steps are:

Tokenize the input
Call the model’s generate function if we’re not streaming the output, otherwise call the stream helper function

model/model.py

    def predict(self, request: Dict):
        stream = request.pop("stream", False)
        prompt = request.pop("prompt")
        generation_args = request.pop("generate_args")
        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.cuda()

        if stream:
            return self.stream(input_ids, generation_args)

        with torch.no_grad():
            output = self.model.generate(input_ids=input_ids, **generation_args)
            generated = output[0][input_ids.shape[-1] :]
            return self.tokenizer.decode(generated, skip_special_tokens=True)

Define the `stream` helper function

In this helper function, we instantiate the TextIteratorStreamer object, which holds the tokens as the model generates them. skip_prompt keeps the request out of the response, so the client only sees the new text, and skip_special_tokens keeps the control tokens out of it.

model/model.py

    def stream(self, input_ids: list, generation_args: dict):
        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)

Build a GenerationConfig from the request’s generation arguments and pass the streamer object alongside it. Every sampling parameter lives in the generation config, so nothing is passed twice.

model/model.py

        generation_config = GenerationConfig(**generation_args)
        generation_kwargs = {
            "input_ids": input_ids,
            "generation_config": generation_config,
            "streamer": streamer,
        }

Spawn a thread to run the generation, so that it does not block the main thread.

model/model.py

        with torch.no_grad():
            # Begin generation in a separate thread
            thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
            thread.start()

In Truss, the way to achieve streaming output is to return a generator that yields content. In this example, we yield the output of the streamer, which produces output and yields it until the generation is complete. We define this inner function to create our generator.

model/model.py

            # Yield generated text as it becomes available
            def inner():
                for text in streamer:
                    yield text
                thread.join()
        return inner()

Set up the `config.yaml`

Running Qwen2.5 7B Instruct requires torch, transformers, and accelerate, which transformers needs to place the model with device_map="auto".

config.yaml

model_name: qwen2-5-7b-instruct
model_metadata:
  example_model_input:
    prompt: What is the meaning of life?
requirements:
  - accelerate==1.15.0
  - torch==2.14.0
  - transformers==5.17.0

Configure resources for Qwen

We will use an L4 to run this model.

config.yaml

resources:
  accelerator: L4
  cpu: "4"
  memory: 16Gi
  use_gpu: true

Deploy the model

Deploy the model like you would other Trusses, with:

baseten model push --dir qwen2-5-7b-instruct

Examples

Models

Engines

Custom Docker servers

Custom Python models

Chains

Stream LLM responses

View the streaming example on GitHub

Why streaming?

Set up the imports

Define the load function

Define the preprocess function

Define the predict function

Define the `stream` helper function

Set up the `config.yaml`

Configure resources for Qwen

Deploy the model

View the streaming example on GitHub

​Why streaming?

​Set up the imports

​Define the load function

​Define the preprocess function

​Define the predict function

​Define the stream helper function

​Set up the config.yaml

​Configure resources for Qwen

​Deploy the model

Why streaming?

Set up the imports

Define the load function

Define the preprocess function

Define the predict function

Define the `stream` helper function

Set up the `config.yaml`

Configure resources for Qwen

Deploy the model