[Model] Let more models to support the score template. (#31335)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
This commit is contained in:
@@ -2,35 +2,70 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
Script to convert Large Language Models (LLMs) to Sequence Classification models.
|
||||
This is particularly useful for converting reranker models that use next-token
|
||||
prediction to a sequence classification format for compatibility with standard
|
||||
classification and rerank pipelines.
|
||||
|
||||
Usage examples:
|
||||
- For BAAI/bge-reranker-v2-gemma:
|
||||
python convert_model_to_seq_cls.py --model_name BAAI/bge-reranker-v2-gemma \
|
||||
--classifier_from_tokens '["Yes"]' --method no_post_processing \
|
||||
--path ./bge-reranker-v2-gemma-seq-cls
|
||||
|
||||
- For mxbai-rerank-v2:
|
||||
python convert_model_to_seq_cls.py --model_name mixedbread-ai/mxbai-rerank-base-v2 \
|
||||
--classifier_from_tokens '["0", "1"]' --method from_2_way_softmax \
|
||||
--path ./mxbai-rerank-base-v2-seq-cls
|
||||
|
||||
- For Qwen3-Reranker:
|
||||
python convert_model_to_seq_cls.py --model_name Qwen/Qwen3-Reranker-0.6B \
|
||||
--classifier_from_tokens '["no", "yes"]' --method from_2_way_softmax \
|
||||
--path ./Qwen3-Reranker-0.6B-seq-cls
|
||||
|
||||
Note: For BAAI/bge-reranker-v2-gemma, "Yes" and "yes" are different tokens.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import torch
|
||||
import transformers
|
||||
|
||||
# Usage:
|
||||
# for BAAI/bge-reranker-v2-gemma
|
||||
# Caution: "Yes" and "yes" are two different tokens
|
||||
# python convert_model_to_seq_cls.py --model_name BAAI/bge-reranker-v2-gemma --classifier_from_tokens '["Yes"]' --method no_post_processing --path ./bge-reranker-v2-gemma-seq-cls
|
||||
# for mxbai-rerank-v2
|
||||
# python convert_model_to_seq_cls.py --model_name mixedbread-ai/mxbai-rerank-base-v2 --classifier_from_tokens '["0", "1"]' --method from_2_way_softmax --path ./mxbai-rerank-base-v2-seq-cls
|
||||
# for Qwen3-Reranker
|
||||
# python convert_model_to_seq_cls.py --model_name Qwen/Qwen3-Reranker-0.6B --classifier_from_tokens '["no", "yes"]' --method from_2_way_softmax --path ./Qwen3-Reranker-0.6B-seq-cls
|
||||
|
||||
|
||||
def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device):
|
||||
# refer to https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
assert len(tokens) == 2
|
||||
"""
|
||||
This method extracts the difference between weights for 'true' and 'false' tokens
|
||||
from the language model head to create a single classification weight vector.
|
||||
|
||||
Args:
|
||||
causal_lm: The original causal language model
|
||||
seq_cls_model: The target sequence classification model
|
||||
tokenizer: Model tokenizer
|
||||
tokens: List of two tokens representing [false_token, true_token]
|
||||
device: Target device (cpu/cuda)
|
||||
|
||||
Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
"""
|
||||
assert len(tokens) == 2, (
|
||||
"Method requires exactly two tokens for binary classification"
|
||||
)
|
||||
|
||||
# Get the language model head weights (vocabulary_size x hidden_size)
|
||||
lm_head_weights = causal_lm.lm_head.weight
|
||||
|
||||
# Convert token strings to their corresponding token IDs
|
||||
false_id = tokenizer.convert_tokens_to_ids(tokens[0])
|
||||
true_id = tokenizer.convert_tokens_to_ids(tokens[1])
|
||||
|
||||
# Compute the classification weight as the difference between true and false token weights
|
||||
# This follows the approach in: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
score_weight = lm_head_weights[true_id].to(device).to(
|
||||
torch.float32
|
||||
) - lm_head_weights[false_id].to(device).to(torch.float32)
|
||||
|
||||
# Copy the computed weights to the sequence classification model
|
||||
with torch.no_grad():
|
||||
seq_cls_model.score.weight.copy_(score_weight.unsqueeze(0))
|
||||
if seq_cls_model.score.bias is not None:
|
||||
@@ -38,12 +73,29 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device):
|
||||
|
||||
|
||||
def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device):
|
||||
"""
|
||||
Directly use token weights from the language model head for classification.
|
||||
|
||||
This method maps each classification label directly to a corresponding token
|
||||
in the vocabulary without additional transformation.
|
||||
|
||||
Args:
|
||||
causal_lm: The original causal language model
|
||||
seq_cls_model: The target sequence classification model
|
||||
tokenizer: Model tokenizer
|
||||
tokens: List of tokens representing class labels
|
||||
device: Target device (cpu/cuda)
|
||||
"""
|
||||
# Get the language model head weights (vocabulary_size x hidden_size)
|
||||
lm_head_weights = causal_lm.lm_head.weight
|
||||
|
||||
# Convert all tokens to their corresponding token IDs
|
||||
token_ids = [tokenizer.convert_tokens_to_ids(t) for t in tokens]
|
||||
|
||||
# Extract weights for the specific tokens (num_tokens x hidden_size)
|
||||
score_weight = lm_head_weights[token_ids].to(device)
|
||||
|
||||
# Copy the weights to the sequence classification model
|
||||
with torch.no_grad():
|
||||
seq_cls_model.score.weight.copy_(score_weight)
|
||||
if seq_cls_model.score.bias is not None:
|
||||
@@ -58,19 +110,33 @@ method_map = {
|
||||
def converting(
|
||||
model_name, classifier_from_tokens, path, method, use_pad_token=False, device="cpu"
|
||||
):
|
||||
assert method in method_map
|
||||
"""
|
||||
Main conversion function to transform a CausalLM model to SequenceClassification.
|
||||
|
||||
Args:
|
||||
model_name: Name or path of the pretrained model
|
||||
classifier_from_tokens: List of tokens used for classification
|
||||
path: Output path to save the converted model
|
||||
method: Conversion method ('from_2_way_softmax' or 'no_post_processing')
|
||||
use_pad_token: Whether to use padding token in the sequence classification model
|
||||
device: Device to load the model on ('cpu' or 'cuda')
|
||||
"""
|
||||
assert method in method_map, f"Unknown method: {method}"
|
||||
|
||||
# Determine number of labels based on conversion method
|
||||
if method == "from_2_way_softmax":
|
||||
assert len(classifier_from_tokens) == 2
|
||||
num_labels = 1
|
||||
else:
|
||||
num_labels = len(classifier_from_tokens)
|
||||
|
||||
# Load tokenizer and original causal language model
|
||||
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
|
||||
causal_lm = transformers.AutoModelForCausalLM.from_pretrained(
|
||||
model_name, device_map=device
|
||||
)
|
||||
|
||||
# Load an empty sequence classification model with the same architecture
|
||||
seq_cls_model = transformers.AutoModelForSequenceClassification.from_pretrained(
|
||||
model_name,
|
||||
num_labels=num_labels,
|
||||
@@ -78,14 +144,17 @@ def converting(
|
||||
device_map=device,
|
||||
)
|
||||
|
||||
# Apply the selected conversion method to transfer weights
|
||||
method_map[method](
|
||||
causal_lm, seq_cls_model, tokenizer, classifier_from_tokens, device
|
||||
)
|
||||
|
||||
# `llm as reranker` defaults to not using pad_token
|
||||
# Configure padding token settings
|
||||
# Note: Reranker models typically don't use padding tokens by default
|
||||
seq_cls_model.config.use_pad_token = use_pad_token
|
||||
seq_cls_model.config.pad_token_id = tokenizer.pad_token_id
|
||||
|
||||
# Save the converted model and tokenizer
|
||||
seq_cls_model.save_pretrained(path)
|
||||
tokenizer.save_pretrained(path)
|
||||
|
||||
@@ -99,25 +168,30 @@ def parse_args():
|
||||
"--model_name",
|
||||
type=str,
|
||||
default="BAAI/bge-reranker-v2-gemma",
|
||||
help="Model name",
|
||||
help="HuggingFace model name or local path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--classifier_from_tokens",
|
||||
type=str,
|
||||
default='["Yes"]',
|
||||
help="classifier from tokens",
|
||||
help="JSON string of tokens used for classification labels",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method", type=str, default="no_post_processing", help="Converting converting"
|
||||
"--method",
|
||||
type=str,
|
||||
default="no_post_processing",
|
||||
help="Conversion method to use",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-pad-token", action="store_true", help="Whether to use pad_token"
|
||||
"--use-pad-token",
|
||||
action="store_true",
|
||||
help="Enable padding token in the sequence classification model",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--path",
|
||||
type=str,
|
||||
default="./bge-reranker-v2-gemma-seq-cls",
|
||||
help="Path to save converted model",
|
||||
help="Output directory to save the converted model",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -1,89 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
from vllm import LLM
|
||||
|
||||
model_name = "Qwen/Qwen3-Reranker-0.6B"
|
||||
|
||||
# What is the difference between the official original version and one
|
||||
# that has been converted into a sequence classification model?
|
||||
# Qwen3-Reranker is a language model that doing reranker by using the
|
||||
# logits of "no" and "yes" tokens.
|
||||
# It needs to computing 151669 tokens logits, making this method extremely
|
||||
# inefficient, not to mention incompatible with the vllm score API.
|
||||
# A method for converting the original model into a sequence classification
|
||||
# model was proposed. See:https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
# Models converted offline using this method can not only be more efficient
|
||||
# and support the vllm score API, but also make the init parameters more
|
||||
# concise, for example.
|
||||
# llm = LLM(model="tomaarsen/Qwen3-Reranker-0.6B-seq-cls", runner="pooling")
|
||||
|
||||
# If you want to load the official original version, the init parameters are
|
||||
# as follows.
|
||||
|
||||
|
||||
def get_llm() -> LLM:
|
||||
"""Initializes and returns the LLM model for Qwen3-Reranker."""
|
||||
return LLM(
|
||||
model=model_name,
|
||||
runner="pooling",
|
||||
hf_overrides={
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# Why do we need hf_overrides for the official original version:
|
||||
# vllm converts it to Qwen3ForSequenceClassification when loaded for
|
||||
# better performance.
|
||||
# - Firstly, we need using `"architectures": ["Qwen3ForSequenceClassification"],`
|
||||
# to manually route to Qwen3ForSequenceClassification.
|
||||
# - Then, we will extract the vector corresponding to classifier_from_token
|
||||
# from lm_head using `"classifier_from_token": ["no", "yes"]`.
|
||||
# - Third, we will convert these two vectors into one vector. The use of
|
||||
# conversion logic is controlled by `using "is_original_qwen3_reranker": True`.
|
||||
|
||||
# Please use the query_template and document_template to format the query and
|
||||
# document for better reranker results.
|
||||
|
||||
prefix = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n'
|
||||
suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
|
||||
|
||||
query_template = "{prefix}<Instruct>: {instruction}\n<Query>: {query}\n"
|
||||
document_template = "<Document>: {doc}{suffix}"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
instruction = (
|
||||
"Given a web search query, retrieve relevant passages that answer the query"
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is the capital of China?",
|
||||
"Explain gravity",
|
||||
]
|
||||
|
||||
documents = [
|
||||
"The capital of China is Beijing.",
|
||||
"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
|
||||
]
|
||||
|
||||
queries = [
|
||||
query_template.format(prefix=prefix, instruction=instruction, query=query)
|
||||
for query in queries
|
||||
]
|
||||
documents = [document_template.format(doc=doc, suffix=suffix) for doc in documents]
|
||||
|
||||
llm = get_llm()
|
||||
outputs = llm.score(queries, documents)
|
||||
|
||||
print("-" * 30)
|
||||
print([output.outputs.score for output in outputs])
|
||||
print("-" * 30)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,27 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
from pathlib import Path
|
||||
|
||||
from vllm import LLM
|
||||
|
||||
model_name = "nvidia/llama-nemotron-rerank-1b-v2"
|
||||
|
||||
# Path to template file
|
||||
template_path = Path(__file__).parent / "template" / "nemotron-rerank.jinja"
|
||||
chat_template = template_path.read_text()
|
||||
|
||||
llm = LLM(model=model_name, runner="pooling", trust_remote_code=True)
|
||||
|
||||
query = "how much protein should a female eat?"
|
||||
documents = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
|
||||
"Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
|
||||
]
|
||||
|
||||
outputs = llm.score(query, documents, chat_template=chat_template)
|
||||
|
||||
print("-" * 30)
|
||||
print([output.outputs.score for output in outputs])
|
||||
print("-" * 30)
|
||||
@@ -1,46 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""
|
||||
Example of using the rerank API with template.
|
||||
|
||||
run:
|
||||
vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
url = "http://127.0.0.1:8000/rerank"
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
query = "how much protein should a female eat?"
|
||||
documents = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
|
||||
"Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
|
||||
]
|
||||
|
||||
data = {
|
||||
"model": "nvidia/llama-nemotron-rerank-1b-v2",
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
response = requests.post(url, headers=headers, json=data)
|
||||
|
||||
# Check the response
|
||||
if response.status_code == 200:
|
||||
print("Request successful!")
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
else:
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
104
examples/pooling/score/qwen3_reranker_offline.py
Normal file
104
examples/pooling/score/qwen3_reranker_offline.py
Normal file
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
What is the difference between the official original version and one
|
||||
that has been converted into a sequence classification model?
|
||||
|
||||
Qwen3-Reranker is a language model that doing reranker by using the
|
||||
logits of "no" and "yes" tokens.
|
||||
This requires computing logits for all 151,669 tokens in the vocabulary,
|
||||
making it inefficient and incompatible with vLLM's score() API.
|
||||
|
||||
A conversion method has been proposed to transform the original model into a
|
||||
sequence classification model. This converted model:
|
||||
1. Is significantly more efficient
|
||||
2. Fully supports vLLM's score() API
|
||||
3. Simplifies initialization parameters
|
||||
Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
|
||||
|
||||
For the converted model, initialization would simply be:
|
||||
llm = LLM(model="tomaarsen/Qwen3-Reranker-0.6B-seq-cls", runner="pooling")
|
||||
|
||||
This example demonstrates loading the ORIGINAL model with special overrides
|
||||
to make it compatible with vLLM's score API.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from vllm import LLM
|
||||
|
||||
model_name = "Qwen/Qwen3-Reranker-0.6B"
|
||||
|
||||
|
||||
def get_llm() -> LLM:
|
||||
"""
|
||||
Initializes and returns the LLM model for Qwen3-Reranker.
|
||||
|
||||
Returns:
|
||||
LLM: Configured vLLM instance for reranking tasks.
|
||||
|
||||
Note:
|
||||
This function loads the ORIGINAL Qwen3-Reranker model with specific
|
||||
overrides to make it compatible with vLLM's score API.
|
||||
"""
|
||||
return LLM(
|
||||
# Specify the original model from HuggingFace
|
||||
model=model_name,
|
||||
# Use pooling runner for score task
|
||||
runner="pooling",
|
||||
# HuggingFace model configuration overrides required for compatibility
|
||||
hf_overrides={
|
||||
# Manually route to sequence classification architecture
|
||||
# This tells vLLM to use Qwen3ForSequenceClassification instead of
|
||||
# the default Qwen3ForCausalLM
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
# Specify which token logits to extract from the language model head
|
||||
# The original reranker uses "no" and "yes" token logits for scoring
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
# Enable special handling for original Qwen3-Reranker models
|
||||
# This flag triggers conversion logic that transforms the two token
|
||||
# vectors into a single classification vector
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
# Load the Jinja template for formatting query-document pairs
|
||||
# The template ensures proper formatting for the reranker model
|
||||
template_home = Path(__file__).parent / "template"
|
||||
template_path = "qwen3_reranker.jinja"
|
||||
chat_template = (template_home / template_path).read_text()
|
||||
|
||||
# Sample queries for testing the reranker
|
||||
queries = [
|
||||
"What is the capital of China?",
|
||||
"Explain gravity",
|
||||
]
|
||||
|
||||
# Corresponding documents to be scored against each query
|
||||
documents = [
|
||||
"The capital of China is Beijing.",
|
||||
"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
|
||||
]
|
||||
|
||||
# Initialize the LLM model with the original Qwen3-Reranker configuration
|
||||
llm = get_llm()
|
||||
|
||||
# Compute relevance scores for each query-document pair
|
||||
# The score() method returns a relevance score for each pair
|
||||
# Higher scores indicate better relevance
|
||||
outputs = llm.score(queries, documents, chat_template=chat_template)
|
||||
|
||||
# Extract and print the relevance scores from the outputs
|
||||
# Each output contains a score representing query-document relevance
|
||||
print("-" * 30)
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
print("-" * 30)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
80
examples/pooling/score/qwen3_reranker_online.py
Normal file
80
examples/pooling/score/qwen3_reranker_online.py
Normal file
@@ -0,0 +1,80 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""
|
||||
What is the difference between the official original version and one
|
||||
that has been converted into a sequence classification model?
|
||||
|
||||
Qwen3-Reranker is a language model that doing reranker by using the
|
||||
logits of "no" and "yes" tokens.
|
||||
This requires computing logits for all 151,669 tokens in the vocabulary,
|
||||
making it inefficient and incompatible with vLLM's score() API.
|
||||
|
||||
A conversion method has been proposed to transform the original model into a
|
||||
sequence classification model. This converted model:
|
||||
1. Is significantly more efficient
|
||||
2. Fully supports vLLM's score() API
|
||||
3. Simplifies initialization parameters
|
||||
Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
|
||||
|
||||
For the converted model, initialization would simply be:
|
||||
vllm serve tomaarsen/Qwen3-Reranker-0.6B-seq-cls --runner pooling --chat-template examples/pooling/score/template/qwen3_reranker.jinja
|
||||
|
||||
This example demonstrates loading the ORIGINAL model with special overrides
|
||||
to make it compatible with vLLM's score API.
|
||||
vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
# URL of the vLLM server's score endpoint
|
||||
# Default vLLM server runs on localhost port 8000
|
||||
url = "http://127.0.0.1:8000/score"
|
||||
|
||||
# HTTP headers for the request
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
# Example queries & documents
|
||||
queries = [
|
||||
"What is the capital of China?",
|
||||
"Explain gravity",
|
||||
]
|
||||
documents = [
|
||||
"The capital of China is Beijing.",
|
||||
"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
|
||||
]
|
||||
|
||||
# Request payload for the score API
|
||||
data = {
|
||||
"model": "Qwen/Qwen3-Reranker-0.6B",
|
||||
"text_1": queries,
|
||||
"text_2": documents,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to send a score request to the vLLM server.
|
||||
|
||||
This function sends a POST request to the /score endpoint with
|
||||
the query and documents, then prints the relevance scores.
|
||||
"""
|
||||
# Send POST request to the vLLM server's score endpoint
|
||||
response = requests.post(url, headers=headers, json=data)
|
||||
|
||||
# Check if the request was successful
|
||||
if response.status_code == 200:
|
||||
print("Request successful!")
|
||||
# Pretty print the JSON response containing relevance scores
|
||||
# The response includes scores for each document's relevance to the query
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
else:
|
||||
# Handle request failure
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,3 @@
|
||||
A: {{ (messages | selectattr("role", "eq", "query") | first).content }}
|
||||
B: {{ (messages | selectattr("role", "eq", "document") | first).content }}
|
||||
Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'.
|
||||
8
examples/pooling/score/template/mxbai_rerank_v2.jinja
Normal file
8
examples/pooling/score/template/mxbai_rerank_v2.jinja
Normal file
@@ -0,0 +1,8 @@
|
||||
<|im_start|>system
|
||||
You are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
|
||||
<|im_start|>user
|
||||
query: {{ (messages | selectattr("role", "eq", "query") | first).content }}
|
||||
document: {{ (messages | selectattr("role", "eq", "document") | first).content }}
|
||||
You are a search relevance expert who evaluates how well documents match search queries. For each query-document pair, carefully analyze the semantic relationship between them, then provide your binary relevance judgment (0 for not relevant, 1 for relevant).
|
||||
Relevance:<|im_end|>
|
||||
<|im_start|>assistant
|
||||
11
examples/pooling/score/template/qwen3_reranker.jinja
Normal file
11
examples/pooling/score/template/qwen3_reranker.jinja
Normal file
@@ -0,0 +1,11 @@
|
||||
<|im_start|>system
|
||||
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".<|im_end|>
|
||||
<|im_start|>user
|
||||
<Instruct>: {{ messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a web search query, retrieve relevant passages that answer the query") }}
|
||||
<Query>: {{ messages | selectattr("role", "eq", "query") | map(attribute="content") | first }}
|
||||
<Document>: {{ messages | selectattr("role", "eq", "document") | map(attribute="content") | first }}<|im_end|>
|
||||
<|im_start|>assistant
|
||||
<think>
|
||||
|
||||
</think>
|
||||
|
||||
159
examples/pooling/score/using_template_offline.py
Normal file
159
examples/pooling/score/using_template_offline.py
Normal file
@@ -0,0 +1,159 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
def parse_args():
|
||||
"""Parse command line arguments for the reranking example.
|
||||
|
||||
This function sets up the argument parser with default values
|
||||
specific to reranking models, including the model name and
|
||||
runner type.
|
||||
"""
|
||||
parser = FlexibleArgumentParser()
|
||||
# Add all EngineArgs command line arguments to the parser
|
||||
parser = EngineArgs.add_cli_args(parser)
|
||||
|
||||
# Set default values specific to this reranking example
|
||||
# These defaults ensure the script works out-of-the-box for reranking tasks
|
||||
parser.set_defaults(
|
||||
model="nvidia/llama-nemotron-rerank-1b-v2", # Default reranking model
|
||||
runner="pooling", # Required for cross-encoder/reranking models
|
||||
trust_remote_code=True, # Allow loading models with custom code
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def get_chat_template(model: str) -> str:
|
||||
"""Load the appropriate chat template for the specified model.
|
||||
|
||||
Reranking models require specific prompt templates to format
|
||||
query-document pairs correctly. This function maps model names
|
||||
to their corresponding template files.
|
||||
"""
|
||||
# Directory containing all chat template files
|
||||
template_home = Path(__file__).parent / "template"
|
||||
|
||||
# Mapping from model names to their corresponding template files
|
||||
# Each reranking model has its own specific prompt format
|
||||
model_name_to_template_path_map = {
|
||||
"BAAI/bge-reranker-v2-gemma": "bge-reranker-v2-gemma.jinja",
|
||||
"Qwen/Qwen3-Reranker-0.6B": "qwen3_reranker.jinja",
|
||||
"Qwen/Qwen3-Reranker-4B": "qwen3_reranker.jinja",
|
||||
"Qwen/Qwen3-Reranker-8B": "qwen3_reranker.jinja",
|
||||
"tomaarsen/Qwen3-Reranker-0.6B-seq-cls": "qwen3_reranker.jinja",
|
||||
"tomaarsen/Qwen3-Reranker-4B-seq-cls": "qwen3_reranker.jinja",
|
||||
"tomaarsen/Qwen3-Reranker-8B-seq-cls": "qwen3_reranker.jinja",
|
||||
"mixedbread-ai/mxbai-rerank-base-v2": "mxbai_rerank_v2.jinja",
|
||||
"mixedbread-ai/mxbai-rerank-large-v2": "mxbai_rerank_v2.jinja",
|
||||
"nvidia/llama-nemotron-rerank-1b-v2": "nemotron-rerank.jinja",
|
||||
}
|
||||
|
||||
# Get the template filename for the specified model
|
||||
template_path = model_name_to_template_path_map.get(model)
|
||||
|
||||
if template_path is None:
|
||||
raise ValueError(f"This demo does not support model name: {model}.")
|
||||
|
||||
# Read and return the template content
|
||||
return (template_home / template_path).read_text()
|
||||
|
||||
|
||||
def get_hf_overrides(model: str) -> dict[str, Any]:
|
||||
"""Convert Large Language Models (LLMs) to Sequence Classification models.
|
||||
|
||||
note:
|
||||
Some reranking models require special configuration overrides to work
|
||||
correctly with vLLM's score API.
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_offline.py
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
|
||||
"""
|
||||
|
||||
model_name_to_hf_overrides_map = {
|
||||
"BAAI/bge-reranker-v2-gemma": {
|
||||
"architectures": ["GemmaForSequenceClassification"],
|
||||
"classifier_from_token": ["Yes"],
|
||||
"method": "no_post_processing",
|
||||
},
|
||||
"Qwen/Qwen3-Reranker-0.6B": {
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
"Qwen/Qwen3-Reranker-4B": {
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
"Qwen/Qwen3-Reranker-8B": {
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
"tomaarsen/Qwen3-Reranker-0.6B-seq-cls": {},
|
||||
"tomaarsen/Qwen3-Reranker-4B-seq-cls": {},
|
||||
"tomaarsen/Qwen3-Reranker-8B-seq-cls": {},
|
||||
"mixedbread-ai/mxbai-rerank-base-v2": {
|
||||
"architectures": ["Qwen2ForSequenceClassification"],
|
||||
"classifier_from_token": ["0", "1"],
|
||||
"method": "from_2_way_softmax",
|
||||
},
|
||||
"mixedbread-ai/mxbai-rerank-large-v2": {
|
||||
"architectures": ["Qwen2ForSequenceClassification"],
|
||||
"classifier_from_token": ["0", "1"],
|
||||
"method": "from_2_way_softmax",
|
||||
},
|
||||
"nvidia/llama-nemotron-rerank-1b-v2": {},
|
||||
}
|
||||
|
||||
hf_overrides = model_name_to_hf_overrides_map.get(model)
|
||||
|
||||
if hf_overrides is None:
|
||||
raise ValueError(f"This demo does not support model name: {model}.")
|
||||
|
||||
return hf_overrides
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
"""Main execution function for the reranking example."""
|
||||
|
||||
# Get the overrides for the specified model
|
||||
args.hf_overrides = get_hf_overrides(args.model)
|
||||
|
||||
# Initialize the LLM with all provided arguments
|
||||
llm = LLM(**vars(args))
|
||||
|
||||
# Example query for demonstration
|
||||
query = "how much protein should a female eat?"
|
||||
|
||||
# Example documents to be reranked based on relevance to the query
|
||||
documents = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
|
||||
"Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
|
||||
]
|
||||
|
||||
# Load the appropriate chat template for the selected model
|
||||
# The template formats query-document pairs for the reranking model
|
||||
chat_template = get_chat_template(args.model)
|
||||
|
||||
# Score documents based on relevance to the query
|
||||
# The score method returns relevance scores for each document
|
||||
outputs = llm.score(query, documents, chat_template=chat_template)
|
||||
|
||||
# Display the relevance scores
|
||||
# Higher scores indicate more relevant documents
|
||||
print("-" * 30)
|
||||
print([output.outputs.score for output in outputs])
|
||||
print("-" * 30)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
75
examples/pooling/score/using_template_online.py
Normal file
75
examples/pooling/score/using_template_online.py
Normal file
@@ -0,0 +1,75 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""
|
||||
Example of using the rerank API with template.
|
||||
|
||||
This script demonstrates how to interact with a vLLM server running
|
||||
a reranking model via the REST API.
|
||||
Before running this script, start the vLLM server with one of the
|
||||
supported reranking models using the commands below.
|
||||
|
||||
note:
|
||||
Some reranking models require special configuration overrides to work correctly
|
||||
with vLLM's score API.
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_online.py
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
|
||||
|
||||
run:
|
||||
vllm serve BAAI/bge-reranker-v2-gemma --hf_overrides '{"architectures": ["GemmaForSequenceClassification"],"classifier_from_token": ["Yes"],"method": "no_post_processing"}' --chat-template examples/pooling/score/template/bge-reranker-v2-gemma.jinja
|
||||
vllm serve tomaarsen/Qwen3-Reranker-0.6B-seq-cls --chat-template examples/pooling/score/template/qwen3_reranker.jinja
|
||||
vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}' --chat-template examples/pooling/score/template/mxbai_rerank_v2.jinja
|
||||
vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja
|
||||
vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
# URL of the vLLM server's rerank endpoint
|
||||
# Default vLLM server runs on localhost port 8000
|
||||
url = "http://127.0.0.1:8000/rerank"
|
||||
|
||||
# HTTP headers for the request
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
# Example query & documents
|
||||
query = "how much protein should a female eat?"
|
||||
documents = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
|
||||
"Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
|
||||
]
|
||||
|
||||
# Request payload for the rerank API
|
||||
data = {
|
||||
"model": "nvidia/llama-nemotron-rerank-1b-v2", # Model to use for reranking
|
||||
"query": query, # The query to score documents against
|
||||
"documents": documents, # List of documents to be scored
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to send a rerank request to the vLLM server.
|
||||
|
||||
This function sends a POST request to the /rerank endpoint with
|
||||
the query and documents, then prints the relevance scores.
|
||||
"""
|
||||
# Send POST request to the vLLM server's rerank endpoint
|
||||
response = requests.post(url, headers=headers, json=data)
|
||||
|
||||
# Check if the request was successful
|
||||
if response.status_code == 200:
|
||||
print("Request successful!")
|
||||
# Pretty print the JSON response containing relevance scores
|
||||
# The response includes scores for each document's relevance to the query
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
else:
|
||||
# Handle request failure
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user