Views
No views yet

Query: What are the health benefits of exercise?
Instruction: Prioritize recent medical research
Documents:
- Regular exercise reduces risk of heart disease and improves mental health.
- A 2024 study shows exercise enhances cognitive function in older adults.
- Ancient Greeks valued physical fitness for military training.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.pip install sentence_transformers1import torch
2from sentence_transformers import CrossEncoder
3
4model = CrossEncoder("ContextualAI/ctxl-rerank-v2-instruct-multilingual-1b", model_kwargs={"dtype": torch.bfloat16})
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
14pairs = [(query, doc) for doc in documents]
15scores = model.predict(pairs, prompt=instruction)
16print(scores)
17# [-0.8515625 0.50390625 -9.375 ]
18
19rankings = model.rank(query, documents, prompt=instruction)
20print(rankings)
21# [{'corpus_id': 1, 'score': np.float32(0.50390625)}, {'corpus_id': 0, 'score': np.float32(-0.8515625)}, {'corpus_id': 2, 'score': np.float32(-9.375)}]prompt argument is optional, you can omit it to score pairs without any custom instruction. Scores are the raw bfloat16 logits at token id 0 at the final position (matching the Transformers path below), so higher means more relevant.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-1b"
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 Query: What are the health benefits of exercise?
55 Instruction: Prioritize recent medical research
56 Score: 0.5039 | Doc: A 2024 study shows exercise enhances cognitive function in older adults.
57 Score: -0.8516 | Doc: Regular exercise reduces risk of heart disease and improves mental health.
58 Score: -9.3750 | Doc: Ancient Greeks valued physical fitness for military training.
59 """
60
61# Example usage
62if __name__ == "__main__":
63 model_path = "ContextualAI/ctxl-rerank-v2-instruct-multilingual-1b"
64 query = "What are the health benefits of exercise?"
65 instruction = "Prioritize recent medical research"
66 documents = [
67 "Regular exercise reduces risk of heart disease and improves mental health.",
68 "A 2024 study shows exercise enhances cognitive function in older adults.",
69 "Ancient Greeks valued physical fitness for military training."
70 ]
71
72 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}