[ChatStream] クイックスタート

こんにちは！ (株)Qualiteg プロダクト開発部 です！

まだまだ暑いですね！

早速、昨日発表しました ChatStream をつかったリアルタイムストリーミングチャットサーバーを作ってみたいと思います。

パッケージのインストール

早速 ChatStream パッケージのインストールをしていきます

pip install chatstream

必要パッケージのインストール

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

ChatStream サーバーの実装

今回は RedPajamaINCITE を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,# 文章生成の最大同時実行数
    max_queue_size=5,# 待ち行列の大きさ
    model=model,
    tokenizer=tokenizer,
    device=device,
    chat_prompt_clazz=ChatPrompt,
)

app = FastAPI()

# ユーザーごとの ChatPrompt を HTTP セッションに保持するため、セッションミドルウェアを指定
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):
    # FastAPI の Request オブジェクトを `handle_chat_stream_request` に渡すだけで自動的にキューイング、同時実行制御します
    response = await chat_stream.handle_chat_stream_request(request)
    return response


@app.on_event("startup")
async def startup():
    # Webサーバー起動と同時に `start_queue_worker` を行い、キューイングシステムを開始します
    await chat_stream.start_queue_worker()

コンストラクターパラメータでパフォーマンス調整

num_of_concurrent_executions

num_of_concurrent_executions=2 は文章生成の最大同時実行数を表しています。この数字が大きいほど、複数の文章生成を並行実行することができます。ただし、大きすぎると、トークン生成速度(トークン/秒)が遅くなるため、GPUの性能にあわせて適切な値を設定してください。たとえば、GPUが A4000 x 1枚 程度の性能の場合は num_of_concurrent_executions=5～10　程度が「快適」と思える生成速度となります。

max_queue_size

max_queue_size=5 は、文章生成の同時実行が埋まっているときに、文章生成待ちをしているリクエストの最大数です。たとえば、上の例のように文章生成の同時実行数の最大数 num_of_concurrent_executions==2 に達しているとき、3番目のリクエストからは待ち行列に入ります。3番目から8番目までは、文章生成を待っている状態で、チャットUI上ではプログレスバーとなります。では、この状態で9番目のリクエストが来たらどうなるでしょうか。

その場合、9番目のリクエストには「現在、文章生成サーバーがBusyです」の旨UIに表示されます。

（ただしこのような状況はサービスとして好ましくないため、このような状況が発生しないよう複数のノードをあらかじめ準備しておくことができます。また、急遽想定以上のアクセスが発生したときに、新しいノードが立ち上がるように設定しておくこともできます。これにより、負荷が多くなったときに、新しいノードが立ち上がり、待ち時間の無いチャットサービスを提供することができます。これら、スケーリングに関する設定方法は別途投稿いたします。）

プロンプト処理クラス "ChatPrompt"を作る

ユーザーによる入力テキストをLLM用のフォーマットに書き換えるためのプロンプト処理クラスはChatPrompt クラスと呼びます。

Qualiteg では、新しい LLM が発表されるたび、そのモデル用の ChatPrompt クラスを作成・バンドルしていますが、自前で作成することもできます。

ここでは今回のサンプルで利用した RedpajamaIncite のモデル用の ChatPromptクラスをご紹介します

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.")

単純なモデルの場合は、LLMに入力するプロンプトテキストは、テンプレートマッチングだけで表現できますが、複雑な処理が必要な場合には、このようにChatPromptクラスとして実装したほうが柔軟な処理ができます。