license: apache-2.0 datasets:

• openbmb/VisRAG-Ret-Train-In-domain-data
• openbmb/VisRAG-Ret-Train-Synthetic-data language:
• en base_model:
• openbmb/MiniCPM-V-2 tags:
• VisRAG pipeline_tag: feature-extraction

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

• 📖 简介 • 🎉 动态 • ✨ 系统架构 • ⚡️ 训练

• 📦 环境要求 • 🔧 使用指南 • 📄 许可协议 • 📑 引用文献 • 📧 联系我们

📖 简介

VisRAG 是一项基于视觉语言模型（VLM）的创新性检索增强生成（RAG）流程。与传统先解析文档获取文本的方案不同，该系统直接将文档作为图像通过VLM进行嵌入表征，继而通过检索增强VLM的生成效果。相较于传统基于文本的RAG方案，VisRAG 能最大限度保留并利用原始文档中的数据信息，规避解析过程导致的信息损耗。

🎉 动态

• 20241104: 在Hugging Face Space发布系统演示空间
• 20241031: 在Colab发布交互式教程
• 20241015: 在Hugging Face发布训练与测试数据集，可通过页面顶部VisRAG集合获取
• 20241014: 在arXiv发布研究论文，在Hugging Face发布模型权重，在GitHub开源项目代码

✨ 系统架构

检索模块

VisRAG-Ret 是基于MiniCPM-V 2.0构建的文档嵌入模型，该视觉语言模型融合了SigLIP视觉编码器与MiniCPM-2B语言模型。

生成模块

研究中我们采用MiniCPM-V 2.0、MiniCPM-V 2.6和GPT-4o作为生成器。实际应用中您可自由选择任意视觉语言模型！

⚡️ 训练

检索模块

我们为VisRAG-Ret构建了包含362,110组查询-文档对的训练数据集，其中34%来自公开学术数据集，66%为基于网络爬取PDF文档页面并通过VLM（GPT-4o）生成伪查询构建的合成数据。该数据集可通过页面顶部Hugging Face的VisRAG集合获取。

生成模块

生成模块未进行微调，直接使用现成的LLMs/VLMs进行生成。

📦 环境要求

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 = ["狗的样貌特征是什么？"]
INSTRUCTION = "为检索相关文档表征该查询："
queries = [INSTRUCTION + query for query in queries]

print("正在下载图像...")
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("图像下载完成")

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

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

📄 许可协议

• 本仓库代码遵循Apache-2.0开源协议
• VisRAG-Ret模型权重使用需严格遵守MiniCPM模型许可条款
• VisRAG-Ret模型及权重完全开放学术研究用途。填写登记问卷后，亦可免费用于商业场景

📑 引用文献

@misc{yu2024visragvisionbasedretrievalaugmentedgeneration,
      title={VisRAG：基于视觉的多模态文档检索增强生成系统}, 
      author={余诗、唐超越、徐博凯、崔峻博、冉俊浩、严宇坤、刘政昊、王硕、韩旭、刘知远、孙茂松},
      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