Update README.md

README.md

@@ -149,16 +149,16 @@ from vllm.inputs.data import TokensPrompt
 def format_instruction(instruction, query, doc):
     text = [
         {"role": "system", "content": "Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\"."},
-        {"role": "user", "content": f"<Instruct>: {instruction}\n\n<Query>: {query}\n\n<Document>: {doc}"}
+        {"role": "user", "content": f"<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {doc}"}
     ]
     return text
 
-def process_inputs(pairs, instruction, max_length, suffix_tokens):
+def process_inputs(pairs, instruction, max_length):
     messages = [format_instruction(instruction, query, doc) for query, doc in pairs]
     messages =  tokenizer.apply_chat_template(
-        messages, tokenize=True, add_generation_prompt=False, enable_thinking=False
+        messages, tokenize=True, add_generation_prompt=True, enable_thinking=False
     )
-    messages = [ele[:max_length] + suffix_tokens for ele in messages]
+    messages = [ele[:max_length] for ele in messages]
     messages = [TokensPrompt(prompt_token_ids=ele) for ele in messages]
     return messages
 
@@ -187,9 +187,8 @@ tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen3-Reranker-0.6B')
 model = LLM(model='Qwen/Qwen3-Reranker-0.6B', tensor_parallel_size=number_of_gpu, max_model_len=10000, enable_prefix_caching=True, gpu_memory_utilization=0.8)
 tokenizer.padding_side = "left"
 tokenizer.pad_token = tokenizer.eos_token
-suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
+
 max_length=8192
-suffix_tokens = tokenizer.encode(suffix, add_special_tokens=False)
 true_token = tokenizer("yes", add_special_tokens=False).input_ids[0]
 false_token = tokenizer("no", add_special_tokens=False).input_ids[0]
 sampling_params = SamplingParams(temperature=0, 
@@ -209,7 +208,7 @@ documents = [
 ]
 
 pairs = list(zip(queries, documents))
-inputs = process_inputs(pairs, task, max_length-len(suffix_tokens), suffix_tokens)
+inputs = process_inputs(pairs, task, max_length)
 scores = compute_logits(model, inputs, sampling_params, true_token, false_token)
 print('scores', scores)