🌌

マルチモーダルな埋め込みモデル「E5-V」をM2 Pro Mac miniで動かしてみた

2024/12/26に公開

はじめに

テキストの埋め込みモデルは、ハイブリッド検索やRAGなどで活用されています。2024年には日本語に特化したRuriという埋め込みモデルが公開されました。(Ruriを使ってハイブリッド検索をやってみたという記事をQiitaに投稿しているので、良かったら読んでくれると嬉しいです)
また、画像の分野ではCNN (Convolutional Neural Networks)をはじめとするモデルで埋め込みベクトルに変換することはよくやられています。

マルチモーダルな埋め込みモデルは、画像とテキストを同じベクトル空間に埋め込むことができるモデルです。今回は、2024年7月に公開された論文「E5-V: Universal Embeddings with Multimodal Large Language Models」で提案されているE5-VをM2 Pro Mac miniで動かしてみました。

E5-Vについて

画像とテキストを同じベクトル空間に埋め込むモデルとしてはCLIPやSigLIPがあります。E5-Vは画像も理解する言語モデルであるLLaVA-NeXT-8Bをファインチューニングしたものです。プロンプトベースで画像とテキストを同じベクトル空間に埋め込むことができるようです。

論文より: E5-Vはモダリティのギャップがなくなっている(右)

M2 Pro Mac miniで試す

今回使用しているMac miniのスペックは以下の通りです。

  • Apple M2 Pro
    • CPU: E-core: 4 / P-core: 8
    • GPU: 19-core
    • Nerual Engine: 16-core
  • RAM: 32GB

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

以下のパッケージをインストールします。transformersは4.45でないと動きませんでした。

pip install transformers==4.45 peft bitsandbytes

今回試したコード

今回はモデルのページにあるサンプルコードをほぼそのまま動かしてみました。Apple Silicon Macで動かしているので、PyTorchのデバイス指定を"cuda"ではなく"mps"に変更しています。
実際に使用したコードはGitHubにあります。: Shakshi3104/e5-v-supplementary

⚠️ E5-Vは8Bのパラメータ数を持っており、ファイルサイズとしては約16GBなので、モデルダウンロードには時間がかかります。

import torch
import torch.nn.functional as F
from transformers import  LlavaNextProcessor, LlavaNextForConditionalGeneration

from PIL import Image


llama3_template = '<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n'

model_name ='royokong/e5-v'

processor = LlavaNextProcessor.from_pretrained(model_name)
model = LlavaNextForConditionalGeneration.from_pretrained(model_name, torch_dtype=torch.float16).to("mps")

processor.patch_size = model.config.vision_config.patch_size
processor.vision_feature_select_strategy = model.config.vision_feature_select_strategy

img_prompt = llama3_template.format('<image>\nSummary above image in one word: ')
text_prompt = llama3_template.format('<sent>\nSummary above sentence in one word: ')

# 試す画像とテキスト
images = [...]
texts = [...]

img_inputs = processor([img_prompt] * len(images), images, return_tensors="pt", padding=True).to('mps')
text_inputs = processor([text_prompt.replace('<sent>', text) for text in texts], return_tensors="pt",
                        padding=True).to('mps')

with torch.no_grad():
    img_embs = model(**img_inputs, output_hidden_states=True, return_dict=True).hidden_states[-1][:, -1, :]
    text_embs = model(**text_inputs, output_hidden_states=True, return_dict=True).hidden_states[-1][:, -1, :]

    text_embs = F.normalize(text_embs, dim=-1)
    img_embs = F.normalize(img_embs, dim=-1)

print(text_embs @ img_embs.t())

試してみた画像とテキスト

以前の記事で紹介したsdxl-emojiで生成した画像と生成する際に使用したプロンプトで試してみました。

画像

使用した画像は以下の3枚です。

テキスト

使用したテキストは以下の3つです。

  • A baby tiger with tiramisu
  • A strawberry parfait
  • A christmas wreath

結果

画像とテキストのコサイン類似度をヒートマップ形式で出してみました。対角線上にある生成画像とプロンプトの類似度が高く、それ以外の類似度が低いということが分かります。このことから、E5-Vで画像とテキストが同じベクトル空間へ埋め込みができそうだと感じました。

メモリ使用量の観察

macOSのアクティビティモニタで実行中のメモリ使用量を観察してみました。実行中は最大でも26GBくらいまで使用量が増えますが、メモリに乗り切っているっぽいので問題なく動作します。また、40秒程度で、3つの画像と3つのテキストを埋め込みベクトルに変換してコサイン類似度を計算するところまで動作できていました。

おわりに

今回は、マルチモーダルな埋め込みモデルであるE5-VをM2 Pro Mac miniで動かしてみました。Apple Silicon MacはUnified Memory Architectureであるため、メモリをフルで活用してそれなりに大きなサイズのモデルを動かすことができるが良いと思いました。

E5-Vでの画像とテキストの類似度を見ると、確かに画像とテキストで同じベクトル空間に埋め込まれている印象です。こういったマルチモーダルなモデルを使うことで、画像も検索できるベクトル検索も実現できそうです。

参考

Discussion