SageMaker Studio Lab と 生成系 AI の OSS で始める画像に対する日本語での対話型問合せ

背景

生成系 AI の応用の幅が広がる技術としてマルチモーダルなモデルがあります。マルチモーダルではモデルの入出力に複数の異なるデータ形式を用いることができます。例えば Amazon Bedrock の Stable Diffusion ではテキストを入力に画像を生成することができます。画像とテキストという異なるデータ形式の入出力をするためマルチモーダルなモデルといえます。他にも入力した画像の説明文章をテキストとして出力できる BLIP-2 (Bootstrapping Language-Image Pre-training) と呼ばれるモデルもあり、AWS では Amazon SageMaker の推論 Endpoint にデプロイすることで利用することができます。詳しくはこちらのブログをご覧ください。さらに、BLIP-2 は画像に対しキャプションをつけるタイプのモデルでしたが、対話形式で画像に対する問い合わせをすることができる MiniGPT-4 が生まれました。MiniGPT-4 は BLIP-2 の ViT (Vision Transformer) 、 Q-Former (Querying Transformer) と LLM (Large Language Models) を線形結合層で繋ぐことで実現されます。LLM を切り替える場合はこの線形結合層のみを再学習するため BLIP-2 や LLM 自体の再学習をしない分少ないパラメータを対象に再学習できるメリットがあります。

これらのように画像を入力にテキストを生成するモデルを Image to Text と呼びます。Image to Text のモデルを使えば、画像に対して説明を求めたり、QA することができるようになります。これを応用して、画像にタグを付与したり、テキストを入力に画像を検索したりできるようになります。従来の物体検出のような Computer Vision 技術で取得できた画像情報とは異なる情報を柔軟に取得できる可能性があります。どれくらいの情報が取得できるのか是非試したいですよね。

画像へ対話形式の問合せを日本語で試したい！と思っても、MiniGPT-4 は英語のモデルのため日本語には難があるかもしれません。そこに、MiniGPT-4 の線形結合層を日本語データで再学習した OSS の 日本語 LLM が発表されました。それが rinna株式会社の bilingual-gpt-neox-4b-minigpt4 です。今回は無料の Jupyter ノートブックサービスである SageMaker Studio Lab から bilingual-gpt-neox-4b-minigpt4 を利用して画像に対する QA に日本語で挑戦してみたいと思います。

制限事項

執筆時点で SageMaker Studio Lab には CPU と GPU いずれも 1 日あたり利用時間上限が決められています。また、閉域接続ではなくインターネットでご利用いただくサービスになります。好評いただいているサービスのため、特に GPU の利用時に一度では Start runtime できず複数回トライいただく場合があります。これらの特徴と上手にお付き合いいただきご利用いただければ嬉しいです。

前準備

SageMaker Studio Lab は AWS アカウント不要の無料のノートブックサービスです。AWS アカウントがなくても利用でき、メールアドレスがあれば登録することができます。こちらのブログをご覧いただき登録を完了いただけると以降の手順がスムーズです。SageMaker Studio Lab は 機械学習帳とも連携しており、SageMaker Studio Lab を使ってすぐに機械学習のスキル獲得を始めていただく事ができます。

sagemaker-distribution を利用して永続化領域を節約

本ブログでは公式の environment.yml ではなく、sagemaker-distribution カーネルを利用する方法をお伝えします。これにより、SageMaker Studio Lab の永続化領域を有効利用することができ、SageMaker Studio への移行も簡単にすることができます。sagemaker-distribution は PyTorch、TensorFlow、Keras などの人気のあるライブラリがプリインストールされている環境を提供します。SageMaker Studio と SageMaker Stdio Lab の両方に互換性のある Pythonライブラリが 18 種類以上付属しています。この環境は永続的であり、SageMaker Studio Lab 利用者に割り当てられた 15 GB の空き容量を使用しません。詳しくはこちらをご覧ください。sagemaker-distribution カーネルは 2023 年 8 月に SageMaker Studio Lab で利用できるようになりました。

bilingual-gpt-neox-4b-minigpt4 のダウンロードと必要なライブラリのインストール

こちらの作業は CPU モードで実行することで GPU 利用時間を節約しましょう。CPU と GPU の切り替えはログインページのラジオボタンからの切り替えになります。ノートブックの実行画面にはないため注意してください。公式のダウンロード手順に従い必要なファイルをダウンロードします。SageMaker Studio Lab から Terminal を起動し、以下を実行します。Terminal は画面左上のプラスボタンから Launch タブを開き、Terminal を選択することで起動できます。Blog 執筆時点では以下のコマンドとなっていました。最新のコマンドはこちらの URLからご確認ください。

git clone https://github.com/Vision-CAIR/MiniGPT-4.git
cd ./MiniGPT-4
git checkout 22d8888 # latest version as of July 31, 2023.
wget https://huggingface.co/rinna/bilingual-gpt-neox-4b-minigpt4/resolve/main/customized_mini_gpt4.py
wget https://huggingface.co/rinna/bilingual-gpt-neox-4b-minigpt4/resolve/main/checkpoint.pth

ダウンロードが終わったことを確認したら、元のモデルファイルを /tmp/model_data に保存するように customized_mini_gpt4.py に変更を加えます。元のままの場合、ユーザの永続化領域にモデルファイルが保存され容量を圧迫してしまうためです。111 行目以降の from_pretreined の引数に cache_dir = ‘/tmp/model_cache/’ を追加し、以下のように変更します。永続化領域を節約する効果はこちらをご覧ください。画面左側のファイル一覧からダブルクリックで customized_mini_gpt4.py の編集画面を開くことができます。

        if self.low_resource:
            self.gpt_neox_model = CustomizedGPTNeoXForCausalLM.from_pretrained(
                gpt_neox_model,
                torch_dtype=torch.float16,
                load_in_8bit=True,
                device_map={'': device_8bit},
                cache_dir = '/tmp/model_cache/'
            )
        else:
            self.gpt_neox_model = CustomizedGPTNeoXForCausalLM.from_pretrained(
                gpt_neox_model,
                torch_dtype=torch.float16,
                cache_dir = '/tmp/model_cache/'
            )

それでは こちらのページを参考にノートブックを作成していきましょう。作成された MiniGPT-4 ディレクトリに sample.ipynb を作成します。以降の作業は全て sample.ipynb に実装します。

画面左上のプラスボタンから sagemaker-distribution: Python を選択いただくとノートブックファイルを作成できます。ファイル名を sample.ipynb に変更してください。

作業中、sagemaker-distribution カーネルが選択されているかを確認する場合は、右上のカーネル選択にてsagemaker-distribution が選択されているかをご覧ください。

次に以下のコマンドをセルで実行してください。

!conda install -y -c conda-forge opencv
!pip install omegaconf
!pip install iopath
!pip install timm
!pip install webdataset
!pip install transformers
!pip install decord
!pip install sentencepiece

これで環境は整いました。

ノートブックで Image QA にトライ

ここからは GPU モードで実行しましょう。移行の実装は公式の I/O Format と How to use the model のソースコードをセルに分割して書き下したものになります。
まずは、必要なモジュールを import します。以下のコードをセルに貼り付けて実行してください。

import torch
import requests
from PIL import Image
from minigpt4.processors.blip_processors import Blip2ImageEvalProcessor
from customized_mini_gpt4 import CustomizedMiniGPT4

GPU が利用可能かを以下のコマンドで確認しておきましょう。以下のコードをセルに貼り付けて実行してください。

torch.cuda.is_available()

GPU モードで SageMaker Studio Lab を起動していても、torch の version と CUDA Driver が合わない場合、上記の結果は False が返ります。この場合、以降のコードは CPU で実行され期待動作しない場合がありますため注意が必要です。sagemaker-distribution カーネルを使えばこの問題は起きません。SageMaker-Distibution カーネルには SageMaker Studio Lab の GPU が使用できるバージョンの torch が予めプリインストールされているためです。もし False が返ってきた場合指定したカーネルが誤っている可能性がありますのでご確認ください。

次に、モデルを準備します。以下のコードをセルに貼り付けて実行してください。CustomizedMiniGPT4 はダウンロード手順で取得した customized_mini_gpt4.py の中に実装があります。興味がある方は是非確認してみてください。checkpoint.pth は再学習した線型結合層のモデルファイルです。こちらも先ほどの手順でダウンロード済です。

model = CustomizedMiniGPT4(gpt_neox_model="rinna/bilingual-gpt-neox-4b")
tokenizer = model.gpt_neox_tokenizer
if torch.cuda.is_available():
    model = model.to("cuda")

ckpt_path = "./checkpoint.pth"
if ckpt_path is not None:
    print("Load BLIP2-LLM Checkpoint: {}".format(ckpt_path))
    ckpt = torch.load(ckpt_path, map_location="cpu")
    model.load_state_dict(ckpt['model'], strict=False)

それでは、問合せ対象の画像を準備しましょう。以下のコードをセルに貼り付けて実行してください。ここでは huggingface にあるサンプル画像を利用します。猫が横たわっている画像が表示されれば成功です。

image_url = "https://huggingface.co/rinna/bilingual-gpt-neox-4b-minigpt4/resolve/main/sample.jpg"
raw_image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
raw_image

BLIP-2 を利用して画像をエンコードします。画像を数字列に変換する Embedding と呼ばれる手法です。この値とテキストを入力することで画像への問合せが可能となります。以下のコードをセルに貼り付けて実行してください。

vis_processor = Blip2ImageEvalProcessor()
image = vis_processor(raw_image).unsqueeze(0).to(model.device)
image_emb = model.encode_img(image)

以下は公式のサンプルプロンプトです。
に先ほどの Embedding した値が入ります。ユーザー/システムで話者を識別しています。以下のコードをセルに貼り付けて実行してください。prompt の中身が表示されれば成功です。

prompt = [
	{
		"speaker": "ユーザー",
		"text": "<Img><ImageHere></Img> 何が写っていますか？"
	},
	{
		"speaker": "システム",
		"text": "a cat on a table with a laptop"
	},
	{
		"speaker": "ユーザー",
		"text": "猫はどんな体勢をしていますか？"
	},
]
prompt = [
	f"{uttr['speaker']}: {uttr['text']}"
	for uttr in prompt
]
prompt = "\n".join(prompt)
prompt = (
    prompt
    + "\n"
    + "システム: "
)
print(prompt)

画像の Embedding に加えて、テキストも Embedding して、それらを組み合わせる処理をします。詳しくは、 customized_mini_gpt4.py の get_context_emb をご覧ください。以下のコードをセルに貼り付けて実行してください。

embs = model.get_context_emb(prompt, [image_emb])

さあ、最後に回答となるテキストを生成するコードです。output に回答テキストが入ります。以下のコードをセルに貼り付けて実行してください。結果はどうなるでしょうか？

output_ids = model.gpt_neox_model.generate(
    inputs_embeds=embs,
    max_new_tokens=512,
    do_sample=True,
    temperature=1.0,
    top_p=0.85,
    pad_token_id=tokenizer.pad_token_id,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id
)

output = tokenizer.decode(output_ids.tolist()[0], skip_special_tokens=True)
print(output)

出力例:

うつ伏せ

以下に、著者が試してみた画像に対する対話型での QA を記載します。是非、皆さんのアイデアで試してみてください。

ユーザー: <Img><ImageHere></Img> 何が写っていますか？
システム: 猫がテーブルの上のパソコンを眺めている
ユーザー: 猫は何色？
システム: 白
ユーザー: 茶色に見えますが、白ですか?
システム: 茶色
ユーザー: 猫可愛いですね。この猫の隣にあるのは何ですか?
システム: パソコン
ユーザー: このシーンの次には何が起きると思いますか?
システム: パソコンが壊れる

最後に

いかがでしたか？本ブログでは、日本語対応の OSS のモデルを利用して、画像に対して対話型で QA することに挑戦してみました。画像を変えてみたり、プロンプトを変えてみたりして結果を確認してみてください。また、bilingual-gpt-neox-4b-minigpt4 の登場によって、日本語 LLM と学習用データを準備すれば線形結合層の再学習のみで日本語対応できる MiniGPT-4 が作れることが示されました。今後、登場する日本語 LLM の OSS を利用して皆さん独自のデータで再学習したい場合にも適用できる方式になります。是非、そちらにも挑戦いただければ嬉しいです。著者もどこかで実験し Blog にできればと考えています。

著者

中島 佑樹 西日本のお客様をメインで担当するソリューションアーキテクト。社会人博士を修了したことをきっかけに AIML を得意分野としている。 システム一般のテーマや Amazon Bedrock を用いた生成系 AI のシステム開発、Amazon SageMaker Studio Lab を用いた AIML への入門まで幅広く活動。