Update README.md

README.md

@@ -3,20 +3,22 @@ language:
 - en
 license: mit
 library_name: Tevatron
+tag:
+- vidore
 datasets:
 - Tevatron/docmatix-ir
 - HuggingFaceM4/Docmatix
 - Tevatron/msmarco-passage-aug
 ---
 
-# DSE-Phi3-Docmatix-V1
+# DSE-Phi3-Docmatix-V2
 
-DSE-Phi3-Docmatix-V1 is a bi-encoder model designed to encode document screenshots into dense vectors for document retrieval. The Document Screenshot Embedding ([DSE](https://arxiv.org/abs/2406.11251)) approach captures documents in their original visual format, preserving all information such as text, images, and layout, thus avoiding tedious parsing and potential information loss.
+DSE-Phi3-Docmatix-V2 is a bi-encoder model designed to encode document screenshots into dense vectors for document retrieval. The Document Screenshot Embedding ([DSE](https://arxiv.org/abs/2406.11251)) approach captures documents in their original visual format, preserving all information such as text, images, and layout, thus avoiding tedious parsing and potential information loss.
 
-The model, `Tevatron/dse-phi3-docmatix-v1`, is trained using 1/10 of the `Tevatron/docmatix-ir` dataset, a variant of `HuggingFaceM4/Docmatix` specifically adapted for training PDF retrievers with Vision Language Models in open-domain question answering scenarios. For more information on dataset filtering and hard negative mining, refer to the [docmatix-ir](https://huggingface.co/datasets/Tevatron/docmatix-ir/blob/main/README.md) dataset page.
+The model, `Tevatron/dse-phi3-docmatix-v2`, is trained using 1/10 of the `Tevatron/docmatix-ir` dataset, a variant of `HuggingFaceM4/Docmatix` specifically adapted for training PDF retrievers with Vision Language Models in open-domain question answering scenarios. For more information on dataset filtering and hard negative mining, refer to the [docmatix-ir](https://huggingface.co/datasets/Tevatron/docmatix-ir/blob/main/README.md) dataset page.
 
 DSE has strong zero-shot effectiveness for document retrieval both with visual input and text input.
-For example, DSE-Phi3-Docmatix-V1 achieves 74.1 nDCG@5 on [ViDoRE](https://huggingface.co/spaces/vidore/vidore-leaderboard) leaderboard in **zero-shot setting** (without finetuning with ViDoRe training data).
+For example, DSE-Phi3-Docmatix-V2 achieves 77.6 nDCG@5 on [ViDoRE](https://huggingface.co/spaces/vidore/vidore-leaderboard) leaderboard in **zero-shot setting** (without finetuning with ViDoRe training data).
 
 
 ## How to Use the Model
@@ -27,8 +29,8 @@ For example, DSE-Phi3-Docmatix-V1 achieves 74.1 nDCG@5 on [ViDoRE](https://huggi
 import torch
 from transformers import AutoProcessor, AutoModelForCausalLM
 
-processor = AutoProcessor.from_pretrained('Tevatron/dse-phi3-docmatix-v1', trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained('Tevatron/dse-phi3-docmatix-v1', trust_remote_code=True, attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, use_cache=False).to('cuda:0')
+processor = AutoProcessor.from_pretrained('Tevatron/dse-phi3-docmatix-v2', trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained('Tevatron/dse-phi3-docmatix-v2', trust_remote_code=True, attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, use_cache=False).to('cuda:0')
 
 def get_embedding(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
     sequence_lengths = attention_mask.sum(dim=1) - 1
@@ -56,15 +58,15 @@ import requests
 from io import BytesIO
 
 # URLs of the images
-url1 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v1/resolve/main/animal-llama.png"
-url2 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v1/resolve/main/meta-llama.png"
+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"
 
 # Download and open images
 response1 = requests.get(url1)
 response2 = requests.get(url2)
 
-passage_image1 = Image.open(BytesIO(response1.content))
-passage_image2 = Image.open(BytesIO(response2.content))
+passage_image1 = Image.open(BytesIO(response1.content)).resize((1344, 1344))
+passage_image2 = Image.open(BytesIO(response2.content)).resize((1344, 1344))
 
 passage_images = [passage_image1, passage_image2]
 passage_prompts = ["<|image_1|>\nWhat is shown in this image?</s>", "<|image_2|>\nWhat is shown in this image?</s>"]