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:
@@ -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.
|
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}
|
```{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
|
```shell
|
||||||
python3 -m sglang.launch_server \
|
python3 -m sglang.launch_server \
|
||||||
@@ -19,7 +37,7 @@ python3 -m sglang.launch_server \
|
|||||||
--port 30000
|
--port 30000
|
||||||
```
|
```
|
||||||
|
|
||||||
## Example Client Request
|
### Example Client Request
|
||||||
|
|
||||||
```python
|
```python
|
||||||
import requests
|
import requests
|
||||||
@@ -32,18 +50,264 @@ payload = {
|
|||||||
"documents": [
|
"documents": [
|
||||||
"hi",
|
"hi",
|
||||||
"The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China."
|
"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 = requests.post(url, json=payload)
|
||||||
response_json = response.json()
|
response_json = response.json()
|
||||||
|
|
||||||
for item in 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 |
|
- `query` (required): The query text to rank documents against
|
||||||
|------------------------------------------------|--------------------------------------|---------------|----------------------------------------------------------------------------------------------------------------------------------|
|
- `documents` (required): List of documents to be ranked
|
||||||
| **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. |
|
- `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.
|
||||||
|
|||||||
7
examples/chat_template/qwen3_reranker.jinja
Normal file
7
examples/chat_template/qwen3_reranker.jinja
Normal 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' }}
|
||||||
32
examples/chat_template/qwen3_vl_reranker.jinja
Normal file
32
examples/chat_template/qwen3_vl_reranker.jinja
Normal 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' }}
|
||||||
185
examples/runtime/qwen3_vl_reranker.py
Normal file
185
examples/runtime/qwen3_vl_reranker.py
Normal 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()
|
||||||
@@ -283,7 +283,7 @@ async def lifespan(fast_api_app: FastAPI):
|
|||||||
_global_state.tokenizer_manager
|
_global_state.tokenizer_manager
|
||||||
)
|
)
|
||||||
fast_api_app.state.openai_serving_rerank = OpenAIServingRerank(
|
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(
|
fast_api_app.state.openai_serving_tokenize = OpenAIServingTokenize(
|
||||||
_global_state.tokenizer_manager
|
_global_state.tokenizer_manager
|
||||||
|
|||||||
@@ -376,6 +376,15 @@ ChatCompletionMessageContentPart = Union[
|
|||||||
ChatCompletionMessageContentAudioPart,
|
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):
|
class FunctionResponse(BaseModel):
|
||||||
"""Function response."""
|
"""Function response."""
|
||||||
@@ -872,16 +881,61 @@ class ScoringResponse(BaseModel):
|
|||||||
|
|
||||||
|
|
||||||
class V1RerankReqInput(BaseModel):
|
class V1RerankReqInput(BaseModel):
|
||||||
query: str
|
query: RerankContent = Field(
|
||||||
documents: List[str]
|
...,
|
||||||
|
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):
|
class RerankResponse(BaseModel):
|
||||||
score: float
|
score: float
|
||||||
document: str
|
document: Optional[str] = None
|
||||||
index: int
|
index: int
|
||||||
meta_info: Optional[dict] = None
|
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):
|
class TokenizeRequest(BaseModel):
|
||||||
"""Request schema for the /tokenize endpoint."""
|
"""Request schema for the /tokenize endpoint."""
|
||||||
|
|||||||
@@ -5,19 +5,217 @@ from fastapi import Request
|
|||||||
from fastapi.responses import ORJSONResponse
|
from fastapi.responses import ORJSONResponse
|
||||||
|
|
||||||
from sglang.srt.entrypoints.openai.protocol import (
|
from sglang.srt.entrypoints.openai.protocol import (
|
||||||
|
ChatCompletionMessageContentImagePart,
|
||||||
|
ChatCompletionMessageContentTextPart,
|
||||||
|
ChatCompletionMessageContentVideoPart,
|
||||||
ErrorResponse,
|
ErrorResponse,
|
||||||
|
RerankContent,
|
||||||
RerankResponse,
|
RerankResponse,
|
||||||
V1RerankReqInput,
|
V1RerankReqInput,
|
||||||
)
|
)
|
||||||
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
|
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__)
|
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):
|
class OpenAIServingRerank(OpenAIServingBase):
|
||||||
"""Handler for /v1/rerank requests"""
|
"""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
|
# NOTE: /v1/rerank is not an official OpenAI endpoint. This module may be moved
|
||||||
# to another module in the future.
|
# to another module in the future.
|
||||||
|
|
||||||
@@ -48,32 +246,67 @@ class OpenAIServingRerank(OpenAIServingBase):
|
|||||||
self,
|
self,
|
||||||
request: V1RerankReqInput,
|
request: V1RerankReqInput,
|
||||||
raw_request: Request = None,
|
raw_request: Request = None,
|
||||||
) -> tuple[EmbeddingReqInput, V1RerankReqInput]:
|
) -> tuple[Union[EmbeddingReqInput, V1RerankReqInput], V1RerankReqInput]:
|
||||||
"""Convert OpenAI rerank request to internal embedding format"""
|
"""
|
||||||
# Create pairs of [query, document] for each document
|
Convert OpenAI rerank request to internal format.
|
||||||
pairs = []
|
|
||||||
for doc in request.documents:
|
|
||||||
pairs.append([request.query, doc])
|
|
||||||
|
|
||||||
adapted_request = EmbeddingReqInput(
|
- For Qwen3-VL reranker (multimodal decoder-only): keep the request.
|
||||||
text=pairs,
|
- For Qwen3 reranker (text-only decoder-only): keep the request and score via
|
||||||
is_cross_encoder_request=True,
|
`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
|
return adapted_request, request
|
||||||
|
|
||||||
async def _handle_non_streaming_request(
|
async def _handle_non_streaming_request(
|
||||||
self,
|
self,
|
||||||
adapted_request: EmbeddingReqInput,
|
adapted_request: Union[EmbeddingReqInput, V1RerankReqInput],
|
||||||
request: V1RerankReqInput,
|
request: V1RerankReqInput,
|
||||||
raw_request: Request,
|
raw_request: Request,
|
||||||
) -> Union[List[RerankResponse], ErrorResponse, ORJSONResponse]:
|
) -> Union[List[RerankResponse], ErrorResponse, ORJSONResponse]:
|
||||||
"""Handle the rerank request"""
|
"""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:
|
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(
|
ret = await self.tokenizer_manager.generate_request(
|
||||||
adapted_request, raw_request
|
adapted_request, raw_request
|
||||||
).__anext__()
|
).__anext__()
|
||||||
|
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
return self.create_error_response(str(e))
|
return self.create_error_response(str(e))
|
||||||
|
|
||||||
@@ -83,22 +316,290 @@ class OpenAIServingRerank(OpenAIServingBase):
|
|||||||
responses = self._build_rerank_response(ret, request)
|
responses = self._build_rerank_response(ret, request)
|
||||||
return responses
|
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(
|
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]:
|
) -> List[RerankResponse]:
|
||||||
"""Build the rerank response from generation results"""
|
"""Build the rerank response from generation results"""
|
||||||
responses = []
|
responses = []
|
||||||
for idx, ret_item in enumerate(ret):
|
for idx, item in enumerate(ret):
|
||||||
responses.append(
|
if isinstance(item, dict):
|
||||||
RerankResponse(
|
score_val = item.get("embedding")
|
||||||
score=ret_item["embedding"],
|
# Some rerank/reward models return scalar score as embedding[0].
|
||||||
document=request.documents[idx],
|
if isinstance(score_val, list):
|
||||||
index=idx,
|
if len(score_val) == 0 or not isinstance(
|
||||||
meta_info=ret_item["meta_info"],
|
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)
|
# Sort by score in descending order (highest relevance first)
|
||||||
responses.sort(key=lambda x: x.score, reverse=True)
|
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
|
return responses
|
||||||
|
|||||||
@@ -8,6 +8,55 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
|
|
||||||
class TokenizerManagerMultiItemMixin:
|
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):
|
def _initialize_multi_item_delimiter_text(self):
|
||||||
"""Initialize multi-item delimiter text from token ID after tokenizer is loaded."""
|
"""Initialize multi-item delimiter text from token ID after tokenizer is loaded."""
|
||||||
if (
|
if (
|
||||||
|
|||||||
309
test/registered/openai_server/basic/test_serving_rerank.py
Normal file
309
test/registered/openai_server/basic/test_serving_rerank.py
Normal 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)
|
||||||
Reference in New Issue
Block a user