main
py 103 lines 3.91 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # ruff: noqa: E501
4
5 """What is the difference between the official original version and one
6 that has been converted into a sequence classification model?
7
8 Qwen3-Reranker is a language model that doing reranker by using the
9 logits of "no" and "yes" tokens.
10 This requires computing logits for all 151,669 tokens in the vocabulary,
11 making it inefficient and incompatible with vLLM's score() API.
12
13 A conversion method has been proposed to transform the original model into a
14 sequence classification model. This converted model:
15 1. Is significantly more efficient
16 2. Fully supports vLLM's score() API
17 3. Simplifies initialization parameters
18 Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
19 Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
20
21 For the converted model, initialization would simply be:
22 llm = LLM(model="tomaarsen/Qwen3-Reranker-0.6B-seq-cls", runner="pooling")
23
24 This example demonstrates loading the ORIGINAL model with special overrides
25 to make it compatible with vLLM's score API.
26 """
27
28 from pathlib import Path
29
30 from vllm import LLM
31
32 model_name = "Qwen/Qwen3-Reranker-0.6B"
33
34
35 def get_llm() -> LLM:
36 """Initializes and returns the LLM model for Qwen3-Reranker.
37
38 Returns:
39 LLM: Configured vLLM instance for reranking tasks.
40
41 Note:
42 This function loads the ORIGINAL Qwen3-Reranker model with specific
43 overrides to make it compatible with vLLM's score API.
44
45 """
46 return LLM(
47 # Specify the original model from HuggingFace
48 model=model_name,
49 # Use pooling runner for score task
50 runner="pooling",
51 # HuggingFace model configuration overrides required for compatibility
52 hf_overrides={
53 # Manually route to sequence classification architecture
54 # This tells vLLM to use Qwen3ForSequenceClassification instead of
55 # the default Qwen3ForCausalLM
56 "architectures": ["Qwen3ForSequenceClassification"],
57 # Specify which token logits to extract from the language model head
58 # The original reranker uses "no" and "yes" token logits for scoring
59 "classifier_from_token": ["no", "yes"],
60 # Enable special handling for original Qwen3-Reranker models
61 # This flag triggers conversion logic that transforms the two token
62 # vectors into a single classification vector
63 "is_original_qwen3_reranker": True,
64 },
65 )
66
67
68 def main() -> None:
69 # Load the Jinja template for formatting query-document pairs
70 # The template ensures proper formatting for the reranker model
71 template_home = Path(__file__).parent / "template"
72 template_path = "qwen3_reranker.jinja"
73 chat_template = (template_home / template_path).read_text()
74
75 # Sample queries for testing the reranker
76 queries = [
77 "What is the capital of China?",
78 "Explain gravity",
79 ]
80
81 # Corresponding documents to be scored against each query
82 documents = [
83 "The capital of China is Beijing.",
84 "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.",
85 ]
86
87 # Initialize the LLM model with the original Qwen3-Reranker configuration
88 llm = get_llm()
89
90 # Compute relevance scores for each query-document pair
91 # The score() method returns a relevance score for each pair
92 # Higher scores indicate better relevance
93 outputs = llm.score(queries, documents, chat_template=chat_template)
94
95 # Extract and print the relevance scores from the outputs
96 # Each output contains a score representing query-document relevance
97 print("-" * 30)
98 print("Relevance scores:", [output.outputs.score for output in outputs])
99 print("-" * 30)
100
101
102 if __name__ == "__main__":
103 main()