pip install sentence_transformers

from sentence_transformers import CrossEncoder
from transformers.image_utils import load_image

model = CrossEncoder("nvidia/llama-nemotron-rerank-vl-1b-v2", trust_remote_code=True)

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

# Text-only reranking
pairs = [(query, doc) for doc in documents]
scores = model.predict(pairs)
print(scores)
# [-0.738  -3.859  -7.75 ]

rankings = model.rank(query, documents)
print(rankings)
# [{'corpus_id': 0, 'score': -0.73828125}, {'corpus_id': 1, 'score': -3.859375}, {'corpus_id': 2, 'score': -7.75}]

# Image-only reranking
pairs = [(query, img) for img in images]
scores = model.predict(pairs)
print(scores)
# [-4.125  -5.156  -5.969]

rankings = model.rank(query, images)
print(rankings)
# [{'corpus_id': 0, 'score': -4.125}, {'corpus_id': 1, 'score': -5.15625}, {'corpus_id': 2, 'score': -5.96875}]

# Image+text reranking
multimodal_docs = [{"image": img, "text": doc} for img, doc in zip(images, documents)]
pairs = [(query, doc) for doc in multimodal_docs]
scores = model.predict(pairs)
print(scores)
# [-2.594  -4.531  -6.219]

rankings = model.rank(query, multimodal_docs)
print(rankings)
# [{'corpus_id': 0, 'score': -2.59375}, {'corpus_id': 1, 'score': -4.53125}, {'corpus_id': 2, 'score': -6.21875}]

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]}")

cat > nemotron-vl-rerank.jinja << 'JINJA'
{%- set query_msg = (messages | selectattr('role', 'equalto', 'query') | list | first) -%}
{%- set doc_msg   = (messages | selectattr('role', 'equalto', 'document') | list | first) -%}
{%- set q = query_msg['content'] -%}
{%- set d = doc_msg['content'] -%}
{%- set has_image = ("<image>" in d) -%}
{%- set d_clean = d | replace("<image>", "") -%}
{%- set q_clean = q | replace("<image>", "") -%}
{%- if has_image -%}<image>{{ " " }}{%- endif -%}
question:{{ q_clean }}{{ " " }}
{{ " " }}
{{ " " }}passage:{{ d_clean }}
JINJA

vllm serve nvidia/llama-nemotron-rerank-vl-1b-v2 \
  --runner pooling \
  --trust-remote-code \
  --max-model-len 10240 \
  --chat-template nemotron-vl-rerank.jinja

import requests

url = "http://localhost:8000/rerank"

# Text-only reranking
response = requests.post(url, json={
    "model": "nvidia/llama-nemotron-rerank-vl-1b-v2",
    "query": "How is AI improving the intelligence and capabilities of robots?",
    "documents": [
        "AI enables robots to perceive, plan, and act autonomously.",
        "A biological foundation model designed to analyze DNA, RNA, and protein sequences.",
    ],
})
print(response.json())

# Image document reranking
response = requests.post(url, json={
    "model": "nvidia/llama-nemotron-rerank-vl-1b-v2",
    "query": "How is AI improving the intelligence and capabilities of robots?",
    "documents": [
        {"content": [{"type": "image_url", "image_url": {"url": "https://example.com/page.png"}}]},
        {"content": [
            {"type": "text", "text": "AI enables robots to perceive, plan, and act autonomously."},
            {"type": "image_url", "image_url": {"url": "https://example.com/page2.png"}},
        ]},
    ],
})
print(response.json())

from vllm import LLM

SCORE_TEMPLATE = """\
{%- set query_msg = (messages | selectattr('role', 'equalto', 'query') | list | first) -%}
{%- set doc_msg   = (messages | selectattr('role', 'equalto', 'document') | list | first) -%}
{%- set q = query_msg['content'] -%}
{%- set d = doc_msg['content'] -%}
{%- set has_image = ("<image>" in d) -%}
{%- set d_clean = d | replace("<image>", "") -%}
{%- set q_clean = q | replace("<image>", "") -%}
{%- if has_image -%}<image>{{ " " }}{%- endif -%}
question:{{ q_clean }}{{ " " }}
{{ " " }}
{{ " " }}passage:{{ d_clean }}"""

llm = LLM(
    model="nvidia/llama-nemotron-rerank-vl-1b-v2",
    runner="pooling",
    max_model_len=10240,
    trust_remote_code=True,
)

query = "How is AI improving the intelligence and capabilities of robots?"
documents = [
    "AI enables robots to perceive, plan, and act autonomously.",
    "A biological foundation model designed to analyze DNA, RNA, and protein sequences.",
]

# Text-only scoring
outputs = llm.score(query, documents, chat_template=SCORE_TEMPLATE)
for doc, output in zip(documents, outputs):
    print(f"Score: {output.outputs.score:.4f} | {doc}")

# Image document scoring
outputs = llm.score(
    query,
    [
        {"content": [{"type": "image_url", "image_url": {"url": "https://example.com/page.png"}}]},
        {"content": [
            {"type": "text", "text": "AI enables robots to perceive, plan, and act autonomously."},
            {"type": "image_url", "image_url": {"url": "https://example.com/page2.png"}},
        ]},
    ],
    chat_template=SCORE_TEMPLATE,
)
for output in outputs:
    print(f"Score: {output.outputs.score:.4f}")

@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}
}