Views
No views yet
1# Choose vLLM (recommended for production) or Transformers (simpler setup)
2# See full implementation in sections below
3
4model_path = "ContextualAI/ctxl-rerank-v2-instruct-multilingual-2b"
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
14# Using vLLM (see full code below):
15infer_w_vllm(model_path, query, instruction, documents)
16
17# OR using Transformers (see full code below):
18infer_w_hf(model_path, query, instruction, documents)Query: What are the health benefits of exercise?
Instruction: Prioritize recent medical research
Score: 0.8398 | Doc: A 2024 study shows exercise enhances cognitive function in older adults.
Score: -2.5469 | Doc: Regular exercise reduces risk of heart disease and improves mental health.
Score: -9.3750 | Doc: Ancient Greeks valued physical fitness for military training.vllm==0.10.0 for NVFP4 or vllm>=0.8.5 for BF16.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-2b"
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)transformers>=4.51.0 for BF16. Not supported for NVFP4.1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4
5def format_prompts(query: str, instruction: str, documents: list[str]) -> list[str]:
6 """Format query and documents into prompts for reranking."""
7 if instruction:
8 instruction = f" {instruction}"
9 prompts = []
10 for doc in documents:
11 prompt = f"Check whether a given document contains information helpful to answer the query.\n<Document> {doc}\n<Query> {query}{instruction} ??"
12 prompts.append(prompt)
13 return prompts
14
15
16def infer_w_hf(model_path: str, query: str, instruction: str, documents: list[str]):
17 device = "cuda" if torch.cuda.is_available() else "cpu"
18 dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
19
20 tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
21 if tokenizer.pad_token is None:
22 tokenizer.pad_token = tokenizer.eos_token
23 tokenizer.padding_side = "left" # so -1 is the real last token for all prompts
24
25 model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype).to(device)
26 model.eval()
27
28 prompts = format_prompts(query, instruction, documents)
29 enc = tokenizer(
30 prompts,
31 return_tensors="pt",
32 padding=True,
33 truncation=True,
34 )
35 input_ids = enc["input_ids"].to(device)
36 attention_mask = enc["attention_mask"].to(device)
37
38 with torch.no_grad():
39 out = model(input_ids=input_ids, attention_mask=attention_mask)
40
41 next_logits = out.logits[:, -1, :] # [batch, vocab]
42
43 scores_bf16 = next_logits[:, 0].to(torch.bfloat16)
44 scores = scores_bf16.float().tolist()
45
46 # Sort by score (descending)
47 results = sorted([(s, i, documents[i]) for i, s in enumerate(scores)], key=lambda x: x[0], reverse=True)
48
49 print(f"Query: {query}")
50 print(f"Instruction: {instruction}")
51 for score, doc_id, doc in results:
52 print(f"Score: {score:.4f} | Doc: {doc}")
53
54
55# Example usage
56if __name__ == "__main__":
57 model_path = "ContextualAI/ctxl-rerank-v2-instruct-multilingual-2b"
58 query = "What are the health benefits of exercise?"
59 instruction = "Prioritize recent medical research"
60 documents = [
61 "Regular exercise reduces risk of heart disease and improves mental health.",
62 "A 2024 study shows exercise enhances cognitive function in older adults.",
63 "Ancient Greeks valued physical fitness for military training."
64 ]
65
66 infer_w_hf(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}