Views
No views yet
1# Requires vLLM==0.10.0 for NVFP4 support
2# See full implementation below
3
4model_path = "ContextualAI/ctxl-rerank-v2-instruct-multilingual-1b-nvfp4"
5
6query = "What are the health benefits of exercise?"
7instruction = "Prioritize recent medical research"
8documents = [
9 "Regular exercise reduces risk of heart disease and improves mental health.",
10 "A 2024 study shows exercise enhances cognitive function in older adults.",
11 "Ancient Greeks valued physical fitness for military training."
12]
13
14infer_w_vllm(model_path, query, instruction, documents)⚠️ Warning: These scores are produced using the BF16 model. If you run the same query with NVFP4, the scores may be slightly different.
Query: What are the health benefits of exercise?
Instruction: Prioritize recent medical research
Score: 0.5039 | Doc: A 2024 study shows exercise enhances cognitive function in older adults.
Score: -0.8398 | Doc: Regular exercise reduces risk of heart disease and improves mental health.
Score: -9.3125 | Doc: Ancient Greeks valued physical fitness for military training.vllm==0.10.0 for NVFP4 support.1import os
2os.environ['VLLM_USE_V1'] = '0' # v1 engine doesn't support logits processor yet
3
4import torch
5from vllm import LLM, SamplingParams
6
7
8def logits_processor(_, scores):
9 """Custom logits processor for vLLM reranking."""
10 index = scores[0].view(torch.uint16)
11 scores = torch.full_like(scores, float("-inf"))
12 scores[index] = 1
13 return scores
14
15
16def format_prompts(query: str, instruction: str, documents: list[str]) -> list[str]:
17 """Format query and documents into prompts for reranking."""
18 if instruction:
19 instruction = f" {instruction}"
20 prompts = []
21 for doc in documents:
22 prompt = f"Check whether a given document contains information helpful to answer the query.\n<Document> {doc}\n<Query> {query}{instruction} ??"
23 prompts.append(prompt)
24 return prompts
25
26
27def infer_w_vllm(model_path: str, query: str, instruction: str, documents: list[str]):
28 model = LLM(
29 model=model_path,
30 gpu_memory_utilization=0.85,
31 max_model_len=8192,
32 dtype="bfloat16",
33 max_logprobs=2,
34 max_num_batched_tokens=262144,
35 )
36 sampling_params = SamplingParams(
37 temperature=0,
38 max_tokens=1,
39 logits_processors=[logits_processor]
40 )
41 prompts = format_prompts(query, instruction, documents)
42
43 outputs = model.generate(prompts, sampling_params, use_tqdm=False)
44
45 # Extract scores and create results
46 results = []
47 for i, output in enumerate(outputs):
48 score = (
49 torch.tensor([output.outputs[0].token_ids[0]], dtype=torch.uint16)
50 .view(torch.bfloat16)
51 .item()
52 )
53 results.append((score, i, documents[i]))
54
55 # Sort by score (descending)
56 results = sorted(results, key=lambda x: x[0], reverse=True)
57
58 print(f"Query: {query}")
59 print(f"Instruction: {instruction}")
60 for score, doc_id, doc in results:
61 print(f"Score: {score:.4f} | Doc: {doc}")
62
63
64# Example usage
65if __name__ == "__main__":
66 model_path = "ContextualAI/ctxl-rerank-v2-instruct-multilingual-1b-nvfp4"
67 query = "What are the health benefits of exercise?"
68 instruction = "Prioritize recent medical research"
69 documents = [
70 "Regular exercise reduces risk of heart disease and improves mental health.",
71 "A 2024 study shows exercise enhances cognitive function in older adults.",
72 "Ancient Greeks valued physical fitness for military training."
73 ]
74
75 infer_w_vllm(model_path, query, instruction, documents)1@misc{ctxl_rerank_v2_instruct_multilingual,
2 title={Contextual AI Reranker v2},
3 author={Halal, George and Agrawal, Sheshansh},
4 year={2025},
5 url={https://contextual.ai/blog/rerank-v2},
6}