284 lines
11 KiB
Markdown
284 lines
11 KiB
Markdown
---
|
|
license: apache-2.0
|
|
language:
|
|
- en
|
|
- it
|
|
- fr
|
|
- de
|
|
- es
|
|
base_model:
|
|
- MrLight/dse-qwen2-2b-mrl-v1
|
|
tags:
|
|
- transformers
|
|
- sentence-transformers
|
|
- Qwen2-VL
|
|
datasets:
|
|
- llamaindex/vdr-multilingual-train
|
|
---
|
|
|
|
# vdr-2b-multi-v1
|
|
|
|
|
|
|
|
vdr-2b-multi-v1 is a multilingual embedding model designed for visual document retrieval across multiple languages and domains. It encodes document page screenshots into dense single-vector representations, this will effectively allow to search and query visually rich multilingual documents without the need for any OCR, data extraction pipelines, chunking...
|
|
|
|
|
|
- **Trained on 🇮🇹 Italian, 🇪🇸 Spanish, 🇬🇧 English, 🇫🇷 French and 🇩🇪 German:** together they form a new large, open-source, multilingual training dataset of 500k high-quality samples.
|
|
|
|
- **Cross-lingual Retrieval**: substantially better on real-world scenarios. For example, this allows for searching german documents with italian queries.
|
|
|
|
- **Matryoshka Representation Learning**: You can reduce the vectors size 3x and still keep 98% of the embeddings quality.
|
|
|
|
# Usage
|
|
|
|
The model uses bf16 tensors and allocates ~4.4GB of VRAM when loaded. You can easily run inference and generate embeddings using 768 image patches and a batch size of 16 even on a cheap NVIDIA T4 GPU. This table reports the memory footprint (GB) under conditions of different batch sizes with HuggingFace Transformers and maximum 768 image patches.
|
|
|
|
| Batch Size | GPU Memory (GB) |
|
|
|------------|-----------------|
|
|
| 4 | 6.9 |
|
|
| 8 | 8.8 |
|
|
| 16 | 11.5 |
|
|
| 32 | 19.7 |
|
|
|
|
You can generate embeddings with this model in many different ways:
|
|
|
|
<details open>
|
|
<summary>
|
|
via LlamaIndex
|
|
</summary>
|
|
|
|
```bash
|
|
pip install -U llama-index-embeddings-huggingface
|
|
```
|
|
|
|
```python
|
|
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
|
|
|
model = HuggingFaceEmbedding(
|
|
model_name="llamaindex/vdr-2b-multi-v1",
|
|
device="cpu", # "mps" for mac, "cuda" for nvidia GPUs
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
image_embedding = model.get_image_embedding("image.png")
|
|
query_embedding = model.get_query_embedding("some query")
|
|
```
|
|
|
|
</details>
|
|
|
|
<details>
|
|
<summary>
|
|
via HuggingFace Transformers
|
|
</summary>
|
|
|
|
```python
|
|
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
|
|
from PIL import Image
|
|
import torch
|
|
import math
|
|
|
|
# more pixels -> better embeddings -> more VRAM -> slower inference
|
|
# From my experience, 768 image patches is the right spot for compute efficient embeddings.
|
|
max_pixels = 768 * 28 * 28
|
|
min_pixels = 1 * 28 * 28
|
|
|
|
# Load the embedding model and processor
|
|
model = Qwen2VLForConditionalGeneration.from_pretrained(
|
|
'llamaindex/vdr-2b-multi-v1',
|
|
# These are the recommended kwargs for the model, but change them as needed
|
|
attn_implementation="flash_attention_2",
|
|
torch_dtype=torch.bfloat16,
|
|
device_map="cuda:0"
|
|
).eval()
|
|
|
|
processor = AutoProcessor.from_pretrained(
|
|
'llamaindex/vdr-2b-multi-v1',
|
|
min_pixels=min_pixels,
|
|
max_pixels=max_pixels
|
|
)
|
|
|
|
model.padding_side = "left"
|
|
processor.tokenizer.padding_side = "left"
|
|
|
|
document_prompt = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is shown in this image?<|im_end|>\n<|endoftext|>"
|
|
|
|
query_prompt = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Query: %s<|im_end|>\n<|endoftext|>"
|
|
```
|
|
|
|
**Encode queries**
|
|
|
|
```python
|
|
def encode_queries(queries: list[str], dimension: int) -> torch.Tensor:
|
|
"""
|
|
Encode a list of queries into a tensor of embeddings.
|
|
|
|
Args:
|
|
queries: A list of strings, each representing a query.
|
|
dimension: The desired dimension of the output embeddings.
|
|
|
|
Returns:
|
|
A tensor of shape (num_queries, dimension) containing the encoded queries.
|
|
"""
|
|
|
|
dummy_image = Image.new('RGB', (56, 56))
|
|
inputs = processor(
|
|
text=[query_prompt % x for x in queries],
|
|
images=[dummy_image for _ in queries],
|
|
videos=None,
|
|
padding='longest',
|
|
return_tensors='pt'
|
|
).to('cuda:0')
|
|
|
|
cache_position = torch.arange(0, len(queries))
|
|
inputs = model.prepare_inputs_for_generation(
|
|
**inputs, cache_position=cache_position, use_cache=False)
|
|
|
|
with torch.no_grad():
|
|
output = self.model(
|
|
**inputs,
|
|
return_dict=True,
|
|
output_hidden_states=True
|
|
)
|
|
|
|
embeddings = output.hidden_states[-1][:, -1]
|
|
return torch.nn.functional.normalize(embeddings[:, :dimension], p=2, dim=-1)
|
|
```
|
|
|
|
**Encode documents**
|
|
|
|
```python
|
|
def round_by_factor(number: float, factor: int) -> int:
|
|
return round(number / factor) * factor
|
|
|
|
def ceil_by_factor(number: float, factor: int) -> int:
|
|
return math.ceil(number / factor) * factor
|
|
|
|
def floor_by_factor(number: float, factor: int) -> int:
|
|
return math.floor(number / factor) * factor
|
|
|
|
def smart_resize(height: int, width: int) -> tuple[int, int]:
|
|
h_bar = max(28, round_by_factor(height, 28))
|
|
w_bar = max(28, round_by_factor(width, 28))
|
|
if h_bar * w_bar > max_pixels:
|
|
beta = math.sqrt((height * width) / max_pixels)
|
|
h_bar = floor_by_factor(height / beta, 28)
|
|
w_bar = floor_by_factor(width / beta, 28)
|
|
elif h_bar * w_bar < min_pixels:
|
|
beta = math.sqrt(min_pixels / (height * width))
|
|
h_bar = ceil_by_factor(height * beta, 28)
|
|
w_bar = ceil_by_factor(width * beta, 28)
|
|
return w_bar, h_bar
|
|
|
|
def resize(image: Image.Image):
|
|
new_size = smart_resize(image.height, image.width)
|
|
return image.resize(new_size)
|
|
|
|
def encode_documents(documents: list[Image.Image], dimension: int):
|
|
"""
|
|
Encode a list of images into a tensor of embeddings.
|
|
|
|
Args:
|
|
documents: A list of PIL Image objects.
|
|
dimension: The desired dimension of the output embeddings.
|
|
|
|
Returns:
|
|
A tensor of shape (num_documents, dimension) containing the encoded images.
|
|
"""
|
|
|
|
inputs = processor(
|
|
text=[document_prompt] * len(documents),
|
|
images=[resize(x) for x in documents],
|
|
videos=None,
|
|
padding='longest',
|
|
return_tensors='pt'
|
|
).to('cuda:0')
|
|
|
|
cache_position = torch.arange(0, len(queries))
|
|
inputs = model.prepare_inputs_for_generation(
|
|
**inputs, cache_position=cache_position, use_cache=False)
|
|
|
|
with torch.no_grad():
|
|
output = self.model(
|
|
**inputs,
|
|
return_dict=True,
|
|
output_hidden_states=True
|
|
)
|
|
|
|
embeddings = output.hidden_states[-1][:, -1]
|
|
return torch.nn.functional.normalize(embeddings[:, :dimension], p=2, dim=-1)
|
|
```
|
|
|
|
</details>
|
|
|
|
|
|
<details>
|
|
<summary>
|
|
via SentenceTransformers
|
|
</summary>
|
|
|
|
```python
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
model = SentenceTransformer(
|
|
model_name_or_path="llamaindex/vdr-2b-multi-v1",
|
|
device="cuda",
|
|
trust_remote_code=True,
|
|
# These are the recommended kwargs for the model, but change them as needed if you don't have CUDA
|
|
model_kwargs={
|
|
"torch_dtype": torch.bfloat16,
|
|
"device_map": "cuda:0",
|
|
"attn_implementation": "flash_attention_2"
|
|
},
|
|
)
|
|
|
|
embeddings = model.encode("image.png")
|
|
```
|
|
|
|
</details>
|
|
|
|
# Training
|
|
|
|
The model is based on [MrLight/dse-qwen2-2b-mrl-v1](https://huggingface.co/MrLight/dse-qwen2-2b-mrl-v1) and it was trained on the new [vdr-multilingual-train](https://huggingface.co/datasets/llamaindex/vdr-multilingual-train) dataset that consinsists of 500k high quality, multilingual query image pairs. It was trained for 1 epoch using the [DSE approach](https://arxiv.org/abs/2406.11251), with a batch size of 128 and hard-mined negatives.
|
|
|
|
# Results
|
|
|
|
|
|
|
|
The model has been evaluated on the Vidore benchmark and on custom-built evaluation sets that allow testing its multilingual capabilities on text-only, visual-only and mixed page screenshots. The evaluation dataset is publicly available [here on HuggingFace](https://huggingface.co/datasets/llamaindex/vdr-multilingual-test).
|
|
|
|
All evaluations are performed by calculating **NDCG@5** scores using **1536 dimensions** vectors and an image resolution that can be represented with **maximum 768 tokens**.
|
|
|
|
| | Avg | Italian (text) | Italian (visual) | Italian (mix) |
|
|
|---------------------|----------|----------------|------------------|---------------|
|
|
| dse-qwen2-2b-mrl-v1 | 95.1 | 95.1 | 94 | 96.2 |
|
|
| vdr-2b-multi-v1 | **97.0** | **96.4** | **96.3** | **98.4** |
|
|
| | **+2%** | | | |
|
|
|
|
| | Avg | French (text) | French (visual) | French (mix) |
|
|
|---------------------|-----------|---------------|-----------------|--------------|
|
|
| dse-qwen2-2b-mrl-v1 | 93.5 | 94.7 | 90.8 | 95.1 |
|
|
| vdr-2b-multi-v1 | **95.6** | **95.6** | **93.3** | **97.9** |
|
|
| | **+2.2%** | | | |
|
|
|
|
| | Avg | Spanish (text) | Spanish (visual) | Spanish (mix) |
|
|
|---------------------|-----------|----------------|------------------|---------------|
|
|
| dse-qwen2-2b-mrl-v1 | 96.7 | 97.2 | 94.7 | 98.2 |
|
|
| vdr-2b-multi-v1 | **98.1** | **98.3** | **96.9** | **99.1** |
|
|
| | **+1.4%** | | | |
|
|
|
|
| | Avg | German (text) | German (visual) | German (mix) |
|
|
|---------------------|-----------|---------------|-----------------|--------------|
|
|
| dse-qwen2-2b-mrl-v1 | 93.0 | 93.4 | 90 | 95.5 |
|
|
| vdr-2b-multi-v1 | **96.2** | **94.8** | **95.7** | **98.1** |
|
|
| | **+3.4%** | | | |
|
|
|
|
| | Avg | English (text) | English (visual) | English (mix) |
|
|
|---------------------|-----------|----------------|------------------|---------------|
|
|
| dse-qwen2-2b-mrl-v1 | 98.0 | **98.3** | 98.5 | 97.1 |
|
|
| vdr-2b-multi-v1 | **98.1** | 97.9 | **99.1** | **97.3** |
|
|
| | **+0.1%** | | | |
|
|
|
|
| | **Avg** | **shiftproject** | **government** | **healthcare** | **energy** | **ai** | **docvqa** | **arxivqa** | **tatdqa** | **infovqa** | **tabfquad** |
|
|
|--------------------:|---------:|-----------------:|---------------:|---------------:|-----------:|-----------:|-----------:|------------:|-----------:|------------:|-------------:|
|
|
| dse-qwen2-2b-mrl-v1 | 83.6 | 79.8 | **95.7** | **96.9** | **92** | 98.2 | 56.3 | **85.2** | **53.9** | **87.5** | 90.3 |
|
|
| vdr-2b-multi-v1 | **84.0** | **82.4** | 95.5 | 96.5 | 91.2 | **98.5** | **58.5** | 84.7 | 53.6 | 87.1 | **92.2** | |