From de94d793ad258e85a254be2e158dd4d37f846c19 Mon Sep 17 00:00:00 2001 From: shuwenn <47200617+alphabetc1@users.noreply.github.com> Date: Thu, 15 Jan 2026 00:45:46 +0800 Subject: [PATCH] feat: support qwen3(-VL) rerank scoring&chat template (#16403) Signed-off-by: Xinyuan Tong Co-authored-by: Xinyuan Tong --- docs/supported_models/rerank_models.md | 282 ++++++++- examples/chat_template/qwen3_reranker.jinja | 7 + .../chat_template/qwen3_vl_reranker.jinja | 32 ++ examples/runtime/qwen3_vl_reranker.py | 185 ++++++ python/sglang/srt/entrypoints/http_server.py | 2 +- .../sglang/srt/entrypoints/openai/protocol.py | 60 +- .../srt/entrypoints/openai/serving_rerank.py | 543 +++++++++++++++++- .../tokenizer_manager_multiitem_mixin.py | 49 ++ .../basic/test_serving_rerank.py | 309 ++++++++++ 9 files changed, 1435 insertions(+), 34 deletions(-) create mode 100644 examples/chat_template/qwen3_reranker.jinja create mode 100644 examples/chat_template/qwen3_vl_reranker.jinja create mode 100644 examples/runtime/qwen3_vl_reranker.py create mode 100644 test/registered/openai_server/basic/test_serving_rerank.py diff --git a/docs/supported_models/rerank_models.md b/docs/supported_models/rerank_models.md index b6f2ffa20..bb989128a 100644 --- a/docs/supported_models/rerank_models.md +++ b/docs/supported_models/rerank_models.md @@ -3,10 +3,28 @@ SGLang offers comprehensive support for rerank models by incorporating optimized serving frameworks with a flexible programming interface. This setup enables efficient processing of cross-encoder reranking tasks, improving the accuracy and relevance of search result ordering. SGLang’s design ensures high throughput and low latency during reranker model deployment, making it ideal for semantic-based result refinement in large-scale retrieval systems. ```{important} -They are executed with `--is-embedding` and some may require `--trust-remote-code` +Rerank models in SGLang fall into two categories: + +- **Cross-encoder rerank models**: run with `--is-embedding` (embedding runner). +- **Decoder-only rerank models**: run **without** `--is-embedding` and use next-token logprob scoring (yes/no). + - Text-only (e.g. Qwen3-Reranker) + - Multimodal (e.g. Qwen3-VL-Reranker): also supports image/video content + +Some models may require `--trust-remote-code`. ``` -## Example Launch Command +## Supported rerank models + +| Model Family (Rerank) | Example HuggingFace Identifier | Chat Template | Description | +|------------------------------------------------|--------------------------------------|---------------|----------------------------------------------------------------------------------------------------------------------------------| +| **BGE-Reranker (BgeRerankModel)** | `BAAI/bge-reranker-v2-m3` | N/A | Currently only support `attention-backend` `triton` and `torch_native`. High-performance cross-encoder reranker model from BAAI. Suitable for reranking search results based on semantic relevance. | +| **Qwen3-Reranker (decoder-only yes/no)** | `Qwen/Qwen3-Reranker-8B` | `examples/chat_template/qwen3_reranker.jinja` | Decoder-only reranker using next-token logprob scoring for labels (yes/no). Launch **without** `--is-embedding`. | +| **Qwen3-VL-Reranker (multimodal yes/no)** | `Qwen/Qwen3-VL-Reranker-2B` | `examples/chat_template/qwen3_vl_reranker.jinja` | Multimodal decoder-only reranker supporting text, images, and videos. Uses yes/no logprob scoring. Launch **without** `--is-embedding`. | + + +## Cross-Encoder Rerank (embedding runner) + +### Launch Command ```shell python3 -m sglang.launch_server \ @@ -19,7 +37,7 @@ python3 -m sglang.launch_server \ --port 30000 ``` -## Example Client Request +### Example Client Request ```python import requests @@ -32,18 +50,264 @@ payload = { "documents": [ "hi", "The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China." - ] + ], + "top_n": 1, + "return_documents": True } response = requests.post(url, json=payload) response_json = response.json() for item in response_json: - print(f"Score: {item['score']:.2f} - Document: '{item['document']}'") + if item.get("document"): + print(f"Score: {item['score']:.2f} - Document: '{item['document']}'") + else: + print(f"Score: {item['score']:.2f} - Index: {item['index']}") ``` -## Supported rerank models +**Request Parameters:** -| Model Family (Rerank) | Example HuggingFace Identifier | Chat Template | Description | -|------------------------------------------------|--------------------------------------|---------------|----------------------------------------------------------------------------------------------------------------------------------| -| **BGE-Reranker (BgeRerankModel)** | `BAAI/bge-reranker-v2-m3` | N/A | Currently only support `attention-backend` `triton` and `torch_native`. high-performance cross-encoder reranker model from BAAI. Suitable for reranking search results based on semantic relevance. | +- `query` (required): The query text to rank documents against +- `documents` (required): List of documents to be ranked +- `model` (required): Model to use for reranking +- `top_n` (optional): Maximum number of documents to return. Defaults to returning all documents. If specified value is greater than the total number of documents, all documents will be returned. +- `return_documents` (optional): Whether to return documents in the response. Defaults to `True`. + +## Qwen3-Reranker (decoder-only yes/no rerank) + +### Launch Command + +```shell +python3 -m sglang.launch_server \ + --model-path Qwen/Qwen3-Reranker-0.6B \ + --trust-remote-code \ + --disable-radix-cache \ + --host 0.0.0.0 \ + --port 8001 \ + --chat-template examples/chat_template/qwen3_reranker.jinja +``` + +```{note} +Qwen3-Reranker uses decoder-only logprob scoring (yes/no). Do NOT launch it with `--is-embedding`. +``` + +### Example Client Request (supports optional instruct, top_n, and return_documents) + +```shell +curl -X POST http://127.0.0.1:8001/v1/rerank \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen3-Reranker-0.6B", + "query": "法国首都是哪里？", + "documents": [ + "法国的首都是巴黎。", + "德国的首都是柏林。", + "香蕉是黄色的水果。" + ], + "instruct": "Given a web search query, retrieve relevant passages that answer the query.", + "top_n": 2, + "return_documents": true + }' +``` + +**Request Parameters:** + +- `query` (required): The query text to rank documents against +- `documents` (required): List of documents to be ranked +- `model` (required): Model to use for reranking +- `instruct` (optional): Instruction text for the reranker +- `top_n` (optional): Maximum number of documents to return. Defaults to returning all documents. If specified value is greater than the total number of documents, all documents will be returned. +- `return_documents` (optional): Whether to return documents in the response. Defaults to `True`. + +### Response Format + +`/v1/rerank` returns a list of objects (sorted by descending score): + +- `score`: float, higher means more relevant +- `document`: the original document string (only included when `return_documents` is `true`) +- `index`: the original index in the input `documents` +- `meta_info`: optional debug/usage info (may be present for some models) + +The number of returned results is controlled by the `top_n` parameter. If `top_n` is not specified or is greater than the total number of documents, all documents are returned. + +Example (with `return_documents: true`): + +```json +[ + {"score": 0.99, "document": "法国的首都是巴黎。", "index": 0}, + {"score": 0.01, "document": "德国的首都是柏林。", "index": 1}, + {"score": 0.00, "document": "香蕉是黄色的水果。", "index": 2} +] +``` + +Example (with `return_documents: false`): + +```json +[ + {"score": 0.99, "index": 0}, + {"score": 0.01, "index": 1}, + {"score": 0.00, "index": 2} +] +``` + +Example (with `top_n: 2`): + +```json +[ + {"score": 0.99, "document": "法国的首都是巴黎。", "index": 0}, + {"score": 0.01, "document": "德国的首都是柏林。", "index": 1} +] +``` + +### Common Pitfalls + +- If you launch Qwen3-Reranker with `--is-embedding`, `/v1/rerank` cannot compute yes/no logprob scores. Relaunch **without** `--is-embedding`. +- If you see a validation error like "score should be a valid number" and the backend returned a list, upgrade to a version that coerces `embedding[0]` into `score` for rerank responses. + +## Qwen3-VL-Reranker (multimodal decoder-only rerank) + +Qwen3-VL-Reranker extends the Qwen3-Reranker to support multimodal content, allowing reranking of documents containing text, images, and videos. + +### Launch Command + +```shell +python3 -m sglang.launch_server \ + --model-path Qwen/Qwen3-VL-Reranker-2B \ + --trust-remote-code \ + --disable-radix-cache \ + --host 0.0.0.0 \ + --port 30000 \ + --chat-template examples/chat_template/qwen3_vl_reranker.jinja +``` + +```{note} +Qwen3-VL-Reranker uses decoder-only logprob scoring (yes/no) like Qwen3-Reranker. Do NOT launch it with `--is-embedding`. +``` + +### Text-Only Reranking (backward compatible) + +```python +import requests + +url = "http://127.0.0.1:30000/v1/rerank" + +payload = { + "model": "Qwen3-VL-Reranker-2B", + "query": "What is machine learning?", + "documents": [ + "Machine learning is a branch of artificial intelligence that enables computers to learn from data.", + "The weather in Paris is usually mild with occasional rain.", + "Deep learning is a subset of machine learning using neural networks with many layers.", + ], + "instruct": "Retrieve passages that answer the question.", + "return_documents": True +} + +response = requests.post(url, json=payload) +results = response.json() + +for item in results: + print(f"Score: {item['score']:.4f} - {item['document'][:60]}...") +``` + +### Image Reranking (text query, image/mixed documents) + +```python +import requests + +url = "http://127.0.0.1:30000/v1/rerank" + +payload = { + "query": "A woman playing with her dog on a beach at sunset.", + "documents": [ + # Document 1: Text description + "A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset.", + # Document 2: Image URL + [ + { + "type": "image_url", + "image_url": { + "url": "https://example.com/beach_dog.jpeg" + } + } + ], + # Document 3: Text + Image (mixed) + [ + {"type": "text", "text": "A joyful scene at the beach:"}, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/beach_dog.jpeg" + } + } + ] + ], + "instruct": "Retrieve images or text relevant to the user's query.", + "return_documents": False +} + +response = requests.post(url, json=payload) +results = response.json() + +for item in results: + print(f"Index: {item['index']}, Score: {item['score']:.4f}") +``` + +### Multimodal Query Reranking (query with image) + +```python +import requests + +url = "http://127.0.0.1:30000/v1/rerank" + +payload = { + # Query with text and image + "query": [ + {"type": "text", "text": "Find similar images to this:"}, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/reference_image.jpeg" + } + } + ], + "documents": [ + "A cat sleeping on a couch.", + "A woman and her dog enjoying the sunset at the beach.", + "A busy city street with cars and pedestrians.", + [ + { + "type": "image_url", + "image_url": { + "url": "https://example.com/similar_image.jpeg" + } + } + ] + ], + "instruct": "Find images or descriptions similar to the query image." +} + +response = requests.post(url, json=payload) +results = response.json() + +for item in results: + print(f"Index: {item['index']}, Score: {item['score']:.4f}") +``` + +### Request Parameters (Multimodal) + +- `query` (required): Can be a string (text-only) or a list of content parts: + - `{"type": "text", "text": "..."}` for text + - `{"type": "image_url", "image_url": {"url": "..."}}` for images + - `{"type": "video_url", "video_url": {"url": "..."}}` for videos +- `documents` (required): List where each document can be a string or list of content parts (same format as query) +- `instruct` (optional): Instruction text for the reranker +- `top_n` (optional): Maximum number of documents to return +- `return_documents` (optional): Whether to return documents in the response (default: `false`) + +### Common Pitfalls + +- Always use `--chat-template examples/chat_template/qwen3_vl_reranker.jinja` for Qwen3-VL-Reranker. +- Do NOT launch with `--is-embedding`. +- For best results, use `--disable-radix-cache` to avoid caching issues with multimodal content. +- **Note**: Currently only `Qwen3-VL-Reranker-2B` is tested and supported. The 8B model may have different behavior and is not guaranteed to work with this template. diff --git a/examples/chat_template/qwen3_reranker.jinja b/examples/chat_template/qwen3_reranker.jinja new file mode 100644 index 000000000..5ab809eea --- /dev/null +++ b/examples/chat_template/qwen3_reranker.jinja @@ -0,0 +1,7 @@ +<|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 | default("Given a web search query, retrieve relevant passages that answer the query.") }} +: {{ messages[0]["content"] }} +: {{ messages[1]["content"] }}<|im_end|> +<|im_start|>assistant{{ '\n' }} diff --git a/examples/chat_template/qwen3_vl_reranker.jinja b/examples/chat_template/qwen3_vl_reranker.jinja new file mode 100644 index 000000000..30447a80f --- /dev/null +++ b/examples/chat_template/qwen3_vl_reranker.jinja @@ -0,0 +1,32 @@ +{#- Qwen3-VL-Reranker chat template for multimodal reranking -#} +{#- This template formats query-document pairs for yes/no relevance judgment -#} +{#- Supports text, images, and videos in both query and documents -#} +<|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 | default("Given a search query, retrieve relevant candidates that answer the query.") }} +{#- Process query content -#} +: {%- for content in query -%} + {%- if content.type == 'image' or 'image' in content or 'image_url' in content -%} + <|vision_start|><|image_pad|><|vision_end|> + {%- elif content.type == 'video' or 'video' in content -%} + <|vision_start|><|video_pad|><|vision_end|> + {%- elif 'text' in content -%} + {{ content.text }} + {%- elif content.type == 'text' -%} + {{ content.text }} + {%- endif -%} +{%- endfor %} +{#- Process document content -#} +{{ '\n' }}: {%- for content in document -%} + {%- if content.type == 'image' or 'image' in content or 'image_url' in content -%} + <|vision_start|><|image_pad|><|vision_end|> + {%- elif content.type == 'video' or 'video' in content -%} + <|vision_start|><|video_pad|><|vision_end|> + {%- elif 'text' in content -%} + {{ content.text }} + {%- elif content.type == 'text' -%} + {{ content.text }} + {%- endif -%} +{%- endfor %}<|im_end|> +<|im_start|>assistant{{ '\n' }} diff --git a/examples/runtime/qwen3_vl_reranker.py b/examples/runtime/qwen3_vl_reranker.py new file mode 100644 index 000000000..09779996f --- /dev/null +++ b/examples/runtime/qwen3_vl_reranker.py @@ -0,0 +1,185 @@ +""" +Example usage of Qwen3-VL-Reranker with SGLang. + +This example demonstrates how to use the Qwen3-VL-Reranker model for multimodal +reranking tasks, supporting text, images, and videos. + +Server Launch: + python -m sglang.launch_server \ + --model-path Qwen/Qwen3-VL-Reranker-2B \ + --served-model-name Qwen3-VL-Reranker-2B \ + --trust-remote-code \ + --disable-radix-cache \ + --chat-template examples/chat_template/qwen3_vl_reranker.jinja + +Client Usage: + python examples/runtime/qwen3_vl_reranker.py +""" + +import requests + +# Server URL +BASE_URL = "http://localhost:30000" + + +def rerank_text_only(): + """Example: Text-only reranking (backward compatible).""" + print("=" * 60) + print("Text-only reranking example") + print("=" * 60) + + request_data = { + "query": "What is machine learning?", + "documents": [ + "Machine learning is a branch of artificial intelligence that enables computers to learn from data.", + "The weather in Paris is usually mild with occasional rain.", + "Deep learning is a subset of machine learning using neural networks with many layers.", + ], + "instruct": "Retrieve passages that answer the question.", + "return_documents": True, + } + + response = requests.post(f"{BASE_URL}/v1/rerank", json=request_data) + results = response.json() + + print("Results (sorted by relevance):") + for i, result in enumerate(results): + print(f" {i+1}. Score: {result['score']:.4f} - {result['document'][:60]}...") + print() + + +def rerank_with_images(): + """Example: Query is text, documents contain images.""" + print("=" * 60) + print("Image reranking example") + print("=" * 60) + + request_data = { + "query": "A woman playing with her dog on a beach at sunset.", + "documents": [ + # Document 1: Text description + "A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset.", + # Document 2: Image URL + [ + { + "type": "image_url", + "image_url": { + "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" + }, + } + ], + # Document 3: Text + Image (mixed) + [ + { + "type": "text", + "text": "A joyful scene at the beach:", + }, + { + "type": "image_url", + "image_url": { + "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" + }, + }, + ], + ], + "instruct": "Retrieve images or text relevant to the user's query.", + "return_documents": False, + } + + response = requests.post(f"{BASE_URL}/v1/rerank", json=request_data) + results = response.json() + + # Debug: print raw response if it's an error + if isinstance(results, dict) and "message" in results: + print(f"Error: {results['message']}") + return + if isinstance(results, str): + print(f"Error: {results}") + return + + print("Results (sorted by relevance):") + for i, result in enumerate(results): + print(f" {i+1}. Index: {result['index']}, Score: {result['score']:.4f}") + print() + + +def rerank_multimodal_query(): + """Example: Query contains both text and image.""" + print("=" * 60) + print("Multimodal query reranking example") + print("=" * 60) + + request_data = { + # Query with text and image + "query": [ + {"type": "text", "text": "Find similar images to this:"}, + { + "type": "image_url", + "image_url": { + "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" + }, + }, + ], + "documents": [ + "A cat sleeping on a couch.", + "A woman and her dog enjoying the sunset at the beach.", + "A busy city street with cars and pedestrians.", + [ + { + "type": "image_url", + "image_url": { + "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" + }, + } + ], + ], + "instruct": "Find images or descriptions similar to the query image.", + } + + response = requests.post(f"{BASE_URL}/v1/rerank", json=request_data) + results = response.json() + + # Debug: print raw response if it's an error + if isinstance(results, dict) and "message" in results: + print(f"Error: {results['message']}") + return + if isinstance(results, str): + print(f"Error: {results}") + return + + print("Results (sorted by relevance):") + for i, result in enumerate(results): + print(f" {i+1}. Index: {result['index']}, Score: {result['score']:.4f}") + print() + + +def main(): + """Run all examples.""" + print("\nQwen3-VL-Reranker Examples") + print("Make sure the server is running with the correct model and template.\n") + + # Check if server is available + try: + response = requests.get(f"{BASE_URL}/health") + if response.status_code != 200: + print(f"Server health check failed: {response.status_code}") + return + except requests.exceptions.ConnectionError: + print(f"Cannot connect to server at {BASE_URL}") + print("Please start the server first with:") + print(" python -m sglang.launch_server \\") + print(" --model-path Qwen/Qwen3-VL-Reranker-2B \\") + print(" --served-model-name Qwen3-VL-Reranker-2B \\") + print(" --trust-remote-code \\") + print(" --disable-radix-cache \\") + print(" --chat-template examples/chat_template/qwen3_vl_reranker.jinja") + return + + # Run examples + rerank_text_only() + rerank_with_images() + rerank_multimodal_query() + + +if __name__ == "__main__": + main() diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 1d44d66bd..652f5b061 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -283,7 +283,7 @@ async def lifespan(fast_api_app: FastAPI): _global_state.tokenizer_manager ) fast_api_app.state.openai_serving_rerank = OpenAIServingRerank( - _global_state.tokenizer_manager + _global_state.tokenizer_manager, _global_state.template_manager ) fast_api_app.state.openai_serving_tokenize = OpenAIServingTokenize( _global_state.tokenizer_manager diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 2204b66db..192852c9a 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -376,6 +376,15 @@ ChatCompletionMessageContentPart = Union[ ChatCompletionMessageContentAudioPart, ] +# Rerank content types for multimodal reranking (e.g., Qwen3-VL-Reranker) +# Can be a simple string (text-only) or a list of multimodal content parts +RerankContentPart = Union[ + ChatCompletionMessageContentTextPart, + ChatCompletionMessageContentImagePart, + ChatCompletionMessageContentVideoPart, +] +RerankContent = Union[str, List[RerankContentPart]] + class FunctionResponse(BaseModel): """Function response.""" @@ -872,16 +881,61 @@ class ScoringResponse(BaseModel): class V1RerankReqInput(BaseModel): - query: str - documents: List[str] + query: RerankContent = Field( + ..., + description="The query to match against documents. Can be a string (text-only) " + "or a list of content parts for multimodal queries (text, image_url, video_url).", + ) + documents: List[RerankContent] = Field( + ..., + description="List of documents to rank. Each document can be a string (text-only) " + "or a list of content parts for multimodal documents (text, image_url, video_url).", + ) + instruct: Optional[str] = Field( + default=None, + description="The instruct to the reranker model.", + ) + top_n: Optional[int] = Field( + default=None, + description="Maximum number of documents to return. Defaults to returning all documents. " + "If specified value is greater than the total number of documents, all documents will be returned.", + ) + return_documents: bool = Field( + default=True, + description="Whether to return documents in the response. Only included when set to true.", + ) + + @field_validator("top_n") + @classmethod + def validate_top_n(cls, v): + if v is not None and v < 1: + raise ValueError("Value error, parameter top_n should be larger than 0.") + return v + + def is_multimodal(self) -> bool: + """Check if the request contains any multimodal content.""" + if isinstance(self.query, list): + return True + for doc in self.documents: + if isinstance(doc, list): + return True + return False class RerankResponse(BaseModel): score: float - document: str + document: Optional[str] = None index: int meta_info: Optional[dict] = None + @model_serializer(mode="wrap") + def _serialize(self, handler): + data = handler(self) + # Exclude document field if it's None + if self.document is None: + data.pop("document", None) + return data + class TokenizeRequest(BaseModel): """Request schema for the /tokenize endpoint.""" diff --git a/python/sglang/srt/entrypoints/openai/serving_rerank.py b/python/sglang/srt/entrypoints/openai/serving_rerank.py index 128215896..c5e80f1d8 100644 --- a/python/sglang/srt/entrypoints/openai/serving_rerank.py +++ b/python/sglang/srt/entrypoints/openai/serving_rerank.py @@ -5,19 +5,217 @@ from fastapi import Request from fastapi.responses import ORJSONResponse from sglang.srt.entrypoints.openai.protocol import ( + ChatCompletionMessageContentImagePart, + ChatCompletionMessageContentTextPart, + ChatCompletionMessageContentVideoPart, ErrorResponse, + RerankContent, RerankResponse, V1RerankReqInput, ) from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase -from sglang.srt.managers.io_struct import EmbeddingReqInput +from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput logger = logging.getLogger(__name__) +def _get_yes_no_token_ids(tokenizer) -> tuple[int, int]: + """Get token IDs for 'yes' and 'no' from the tokenizer. + + Different model sizes may have different token IDs, so we look them up dynamically. + """ + # Try to encode 'yes' and 'no' to get their token IDs + # The tokenizer should return a single token for these common words + try: + yes_tokens = tokenizer.encode("yes", add_special_tokens=False) + no_tokens = tokenizer.encode("no", add_special_tokens=False) + + if len(yes_tokens) == 1 and len(no_tokens) == 1: + return yes_tokens[0], no_tokens[0] + + # Fallback: try convert_tokens_to_ids + yes_id = tokenizer.convert_tokens_to_ids("yes") + no_id = tokenizer.convert_tokens_to_ids("no") + if yes_id is not None and no_id is not None: + return yes_id, no_id + + except Exception as e: + logger.warning(f"Failed to get yes/no token IDs dynamically: {e}") + + # Fallback to known Qwen3 token IDs (may not work for all model sizes) + logger.warning("Using fallback token IDs for yes/no (9693/2152)") + return 9693, 2152 + + +def _is_qwen3_reranker_template(chat_template: str) -> bool: + """Detect if the chat template is for Qwen3 text-only reranker.""" + if not chat_template: + return False + t = chat_template.lower() + return ('answer can only be "yes" or "no"' in t) or ( + "answer can only be" in t and '"yes"' in t and '"no"' in t + ) + + +def _is_qwen3_vl_reranker_template(chat_template: str) -> bool: + """Detect if the chat template is for Qwen3-VL multimodal reranker. + + VL reranker templates use `query` and `document` as jinja variables + and include vision token placeholders for image/video support. + """ + if not chat_template: + return False + t = chat_template.lower() + # Check for reranker phrase (yes/no judgment) + has_reranker_phrase = ('answer can only be "yes" or "no"' in t) or ( + "answer can only be" in t and '"yes"' in t and '"no"' in t + ) + # Check for vision token placeholders (unique to VL templates) + has_vision_tokens = "<|vision_start|>" in t or "<|image_pad|>" in t + return has_reranker_phrase and has_vision_tokens + + +def _is_qwen3_vl_model(model_path: str) -> bool: + """Check if the model is a Qwen3-VL model based on model path.""" + if not model_path: + return False + model_lower = model_path.lower() + return "qwen3-vl" in model_lower or "qwen3vl" in model_lower + + +def _detect_rerank_backend( + *, + request: V1RerankReqInput, + chat_template: Optional[str], + model_path: str, +) -> str: + """ + Unify rerank routing decisions used by both `_convert_to_internal_request` and + `_handle_non_streaming_request`. + + Returns: + "vl_decoder" | "text_decoder" | "cross_encoder" + """ + is_multimodal = request.is_multimodal() + is_vl_model = _is_qwen3_vl_model(model_path) + is_vl_template = _is_qwen3_vl_reranker_template(chat_template) + is_text_template = _is_qwen3_reranker_template(chat_template) + + # Prefer VL when template/model indicates VL, or request is multimodal with reranker template. + if is_vl_template or is_vl_model or (is_multimodal and is_text_template): + return "vl_decoder" + if is_text_template: + return "text_decoder" + return "cross_encoder" + + +def _qwen3_rerank_score(p_yes: float, p_no: float) -> float: + denom = p_yes + p_no + if denom <= 0.0: + return 0.0 + return p_yes / denom + + +def _get_jinja_env(): + try: + import jinja2 # Lazy import: server env should provide this dependency. + except ModuleNotFoundError as e: + raise ValueError( + "Rendering Qwen3 reranker prompts requires `jinja2`. " + "Please install it in your runtime environment (e.g., `pip install jinja2`)." + ) from e + + return jinja2.Environment( + loader=jinja2.BaseLoader(), + autoescape=False, + undefined=jinja2.Undefined, + ) + + +def _render_jinja_chat_template( + chat_template: str, + *, + query: RerankContent, + document: RerankContent, + instruct: Optional[str], +) -> str: + """Render a loaded Jinja chat template for Qwen3 reranker prompts (text-only).""" + env = _get_jinja_env() + template = env.from_string(chat_template) + + # For text-only template, extract text content + query_text = query if isinstance(query, str) else _extract_text_from_content(query) + doc_text = ( + document if isinstance(document, str) else _extract_text_from_content(document) + ) + + render_kwargs = { + "messages": [ + {"role": "user", "content": query_text}, + {"role": "user", "content": doc_text}, + ] + } + # Only pass instruct when explicitly provided; template uses `default(...)` + # which works only when the variable is undefined (not None). + if instruct: + render_kwargs["instruct"] = instruct + return template.render(**render_kwargs) + + +def _render_vl_jinja_template( + chat_template: str, + *, + query: List[Dict[str, Any]], + document: List[Dict[str, Any]], + instruct: Optional[str], +) -> str: + """Render a loaded Jinja chat template for Qwen3-VL reranker prompts (multimodal). + + The template expects `query` and `document` as lists of content parts, + where each part has a `type` field (text, image, video) and corresponding data. + """ + env = _get_jinja_env() + template = env.from_string(chat_template) + + render_kwargs = { + "query": query, + "document": document, + } + if instruct: + render_kwargs["instruct"] = instruct + return template.render(**render_kwargs) + + +def _extract_text_from_content(content: RerankContent) -> str: + """Extract text from multimodal content.""" + if isinstance(content, str): + return content + texts = [] + for part in content: + if isinstance(part, ChatCompletionMessageContentTextPart): + texts.append(part.text) + elif isinstance(part, dict) and part.get("type") == "text": + texts.append(part.get("text", "")) + return " ".join(texts) + + class OpenAIServingRerank(OpenAIServingBase): """Handler for /v1/rerank requests""" + def __init__(self, tokenizer_manager, template_manager=None): + super().__init__(tokenizer_manager) + # TemplateManager is optional; rerank uses tokenizer.chat_template today. + # Keeping this explicit makes the dependency clear and supports future extensions. + self.template_manager = template_manager + + # Cache yes/no token IDs for Qwen3 reranker scoring + self._yes_token_id, self._no_token_id = _get_yes_no_token_ids( + tokenizer_manager.tokenizer + ) + logger.info( + f"Reranker yes/no token IDs: yes={self._yes_token_id}, no={self._no_token_id}" + ) + # NOTE: /v1/rerank is not an official OpenAI endpoint. This module may be moved # to another module in the future. @@ -48,32 +246,67 @@ class OpenAIServingRerank(OpenAIServingBase): self, request: V1RerankReqInput, raw_request: Request = None, - ) -> tuple[EmbeddingReqInput, V1RerankReqInput]: - """Convert OpenAI rerank request to internal embedding format""" - # Create pairs of [query, document] for each document - pairs = [] - for doc in request.documents: - pairs.append([request.query, doc]) + ) -> tuple[Union[EmbeddingReqInput, V1RerankReqInput], V1RerankReqInput]: + """ + Convert OpenAI rerank request to internal format. - adapted_request = EmbeddingReqInput( - text=pairs, - is_cross_encoder_request=True, + - For Qwen3-VL reranker (multimodal decoder-only): keep the request. + - For Qwen3 reranker (text-only decoder-only): keep the request and score via + `tokenizer_manager.score_prompts(...)` in the handler. + - For cross-encoder rerank models: adapt into `EmbeddingReqInput` pairs. + """ + chat_template = self.tokenizer_manager.tokenizer.chat_template + model_path = getattr(self.tokenizer_manager.model_config, "model_path", "") + backend = _detect_rerank_backend( + request=request, + chat_template=chat_template if isinstance(chat_template, str) else None, + model_path=model_path, ) + if backend in ("vl_decoder", "text_decoder"): + return request, request + # Cross-encoder rerank: Create pairs of [query, document] for each document. + # Note: Cross-encoder only supports text-only content + if request.is_multimodal(): + # Extract text for cross-encoder (multimodal not supported) + query_text = _extract_text_from_content(request.query) + doc_texts = [_extract_text_from_content(doc) for doc in request.documents] + pairs = [[query_text, doc] for doc in doc_texts] + else: + pairs = [[request.query, doc] for doc in request.documents] + + adapted_request = EmbeddingReqInput(text=pairs, is_cross_encoder_request=True) return adapted_request, request async def _handle_non_streaming_request( self, - adapted_request: EmbeddingReqInput, + adapted_request: Union[EmbeddingReqInput, V1RerankReqInput], request: V1RerankReqInput, raw_request: Request, ) -> Union[List[RerankResponse], ErrorResponse, ORJSONResponse]: """Handle the rerank request""" + chat_template = getattr(self.tokenizer_manager.tokenizer, "chat_template", None) + model_path = getattr(self.tokenizer_manager.model_config, "model_path", "") + rerank_ret = await self._handle_rerank_paths( + request=request, + raw_request=raw_request, + chat_template=chat_template, + model_path=model_path, + ) + if rerank_ret is not None: + return rerank_ret + + # Default cross-encoder rerank path (existing behavior). try: + if not isinstance(adapted_request, EmbeddingReqInput): + raise ValueError( + "Invalid rerank request adaptation. " + "If you are serving a decoder-only reranker (e.g., Qwen3-Reranker), " + "please provide the corresponding --chat-template and launch without --is-embedding." + ) ret = await self.tokenizer_manager.generate_request( adapted_request, raw_request ).__anext__() - except ValueError as e: return self.create_error_response(str(e)) @@ -83,22 +316,290 @@ class OpenAIServingRerank(OpenAIServingBase): responses = self._build_rerank_response(ret, request) return responses + async def _handle_rerank_paths( + self, + *, + request: V1RerankReqInput, + raw_request: Request, + chat_template: Optional[str], + model_path: str, + ) -> Optional[Union[List[RerankResponse], ErrorResponse, ORJSONResponse]]: + """ + Handle decoder-only rerank paths (VL/text) and return a response if matched. + + Returns None if the request should fall back to cross-encoder rerank. + """ + backend = _detect_rerank_backend( + request=request, + chat_template=chat_template, + model_path=model_path, + ) + + # Qwen3-VL reranker path (decoder-only scoring with query/document template format) + if backend == "vl_decoder": + return await self._handle_vl_reranker_request( + request, raw_request, chat_template or "" + ) + + # Qwen3 text-only reranker path (decoder-only scoring). + if backend == "text_decoder": + return await self._handle_text_reranker_request( + request=request, + raw_request=raw_request, + chat_template=chat_template or "", + ) + + return None + + async def _handle_text_reranker_request( + self, + *, + request: V1RerankReqInput, + raw_request: Request, + chat_template: str, + ) -> Union[List[RerankResponse], ErrorResponse]: + """Handle text-only decoder reranker request via score_prompts().""" + # Qwen3 reranker relies on decoder-only logprobs. If the server is launched + # with --is-embedding, model_config.is_generation is typically False and + # logprob scoring is not supported. + if not self.tokenizer_manager.model_config.is_generation: + return self.create_error_response( + "Detected Qwen3 reranker chat template, but the server is not in generation mode. " + "Please relaunch without --is-embedding for Qwen3-Reranker models." + ) + + try: + prompts = [ + _render_jinja_chat_template( + chat_template, + query=request.query, + document=doc, + instruct=getattr(request, "instruct", None), + ) + for doc in request.documents + ] + + probs = await self.tokenizer_manager.score_prompts( + prompts, + label_token_ids=[self._yes_token_id, self._no_token_id], + apply_softmax=False, + request=raw_request, + ) + scores = [_qwen3_rerank_score(p[0], p[1]) for p in probs] + except ValueError as e: + return self.create_error_response(str(e)) + except Exception as e: + # Includes template rendering errors from jinja2. + return self.create_error_response(str(e)) + + responses = self._build_rerank_response(scores, request) + return responses + + async def _handle_vl_reranker_request( + self, + request: V1RerankReqInput, + raw_request: Request, + _chat_template: str, + ) -> Union[List[RerankResponse], ErrorResponse]: + """Handle multimodal VL reranker request using chat completion with logprobs.""" + if not self.tokenizer_manager.model_config.is_generation: + return self.create_error_response( + "Detected Qwen3-VL reranker, but the server is not in generation mode. " + "Please relaunch without --is-embedding for Qwen3-VL-Reranker models." + ) + + try: + scores = [] + instruct = getattr(request, "instruct", None) + + for doc in request.documents: + # Build multimodal content lists and render prompt using jinja template + query_content, doc_content, image_data, video_data = ( + self._build_vl_reranker_content( + query=request.query, + document=doc, + ) + ) + + # Render the chat template directly with query/document variables + prompt = _render_vl_jinja_template( + chat_template=_chat_template, + query=query_content, + document=doc_content, + instruct=instruct, + ) + + # Create generate request with logprobs + gen_request = GenerateReqInput( + text=prompt, + image_data=image_data if image_data else None, + video_data=video_data if video_data else None, + sampling_params={ + "max_new_tokens": 1, + "temperature": 0, + }, + return_logprob=True, + top_logprobs_num=50, # Get enough logprobs to find yes/no tokens + logprob_start_len=0, + ) + + # Execute generation request + ret = await self.tokenizer_manager.generate_request( + gen_request, raw_request + ).__anext__() + + # Extract yes/no probabilities from logprobs + score = self._extract_score_from_logprobs(ret) + scores.append(score) + + responses = self._build_rerank_response(scores, request) + return responses + + except ValueError as e: + return self.create_error_response(str(e)) + except Exception as e: + logger.exception("Error handling VL reranker request") + return self.create_error_response(str(e)) + + def _build_vl_reranker_content( + self, + query: RerankContent, + document: RerankContent, + ) -> tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[str], List[str]]: + """Build content lists for VL reranker request. + + Returns: + Tuple of (query_content, document_content, image_data, video_data) + where query_content and document_content are lists suitable for jinja template. + """ + image_data = [] + video_data = [] + + # Build query content list + query_content = self._content_to_template_list(query, image_data, video_data) + + # Build document content list + doc_content = self._content_to_template_list(document, image_data, video_data) + + return query_content, doc_content, image_data, video_data + + def _content_to_template_list( + self, + content: RerankContent, + image_data: List[str], + video_data: List[str], + ) -> List[Dict[str, Any]]: + """Convert RerankContent to a list format suitable for jinja template.""" + result = [] + + if isinstance(content, str): + result.append({"type": "text", "text": content}) + return result + + for part in content: + if isinstance(part, ChatCompletionMessageContentTextPart): + result.append({"type": "text", "text": part.text}) + elif isinstance(part, ChatCompletionMessageContentImagePart): + if part.image_url: + image_data.append(part.image_url.url) + result.append({"type": "image"}) + elif isinstance(part, ChatCompletionMessageContentVideoPart): + if part.video_url: + video_data.append(part.video_url.url) + result.append({"type": "video"}) + elif isinstance(part, dict): + part_type = part.get("type") + if part_type == "text": + result.append({"type": "text", "text": part.get("text", "")}) + elif part_type == "image_url": + image_url = part.get("image_url", {}) + if isinstance(image_url, dict): + url = image_url.get("url") + else: + url = image_url + if url: + image_data.append(url) + result.append({"type": "image"}) + elif part_type == "video_url": + video_url = part.get("video_url", {}) + if isinstance(video_url, dict): + url = video_url.get("url") + else: + url = video_url + if url: + video_data.append(url) + result.append({"type": "video"}) + + return result + + def _extract_score_from_logprobs(self, ret: Dict[str, Any]) -> float: + """Extract reranking score from generation response with logprobs.""" + import math + + # Get logprobs from the response + meta_info = ret.get("meta_info", {}) + output_top_logprobs = meta_info.get("output_top_logprobs", []) + + # Use output_top_logprobs[0] - the model's prediction for the first generated token + top_logprobs = output_top_logprobs[0] if output_top_logprobs else [] + + # Find yes and no token probabilities + # Format: list of tuples (logprob, token_id, token_text) + p_yes = 0.0 + p_no = 0.0 + + for item in top_logprobs: + logprob, token_id = item[0], item[1] + if token_id == self._yes_token_id: + p_yes = math.exp(logprob) + elif token_id == self._no_token_id: + p_no = math.exp(logprob) + + return _qwen3_rerank_score(p_yes, p_no) + def _build_rerank_response( - self, ret: List[Dict[str, Any]], request: V1RerankReqInput + self, ret: Union[List[Dict[str, Any]], List[float]], request: V1RerankReqInput ) -> List[RerankResponse]: """Build the rerank response from generation results""" responses = [] - for idx, ret_item in enumerate(ret): - responses.append( - RerankResponse( - score=ret_item["embedding"], - document=request.documents[idx], - index=idx, - meta_info=ret_item["meta_info"], + for idx, item in enumerate(ret): + if isinstance(item, dict): + score_val = item.get("embedding") + # Some rerank/reward models return scalar score as embedding[0]. + if isinstance(score_val, list): + if len(score_val) == 0 or not isinstance( + score_val[0], (int, float) + ): + raise ValueError( + f"Invalid embedding score for rerank at index {idx}: {score_val!r}" + ) + score_val = float(score_val[0]) + responses.append( + RerankResponse( + score=float(score_val), + document=( + request.documents[idx] if request.return_documents else None + ), + index=idx, + meta_info=item.get("meta_info"), + ) + ) + else: + responses.append( + RerankResponse( + score=float(item), + document=( + request.documents[idx] if request.return_documents else None + ), + index=idx, + ) ) - ) # Sort by score in descending order (highest relevance first) responses.sort(key=lambda x: x.score, reverse=True) + # Apply top_n limit if specified + if request.top_n is not None and request.top_n > 0: + responses = responses[: request.top_n] + return responses diff --git a/python/sglang/srt/managers/tokenizer_manager_multiitem_mixin.py b/python/sglang/srt/managers/tokenizer_manager_multiitem_mixin.py index 55f9442b2..2ab5dd11c 100644 --- a/python/sglang/srt/managers/tokenizer_manager_multiitem_mixin.py +++ b/python/sglang/srt/managers/tokenizer_manager_multiitem_mixin.py @@ -8,6 +8,55 @@ logger = logging.getLogger(__name__) class TokenizerManagerMultiItemMixin: + async def score_prompts( + self, + prompts: Union[str, List[str], List[List[int]]], + label_token_ids: List[int], + apply_softmax: bool = False, + request: Optional[Any] = None, + ) -> List[List[float]]: + """ + Score probabilities of specified token IDs after each *full prompt*. + + This is a thin wrapper over `score_request` that treats `prompts` as + already-composed inputs (i.e., no query/item concatenation needed). + + Args: + prompts: A single prompt string, a list of prompt strings, or a list of + pre-tokenized prompt token ID sequences. + label_token_ids: Token IDs to compute probabilities for. + apply_softmax: Whether to normalize probabilities using softmax. + request: Optional FastAPI request object. + + Returns: + List of score lists, one for each prompt, each in the order of label_token_ids. + """ + # Text prompts + if isinstance(prompts, str) or ( + isinstance(prompts, list) and (not prompts or isinstance(prompts[0], str)) + ): + return await self.score_request( + query="", + items=prompts, # type: ignore[arg-type] + label_token_ids=label_token_ids, + apply_softmax=apply_softmax, + item_first=False, + request=request, + ) + + # Tokenized prompts + if isinstance(prompts, list) and (not prompts or isinstance(prompts[0], list)): + return await self.score_request( + query=[], + items=prompts, + label_token_ids=label_token_ids, + apply_softmax=apply_softmax, + item_first=False, + request=request, + ) + + raise ValueError("Invalid prompts type for score_prompts.") + def _initialize_multi_item_delimiter_text(self): """Initialize multi-item delimiter text from token ID after tokenizer is loaded.""" if ( diff --git a/test/registered/openai_server/basic/test_serving_rerank.py b/test/registered/openai_server/basic/test_serving_rerank.py new file mode 100644 index 000000000..51d6708d0 --- /dev/null +++ b/test/registered/openai_server/basic/test_serving_rerank.py @@ -0,0 +1,309 @@ +import asyncio +import unittest +from unittest.mock import Mock + +from sglang.srt.entrypoints.openai.protocol import V1RerankReqInput +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +# Keep consistent with other openai_server/basic unit tests. +register_cuda_ci(est_time=10, suite="stage-b-test-small-1-gpu") +register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu") + +try: + from sglang.srt.entrypoints.openai.serving_rerank import ( + OpenAIServingRerank, + _is_qwen3_reranker_template, + _qwen3_rerank_score, + _render_jinja_chat_template, + ) +except ModuleNotFoundError as e: + # Some minimal environments used for unit tests may not have FastAPI/torch installed. + # Skip this test in that case. + if e.name in ("fastapi", "torch"): + OpenAIServingRerank = None # type: ignore[assignment] + else: + raise + + +class _DummyModelConfig: + # Keep consistent with TokenizerManager.model_config usage + is_generation = False + + +class _DummyTokenizer: + chat_template = "" + + +class _DummyTokenizerManager: + # Minimal surface required by OpenAIServingBase/OpenAIServingRerank + server_args = object() + model_config = _DummyModelConfig() + tokenizer = _DummyTokenizer() + + async def generate_request(self, *_args, **_kwargs): + raise AssertionError("generate_request should not be called in this unit test") + + +@unittest.skipIf(OpenAIServingRerank is None, "fastapi/torch is not installed") +class TestOpenAIServingRerankUnit(unittest.TestCase): + def setUp(self): + self.handler = OpenAIServingRerank(_DummyTokenizerManager()) + + def test_convert_to_internal_request_cross_encoder_pairs(self): + req = V1RerankReqInput( + query="q", + documents=["doc-a", "doc-b"], + instruct="Retrieve semantically similar text.", + ) + + adapted, processed = self.handler._convert_to_internal_request(req) + + # Avoid importing EmbeddingReqInput (requires torch). Use duck-typing checks instead. + self.assertTrue(hasattr(adapted, "is_cross_encoder_request")) + self.assertTrue(adapted.is_cross_encoder_request) + self.assertEqual(getattr(adapted, "text"), [["q", "doc-a"], ["q", "doc-b"]]) + self.assertEqual(processed, req) + + def test_convert_to_internal_request_qwen3_template_returns_request(self): + tm = _DummyTokenizerManager() + tm.tokenizer.chat_template = ( + '... Note that the answer can only be "yes" or "no". ...' + ) + handler = OpenAIServingRerank(tm) + req = V1RerankReqInput(query="q", documents=["d1"]) + adapted, processed = handler._convert_to_internal_request(req) + self.assertIs(adapted, req) + self.assertIs(processed, req) + + def test_build_rerank_response_embedding_list_uses_first_scalar(self): + req = V1RerankReqInput( + query="q", + documents=["doc-a", "doc-b"], + return_documents=True, + ) + # Two results with embedding as list, should coerce embedding[0] to float. + # Also verifies sorting (doc-b > doc-a). + ret = [ + {"embedding": [0.1, 0.2], "meta_info": {"id": "a"}}, + {"embedding": [0.9, -1.0], "meta_info": {"id": "b"}}, + ] + + res = self.handler._build_rerank_response(ret, req) + + self.assertEqual(len(res), 2) + + # Sorted descending by score, so doc-b first. + self.assertEqual(res[0].document, "doc-b") + self.assertEqual(res[0].index, 1) + self.assertAlmostEqual(res[0].score, 0.9) + self.assertEqual(res[0].meta_info, {"id": "b"}) + + self.assertEqual(res[1].document, "doc-a") + self.assertEqual(res[1].index, 0) + self.assertAlmostEqual(res[1].score, 0.1) + self.assertEqual(res[1].meta_info, {"id": "a"}) + + def test_build_rerank_response_float_list(self): + req = V1RerankReqInput( + query="q", documents=["a", "b", "c"], return_documents=True + ) + scores = [0.2, 0.9, 0.1] + res = self.handler._build_rerank_response(scores, req) + self.assertEqual([r.document for r in res], ["b", "a", "c"]) + self.assertEqual([r.index for r in res], [1, 0, 2]) + self.assertAlmostEqual(res[0].score, 0.9) + self.assertAlmostEqual(res[1].score, 0.2) + self.assertAlmostEqual(res[2].score, 0.1) + + def test_helper_is_qwen3_reranker_template(self): + self.assertTrue( + _is_qwen3_reranker_template( + 'Note that the answer can only be "yes" or "no".' + ) + ) + self.assertFalse(_is_qwen3_reranker_template("plain template")) + + def test_helper_qwen3_rerank_score(self): + self.assertAlmostEqual(_qwen3_rerank_score(0.9, 0.1), 0.9) + self.assertAlmostEqual(_qwen3_rerank_score(0.0, 0.0), 0.0) + + def test_helper_render_jinja_chat_template(self): + # Skip if jinja2 isn't installed in this environment. + try: + import jinja2 # noqa: F401 + except ModuleNotFoundError: + self.skipTest("jinja2 is not installed") + + tpl = "{{ instruct | default('DEF') }}|{{ messages[0]['content'] }}|{{ messages[1]['content'] }}" + self.assertEqual( + _render_jinja_chat_template(tpl, query="Q", document="D", instruct=None), + "DEF|Q|D", + ) + self.assertEqual( + _render_jinja_chat_template(tpl, query="Q", document="D", instruct="I"), + "I|Q|D", + ) + + def test_handle_non_streaming_request_qwen3_path_uses_score_prompts(self): + class _TM(_DummyTokenizerManager): + def __init__(self): + self.server_args = object() + self.model_config = Mock() + self.model_config.is_generation = True + self.model_config.model_path = "qwen/qwen3" + self.tokenizer = Mock() + self.tokenizer.chat_template = ( + 'Note that the answer can only be "yes" or "no". ' + "{{ messages[0]['content'] }} {{ messages[1]['content'] }}" + ) + + async def score_prompts( + self, prompts, label_token_ids, apply_softmax, request + ): + # Return [p_yes, p_no] for each prompt + assert len(prompts) == 2 + assert label_token_ids and len(label_token_ids) == 2 + return [[0.9, 0.1], [0.2, 0.8]] + + handler = OpenAIServingRerank(_TM()) + req = V1RerankReqInput(query="q", documents=["d1", "d2"], return_documents=True) + adapted, _ = handler._convert_to_internal_request(req) + raw_request = Mock() + + res = asyncio.run( + handler._handle_non_streaming_request(adapted, req, raw_request) + ) + self.assertEqual([r.document for r in res], ["d1", "d2"]) + self.assertAlmostEqual(res[0].score, 0.9 / (0.9 + 0.1)) + self.assertAlmostEqual(res[1].score, 0.2 / (0.2 + 0.8)) + + def test_build_rerank_response_return_documents_false(self): + """Test that document field is None when return_documents=False""" + req = V1RerankReqInput( + query="q", documents=["a", "b", "c"], return_documents=False + ) + scores = [0.2, 0.9, 0.1] + res = self.handler._build_rerank_response(scores, req) + # All documents should be None + self.assertEqual([r.document for r in res], [None, None, None]) + # But scores and indices should still be correct + self.assertEqual([r.index for r in res], [1, 0, 2]) + self.assertAlmostEqual(res[0].score, 0.9) + + def test_build_rerank_response_top_n(self): + """Test that top_n limits the number of returned results""" + req = V1RerankReqInput( + query="q", documents=["a", "b", "c"], return_documents=True, top_n=2 + ) + scores = [0.2, 0.9, 0.1] + res = self.handler._build_rerank_response(scores, req) + # Should only return top 2 results + self.assertEqual(len(res), 2) + self.assertEqual([r.document for r in res], ["b", "a"]) + self.assertEqual([r.index for r in res], [1, 0]) + self.assertAlmostEqual(res[0].score, 0.9) + self.assertAlmostEqual(res[1].score, 0.2) + + def test_build_rerank_response_top_n_greater_than_total(self): + """Test that top_n greater than total documents returns all documents""" + req = V1RerankReqInput( + query="q", documents=["a", "b"], return_documents=True, top_n=10 + ) + scores = [0.2, 0.9] + res = self.handler._build_rerank_response(scores, req) + # Should return all 2 documents even though top_n=10 + self.assertEqual(len(res), 2) + self.assertEqual([r.document for r in res], ["b", "a"]) + + def test_build_rerank_response_top_n_with_return_documents_false(self): + """Test top_n works correctly with return_documents=False""" + req = V1RerankReqInput( + query="q", documents=["a", "b", "c"], return_documents=False, top_n=1 + ) + scores = [0.2, 0.9, 0.1] + res = self.handler._build_rerank_response(scores, req) + # Should only return top 1 result, and document should be None + self.assertEqual(len(res), 1) + self.assertIsNone(res[0].document) + self.assertEqual(res[0].index, 1) + self.assertAlmostEqual(res[0].score, 0.9) + + def test_handle_vl_reranker_request(self): + """Test the Qwen3-VL reranker path with mocked logprobs.""" + import math + + # Mock tokenizer manager that supports generate_request + class _AsyncGen: + def __init__(self, val): + self.val = val + + def __aiter__(self): + return self + + async def __anext__(self): + return self.val + + class _TM(_DummyTokenizerManager): + def __init__(self): + self.server_args = object() + self.model_config = Mock() + self.model_config.is_generation = True + self.model_config.model_path = "qwen/qwen3-vl" + self.tokenizer = Mock() + # Mock VL template detection + self.tokenizer.chat_template = ( + "{% for x in query %}{{ x.text }}{% endfor %}" + "{% for x in document %}{{ x.text }}{% endfor %}" + 'answer can only be "yes" or "no" <|vision_start|>' + ) + + async def generate_request(self, req, _raw): + # Return logprobs for yes/no + # Mock logprobs: P(yes) > P(no) for first doc, P(no) > P(yes) for second + + if not hasattr(self, "call_count"): + self.call_count = 0 + + if self.call_count == 0: + # First doc: yes is likely + yes_logprob = math.log(0.8) + no_logprob = math.log(0.2) + else: + # Second doc: no is likely + yes_logprob = math.log(0.3) + no_logprob = math.log(0.7) + + self.call_count += 1 + + # Qwen3 token IDs: YES=9693, NO=2152 + top_logprobs = [ + (yes_logprob, 9693, "yes"), + (no_logprob, 2152, "no"), + ] + + # The rerank handler checks output_top_logprobs[0] for the first generated token + meta_info = {"output_top_logprobs": [top_logprobs]} + + yield {"meta_info": meta_info, "embedding": None} + + handler = OpenAIServingRerank(_TM()) + req = V1RerankReqInput( + query="query", documents=["doc1", "doc2"], return_documents=True + ) + # Force VL path is handled by detection logic inside handler + # We mocked chat_template to satisfy _is_qwen3_vl_reranker_template + + raw_request = Mock() + res = asyncio.run(handler._handle_non_streaming_request(req, req, raw_request)) + + self.assertEqual(len(res), 2) + # First doc should have higher score + self.assertEqual(res[0].document, "doc1") + self.assertAlmostEqual(res[0].score, 0.8) # 0.8 / (0.8+0.2) = 0.8 + + self.assertEqual(res[1].document, "doc2") + self.assertAlmostEqual(res[1].score, 0.3) # 0.3 / (0.3+0.7) = 0.3 + + +if __name__ == "__main__": + unittest.main(verbosity=2)