Views
No views yet
1from src.models import Ret2Model
2import requests
3from PIL import Image
4from io import BytesIO
5import torch
6import torch.nn.functional as F
7
8device = 'cuda' if torch.cuda.is_available() else 'cpu'
9
10headers = {
11 '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'
12}
13
14query_img_url = 'https://upload.wikimedia.org/wikipedia/commons/8/84/Ghirlandina_%28Modena%29.jpg'
15response = requests.get(query_img_url, headers=headers)
16query_image = Image.open(BytesIO(response.content)).convert('RGB')
17query_text = 'Where is this building located?'
18
19passage_img_url = 'https://upload.wikimedia.org/wikipedia/commons/0/09/Absidi_e_Ghirlandina.jpg'
20response = requests.get(query_img_url, headers=headers)
21passage_image = Image.open(BytesIO(response.content)).convert('RGB')
22passage_text = (
23 "The Ghirlandina is the bell tower of the Cathedral of Modena, in Modena, Italy. "
24 "It is 86.12 metres (282.7 ft) high and is the symbol of the city. "
25 "It was built in Romanesque style in the 12th century and is part of a UNESCO World Heritage Site."
26)
27
28model = Ret2Model.from_pretrained('aimagelab/ReT2-M2KR-ColBERT-SigLIP2-ViT-L', device_map=device)
29
30query_txt_inputs = model.tokenizer([query_text], return_tensors='pt').to(device)
31query_img_inputs = model.image_processor([query_image], return_tensors='pt').to(device)
32passage_txt_inputs = model.tokenizer([passage_text], return_tensors='pt').to(device)
33passage_img_inputs = model.image_processor([passage_image], return_tensors='pt').to(device)
34
35with torch.inference_mode():
36 query_feats = model.get_ret_features(
37 input_ids=query_txt_inputs.input_ids,
38 attention_mask=query_txt_inputs.attention_mask,
39 pixel_values=query_img_inputs.pixel_values
40 )
41
42 passage_feats = model.get_ret_features(
43 input_ids=passage_txt_inputs.input_ids,
44 attention_mask=passage_txt_inputs.attention_mask,
45 pixel_values=passage_img_inputs.pixel_values
46 )
47
48 sim = F.normalize(query_feats, p=2, dim=-1) @ F.normalize(passage_feats, p=2, dim=-1).T
49
50print(f"query-passage similarity: {sim.item():.3f}")@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}
}