Visrag Ret

🚀 VisRAG：多模态文档上基于视觉的检索增强生成

VisRAG 是一种基于新型视觉语言模型（VLM）的检索增强生成（RAG）管道。它直接将文档作为图像进行嵌入，避免了传统文本解析过程中的信息损失，能最大程度保留和利用原始文档中的数据信息。

🚀 快速开始

VisRAG 是一个创新的基于视觉语言模型（VLM）的 RAG 管道。在这个管道中，它直接将文档作为图像使用 VLM 进行嵌入，然后进行检索以增强 VLM 的生成能力，而不是先解析文档以获取文本。与传统的基于文本的 RAG 相比，VisRAG 最大程度地保留和利用了原始文档中的数据信息，消除了解析过程中引入的信息损失。

✨ 主要特性

VisRAG-Ret

VisRAG-Ret 是一个基于 MiniCPM-V 2.0 构建的文档嵌入模型。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%）和一个合成数据集（66%）组成。合成数据集由网络爬取的 PDF 文档页面组成，并通过 VLM 生成（GPT - 4o）的伪查询进行增强。你可以在 Hugging Face 上的 VisRAG 集合中找到它，该集合在本页面开头有引用。

VisRAG-Gen

生成部分不使用任何微调；我们直接使用现成的大语言模型/视觉语言模型进行生成。

📄 许可证

• 本仓库中的代码遵循 Apache - 2.0 许可证发布。
• VisRAG - Ret 模型权重的使用必须严格遵循 MiniCPM 模型许可证.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}, 
}

📧 联系信息

• 喻石：yus21@mails.tsinghua.edu.cn
• 唐超越：tcy006@gmail.com

🎉 最新消息

• 2024 年 11 月 4 日：我们在 Hugging Face Space 上发布了 VisRAG 管道。
• 2024 年 10 月 31 日：我们在 Colab 上发布了 VisRAG 管道。
• 2024 年 10 月 15 日：我们在 Hugging Face 上发布了训练数据和测试数据，你可以在 Hugging Face 上的 VisRAG 集合中找到它们。该集合在本页面开头有引用。
• 2024 年 10 月 14 日：我们在 arXiv 上发布了 论文。在 Hugging Face 上发布了 模型。在 GitHub 上发布了 代码。

• 📖 介绍 • 🎉 最新消息 • ✨ VisRAG 管道 • ⚡️ 训练

• 📦 依赖项 • 🔧 使用方法 • 📄 许可证 • 📑 引用 • 📧 联系信息