Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,188 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+language:
+- en
+base_model:
+- MrLight/dse-qwen2-2b-mrl-v1
+tags:
+- transformers
+- Qwen2-VL
+---
+
+# vdr-2b-v1
+
+![](cover.png)
+
+vdr-2b-v1 is an english only embedding model designed for visual document retrieval. This model is designed to encode document page screenshots into dense single-vector representations, this will effectively allow to search and query visually rich documents without the need for any OCR, data extraction pipelines, chunking...
+
+- **Trained on the 🇬🇧 English vdr-multi-train subset:** extensive training dataset of 100k high-quality english samples.
+
+- **Low VRAM and Faster Inference**: achieves better results on synthetic Vidore benchmarks with just 30% of the base model image resolution. This results in 3x faster inference and much lower VRAM usage.
+
+- **Matryoshka Representation Learning**: You can reduce the vectors size 3x and still keep 98% of the embeddings quality.
+
+The multilingual version is available [here](https://huggingface.co/llamaindex/vdr-2b-multi-v1). To know more about both models, read the [announcement blogpost](https://huggingface.co/blog/marco/vdr-2b-multilingual).
+
+# Usage
+
+**Initialize model and processor**
+
+```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-v1',
+    attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16,
+    device_map="cuda:0"
+).eval()
+
+processor = AutoProcessor.from_pretrained(
+    'llamaindex/vdr-2b-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)
+```
+
+# 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) english subset that consinsists of 100k high quality samples. 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. All evaluations are performed by calculating **NDCG@5** scores using an image resolution that can be represented with **maximum 768 tokens**.
+
+On the full Vidore benchmark (evaluated with 768 image tokens), both the multilingual and the english-only version performs better than the base model.
+
+|                     | **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**     |
+| vdr-2b-v1           | **84.3** | **83.4**         | **96.9**       | **97.2**       | **92.6**   | 96.8     | 57.4       | 85.1        | **54.1**   | **87.9**    | 91.3         |
+
+![](chart.png)
+
+|                                         | Avg      | shiftproject | government | healthcare | energy   | ai       |
+|-----------------------------------------|----------|--------------|------------|------------|----------|----------|
+| dse-qwen2-2b-mrl-v1 (2560 image tokens) | 93.0     | 82           | 96         | 96.4       | **92.9** | **97.5** |
+| vdr-2b-v1 (768 image tokens)            | **93.4** | **83.4**     | **96.9**   | **97.2**   | 92.6     | 96.8     |
+
+vdr-2b-v1 matches the performance of the base model on vidore synthetic datasets, while only using 30% of the image tokens (768 vs. 2560). This results in 3x faster inference and much lower VRAM usage.

added_tokens.json

@@ -0,0 +1,16 @@
+{
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.json

@@ -0,0 +1,3 @@
+{
+  "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
+}

config.json

@@ -0,0 +1,48 @@
+{
+  "_name_or_path": "MrLight/dse-qwen2-2b-mrl-v1",
+  "architectures": [
+    "Qwen2VLForConditionalGeneration"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "image_token_id": 151655,
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 28,
+  "model_type": "qwen2_vl",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": {
+    "mrope_section": [
+      16,
+      24,
+      24
+    ],
+    "rope_type": "default",
+    "type": "default"
+  },
+  "rope_theta": 1000000.0,
+  "sliding_window": 32768,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.47.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "video_token_id": 151656,
+  "vision_config": {
+    "hidden_size": 1536,
+    "in_chans": 3,
+    "model_type": "qwen2_vl",
+    "spatial_patch_size": 14
+  },
+  "vision_end_token_id": 151653,
+  "vision_start_token_id": 151652,
+  "vision_token_id": 151654,
+  "vocab_size": 151936
+}

generation_config.json

@@ -0,0 +1,14 @@
+{
+  "attn_implementation": "flash_attention_2",
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "temperature": 0.01,
+  "top_k": 1,
+  "top_p": 0.001,
+  "transformers_version": "4.47.1"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1e1bfa7e70273b9532c657209075fc9ddd34673bb90e3a13ace3bd9d1cd0838
+size 4418050848

preprocessor_config.json

@@ -0,0 +1,29 @@
+{
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "Qwen2VLImageProcessor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "max_pixels": 602112,
+  "merge_size": 2,
+  "min_pixels": 784,
+  "patch_size": 14,
+  "processor_class": "Qwen2VLProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "max_pixels": 12845056,
+    "min_pixels": 3136
+  },
+  "temporal_patch_size": 2
+}

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:091aa7594dc2fcfbfa06b9e3c22a5f0562ac14f30375c13af7309407a0e67b8a
+size 11420371

tokenizer_config.json

@@ -0,0 +1,147 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "max_pixels": 602112,
+  "min_pixels": 784,
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "padding_side": "left",
+  "processor_class": "Qwen2VLProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}