Visrag Ret

🚀 VisRAG: マルチモーダルドキュメントにおけるビジョンベースの検索拡張生成

VisRAGは、新しいビジョン言語モデル（VLM）ベースの検索拡張生成（RAG）パイプラインです。このパイプラインでは、ドキュメントを最初に解析してテキストを取得するのではなく、ドキュメントを画像としてVLMを使用して直接埋め込み、その後検索してVLMの生成を強化します。従来のテキストベースのRAGと比較して、VisRAGは元のドキュメント内のデータ情報を最大限に保持し、解析プロセス中に導入される情報損失を排除します。

• 📖 紹介 • 🎉 ニュース • ✨ VisRAGパイプライン • ⚡️ トレーニング

• 📦 必要条件 • 🔧 使用方法 • 📄 ライセンス • 📑 引用 • 📧 連絡先

🚀 クイックスタート

VisRAGは、新しいビジョン言語モデル（VLM）ベースのRAGパイプラインです。このパイプラインでは、ドキュメントを画像として直接埋め込み、その後検索してVLMの生成を強化します。

✨ 主な機能

VisRAG-Ret

VisRAG-Retは、MiniCPM-V 2.0をベースに構築されたドキュメント埋め込みモデルです。このビジョン言語モデルは、SigLIPをビジョンエンコーダとして、MiniCPM-2Bを言語モデルとして統合しています。

VisRAG-Gen

論文では、MiniCPM-V 2.0、MiniCPM-V 2.6、およびGPT-4oを生成器として使用しています。実際には、好きなVLMを使用することができます！

📦 インストール

必要なライブラリをインストールするには、以下のコマンドを実行します。

torch==2.1.2
torchvision==0.16.2
transformers==4.40.2
sentencepiece==0.1.99
decord==0.6.0
Pillow==10.1.0

💻 使用例

基本的な使用法

from transformers import AutoModel, AutoTokenizer
import torch
import torch.nn.functional as F
from PIL import Image
import requests
from io import BytesIO

def weighted_mean_pooling(hidden, attention_mask):
    attention_mask_ = attention_mask * attention_mask.cumsum(dim=1)
    s = torch.sum(hidden * attention_mask_.unsqueeze(-1).float(), dim=1)
    d = attention_mask_.sum(dim=1, keepdim=True).float()
    reps = s / d
    return reps

@torch.no_grad()
def encode(text_or_image_list):
    
    if (isinstance(text_or_image_list[0], str)):
        inputs = {
            "text": text_or_image_list,
            'image': [None] * len(text_or_image_list),
            'tokenizer': tokenizer
        }
    else:
        inputs = {
            "text": [''] * len(text_or_image_list),
            'image': text_or_image_list,
            'tokenizer': tokenizer
        }
    outputs = model(**inputs)
    attention_mask = outputs.attention_mask
    hidden = outputs.last_hidden_state

    reps = weighted_mean_pooling(hidden, attention_mask)   
    embeddings = F.normalize(reps, p=2, dim=1).detach().cpu().numpy()
    return embeddings

model_name_or_path = "openbmb/VisRAG-Ret"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name_or_path, torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
model.eval()

queries = ["What does a dog look like?"]
INSTRUCTION = "Represent this query for retrieving relevant documents: "
queries = [INSTRUCTION + query for query in queries]

print("Downloading images...")
passages = [
    Image.open(BytesIO(requests.get(
        'https://github.com/OpenBMB/VisRAG/raw/refs/heads/master/scripts/demo/retriever/test_image/cat.jpeg'
    ).content)).convert('RGB'),
    Image.open(BytesIO(requests.get(
        'https://github.com/OpenBMB/VisRAG/raw/refs/heads/master/scripts/demo/retriever/test_image/dog.jpg'
    ).content)).convert('RGB')
]
print("Images downloaded.")

embeddings_query = encode(queries)
embeddings_doc = encode(passages)

scores = (embeddings_query @ embeddings_doc.T)
print(scores.tolist())

🔧 技術詳細

VisRAG-Ret

VisRAG-Retのトレーニングデータセットは、362,110のクエリ - ドキュメント（Q - D）ペアで構成されています。これは、公開されている学術データセットのトレーニングセット（34％）と、ウェブクロールしたPDFドキュメントのページから構成され、VLM生成（GPT-4o）の疑似クエリで拡張された合成データセット（66％）で構成されています。これは、このページの冒頭で参照されているHugging FaceのVisRAGコレクションにあります。

VisRAG-Gen

生成部分では、微調整は行わず、既存の大規模言語モデル（LLM）/VLMを直接使用して生成を行います。

📄 ライセンス

• このリポジトリ内のコードは、Apache-2.0ライセンスの下で公開されています。
• VisRAG-Retモデルの重みの使用は、MiniCPM Model License.mdに厳密に従う必要があります。
• VisRAG-Retのモデルと重みは、学術研究には完全に無料です。"アンケート"に記入して登録すると、VisRAG-Retの重みは商用利用にも無料で利用できます。

📑 引用

@misc{yu2024visragvisionbasedretrievalaugmentedgeneration,
      title={VisRAG: Vision-based Retrieval-augmented Generation on Multi-modality Documents}, 
      author={Shi Yu and Chaoyue Tang and Bokai Xu and Junbo Cui and Junhao Ran and Yukun Yan and Zhenghao Liu and Shuo Wang and Xu Han and Zhiyuan Liu and Maosong Sun},
      year={2024},
      eprint={2410.10594},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2410.10594}, 
}

📧 連絡先

• Shi Yu: yus21@mails.tsinghua.edu.cn
• Chaoyue Tang: tcy006@gmail.com