ChatStream

English | 日本語

ChatStream is a chat toolkit for pre-trained large language models.

It can be embedded in FastAPI/Starlette based web applications/web APIs to perform sequential sentence generation with pre-trained language models under load control.

Installation

pip install chatstream

Quick Start

Install required packages

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
pip install transformers
pip install "uvicorn[standard]" gunicorn 

Implementing a ChatStream server

Implement a streaming chat server for pre-trained models.

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,  # max_concurrent_executions for sentence generation
    max_queue_size=5,  # size of queue
    model=model,
    tokenizer=tokenizer,
    device=device,
    chat_prompt_clazz=ChatPrompt,
)

app = FastAPI()

# Specify session middleware to keep per-user ChatPrompt 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):
    # Just pass a FastAPI Request object to `handle_chat_stream_request` to automatically queue and control concurrency
    response = await chat_stream.handle_chat_stream_request(request)
    return response


@app.on_event("startup")
async def startup():
    # start the queueing system by doing `start_queue_worker` at the same time the web server starts up
    await chat_stream.start_queue_worker()

Table of Contents

• What is ChatStream

• - Importing Prompt Class ChatPrompt

• - Loading model classes

• -- HTTP session middleware configuration

• - Create and initialize ChatStream

• Implementation of Web API Endpoints

  • - Endpoint Implementation
  • - Receive streaming transmission completion callback
  • - Read requests from users
  • - How to set up an HTTP session

• Queueing System and Concurrency Limit

  • - What is the queueing system
  • - Starting the Queueing System

• Start the Web server (ASGI server)

  • - uvicorn (start from inside)
  • - uvicorn (start from outside)
  • - gunicorn

• Console chat implementation

  • - Run a simple console chat to check the model behavior

• Configuration during development

  • - CORS middleware settings
  • - Using Mock Response (Fast Startup)
  • - Logging Settings
  • - Reading requests from users
  • - Attach progress bar to time-consuming model loading

• Advanced Settings

  • Chat History Persistence
    • - Implement custom request handler
  • Configuration for large scale access
    • Interfacing with login authentication using OAuth
    • - Load Balancing on Multi-GPU
    • - Load Balancing with Multi-GPU Server

License

LICENSE.md

Citing ChatStream

@software{chatstream,
  title = {{ChatStream: A streaming chat toolkit for pre-trained large language models(LLM)}},
  author = {Qualiteg Inc.(https://qualiteg.com) },
  url = {https://github.com/qualiteg/ChatStream}
  month = {5},
  year = {2023},
  version = {0.15},
}