llama-nemotron-rerank-vl-1b-v2

import torch
from transformers import AutoModelForSequenceClassification, AutoProcessor
from transformers.image_utils import load_image

modality = "image"

# Load model
model_path = "nvidia/llama-nemotron-rerank-vl-1b-v2"
device = "cuda" if torch.cuda.is_available() else "cpu"

model = AutoModelForSequenceClassification.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    attn_implementation="flash_attention_2",
    device_map="auto"
).eval()

# Build processor kwargs (base settings)
processor_kwargs = {
    "trust_remote_code": True,
    "max_input_tiles": 6,
    "use_thumbnail": True
}

# Set rerank_max_length based on modality
if modality == "image":
    processor_kwargs["rerank_max_length"] = 2048
elif modality == "image_text":
    processor_kwargs["rerank_max_length"] = 10240
elif modality == "text":
    processor_kwargs["rerank_max_length"] = 8192

# Load processor with modality-specific kwargs
processor = AutoProcessor.from_pretrained(
    model_path,
    **processor_kwargs
)

query = "How is AI improving the intelligence and capabilities of robots?"
image_paths = [
    "https://developer.download.nvidia.com/images/isaac/nvidia-isaac-lab-1920x1080.jpg",
    "https://blogs.nvidia.com/wp-content/uploads/2018/01/automotive-key-visual-corp-blog-level4-av-og-1280x680-1.png",
    "https://developer-blogs.nvidia.com/wp-content/uploads/2025/02/hc-press-evo2-nim-25-featured-b.jpg"
]

# Load all images
images = [load_image(img_path) for img_path in image_paths]

# Text descriptions corresponding to each image/document
document_texts = [
    "AI enables robots to perceive, plan, and act autonomously.",
    "AI is transforming autonomous vehicles by enabling safer, smarter, and more reliable decision-making on the road.",
    "A biological foundation model designed to analyze and generate DNA, RNA, and protein sequences.",
]

if modality == "image":
    # Prepare inputs: same query, different images
    examples = [{
        "question": query,
        "doc_text": "",
        "doc_image": image
    } for image in images]

elif modality == "image_text":
    examples = [{
        "question": query,
        "doc_text": doc_text,
        "doc_image": image
    } for image, doc_text in zip(images, document_texts)]

elif modality == "text":
    # Prepare inputs: same query, different texts
    examples = [{
        "question": query,
        "doc_text": doc_text,
        "doc_image": ""
    } for doc_text in document_texts]

else:
    raise ValueError(f"Invalid modality: {modality}. Must be 'image', 'image_text', or 'text'")

# Process with processor
batch_dict = processor.process_queries_documents_crossencoder(examples)

# Move to device
batch_dict = {
    k: v.to(device) if isinstance(v, torch.Tensor) else v
    for k, v in batch_dict.items()
}

# Run inference
with torch.no_grad():
    outputs = model(**batch_dict, return_dict=True)

# Get logits
logits = outputs.logits
logits_flat = logits.squeeze(-1)
    
# Get sorted indices (highest to lowest)
sorted_indices = torch.argsort(logits_flat, descending=True)

print(f"\nRanking (highest to lowest relevance for the modality {modality}):")
for rank, idx in enumerate(sorted_indices, 1):
    doc_idx = idx.item()
    logit_val = logits_flat[doc_idx].item()
    if modality == "text":
        print(f"  Rank {rank}: logit={logit_val:.4f} | Text: {document_texts[doc_idx]}")
    else:  # image or image_text modality
        print(f"  Rank {rank}: logit={logit_val:.4f} | Image: {image_paths[doc_idx]}")

@inproceedings{moreira2025_nvretriever,
author = {Moreira, Gabriel de Souza P. and Osmulski, Radek and Xu, Mengyao and Ak, Ronay and Schifferer, Benedikt and Oldridge, Even},
title = {Improving Text Embedding Models with Positive-aware Hard-negative Mining},
year = {2025},
isbn = {9798400720406},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3746252.3761254},
doi = {10.1145/3746252.3761254},
pages = {2169–2178},
numpages = {10},
keywords = {contrastive learning, distillation, embedding models, hard-negative mining, rag, text retrieval, transformers},
location = {Seoul, Republic of Korea},
series = {CIKM '25}
}