1from vllm import LLM, SamplingParams
2import math
3
4# Initialize the model with vLLM
5model = LLM(
6 model="jhu-clsp/rank1-32b",
7 tensor_parallel_size=1, # Number of GPUs
8 trust_remote_code=True,
9 max_model_len=16000, # Context length
10 gpu_memory_utilization=0.9,
11 dtype="float16",
12)
13
14# Set up sampling parameters
15sampling_params = SamplingParams(
16 temperature=0,
17 max_tokens=8192,
18 logprobs=20,
19 stop=["</think> true", "</think> false"],
20 skip_special_tokens=False
21)
22
23# Prepare the prompt
24def create_prompt(query, document):
25 return (
26 "Determine if the following passage is relevant to the query. "
27 "Answer only with 'true' or 'false'.\n"
28 f"Query: {query}\n"
29 f"Passage: {document}\n"
30 "<think>"
31 )
32
33# Example usage
34query = "What are the effects of climate change?"
35document = "Climate change leads to rising sea levels, extreme weather events, and disruptions to ecosystems. These effects are caused by increasing greenhouse gas concentrations in the atmosphere due to human activities."
36
37# Generate prediction
38prompt = create_prompt(query, document)
39outputs = model.generate([prompt], sampling_params)
40
41# Extract score
42output = outputs[0].outputs[0]
43text = output.text
44final_logits = output.logprobs[-1]
45
46# Get token IDs for "true" and "false" tokens
47from transformers import AutoTokenizer
48tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/rank1-32b")
49true_token = tokenizer(" true", add_special_tokens=False).input_ids[0]
50false_token = tokenizer(" false", add_special_tokens=False).input_ids[0]
51
52# Calculate relevance score (probability of "true")
53true_logit = final_logits[true_token].logprob
54false_logit = final_logits[false_token].logprob
55true_score = math.exp(true_logit)
56false_score = math.exp(false_logit)
57relevance_score = true_score / (true_score + false_score)
58
59print(f"Reasoning chain: {text}")
60print(f"Relevance score: {relevance_score}")