| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # ruff: noqa: E501 |
| 4 | from argparse import Namespace |
| 5 | from pathlib import Path |
| 6 | from typing import Any |
| 7 | |
| 8 | from vllm import LLM, EngineArgs |
| 9 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 10 | |
| 11 | |
| 12 | def parse_args(): |
| 13 | """Parse command line arguments for the reranking example. |
| 14 | |
| 15 | This function sets up the argument parser with default values |
| 16 | specific to reranking models, including the model name and |
| 17 | runner type. |
| 18 | """ |
| 19 | parser = FlexibleArgumentParser() |
| 20 | # Add all EngineArgs command line arguments to the parser |
| 21 | parser = EngineArgs.add_cli_args(parser) |
| 22 | |
| 23 | # Set default values specific to this reranking example |
| 24 | # These defaults ensure the script works out-of-the-box for reranking tasks |
| 25 | parser.set_defaults( |
| 26 | model="nvidia/llama-nemotron-rerank-1b-v2", # Default reranking model |
| 27 | runner="pooling", # Required for cross-encoder/reranking models |
| 28 | trust_remote_code=True, # Allow loading models with custom code |
| 29 | ) |
| 30 | return parser.parse_args() |
| 31 | |
| 32 | |
| 33 | def get_chat_template(model: str) -> str: |
| 34 | """Load the appropriate chat template for the specified model. |
| 35 | |
| 36 | Reranking models require specific prompt templates to format |
| 37 | query-document pairs correctly. This function maps model names |
| 38 | to their corresponding template files. |
| 39 | """ |
| 40 | # Directory containing all chat template files |
| 41 | template_home = Path(__file__).parent / "template" |
| 42 | |
| 43 | # Mapping from model names to their corresponding template files |
| 44 | # Each reranking model has its own specific prompt format |
| 45 | model_name_to_template_path_map = { |
| 46 | "BAAI/bge-reranker-v2-gemma": "bge-reranker-v2-gemma.jinja", |
| 47 | "Qwen/Qwen3-Reranker-0.6B": "qwen3_reranker.jinja", |
| 48 | "Qwen/Qwen3-Reranker-4B": "qwen3_reranker.jinja", |
| 49 | "Qwen/Qwen3-Reranker-8B": "qwen3_reranker.jinja", |
| 50 | "tomaarsen/Qwen3-Reranker-0.6B-seq-cls": "qwen3_reranker.jinja", |
| 51 | "tomaarsen/Qwen3-Reranker-4B-seq-cls": "qwen3_reranker.jinja", |
| 52 | "tomaarsen/Qwen3-Reranker-8B-seq-cls": "qwen3_reranker.jinja", |
| 53 | "mixedbread-ai/mxbai-rerank-base-v2": "mxbai_rerank_v2.jinja", |
| 54 | "mixedbread-ai/mxbai-rerank-large-v2": "mxbai_rerank_v2.jinja", |
| 55 | "nvidia/llama-nemotron-rerank-1b-v2": "nemotron-rerank.jinja", |
| 56 | } |
| 57 | |
| 58 | # Get the template filename for the specified model |
| 59 | template_path = model_name_to_template_path_map.get(model) |
| 60 | |
| 61 | if template_path is None: |
| 62 | raise ValueError(f"This demo does not support model name: {model}.") |
| 63 | |
| 64 | # Read and return the template content |
| 65 | return (template_home / template_path).read_text() |
| 66 | |
| 67 | |
| 68 | def get_hf_overrides(model: str) -> dict[str, Any]: |
| 69 | """Convert Large Language Models (LLMs) to Sequence Classification models. |
| 70 | |
| 71 | Note: |
| 72 | Some reranking models require special configuration overrides to work |
| 73 | correctly with vLLM's score API. |
| 74 | Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_offline.py |
| 75 | Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py |
| 76 | |
| 77 | """ |
| 78 | model_name_to_hf_overrides_map = { |
| 79 | "BAAI/bge-reranker-v2-gemma": { |
| 80 | "architectures": ["GemmaForSequenceClassification"], |
| 81 | "classifier_from_token": ["Yes"], |
| 82 | "method": "no_post_processing", |
| 83 | }, |
| 84 | "Qwen/Qwen3-Reranker-0.6B": { |
| 85 | "architectures": ["Qwen3ForSequenceClassification"], |
| 86 | "classifier_from_token": ["no", "yes"], |
| 87 | "is_original_qwen3_reranker": True, |
| 88 | }, |
| 89 | "Qwen/Qwen3-Reranker-4B": { |
| 90 | "architectures": ["Qwen3ForSequenceClassification"], |
| 91 | "classifier_from_token": ["no", "yes"], |
| 92 | "is_original_qwen3_reranker": True, |
| 93 | }, |
| 94 | "Qwen/Qwen3-Reranker-8B": { |
| 95 | "architectures": ["Qwen3ForSequenceClassification"], |
| 96 | "classifier_from_token": ["no", "yes"], |
| 97 | "is_original_qwen3_reranker": True, |
| 98 | }, |
| 99 | "tomaarsen/Qwen3-Reranker-0.6B-seq-cls": {}, |
| 100 | "tomaarsen/Qwen3-Reranker-4B-seq-cls": {}, |
| 101 | "tomaarsen/Qwen3-Reranker-8B-seq-cls": {}, |
| 102 | "mixedbread-ai/mxbai-rerank-base-v2": { |
| 103 | "architectures": ["Qwen2ForSequenceClassification"], |
| 104 | "classifier_from_token": ["0", "1"], |
| 105 | "method": "from_2_way_softmax", |
| 106 | }, |
| 107 | "mixedbread-ai/mxbai-rerank-large-v2": { |
| 108 | "architectures": ["Qwen2ForSequenceClassification"], |
| 109 | "classifier_from_token": ["0", "1"], |
| 110 | "method": "from_2_way_softmax", |
| 111 | }, |
| 112 | "nvidia/llama-nemotron-rerank-1b-v2": {}, |
| 113 | } |
| 114 | |
| 115 | hf_overrides = model_name_to_hf_overrides_map.get(model) |
| 116 | |
| 117 | if hf_overrides is None: |
| 118 | raise ValueError(f"This demo does not support model name: {model}.") |
| 119 | |
| 120 | return hf_overrides |
| 121 | |
| 122 | |
| 123 | def main(args: Namespace): |
| 124 | """Main execution function for the reranking example.""" |
| 125 | # Get the overrides for the specified model |
| 126 | args.hf_overrides = get_hf_overrides(args.model) |
| 127 | |
| 128 | # Initialize the LLM with all provided arguments |
| 129 | llm = LLM(**vars(args)) |
| 130 | |
| 131 | # Example query for demonstration |
| 132 | query = "how much protein should a female eat?" |
| 133 | |
| 134 | # Example documents to be reranked based on relevance to the query |
| 135 | documents = [ |
| 136 | "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.", |
| 137 | "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.", |
| 138 | "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.", |
| 139 | ] |
| 140 | |
| 141 | # Load the appropriate chat template for the selected model |
| 142 | # The template formats query-document pairs for the reranking model |
| 143 | chat_template = get_chat_template(args.model) |
| 144 | |
| 145 | # Score documents based on relevance to the query |
| 146 | # The score method returns relevance scores for each document |
| 147 | outputs = llm.score(query, documents, chat_template=chat_template) |
| 148 | |
| 149 | # Display the relevance scores |
| 150 | # Higher scores indicate more relevant documents |
| 151 | print("-" * 30) |
| 152 | print([output.outputs.score for output in outputs]) |
| 153 | print("-" * 30) |
| 154 | |
| 155 | |
| 156 | if __name__ == "__main__": |
| 157 | args = parse_args() |
| 158 | main(args) |