import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ColModernVBertForRetrieval, ColModernVBertProcessor


processor = ColModernVBertProcessor.from_pretrained("ModernVBERT/colmodernvbert-hf")
model = ColModernVBertForRetrieval.from_pretrained("ModernVBERT/colmodernvbert-hf", device_map="auto")

# Load the test dataset
queries = [
    "A paint on the wall",
    "ColModernVBERT matches the performance of models nearly 10x larger on visual document benchmarks."
]

images = [
    Image.open(hf_hub_download("HuggingFaceTB/SmolVLM", "example_images/rococo.jpg", repo_type="space")),
    Image.open(hf_hub_download("ModernVBERT/colmodernvbert", "table.png", repo_type="model"))
]

# Preprocess the examples
batch_images = processor(images=images).to(model.device)
batch_queries = processor(text=queries).to(model.device)

# Run inference
with torch.inference_mode():
    image_embeddings = model(**batch_images).embeddings
    query_embeddings = model(**batch_queries).embeddings

# Compute retrieval scores
scores = processor.score_retrieval(
    query_embeddings=query_embeddings,
    passage_embeddings=image_embeddings,
)

scores = torch.softmax(scores, dim=-1)

print(scores)    # [[0.9350, 0.0650], [0.0015, 0.9985]]