feat: support qwen3(-VL) rerank scoring&chat template (#16403)

Signed-off-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
This commit is contained in:
shuwenn
2026-01-15 00:45:46 +08:00
committed by GitHub
parent 0d904ef44c
commit de94d793ad
9 changed files with 1435 additions and 34 deletions

View File

@@ -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. SGLangs 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.

View File

@@ -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>: {{ instruct | default("Given a web search query, retrieve relevant passages that answer the query.") }}
<Query>: {{ messages[0]["content"] }}
<Document>: {{ messages[1]["content"] }}<|im_end|>
<|im_start|>assistant{{ '\n' }}

View File

@@ -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>: {{ instruct | default("Given a search query, retrieve relevant candidates that answer the query.") }}
{#- Process query content -#}
<Query>: {%- 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' }}<Document>: {%- 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' }}

View File

@@ -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()

View File

@@ -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

View File

@@ -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."""

View File

@@ -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

View File

@@ -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 (

View File

@@ -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)