[ChatStream] Quickstart

Hello from the Product Development Division at Qualiteg Inc.

The summer heat just keeps going, doesn't it?

Let's jump right in and build a real-time streaming chat server using ChatStream, which we announced yesterday.

Installing the package

First, install the ChatStream package.

pip install chatstream

Installing the required packages

pip install torch torchvision torchaudio
pip install transformers
pip install "uvicorn[standard]" gunicorn 

Implementing the ChatStream server

In this example, we will implement a streaming chat server that uses RedPajama-INCITE as the LLM.

chatstream_server.py

import torch
from fastapi import FastAPI, Request
from fastsession import FastSessionMiddleware, MemoryStore
from transformers import AutoTokenizer, AutoModelForCausalLM

from chatstream import ChatStream,ChatPromptTogetherRedPajamaINCITEChat as ChatPrompt

model_path = "togethercomputer/RedPajama-INCITE-Chat-3B-v1"
device = "cuda"  # "cuda" / "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)
model.to(device)

chat_stream = ChatStream(
    num_of_concurrent_executions=2,# maximum number of concurrent text generations
    max_queue_size=5,# size of the wait queue
    model=model,
    tokenizer=tokenizer,
    device=device,
    chat_prompt_clazz=ChatPrompt,
)

app = FastAPI()

# Register the session middleware so that each user's ChatPrompt is kept in the HTTP session
app.add_middleware(FastSessionMiddleware,
                   secret_key="your-session-secret-key",
                   store=MemoryStore(),
                   http_only=True,
                   secure=False,
                   )


@app.post("/chat_stream")
async def stream_api(request: Request):
    # Simply pass the FastAPI Request object to `handle_chat_stream_request`; queuing and concurrency control are handled automatically
    response = await chat_stream.handle_chat_stream_request(request)
    return response


@app.on_event("startup")
async def startup():
    # Call `start_queue_worker` when the web server starts to launch the queuing system
    await chat_stream.start_queue_worker()

Tuning performance with constructor parameters

num_of_concurrent_executions

num_of_concurrent_executions=2 is the maximum number of text generations that run concurrently. The larger this number, the more text generations can run in parallel. However, if it is too large, the token generation speed (tokens per second) drops, so set an appropriate value for the performance of your GPU. For example, with a single A4000-class GPU, num_of_concurrent_executions=5 to 10 gives a generation speed that feels comfortable.

max_queue_size

max_queue_size=5 is the maximum number of requests that can wait for text generation when all concurrent generation slots are occupied. For example, as in the code above, once the maximum number of concurrent generations num_of_concurrent_executions==2 has been reached, the third and subsequent requests go into the wait queue. Requests three through eight are waiting for text generation, and the chat UI shows a progress bar for them. So what happens if a ninth request arrives in this state?

In that case, the ninth request receives a message in the UI to the effect that the text generation server is currently busy.

(Since this situation is undesirable for a service, you can prepare multiple nodes in advance so that it does not occur. You can also configure new nodes to spin up when traffic unexpectedly exceeds expectations. That way, when the load grows, new nodes come online and you can provide a chat service with no waiting time. We will cover these scaling settings in a separate post.)

Creating a prompt-handling class, "ChatPrompt"

The prompt-handling class that rewrites the user's input text into the format expected by the LLM is called the ChatPrompt class.

At Qualiteg, every time a new LLM is released we create and bundle a ChatPrompt class for that model, but you can also write your own.

Here is the ChatPrompt class for the RedPajama-INCITE model used in this example.

from chatstream import AbstractChatPrompt


class ChatPromptTogetherRedPajamaINCITEChat(AbstractChatPrompt):
    """
    togethercomputer/RedPajama-INCITE-7B-Chat
    """

    def __init__(self):
        super().__init__()  # Call the initialization of the base class
        self.set_requester("<human>")
        self.set_responder("<bot>")
        self.set_prefix_as_stop_str_enabled(True)  # enable requester's prompt suffix as stop str

    def get_stop_strs(self):
        return ['<|endoftext|>']

    def create_prompt(self, opts={}):
        """
        Build prompts according to the characteristics of each language model
        :return:
        """
        if self.chat_mode == False:
            return self.get_requester_last_msg()

        ret = self.system
        for chat_content in self.get_contents(opts):
            chat_content_role = chat_content.get_role()
            chat_content_message = chat_content.get_message()

            if chat_content_role:
                if chat_content_message:
                    merged_message = chat_content_role + ": " + chat_content_message + "\n"
                else:
                    merged_message = chat_content_role + ":"

                ret += merged_message

        return ret

    async def build_initial_prompt(self, chat_prompt):
        pass
        # If you want a common initial prompt for instructions, override this method and implement
        # chat_prompt.add_requester_msg("Do you know about the Titanic movie?")
        # chat_prompt.add_responder_msg("Yes, I am familiar with it.")
        # chat_prompt.add_requester_msg("Who starred in the movie?")
        # chat_prompt.add_responder_msg("Leonardo DiCaprio and Kate Winslet.")

For simple models, the prompt text fed to the LLM can be expressed with template matching alone, but when more complex processing is required, implementing it as a ChatPrompt class like this gives you far more flexibility.