language:

• 英文
• 法文 license: apache-2.0 library_name: Tevatron tags:
• vidore datasets:
• Tevatron/docmatix-ir
• HuggingFaceM4/Docmatix
• Tevatron/msmarco-passage-aug
• vidore/colpali_train_set
• Tevatron/wiki-ss-nq base_model:
• Qwen/Qwen2-VL-2B-Instruct pipeline_tag: 视觉文档检索

DSE-QWen2-2b-MRL-V1

DSE-QWen2-2b-MRL-V1 是一个双编码器模型，专为将文档截图编码为密集向量以实现文档检索而设计。 文档截图嵌入（DSE）方法以原始视觉格式捕获文档，保留所有信息，如文本、图像和布局，从而避免繁琐的解析和潜在的信息丢失。 DSE 旨在为文本、PDF 文档、网页和幻灯片检索提供一个通用的嵌入模型。

例如，DSE-QWen2-2b-MRL-V1 在 ViDoRE 排行榜上实现了 85.8 的 nDCG@5。

注意：

如果输入图像较大，QWen 视觉编码器可能会占用较高的 GPU 内存。根据 GPU 资源调整 'resized_height':680 , 'resized_width':680（见下文）以适应显存。

如何使用模型

为了支持更好的效果与效率权衡，此检查点训练支持：

1. 灵活的表示维度。
2. 灵活的输入图像尺寸。

加载模型和处理器

import torch
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
from qwen_vl_utils import process_vision_info

min_pixels = 1*28*28
max_pixels = 2560*28*28

processor = AutoProcessor.from_pretrained("MrLight/dse-qwen2-2b-mrl-v1", min_pixels=min_pixels, max_pixels=max_pixels)
model = Qwen2VLForConditionalGeneration.from_pretrained('MrLight/dse-qwen2-2b-mrl-v1', attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16).to('cuda:0').eval()
processor.tokenizer.padding_side = "left"
model.padding_side = "left"

def get_embedding(last_hidden_state: torch.Tensor, dimension: int) -> torch.Tensor:
    reps = last_hidden_state[:, -1]
    reps = torch.nn.functional.normalize(reps[:, :dimension], p=2, dim=-1)
    return reps

编码文本查询

from PIL import Image
queries = ["我们在哪里可以看到羊驼？", "LLaMA AI 模型是什么？"]
query_messages = []
for query in queries:
    message = [
        {
            'role': 'user',
            'content': [
                {'type': 'image', 'image': Image.new('RGB', (28, 28)), 'resized_height':1 , 'resized_width':1}, # 需要一个虚拟图像以便处理。
                {'type': 'text', 'text': f'查询: {query}'},
            ]
        }
    ]
    query_messages.append(message)
query_texts = [
    processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) + "<|endoftext|>"
    for msg in query_messages
]
query_image_inputs, query_video_inputs = process_vision_info(query_messages)
query_inputs = processor(text=query_texts, images=query_image_inputs, videos=query_video_inputs, padding='longest', return_tensors='pt').to('cuda:0')
cache_position = torch.arange(0, len(query_texts))
query_inputs = model.prepare_inputs_for_generation(**query_inputs, cache_position=cache_position, use_cache=False)
with torch.no_grad():
  output = model(**query_inputs, return_dict=True, output_hidden_states=True)
query_embeddings = get_embedding(output.hidden_states[-1], 1536) # 调整维度以权衡效率，例如 512

编码文档截图

import requests
from io import BytesIO

# 图像 URL
url1 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v2/resolve/main/animal-llama.png"
url2 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v2/resolve/main/meta-llama.png"

# 下载并打开图像
response1 = requests.get(url1)
response2 = requests.get(url2)

doc_image1 = Image.open(BytesIO(response1.content))
doc_image2 = Image.open(BytesIO(response2.content))

doc_images = [doc_image1, doc_image2]
doc_messages = []
for doc in doc_images:
    message = [
        {
            'role': 'user',
            'content': [
                {'type': 'image', 'image': doc}, #'resized_height':680 , 'resized_width':680} # 调整图像尺寸以权衡效率
                {'type': 'text', 'text': '这张图片展示了什么？'}
            ]
        }
    ]
    doc_messages.append(message)
doc_texts = [
    processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) + "<|endoftext|>"
    for msg in doc_messages
]
doc_image_inputs, doc_video_inputs = process_vision_info(doc_messages)
doc_inputs = processor(text=doc_texts, images=doc_image_inputs, videos=doc_video_inputs, padding='longest', return_tensors='pt').to('cuda:0')
cache_position = torch.arange(0, len(doc_texts))
doc_inputs = model.prepare_inputs_for_generation(**doc_inputs, cache_position=cache_position, use_cache=False)
with torch.no_grad():
    output = model(**doc_inputs, return_dict=True, output_hidden_states=True)
doc_embeddings = get_embedding(output.hidden_states[-1], 1536) # 调整维度以权衡效率，例如 512

计算相似度

from torch.nn.functional import cosine_similarity
num_queries = query_embeddings.size(0)
num_passages = doc_embeddings.size(0)

for i in range(num_queries):
    query_embedding = query_embeddings[i].unsqueeze(0)
    similarities = cosine_similarity(query_embedding, doc_embeddings)
    print(f"查询 {i+1} 的相似度: {similarities.cpu().float().numpy()}")

编码文档文本

此 DSE 检查点通过 Tevatron/msmarco-passage-aug 进行了预热，因此模型也可以有效地将文档编码为文本输入。

doc_texts = [
  "羊驼（学名：Lama glama）是一种驯化的南美洲骆驼科动物，自前哥伦布时代以来被安第斯文化广泛用作肉用和驮畜。",
  "Llama（Large Language Model Meta AI 的缩写，曾用名 LLaMA）是 Meta AI 自 2023 年 2 月起发布的一系列自回归大语言模型（LLMs）。最新版本是 2024 年 7 月发布的 Llama 3.1。"
]
doc_messages = []
for doc in doc_texts:
    message = [
        {
            'role': 'user',
            'content': [
                {'type': 'image', 'image': Image.new('RGB', (28, 28)), 'resized_height':1 , 'resized_width':1}, # 需要一个虚拟图像以便处理。
                {'type': 'text', 'text': f'文档: {doc}'}
            ]
        }
    ]
    doc_messages.append(message)
doc_texts = [
    processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) + "<|endoftext|>"
    for msg in doc_messages
]
doc_image_inputs, doc_video_inputs = process_vision_info(doc_messages)
doc_inputs = processor(text=doc_texts, images=doc_image_inputs, videos=doc_video_inputs, padding='longest', return_tensors='pt').to('cuda:0')
cache_position = torch.arange(0, len(doc_texts))
doc_inputs = model.prepare_inputs_for_generation(**doc_inputs, cache_position=cache_position, use_cache=False)
with torch.no_grad():
    output = model(**doc_inputs, return_dict=True, output_hidden_states=True)
doc_embeddings = get_embedding(output.hidden_states[-1], 1536) # 调整维度以权衡效率，例如 512

for i in range(num_queries):
    query_embedding = query_embeddings[i].unsqueeze(0)
    similarities = cosine_similarity(query_embedding, doc_embeddings)
    print(f"查询 {i+1} 的相似度: {similarities.cpu().float().numpy()}")

引用

如果您发现此检查点有帮助，请考虑引用 QWen2、Docmatix、ViDoRE 和我们的 DSE 工作。