Taste of Tech Topics

Acroquest Technology株式会社のエンジニアが書く技術ブログ

PDFドキュメントを画像のまま検索できるColQwen2でマルチモーダル検索を試す

こんにちは。データサイエンスチームYAMALEXの@Ssk1029Takashiです。
完全に年末ですが、最後まで油断せずに年を越したいですね。

このブログは、LLM・LLM活用 Advent Calendar 2024の25日目になります。
qiita.com

今年は1年間RAGを扱う多くの事例に触れさせていただきましたが、どこでも難しいのが図や画像が使われているドキュメントですね。
ただのPDFであれば、pdfminerなどのライブラリを使用して文字を読み取り検索すればよいですが、画像が使われるとそう簡単にもいきません。
画像部分からはOCRを使って文字を取り出すなど、工夫が必要になります。

今回はドキュメントを画像にした状態で検索可能にするColQwen2というモデルを使って日本語ドキュメントの検索を試してみます。

ColQwen2とは

手法自体はColPaliというモデルがもとになっています。
ColPaliは画像も扱えるVMLとColBERTを組み合わせたモデルになります。
arxiv.org
ColBERTとは詳細は省きますが、BERTを使って効率的に精度の良い文章のEmbeddingを作成する手法です。
arxiv.org

ColPaliはEmbeddingを計算するモデルに、画像に対応したVision LLMを使用することで、LLMの画像解釈能力を生かしつつ、テキストと画像を同じベクトル空間に表現できます。
ColPaliの段階ではモデルとしてPaliGemmaが使用されていたので、日本語には対応していませんでしたが、ColQwen2ではQwen2-VLが使用されているため、多言語対応の一部として日本語でも使用できます。

https://huggingface.co/vidore/colpali

実験

今回は製造業を巡る現状と課題 今後の政策の方向性という経済産業省の一般公開されている資料を基に検索を試してみます。
流れとしては以下になります。

  1. pdfをページごとの画像に変換する
  2. 各画像をColQwen2を使ってベクトルに変換する
  3. テキストも同じくベクトルに変換して、類似度が高いページを検索する

また、基本的なサンプルコードは以下のページで紹介されているため本記事では部分的に重要なコードを紹介します。
huggingface.co

環境

今回は以下の環境で実行しています。

項目
OS Ubuntu24.04
Python 3.11.9
GPU RTX4090

GPUはRTX4090を使っていますが、推論だけであればVRAMは8GB程度あれば実行可能です

PDFドキュメントのEmbeddingを作成する

それでは実際に実装に入っていきます。
colpaliやcolqwen2を使うにはcolpali_engineというライブラリが便利です。
PDFの画像変換、ベクトル作成は以下のコードで実行できます。

# PDFを画像に変換
images = convert_from_path(pdf_path)

# プロセッサ作成
processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v0.1")

# モデルの読み込み
model = ColQwen2.from_pretrained(
    "vidore/colqwen2-v0.1", 
    torch_dtype=torch.bfloat16, 
    device_map="cuda"
)

# 2ページずつ読み込んで変換する
dataloader = DataLoader(
    images,
    batch_size=2, 
    shuffle=False, 
    collate_fn=lambda x: processor.process_images(x)
)

embedding = []
for batch_doc in tqdm(dataloader):
    with torch.no_grad(): 
        batch_doc = {k: v.to("cuda") for k, v in batch_doc.items()} 
        embeddings_doc = model(**batch_doc) 
    embedding.extend(list(torch.unbind(embeddings_doc))) 

# ベクトル情報を保存
torch.save(embedding, "embedding.pt")

この結果としてembedding.ptというファイルが生成されれば完了です。

検索を試してみる

それでは例として「国製造業の売上高は?」というクエリで検索を試してみます。
流れとしては、事前に作ったベクトルをロードした後、クエリテキストから作成したベクトルと最も類似度が高いページを計算しています。

query = "国製造業の売上高は?",

# モデルの読み込み
model = ColQwen2.from_pretrained(
    "vidore/colqwen2-v0.1", 
    torch_dtype=torch.bfloat16, 
    device_map="cuda"
)

# プロセッサの作成
processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v0.1")

# クエリのベクトルを作成
processed_query = processor.process_queries([query]).to(model.device)
with torch.no_grad():
    query_embedding = model(**processed_query)

# 事前に作成したPDFのベクトルをロード
embedding = torch.load("./embedding.pt", weights_only=True)

# クエリのベクトルとPDFのベクトルで類似度を計算
scores = processor.score_multi_vector(query_embedding, embedding)[0]

# トップ5のページを取得
scores_indices = scores.argsort().tolist()[-5:][::-1]
print([x + 1 for x in scores_indices])

上記を実行すると以下の結果を得られます。

[5, 7, 10, 6, 12]

トップはP5となっているので、P5の内容は以下のようになっています。

これを見ると売上高について言及しているページを取れているようです。
また、上記の場合は文字情報から類似度が高いページを選んでいましたが、それ以外の情報ではどうでしょうか?
例えば、「世界地図」というワードで検索してみましょう。
結果としては以下のページが出力されます。

[66, 65, 56, 116, 69]

P66は以下のようになっています。

「世界地図」というワードはページに含まれていませんが、世界地図が載っているページを取得できています。
つまり、Qwen2-VLモデルがページに世界地図が載っていることを認識できるため、ベクトルに反映されているようです。

このように、ColQwen2を使うことでドキュメントに載っているテキストやグラフの情報をもとに、複雑な処理をせずとも同じベクトル空間で検索可能になります。

まとめ

今回はColQwen2を使ってPDFドキュメントを画像として検索する手法を試しました。
複雑なデータ前処理なく載っているテキストや図表の情報をまとめて検索できるのはかなり便利そうです。
テキストと画像を同じベクトル空間で検索する手法はCLIPなどもあるのでそれらとの比較してみるのも面白そうです。
マルチモーダル検索は今後も注目の領域ですね。
それではまた。

Acroquest Technologyでは、キャリア採用を行っています。

  • Azure OpenAI/Amazon Bedrock等を使った生成AIソリューションの開発
  • ディープラーニング等を使った自然言語/画像/音声/動画解析の研究開発
  • マイクロサービス、DevOps、最新のOSSクラウドサービスを利用する開発プロジェクト
  • 書籍・雑誌等の執筆や、社内外での技術の発信・共有によるエンジニアとしての成長

少しでも上記に興味を持たれた方は、是非以下のページをご覧ください。

www.wantedly.com