何をやろうとしているのか？

rinna Youri7Bをdockerを使って自分のPCで動作させる

自分のノートPCで最新の日本語LLMでチャットができるかどうかを確認したかった。

rinna、Llama 2の日本語継続事前学習モデル「Youri 7B」を公開

以下の記事に、google collaboratory で動かす記事が紹介されていたので、早速自分でも実行してみた。

そもそも、チャットのアプリとしては、これまでの会話の内容をその都度入力しないといけないので、WebUIにしたほうが試しやすいということで改良を試みようと始めたが・・・

collaboratory はPro版ではなくノーマル版（お金払えていない）で動作させていると、T4 GPU と、AutoGPTQForCausalLMが想定しているGPU は、異なるので動作できなくなっていた。cuda ライブラリを読み込むところでエラーとなってしまった。

準備

適当な場所に作業用フォルダを作成します

mkdir Youri7B
cd Youri7B

Dockerの準備

cuda が動作できる docker を準備してください

自分の環境は、RTX3060 6GBのノートパソコンで、ubuntu 22.04の環境を使っています。

次に、Dockerfileというファイル名で、以下の内容を保存します

FROM ubuntu:22.04

RUN apt-get update && apt-get install -y python3-pip
RUN python3 -m pip install pip --upgrade
RUN python3 -m pip install transformers accelerate bitsandbytes auto_gptq scipy gradio
RUN python3 -m pip install huggingface-hub

COPY chat.py /root/
# set your GPU device id
ENV CUDA_VISIBLE_DEVICES 0
# prepare download youri7b files
RUN python3 /root/chat.py download

CMD ["/usr/bin/python3","/root/chat.py"]

変数のCUDA_VISIBLE_DEVICES には、使用するGPUの番号を入力します。1枚しか使っていない場合には、０で良いですが、２枚挿ししている場合などは変更できます。

RUN python3 /root/chat.py download

は、事前にhuggingfaceからモデルをダウンロードして docker のイメージに格納するための処理をしています。

チャットのアプリ

chat のプログラムとして、以下の内容を chat.py として保存します。

#!/usr/bin/env python3

import torch
from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM
import gradio as gr
import sys

class ChatUI:
    def __init__(self, modelname ):
        self.modelname = modelname
        self.default_max_token = 200
        self.default_temperature = 0.5

    def prepare_pretrained(self):
        # トークナイザーとモデルの準備
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.modelname
        )
        print('prepare model is loaded.')

    def prepare_model(self):
        self.model = AutoGPTQForCausalLM.from_quantized(
            self.modelname,
            use_safetensors=True
        )

    # 推論の実行
    def generate(self,prompt,max_token = -1, temp = -1):
        if max_token < 0:
            max_token = self.default_max_token
        if temp < 0:
            temp = self.default_temperature

        token_ids = self.tokenizer.encode(
            prompt,
            add_special_tokens=False,
            return_tensors="pt")
        with torch.no_grad():
            output_ids = self.model.generate(
                input_ids=token_ids.to(self.model.device),
                max_new_tokens=max_token, #200,
                do_sample=True,
                temperature=temp, #0.5,
                pad_token_id=self.tokenizer.pad_token_id,
                bos_token_id=self.tokenizer.bos_token_id,
                eos_token_id=self.tokenizer.eos_token_id
            )
        return self.tokenizer.decode(
            output_ids[0][token_ids.size(1) :],
            skip_special_tokens=True
        )

    def make_prompt(self, message, chat_history, actor, max_context_size: int = 10, max_token_length = 128, temp = 0.5):
        prompt = ""
        if len(actor) > 0:
            prompt = "設定:" + actor + '\n'

        for u,s in chat_history:
            prompt += "ユーザー:" + u + '\n'
            prompt += "システム:" + s + '\n'

        prompt += "ユーザー:" + message
        prompt += "システム:"

        print("[prompt]" + prompt)

        generated = self.generate(prompt,max_token_length, temp)

        print('[generated]' + generated)

        return generated


    def interact_func(self, message, chat_history, actor, max_context_size, max_length, sampling_temperature):
        generated = self.make_prompt(message,chat_history,actor,max_context_size,max_length, sampling_temperature)
        chat_history.append((message,generated))
        return "", chat_history
    
    def run(self):
        # load model
        self.prepare_pretrained()
        self.prepare_model()

        with gr.Blocks() as demo:
            with gr.Accordion("Configs", open=False):
                # max_context_size = the number of turns * 2
                max_context_size = gr.Number(value=5, label="max_context_size", precision=0)
                max_length = gr.Number(value=128, label="max_length", precision=0)
                sampling_temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="temperature")
            actor = gr.Textbox(label="設定")
            chatbot = gr.Chatbot()
            msg = gr.Textbox(label='入力')
            print(type(msg))
            clear = gr.Button("Clear")
            msg.submit(
                self.interact_func,
                [msg, chatbot, actor, max_context_size, max_length, sampling_temperature],
                [msg, chatbot],
            )
            clear.click(lambda: None, None, chatbot, queue=False)

        demo.launch(debug=True, share=True)


    def test(self):
        prompt = """設定: 次の日本語を英語に翻訳してください。
    ユーザー: 自然言語による指示に基づきタスクが解けるよう学習させることを Instruction tuning と呼びます。
    システム:"""
        print(self.generate(prompt))

if __name__ == '__main__':
    modelname = "rinna/youri-7b-chat-gptq"

    if len(sys.argv) > 1:
        if sys.argv[1] == "download":
            from huggingface_hub import snapshot_download
            snapshot_download(repo_id=modelname,revision="main")

    else:
        chatbot = ChatUI(modelname=modelname)
        chatbot.run()

プログラムの解説

dockerのイメージに含めるために、huggingface_hub.snapshot_download を使ってビルド時にダウンロードさせておきます。ファイルは、/root/.cache/huggingface/hub に保存されます。
make_prompt() で、チャット用に変換します。「設定：」は、事前に入力されている状態。「ユーザー：」は、ユーザーからの入力。「システム：」はLLMからの回答

docker ビルド

以下のコマンドで、dockerのイメージを作成します

docker build -t youri7b:latest .

ビルド名は適宜変えてください。実行をすると、huggingfaceからsnapshotの4GB程度をダウンロードするので、接続によっては時間がかかります。

docker の実行

イメージを作成が完了すると、実行できるようになります

docker run -d --rm --gpus all -p 7860:7860 --net host youri7b:latest

引数の説明

-d デーモンとして起動します
--rm 実行後に使用したコンテナを廃棄します
--gpus all nvidiaを実行するコンテナを呼び出します.all は、すべてのGPUを使う設定
-p 7860:7860 gradioのデフォルトのポートで、ポートフォワードの設定
--net host ネットワークとしてホストと同じリソースを利用
youri7b:lastest は、ビルド時に指定したビルド名と同じにします

これで起動すると、Dockerfile の最終行の python3 chat.py が実行されます。

動作確認

ブラウザを立ち上げて、http://localhost:7860 にアクセスします。

設定は、事前の設定を書いておきます。例では、「次の文書を英語に翻訳しなさい。」と役割を指定しています。
入力は、ここでは翻訳したい文書を入力して、最後にリターンを押すと入力されます。
結果がchat の欄に表示されます。

例２

設定に、役割としてユーザーサポートと、商品に関する事前知識を与えておきます。
質問として、記載に関連することを「入力」に入れています

終了する場合には、デーモンモードで動作しているので、docker psで番号を調べて

docker ps

調べてCONTAINER IDの番号を docker kill に指定

docker kill <CONTAINER ID>

既存の問題

しばらく遊んでいるとGPUメモリが足りなくなったとエラーが出て止まります
Win11 + wsl2 + docker on ubuntu + nvidia-container-toolkit で環境を作って、--runtime=nvidia オプションで動作させることができました。

docker run -d --rm --runtime=nvidia --gpus all -p 7860:7860 --net host youri7b:latest