Visual Document Retrieval
Transformers
Safetensors
ret2
nielsr's picture
nielsr HF Staff
Add sample usage section to model card
7f6f02d verified
|
raw
history blame
3.49 kB
metadata
base_model:
  - google/siglip2-large-patch16-256
  - colbert-ir/colbertv2.0
datasets:
  - aimagelab/ReT-M2KR
library_name: transformers
license: apache-2.0
pipeline_tag: visual-document-retrieval

Model Card: ReT-2

Official implementation of ReT-2: Recurrence Meets Transformers for Universal Multimodal Retrieval.

This model features a visual backbone based on google/siglip2-large-patch16-256 and a textual backbone based on colbert-ir/colbertv2.0.
The backbones have been fine-tuned on the M2KR dataset.

Model Sources

Training Data

aimagelab/ReT-M2KR

Use with 🤗's Transformers

from src.models import Ret2Model
import requests
from PIL import Image
from io import BytesIO
import torch
import torch.nn.functional as F

device = 'cuda' if torch.cuda.is_available() else 'cpu'

headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
}

query_img_url = 'https://upload.wikimedia.org/wikipedia/commons/8/84/Ghirlandina_%28Modena%29.jpg'
response = requests.get(query_img_url, headers=headers)
query_image = Image.open(BytesIO(response.content)).convert('RGB')
query_text = 'Where is this building located?'

passage_img_url = 'https://upload.wikimedia.org/wikipedia/commons/0/09/Absidi_e_Ghirlandina.jpg'
response = requests.get(query_img_url, headers=headers)
passage_image = Image.open(BytesIO(response.content)).convert('RGB')
passage_text = (
    "The Ghirlandina is the bell tower of the Cathedral of Modena, in Modena, Italy. "
    "It is 86.12 metres (282.7 ft) high and is the symbol of the city. "
    "It was built in Romanesque style in the 12th century and is part of a UNESCO World Heritage Site."
)

model = Ret2Model.from_pretrained('aimagelab/ReT2-M2KR-ColBERT-SigLIP2-ViT-L', device_map=device)

query_txt_inputs = model.tokenizer([query_text], return_tensors='pt').to(device)
query_img_inputs = model.image_processor([query_image], return_tensors='pt').to(device)
passage_txt_inputs = model.tokenizer([passage_text], return_tensors='pt').to(device)
passage_img_inputs = model.image_processor([passage_image], return_tensors='pt').to(device)

with torch.inference_mode():
    query_feats = model.get_ret_features(
        input_ids=query_txt_inputs.input_ids,
        attention_mask=query_txt_inputs.attention_mask,
        pixel_values=query_img_inputs.pixel_values
    )

    passage_feats = model.get_ret_features(
        input_ids=passage_txt_inputs.input_ids,
        attention_mask=passage_txt_inputs.attention_mask,
        pixel_values=passage_img_inputs.pixel_values
    )

    sim = F.normalize(query_feats, p=2, dim=-1) @ F.normalize(passage_feats, p=2, dim=-1).T

print(f"query-passage similarity: {sim.item():.3f}")

Citation

@article{caffagni2025recurrencemeetstransformers,
      title={{Recurrence Meets Transformers for Universal Multimodal Retrieval}},
      author={Davide Caffagni and Sara Sarto and Marcella Cornia and Lorenzo Baraldi and Rita Cucchiara},
      journal={arXiv preprint arXiv:2509.08897},
      year={2025}
}