ローカルLLMの生成文字をWebSocket経由で順次出力させる

## はじめに

最近ローカルLLMをAPI経由で飛ばすアプリケーションを開発しているのですが、APIを叩いた際に、出力文字が一括で出力されるからか、待ち時間が長くストレスになっていました。
ChatGPTはもちろんですが、FastChatなどローカルLLMでも同じように文字が順次出力されており、同じことができないかと調査したところ、transformersの場合streamerというストリームミングを実現するオプションの存在を知りました。
この記事はstreamerとWebSocketでLLMの出力を順次出力させる内容になります。

この記事のコードは以下リポジトリに置いてあります。
[WSLLMComm](https://github.com/yamichi77/WSLLMComm/tree/v1.0?tab=readme-ov-file)

## 動作確認環境
OS：Windows 11（WSL+Docker）
GPU：RTX 2070 Super（VRAM：8GB）
バックエンド言語：Python
Webフレームワーク：FastAPI
フロントエンド：Vite+React+TypeScript

## 技術概要
### WebSocketとは

[Wikipediaより引用](https://ja.wikipedia.org/wiki/WebSocket)

> WebSocket（ウェブソケット）は、単一のTCPコネクション上に双方向通信のチャンネルを提供する、コンピュータの通信プロトコルの1つである。

> WebSocketはHTTPとは異なるプロトコルである。ともにOSI参照モデルのレイヤー7に位置し、レイヤー4のTCPに依存している。両者は異なるプロトコルであるが、RFC 6455では、WebSocketは「HTTPプロキシと仲介者をサポートするために、HTTPの443番および80番ポート上で動作するように設計されている」と述べられているように、HTTPプロトコルと互換性がある。互換性を実現するために、WebSocketのハンドシェイクはHTTP/1.1 Upgradeヘッダーを使用し、HTTPプロトコルをWebSocketプロトコルに変更するように実現されている。

つまり、Socket通信のようなことをHTTPポートで行えるプロトコルということです。

### transformersのstreamerについて

transformersの出力をストリームミングでテキスト出力してくれるものになります。使用方法はtokenizerを引数として入れ、generateの引数に渡します。
```Python
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct")
model = AutoModelForCausalLM.from_pretrained(
            "microsoft/Phi-3-mini-128k-instruct",
            device_map="cuda",
            torch_dtype="auto",
            trust_remote_code=True,
        )
inputs = tokenizer("Can you provide ways to eat combinations of bananas and dragonfruits?", return_tensors="pt")
streamer = TextStreamer(tokenizer)

_ = model.generate(**inputs, streamer=streamer, max_new_tokens=20)

for text in streamer:
    print(text, end="")
```

## 実装コード

### バックエンド

FastAPIの実装になります。
まずmainとrouterから
```Python:main
from api.router import llm
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(llm.router)
```

```Python:router
from api.services.llm import LlmService
from fastapi import APIRouter, HTTPException, WebSocket

router = APIRouter()


@router.on_event("startup")
async def startup_event():
    global service
    service = LlmService()

@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    global service
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        print("WebSocket Data: ", data)
        await service.get_inference_result_as_ws(data, websocket)
        print("WebSocket Success!")
```

mainでrouterを呼んでいます。
routerの特徴は以下の通りです。
* ロードしたモデルを保持するため、startupでサービスを初期化してグローバル変数にしています
* `@rouer.websocket("/ws")`でWebSocketを定義しています

次はserviceです。
```Python:service
import asyncio
from threading import Thread

import torch
from fastapi import WebSocket
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer


class LlmService:
    def __init__(self):
        self.tokenizer = None
        self.model = None

    def get_inference_result(self, str: str) -> TextIteratorStreamer:
        if self.model is None:
            self.set_llm()
        chat = [
            {"role": "user", "content": str},
        ]
        prompt = self.tokenizer.apply_chat_template(
            chat, tokenize=False, add_generation_prompt=True
        )
        token_ids = self.tokenizer.encode(
            prompt, add_special_tokens=False, return_tensors="pt"
        )
        streamer = TextIteratorStreamer(
            self.tokenizer, skip_prompt=True, skip_special_tokens=False
        )
        generation_kwargs = dict(
            input_ids=token_ids.to(self.model.device),
            do_sample=True,
            temperature=0.6,
            max_new_tokens=256,
            streamer=streamer,
        )
        with torch.no_grad():
            thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
            thread.start()
        return streamer

    def set_llm(self):
        self.model = AutoModelForCausalLM.from_pretrained(
            "microsoft/Phi-3-mini-128k-instruct",
            device_map="cuda",
            torch_dtype="auto",
            trust_remote_code=True,
            load_in_4bit=True,
        )
        self.tokenizer = AutoTokenizer.from_pretrained(
            "microsoft/Phi-3-mini-128k-instruct"
        )

    async def get_inference_result_as_ws(self, input: str, ws: WebSocket) -> None:
        streamer = self.get_inference_result(input)
        for text in streamer:
            await ws.send_text(text)
            print(text, end="")
            await asyncio.sleep(0)
        await ws.send_text("
")
```

serviceではモデルのロードおよび文字の生成を行っています。
今回は`microsoft/Phi-3-mini-128k-instruct`を使用させていただきました。
* `get_inference_result`で適切なパラメータをセットし、generateをしたstreamerを返しています
* `get_inference_result_as_ws`でWebSocketへSendを行っています
    * asyncio.sleep(0)を入れないとSendパラメータ送信タイミングが処理後になるので注意
* 最後に終わりを検知するテキスト`
`を送信します

### フロントエンド

バックエンドで定義したWebSocketを受け取るフロントエンドの実装です。
今回はReact+TypeScriptの環境想定です。

```TypeScript:container.tsx
import { SelectChangeEvent } from "@mui/material";
import { useEffect, useRef, useState } from "react";
import { Layout } from "../../components/layout/container";
import { useAPI } from "../../hooks/useAPI";
import { useWebSocket } from "../../hooks/useWebSocket";
import { LlmTextDto } from "../../types/LlmTextDto";
import { IndexPagePresenter } from "./presenter";

type LlmTextDto = {
  speaker: string;
  text: string;
};

export const IndexPage = () => {
  const [ws, setWs] = useState<WebSocket | null>(null);
  const [llmText, setLlmText] = useState<string>("");
  const [responseText, setResponseText] = useState<string>("");
  const responseTextRef = useRef<string>("");
  const [llmTextList, setLlmList] = useState<LlmTextDto[]>([]);
  const handleGetLlmText = async () => {
    setLlmText("");
    setIsSendButtonDisabled(true);
    setLlmList((llmTextList) => [
      ...llmTextList,
      { speaker: "user", text: llmText },
    ]);
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws?.send(llmText);
    }
  };
  useEffect(() => {
    const ws = new WebSocket("ws://localhost:8000/ws");
    ws.onopen = () => {
      console.log("connected");
    };
    ws.onclose = () => {
      console.log("disconnected");
    };
    ws.onmessage = (event) => {
      const message = event.data;
      if (message === "
") {
        const text = responseTextRef.current;
        setLlmList((llmTextList) => [
          ...llmTextList,
          { speaker: "bot", text: text },
        ]);
        responseTextRef.current = "";
        setResponseText("");
        return;
      }
      responseTextRef.current += message;
      setResponseText(responseTextRef.current);
    };
    setWs(ws);
    return () => {
      ws.close();
    };
  }, [getWs]);
};
```

上記プログラムのコードは以下の通りです。
* `llmText`はユーザー入力のテキスト、`responseText`は受信したテキスト、`llmTextList`は受信、送信両方のテキストをListとして持つ想定です
* useEffectでWebSocketの定義を行っています
    * メッセージ受信時にResponseTextへ受信したテキストを追加しています
    * `
`が来た際に`llmTextList`へテキストを入れ、受信したテキストの初期化をおこなっております
* `handleGetLlmText`はボタンなどが押された際にテキストを送信するメソッドです

## 動作比較

通常の生成速度と比較をしてみます。

### 通常
![rest20240602141717.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3739636/60269a75-e207-92a8-307b-6bdf671009db.gif)
生成時間：25.6秒

### WebSocket
![websocket20240602142014.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3739636/90b0e329-f986-8bde-2e5a-88af79d59e04.gif)
生成開始時間：0.77秒
生成終了時間：24.7秒

どちらもほぼ同じ時間ですが、WebSocketは待ち時間が少なく、ストレスを感じにくいですね！

# おわりに

今回の記事をまとめると以下の通りです。

* WebSocketとstreamerを使うことでLLMの生成物を順次出力することができる
* 順次出力のほうが、待ち時間が少なくユーザーにストレスを感じさせない

この記事が、ローカルLLM生成時間の長さで悩んでいる方の助けになれば幸いです。

# 参考・出展リンク
- [WebSocket - Wikipedia](https://ja.wikipedia.org/wiki/WebSocket)
- [Utilities for Generation - Hugging Face](https://huggingface.co/docs/transformers/main/en/internal/generation_utils#transformers.TextStreamer)
- [microsoft/Phi-3-mini-128k-instruct - Hugging Face](https://huggingface.co/microsoft/Phi-3-mini-128k-instruct)
- [huggingface/transformersでストリーミング](https://cockscomb.hatenablog.com/entry/streaming-with-huggingface-transformers)
- [ご家庭用LLMでストリーミングする方法](https://note.com/shi3zblog/n/n66ae41af7c64)