Add vLLM v0.18.1 source tree with KV transfer abort fix
third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:
vllm/v1/core/sched/scheduler.py:
Replace fatal assert with graceful skip when KV transfer callback
arrives for an already-aborted request during PD disaggregated serving.
Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
67
third_party/vllm/examples/pooling/classify/classification_online.py
vendored
Normal file
67
third_party/vllm/examples/pooling/classify/classification_online.py
vendored
Normal file
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Example Python client for classification API using vLLM API server
|
||||
NOTE:
|
||||
start a supported classification model server with `vllm serve`, e.g.
|
||||
vllm serve jason9693/Qwen2.5-1.5B-apeach
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parse = argparse.ArgumentParser()
|
||||
parse.add_argument("--host", type=str, default="localhost")
|
||||
parse.add_argument("--port", type=int, default=8000)
|
||||
return parse.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
models_url = base_url + "/v1/models"
|
||||
classify_url = base_url + "/classify"
|
||||
tokenize_url = base_url + "/tokenize"
|
||||
|
||||
response = requests.get(models_url, headers=headers)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
# /classify can accept str as input
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"input": prompts,
|
||||
}
|
||||
response = requests.post(classify_url, headers=headers, json=payload)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
# /classify can accept token ids as input
|
||||
token_ids = []
|
||||
for prompt in prompts:
|
||||
response = requests.post(
|
||||
tokenize_url,
|
||||
json={"model": model, "prompt": prompt},
|
||||
)
|
||||
token_ids.append(response.json()["tokens"])
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"input": token_ids,
|
||||
}
|
||||
response = requests.post(classify_url, headers=headers, json=payload)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
110
third_party/vllm/examples/pooling/classify/vision_classification_online.py
vendored
Normal file
110
third_party/vllm/examples/pooling/classify/vision_classification_online.py
vendored
Normal file
@@ -0,0 +1,110 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""Example Python client for multimodal classification API using vLLM API server
|
||||
NOTE:
|
||||
start a supported multimodal classification model server with `vllm serve`, e.g.
|
||||
vllm serve muziyongshixin/Qwen2.5-VL-7B-for-VideoCls \
|
||||
--runner pooling \
|
||||
--max-model-len 5000 \
|
||||
--limit-mm-per-prompt.video 1 \
|
||||
--hf-overrides '{"architectures": ["Qwen2_5_VLForSequenceClassification"]}'
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
input_text = "This product was excellent and exceeded my expectations"
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
image_base64 = {"url": encode_image_url(fetch_image(image_url))}
|
||||
video_url = "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4"
|
||||
|
||||
|
||||
def parse_args():
|
||||
parse = argparse.ArgumentParser()
|
||||
parse.add_argument("--host", type=str, default="localhost")
|
||||
parse.add_argument("--port", type=int, default=8000)
|
||||
return parse.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
models_url = base_url + "/v1/models"
|
||||
classify_url = base_url + "/classify"
|
||||
|
||||
response = requests.get(models_url)
|
||||
model_name = response.json()["data"][0]["id"]
|
||||
|
||||
print("Text classification output:")
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Please classify this text request.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": input_text,
|
||||
},
|
||||
]
|
||||
response = requests.post(
|
||||
classify_url,
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Image url classification output:")
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Please classify this image."},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
]
|
||||
response = requests.post(
|
||||
classify_url,
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Image base64 classification output:")
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Please classify this image."},
|
||||
{"type": "image_url", "image_url": image_base64},
|
||||
],
|
||||
}
|
||||
]
|
||||
response = requests.post(
|
||||
classify_url,
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Video url classification output:")
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Please classify this video."},
|
||||
{"type": "video_url", "video_url": {"url": video_url}},
|
||||
],
|
||||
}
|
||||
]
|
||||
response = requests.post(
|
||||
classify_url,
|
||||
json={"model": model_name, "messages": messages},
|
||||
)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
60
third_party/vllm/examples/pooling/embed/embed_jina_embeddings_v3_offline.py
vendored
Normal file
60
third_party/vllm/examples/pooling/embed/embed_jina_embeddings_v3_offline.py
vendored
Normal file
@@ -0,0 +1,60 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser()
|
||||
parser = EngineArgs.add_cli_args(parser)
|
||||
# Set example specific arguments
|
||||
parser.set_defaults(
|
||||
model="jinaai/jina-embeddings-v3",
|
||||
runner="pooling",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
# Sample prompts.
|
||||
prompts = [
|
||||
"Follow the white rabbit.", # English
|
||||
"Sigue al conejo blanco.", # Spanish
|
||||
"Suis le lapin blanc.", # French
|
||||
"跟着白兔走。", # Chinese
|
||||
"اتبع الأرنب الأبيض.", # Arabic
|
||||
"Folge dem weißen Kaninchen.", # German
|
||||
]
|
||||
|
||||
# Create an LLM.
|
||||
# You should pass runner="pooling" for embedding models
|
||||
llm = LLM(**vars(args))
|
||||
|
||||
# Generate embedding. The output is a list of EmbeddingRequestOutputs.
|
||||
# Only text matching task is supported for now. See #16120
|
||||
outputs = llm.embed(prompts)
|
||||
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:")
|
||||
print("Only text matching task is supported for now. See #16120")
|
||||
print("-" * 60)
|
||||
for prompt, output in zip(prompts, outputs):
|
||||
embeds = output.outputs.embedding
|
||||
embeds_trimmed = (
|
||||
(str(embeds[:16])[:-1] + ", ...]") if len(embeds) > 16 else embeds
|
||||
)
|
||||
print(
|
||||
f"Prompt: {prompt!r} \n"
|
||||
f"Embeddings for text matching: {embeds_trimmed} "
|
||||
f"(size={len(embeds)})"
|
||||
)
|
||||
print("-" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
54
third_party/vllm/examples/pooling/embed/embed_matryoshka_fy_offline.py
vendored
Normal file
54
third_party/vllm/examples/pooling/embed/embed_matryoshka_fy_offline.py
vendored
Normal file
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs, PoolingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser()
|
||||
parser = EngineArgs.add_cli_args(parser)
|
||||
# Set example specific arguments
|
||||
parser.set_defaults(
|
||||
model="jinaai/jina-embeddings-v3",
|
||||
runner="pooling",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
# Sample prompts.
|
||||
prompts = [
|
||||
"Follow the white rabbit.", # English
|
||||
"Sigue al conejo blanco.", # Spanish
|
||||
"Suis le lapin blanc.", # French
|
||||
"跟着白兔走。", # Chinese
|
||||
"اتبع الأرنب الأبيض.", # Arabic
|
||||
"Folge dem weißen Kaninchen.", # German
|
||||
]
|
||||
|
||||
# Create an LLM.
|
||||
# You should pass runner="pooling" for embedding models
|
||||
llm = LLM(**vars(args))
|
||||
|
||||
# Generate embedding. The output is a list of EmbeddingRequestOutputs.
|
||||
outputs = llm.embed(prompts, pooling_params=PoolingParams(dimensions=32))
|
||||
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:")
|
||||
print("-" * 60)
|
||||
for prompt, output in zip(prompts, outputs):
|
||||
embeds = output.outputs.embedding
|
||||
embeds_trimmed = (
|
||||
(str(embeds[:16])[:-1] + ", ...]") if len(embeds) > 16 else embeds
|
||||
)
|
||||
print(f"Prompt: {prompt!r} \nEmbeddings: {embeds_trimmed} (size={len(embeds)})")
|
||||
print("-" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
66
third_party/vllm/examples/pooling/embed/embedding_requests_base64_online.py
vendored
Normal file
66
third_party/vllm/examples/pooling/embed/embedding_requests_base64_online.py
vendored
Normal file
@@ -0,0 +1,66 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Example Python client for embedding API using vLLM API server
|
||||
NOTE:
|
||||
start a supported embeddings model server with `vllm serve`, e.g.
|
||||
vllm serve intfloat/e5-small
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
|
||||
import requests
|
||||
import torch
|
||||
|
||||
from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor
|
||||
|
||||
|
||||
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
response = requests.post(api_url, headers=headers, json=prompt)
|
||||
return response
|
||||
|
||||
|
||||
def parse_args():
|
||||
parse = argparse.ArgumentParser()
|
||||
parse.add_argument("--host", type=str, default="localhost")
|
||||
parse.add_argument("--port", type=int, default=8000)
|
||||
return parse.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
models_url = base_url + "/v1/models"
|
||||
embeddings_url = base_url + "/v1/embeddings"
|
||||
|
||||
response = requests.get(models_url)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
input_texts = [
|
||||
"The best thing about vLLM is that it supports many different models",
|
||||
] * 2
|
||||
|
||||
# The OpenAI client does not support the embed_dtype and endianness parameters.
|
||||
for embed_dtype in EMBED_DTYPES:
|
||||
for endianness in ENDIANNESS:
|
||||
prompt = {
|
||||
"model": model,
|
||||
"input": input_texts,
|
||||
"encoding_format": "base64",
|
||||
"embed_dtype": embed_dtype,
|
||||
"endianness": endianness,
|
||||
}
|
||||
response = post_http_request(prompt=prompt, api_url=embeddings_url)
|
||||
|
||||
embedding = []
|
||||
for data in response.json()["data"]:
|
||||
binary = base64.b64decode(data["embedding"])
|
||||
tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
|
||||
embedding.append(tensor.to(torch.float32))
|
||||
embedding = torch.stack(embedding)
|
||||
print(embed_dtype, endianness, embedding.shape)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
102
third_party/vllm/examples/pooling/embed/embedding_requests_bytes_online.py
vendored
Normal file
102
third_party/vllm/examples/pooling/embed/embedding_requests_bytes_online.py
vendored
Normal file
@@ -0,0 +1,102 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Example Python client for embedding API using vLLM API server
|
||||
NOTE:
|
||||
start a supported embeddings model server with `vllm serve`, e.g.
|
||||
vllm serve intfloat/e5-small
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import requests
|
||||
import torch
|
||||
|
||||
from vllm.entrypoints.pooling.utils import (
|
||||
MetadataItem,
|
||||
build_metadata_items,
|
||||
decode_pooling_output,
|
||||
)
|
||||
from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS
|
||||
|
||||
|
||||
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
response = requests.post(api_url, headers=headers, json=prompt)
|
||||
return response
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
models_url = base_url + "/v1/models"
|
||||
embeddings_url = base_url + "/v1/embeddings"
|
||||
|
||||
response = requests.get(models_url)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
embedding_size = 0
|
||||
|
||||
input_texts = [
|
||||
"The best thing about vLLM is that it supports many different models",
|
||||
] * 2
|
||||
|
||||
# The OpenAI client does not support the bytes encoding_format.
|
||||
# The OpenAI client does not support the embed_dtype and endianness parameters.
|
||||
for embed_dtype in EMBED_DTYPES:
|
||||
for endianness in ENDIANNESS:
|
||||
prompt = {
|
||||
"model": model,
|
||||
"input": input_texts,
|
||||
"encoding_format": "bytes",
|
||||
"embed_dtype": embed_dtype,
|
||||
"endianness": endianness,
|
||||
}
|
||||
response = post_http_request(prompt=prompt, api_url=embeddings_url)
|
||||
metadata = json.loads(response.headers["metadata"])
|
||||
body = response.content
|
||||
items = [MetadataItem(**x) for x in metadata["data"]]
|
||||
|
||||
embedding = decode_pooling_output(items=items, body=body)
|
||||
embedding = [x.to(torch.float32) for x in embedding]
|
||||
embedding = torch.stack(embedding)
|
||||
embedding_size = embedding.shape[-1]
|
||||
print(embed_dtype, endianness, embedding.shape)
|
||||
|
||||
# The vllm server always sorts the returned embeddings in the order of input. So
|
||||
# returning metadata is not necessary. You can set encoding_format to bytes_only
|
||||
# to let the server not return metadata.
|
||||
for embed_dtype in EMBED_DTYPES:
|
||||
for endianness in ENDIANNESS:
|
||||
prompt = {
|
||||
"model": model,
|
||||
"input": input_texts,
|
||||
"encoding_format": "bytes_only",
|
||||
"embed_dtype": embed_dtype,
|
||||
"endianness": endianness,
|
||||
}
|
||||
response = post_http_request(prompt=prompt, api_url=embeddings_url)
|
||||
body = response.content
|
||||
|
||||
items = build_metadata_items(
|
||||
embed_dtype=embed_dtype,
|
||||
endianness=endianness,
|
||||
shape=(embedding_size,),
|
||||
n_request=len(input_texts),
|
||||
)
|
||||
embedding = decode_pooling_output(items=items, body=body)
|
||||
embedding = [x.to(torch.float32) for x in embedding]
|
||||
embedding = torch.stack(embedding)
|
||||
print(embed_dtype, endianness, embedding.shape)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
40
third_party/vllm/examples/pooling/embed/openai_embedding_client.py
vendored
Normal file
40
third_party/vllm/examples/pooling/embed/openai_embedding_client.py
vendored
Normal file
@@ -0,0 +1,40 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Example Python client for embedding API using vLLM API server
|
||||
NOTE:
|
||||
start a supported embeddings model server with `vllm serve`, e.g.
|
||||
vllm serve intfloat/e5-small
|
||||
"""
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||
openai_api_key = "EMPTY"
|
||||
openai_api_base = "http://localhost:8000/v1"
|
||||
|
||||
|
||||
def main():
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
models = client.models.list()
|
||||
model = models.data[0].id
|
||||
|
||||
responses = client.embeddings.create(
|
||||
# ruff: noqa: E501
|
||||
input=[
|
||||
"Hello my name is",
|
||||
"The best thing about vLLM is that it supports many different models",
|
||||
],
|
||||
model=model,
|
||||
)
|
||||
|
||||
for data in responses.data:
|
||||
print(data.embedding) # List of float of len 4096
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
186
third_party/vllm/examples/pooling/embed/openai_embedding_long_text/README.md
vendored
Normal file
186
third_party/vllm/examples/pooling/embed/openai_embedding_long_text/README.md
vendored
Normal file
@@ -0,0 +1,186 @@
|
||||
# Long Text Embedding with Chunked Processing
|
||||
|
||||
This directory contains examples for using vLLM's **chunked processing** feature to handle long text embedding that exceeds the model's maximum context length.
|
||||
|
||||
## 🚀 Quick Start
|
||||
|
||||
### Start the Server
|
||||
|
||||
Use the provided script to start a vLLM server with chunked processing enabled:
|
||||
|
||||
```bash
|
||||
# Basic usage (supports very long texts up to ~3M tokens)
|
||||
./service.sh
|
||||
|
||||
# Custom configuration with different models
|
||||
MODEL_NAME="jinaai/jina-embeddings-v3" \
|
||||
MAX_EMBED_LEN=1048576 \
|
||||
./service.sh
|
||||
|
||||
# For extremely long documents
|
||||
MODEL_NAME="intfloat/multilingual-e5-large" \
|
||||
MAX_EMBED_LEN=3072000 \
|
||||
./service.sh
|
||||
```
|
||||
|
||||
### Test Long Text Embedding
|
||||
|
||||
Run the comprehensive test client:
|
||||
|
||||
```bash
|
||||
python client.py
|
||||
```
|
||||
|
||||
## 📁 Files
|
||||
|
||||
| File | Description |
|
||||
| ---- | ----------- |
|
||||
| `service.sh` | Server startup script with chunked processing enabled |
|
||||
| `client.py` | Comprehensive test client for long text embedding |
|
||||
|
||||
## ⚙️ Configuration
|
||||
|
||||
### Server Configuration
|
||||
|
||||
The key parameters for chunked processing are in the `--pooler-config`:
|
||||
|
||||
```json
|
||||
{
|
||||
"pooling_type": "auto",
|
||||
"use_activation": true,
|
||||
"enable_chunked_processing": true,
|
||||
"max_embed_len": 3072000
|
||||
}
|
||||
```
|
||||
|
||||
!!! note
|
||||
`pooling_type` sets the model's own pooling strategy for processing within each chunk. The cross-chunk aggregation automatically uses MEAN strategy when input exceeds the model's native maximum length.
|
||||
|
||||
#### Chunked Processing Behavior
|
||||
|
||||
Chunked processing uses **MEAN aggregation** for cross-chunk combination when input exceeds the model's native maximum length:
|
||||
|
||||
| Component | Behavior | Description |
|
||||
| --------- | -------- | ----------- |
|
||||
| **Within chunks** | Model's native pooling | Uses the model's configured pooling strategy |
|
||||
| **Cross-chunk aggregation** | Always MEAN | Weighted averaging based on chunk token counts |
|
||||
| **Performance** | Optimal | All chunks processed for complete semantic coverage |
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Default | Description |
|
||||
| -------- | ------- | ----------- |
|
||||
| `MODEL_NAME` | `intfloat/multilingual-e5-large` | Embedding model to use (supports multiple models) |
|
||||
| `PORT` | `31090` | Server port |
|
||||
| `GPU_COUNT` | `1` | Number of GPUs to use |
|
||||
| `MAX_EMBED_LEN` | `3072000` | Maximum embedding input length (supports very long documents) |
|
||||
| `POOLING_TYPE` | `auto` | Model's native pooling type: `auto`, `MEAN`, `CLS`, `LAST` (only affects within-chunk pooling, not cross-chunk aggregation) |
|
||||
| `API_KEY` | `EMPTY` | API key for authentication |
|
||||
|
||||
## 🔧 How It Works
|
||||
|
||||
1. **Enhanced Input Validation**: `max_embed_len` allows accepting inputs longer than `max_model_len` without environment variables
|
||||
2. **Smart Chunking**: Text is split based on `max_position_embeddings` to maintain semantic integrity
|
||||
3. **Unified Processing**: All chunks processed separately through the model using its configured pooling strategy
|
||||
4. **MEAN Aggregation**: When input exceeds model's native length, results combined using token count-based weighted averaging across all chunks
|
||||
5. **Consistent Output**: Final embeddings maintain the same dimensionality as standard processing
|
||||
|
||||
### Input Length Handling
|
||||
|
||||
- **Within max_embed_len**: Input is accepted and processed (up to 3M+ tokens)
|
||||
- **Exceeds max_position_embeddings**: Chunked processing is automatically triggered
|
||||
- **Exceeds max_embed_len**: Input is rejected with clear error message
|
||||
- **No environment variables required**: Works without `VLLM_ALLOW_LONG_MAX_MODEL_LEN`
|
||||
|
||||
### Extreme Long Text Support
|
||||
|
||||
With `MAX_EMBED_LEN=3072000`, you can process:
|
||||
|
||||
- **Academic papers**: Full research papers with references
|
||||
- **Legal documents**: Complete contracts and legal texts
|
||||
- **Books**: Entire chapters or small books
|
||||
- **Code repositories**: Large codebases and documentation
|
||||
|
||||
## 📊 Performance Characteristics
|
||||
|
||||
### Chunked Processing Performance
|
||||
|
||||
| Aspect | Behavior | Performance |
|
||||
| ------ | -------- | ----------- |
|
||||
| **Chunk Processing** | All chunks processed with native pooling | Consistent with input length |
|
||||
| **Cross-chunk Aggregation** | MEAN weighted averaging | Minimal overhead |
|
||||
| **Memory Usage** | Proportional to number of chunks | Moderate, scalable |
|
||||
| **Semantic Quality** | Complete text coverage | Optimal for long documents |
|
||||
|
||||
## 🧪 Test Cases
|
||||
|
||||
The test client demonstrates:
|
||||
|
||||
- ✅ **Short text**: Normal processing (baseline)
|
||||
- ✅ **Medium text**: Single chunk processing
|
||||
- ✅ **Long text**: Multi-chunk processing with aggregation
|
||||
- ✅ **Very long text**: Many chunks processing
|
||||
- ✅ **Extreme long text**: Document-level processing (100K+ tokens)
|
||||
- ✅ **Batch processing**: Mixed-length inputs in one request
|
||||
- ✅ **Consistency**: Reproducible results across runs
|
||||
|
||||
## 🐛 Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
1. **Chunked processing not enabled**:
|
||||
|
||||
```log
|
||||
ValueError: This model's maximum position embeddings length is 4096 tokens...
|
||||
```
|
||||
|
||||
**Solution**: Ensure `enable_chunked_processing: true` in pooler config
|
||||
|
||||
2. **Input exceeds max_embed_len**:
|
||||
|
||||
```log
|
||||
ValueError: This model's maximum embedding input length is 3072000 tokens...
|
||||
```
|
||||
|
||||
**Solution**: Increase `max_embed_len` in pooler config or reduce input length
|
||||
|
||||
3. **Memory errors**:
|
||||
|
||||
```log
|
||||
RuntimeError: CUDA out of memory
|
||||
```
|
||||
|
||||
**Solution**: Reduce chunk size by adjusting model's `max_position_embeddings` or use fewer GPUs
|
||||
|
||||
4. **Slow processing**:
|
||||
**Expected**: Long text takes more time due to multiple inference calls
|
||||
|
||||
### Debug Information
|
||||
|
||||
Server logs show chunked processing activity:
|
||||
|
||||
```log
|
||||
INFO: Input length 150000 exceeds max_position_embeddings 4096, will use chunked processing
|
||||
INFO: Split input of 150000 tokens into 37 chunks (max_chunk_size: 4096)
|
||||
```
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
To extend chunked processing support to other embedding models:
|
||||
|
||||
1. Check model compatibility with the pooling architecture
|
||||
2. Test with various text lengths
|
||||
3. Validate embedding quality compared to single-chunk processing
|
||||
4. Submit PR with test cases and documentation updates
|
||||
|
||||
## 🆕 Enhanced Features
|
||||
|
||||
### max_embed_len Parameter
|
||||
|
||||
The new `max_embed_len` parameter provides:
|
||||
|
||||
- **Simplified Configuration**: No need for `VLLM_ALLOW_LONG_MAX_MODEL_LEN` environment variable
|
||||
- **Flexible Input Validation**: Accept inputs longer than `max_model_len` up to `max_embed_len`
|
||||
- **Extreme Length Support**: Process documents with millions of tokens
|
||||
- **Clear Error Messages**: Better feedback when inputs exceed limits
|
||||
- **Backward Compatibility**: Existing configurations continue to work
|
||||
366
third_party/vllm/examples/pooling/embed/openai_embedding_long_text/client.py
vendored
Normal file
366
third_party/vllm/examples/pooling/embed/openai_embedding_long_text/client.py
vendored
Normal file
@@ -0,0 +1,366 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Example script demonstrating long text embedding with chunked processing in vLLM.
|
||||
|
||||
This example shows how to use vLLM's chunked processing feature to handle text
|
||||
inputs that exceed the model's maximum token length. The feature automatically
|
||||
splits long text into chunks and handles different pooling types optimally.
|
||||
|
||||
Prerequisites:
|
||||
1. Start vLLM server with chunked processing enabled:
|
||||
|
||||
# MEAN pooling (processes all chunks, recommended for complete coverage)
|
||||
vllm serve intfloat/multilingual-e5-large \
|
||||
--pooler-config \
|
||||
'{"pooling_type": "MEAN", "use_activation": true, ' \
|
||||
'"enable_chunked_processing": true, "max_embed_len": 3072000}' \
|
||||
--served-model-name multilingual-e5-large \
|
||||
--trust-remote-code \
|
||||
--port 31090 \
|
||||
--api-key your-api-key
|
||||
|
||||
# OR CLS pooling (native CLS within chunks, MEAN aggregation across chunks)
|
||||
vllm serve BAAI/bge-large-en-v1.5 \
|
||||
--pooler-config \
|
||||
'{"pooling_type": "CLS", "use_activation": true, ' \
|
||||
'"enable_chunked_processing": true, "max_embed_len": 1048576}' \
|
||||
--served-model-name bge-large-en-v1.5 \
|
||||
--trust-remote-code \
|
||||
--port 31090 \
|
||||
--api-key your-api-key
|
||||
|
||||
2. Install required dependencies:
|
||||
pip install openai requests
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
from openai import OpenAI
|
||||
|
||||
# Configuration
|
||||
API_KEY = "your-api-key" # Replace with your actual API key
|
||||
BASE_URL = "http://localhost:31090/v1"
|
||||
MODEL_NAME = "multilingual-e5-large"
|
||||
|
||||
|
||||
def generate_long_text(base_text: str, repeat_count: int) -> str:
|
||||
"""Generate long text by repeating base text."""
|
||||
return base_text * repeat_count
|
||||
|
||||
|
||||
def test_embedding_with_different_lengths():
|
||||
"""Test embedding generation with different text lengths."""
|
||||
client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
|
||||
|
||||
# Test cases with different text lengths
|
||||
test_cases = [
|
||||
{
|
||||
"name": "Short Text",
|
||||
"text": "Hello, this is a short text for embedding.",
|
||||
"expected_chunks": 1,
|
||||
},
|
||||
{
|
||||
"name": "Medium Text",
|
||||
"text": generate_long_text(
|
||||
"This is a medium-length text that should fit within the "
|
||||
"model's context window. " * 20,
|
||||
2,
|
||||
),
|
||||
"expected_chunks": 1,
|
||||
},
|
||||
{
|
||||
"name": "Long Text (2 chunks)",
|
||||
"text": generate_long_text(
|
||||
"This is a very long text that will exceed the model's "
|
||||
"maximum context length and trigger chunked processing. " * 50,
|
||||
5,
|
||||
),
|
||||
"expected_chunks": 2,
|
||||
},
|
||||
{
|
||||
"name": "Very Long Text (3+ chunks)",
|
||||
"text": generate_long_text(
|
||||
"This text is extremely long and will definitely "
|
||||
"require multiple chunks for processing. " * 100,
|
||||
10,
|
||||
),
|
||||
"expected_chunks": 3,
|
||||
},
|
||||
]
|
||||
|
||||
print("🧪 Testing vLLM Long Text Embedding with Chunked Processing")
|
||||
print("=" * 70)
|
||||
|
||||
for i, test_case in enumerate(test_cases, 1):
|
||||
print(f"\n📝 Test {i}: {test_case['name']}")
|
||||
print(f"Text length: {len(test_case['text'])} characters")
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
|
||||
response = client.embeddings.create(
|
||||
input=test_case["text"], model=MODEL_NAME, encoding_format="float"
|
||||
)
|
||||
|
||||
end_time = time.time()
|
||||
processing_time = end_time - start_time
|
||||
|
||||
# Extract embedding data
|
||||
embedding = response.data[0].embedding
|
||||
embedding_dim = len(embedding)
|
||||
|
||||
print("✅ Success!")
|
||||
print(f" - Embedding dimension: {embedding_dim}")
|
||||
print(f" - Processing time: {processing_time:.2f}s")
|
||||
print(f" - Expected chunks: ~{test_case['expected_chunks']}")
|
||||
print(f" - First 5 values: {embedding[:5]}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Failed: {str(e)}")
|
||||
|
||||
|
||||
def test_batch_embedding():
|
||||
"""Test batch embedding with mixed-length inputs."""
|
||||
client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
|
||||
|
||||
print("\n🔄 Testing Batch Embedding with Mixed Lengths")
|
||||
print("=" * 50)
|
||||
|
||||
# Mix of short and long texts
|
||||
batch_inputs = [
|
||||
"Short text 1",
|
||||
generate_long_text("Medium length text that fits in one chunk. " * 20, 1),
|
||||
"Another short text",
|
||||
generate_long_text("Long text requiring chunked processing. " * 100, 5),
|
||||
]
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
|
||||
response = client.embeddings.create(
|
||||
input=batch_inputs, model=MODEL_NAME, encoding_format="float"
|
||||
)
|
||||
|
||||
end_time = time.time()
|
||||
processing_time = end_time - start_time
|
||||
|
||||
print("✅ Batch processing successful!")
|
||||
print(f" - Number of inputs: {len(batch_inputs)}")
|
||||
print(f" - Number of embeddings: {len(response.data)}")
|
||||
print(f" - Total processing time: {processing_time:.2f}s")
|
||||
print(
|
||||
f" - Average time per input: {processing_time / len(batch_inputs):.2f}s"
|
||||
)
|
||||
|
||||
for i, data in enumerate(response.data):
|
||||
input_length = len(batch_inputs[i])
|
||||
embedding_dim = len(data.embedding)
|
||||
print(
|
||||
f" - Input {i + 1}: {input_length} chars → {embedding_dim}D embedding"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Batch processing failed: {str(e)}")
|
||||
|
||||
|
||||
def test_multiple_long_texts_batch():
|
||||
"""Test batch processing with multiple long texts to verify chunk ID uniqueness."""
|
||||
client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
|
||||
|
||||
print("\n🔧 Testing Multiple Long Texts in Batch (Chunk ID Fix Verification)")
|
||||
print("=" * 70)
|
||||
|
||||
# Create multiple distinct long texts that will all require chunking
|
||||
# Note: All pooling types now use MEAN aggregation across chunks:
|
||||
# - Native pooling (MEAN/CLS/LAST) is used within each chunk
|
||||
# - MEAN aggregation combines results across all chunks
|
||||
# - Full semantic coverage for all pooling types
|
||||
long_texts = [
|
||||
generate_long_text(
|
||||
"First long document about artificial intelligence and machine learning. "
|
||||
* 80,
|
||||
6,
|
||||
),
|
||||
generate_long_text(
|
||||
"Second long document about natural language processing and transformers. "
|
||||
* 80,
|
||||
6,
|
||||
),
|
||||
generate_long_text(
|
||||
"Third long document about computer vision and neural networks. " * 80, 6
|
||||
),
|
||||
]
|
||||
|
||||
# Add some short texts to mix things up
|
||||
batch_inputs = [
|
||||
"Short text before long texts",
|
||||
long_texts[0],
|
||||
"Short text between long texts",
|
||||
long_texts[1],
|
||||
long_texts[2],
|
||||
"Short text after long texts",
|
||||
]
|
||||
|
||||
print("📊 Batch composition:")
|
||||
for i, text in enumerate(batch_inputs):
|
||||
length = len(text)
|
||||
text_type = "Long (will be chunked)" if length > 5000 else "Short"
|
||||
print(f" - Input {i + 1}: {length} chars ({text_type})")
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
|
||||
response = client.embeddings.create(
|
||||
input=batch_inputs, model=MODEL_NAME, encoding_format="float"
|
||||
)
|
||||
|
||||
end_time = time.time()
|
||||
processing_time = end_time - start_time
|
||||
|
||||
print("\n✅ Multiple long texts batch processing successful!")
|
||||
print(f" - Number of inputs: {len(batch_inputs)}")
|
||||
print(f" - Number of embeddings returned: {len(response.data)}")
|
||||
print(f" - Total processing time: {processing_time:.2f}s")
|
||||
|
||||
# Verify each embedding is different (no incorrect aggregation)
|
||||
embeddings = [data.embedding for data in response.data]
|
||||
|
||||
if len(embeddings) >= 3:
|
||||
import numpy as np
|
||||
|
||||
# Compare embeddings of the long texts (indices 1, 3, 4)
|
||||
long_embeddings = [
|
||||
np.array(embeddings[1]), # First long text
|
||||
np.array(embeddings[3]), # Second long text
|
||||
np.array(embeddings[4]), # Third long text
|
||||
]
|
||||
|
||||
print("\n🔍 Verifying embedding uniqueness:")
|
||||
for i in range(len(long_embeddings)):
|
||||
for j in range(i + 1, len(long_embeddings)):
|
||||
cosine_sim = np.dot(long_embeddings[i], long_embeddings[j]) / (
|
||||
np.linalg.norm(long_embeddings[i])
|
||||
* np.linalg.norm(long_embeddings[j])
|
||||
)
|
||||
print(
|
||||
f" - Similarity between long text {i + 1} and {j + 1}: "
|
||||
f"{cosine_sim:.4f}"
|
||||
)
|
||||
|
||||
if (
|
||||
cosine_sim < 0.9
|
||||
): # Different content should have lower similarity
|
||||
print(" ✅ Good: Embeddings are appropriately different")
|
||||
else:
|
||||
print(
|
||||
" ⚠️ High similarity - may indicate chunk "
|
||||
"aggregation issue"
|
||||
)
|
||||
|
||||
print("\n📋 Per-input results:")
|
||||
for i, data in enumerate(response.data):
|
||||
input_length = len(batch_inputs[i])
|
||||
embedding_dim = len(data.embedding)
|
||||
embedding_norm = np.linalg.norm(data.embedding)
|
||||
print(
|
||||
f" - Input {i + 1}: {input_length} chars → {embedding_dim}D "
|
||||
f"embedding (norm: {embedding_norm:.4f})"
|
||||
)
|
||||
|
||||
print(
|
||||
"\n✅ This test verifies the fix for chunk ID collisions in "
|
||||
"batch processing"
|
||||
)
|
||||
print(" - Before fix: Multiple long texts would have conflicting chunk IDs")
|
||||
print(" - After fix: Each prompt's chunks have unique IDs with prompt index")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Multiple long texts batch test failed: {str(e)}")
|
||||
print(" This might indicate the chunk ID collision bug is present!")
|
||||
|
||||
|
||||
def test_embedding_consistency():
|
||||
"""Test that chunked processing produces consistent results."""
|
||||
client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
|
||||
|
||||
print("\n🔍 Testing Embedding Consistency")
|
||||
print("=" * 40)
|
||||
|
||||
# Use the same long text multiple times
|
||||
long_text = generate_long_text(
|
||||
"Consistency test text for chunked processing validation. " * 50, 3
|
||||
)
|
||||
|
||||
embeddings = []
|
||||
|
||||
try:
|
||||
for i in range(3):
|
||||
response = client.embeddings.create(
|
||||
input=long_text, model=MODEL_NAME, encoding_format="float"
|
||||
)
|
||||
embeddings.append(response.data[0].embedding)
|
||||
print(f" - Generated embedding {i + 1}")
|
||||
|
||||
# Check consistency (embeddings should be identical)
|
||||
if len(embeddings) >= 2:
|
||||
# Calculate similarity between first two embeddings
|
||||
|
||||
emb1 = np.array(embeddings[0])
|
||||
emb2 = np.array(embeddings[1])
|
||||
|
||||
# Cosine similarity
|
||||
cosine_sim = np.dot(emb1, emb2) / (
|
||||
np.linalg.norm(emb1) * np.linalg.norm(emb2)
|
||||
)
|
||||
|
||||
print("✅ Consistency test completed!")
|
||||
print(f" - Cosine similarity between runs: {cosine_sim:.6f}")
|
||||
print(" - Expected: ~1.0 (identical embeddings)")
|
||||
|
||||
if cosine_sim > 0.999:
|
||||
print(" - ✅ High consistency achieved!")
|
||||
else:
|
||||
print(" - ⚠️ Consistency may vary due to numerical precision")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Consistency test failed: {str(e)}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to run all tests."""
|
||||
print("🚀 vLLM Long Text Embedding Client")
|
||||
print(f"📡 Connecting to: {BASE_URL}")
|
||||
print(f"🤖 Model: {MODEL_NAME}")
|
||||
masked_key = "*" * (len(API_KEY) - 4) + API_KEY[-4:] if len(API_KEY) > 4 else "****"
|
||||
print(f"🔑 API Key: {masked_key}")
|
||||
|
||||
# Run all test cases
|
||||
test_embedding_with_different_lengths()
|
||||
test_batch_embedding()
|
||||
test_multiple_long_texts_batch()
|
||||
test_embedding_consistency()
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print("🎉 All tests completed!")
|
||||
print("\n💡 Key Features Demonstrated:")
|
||||
print(" - ✅ Automatic chunked processing for long text")
|
||||
print(" - ✅ Seamless handling of mixed-length batches")
|
||||
print(" - ✅ Multiple long texts in single batch (chunk ID fix)")
|
||||
print(" - ✅ Unified chunked processing:")
|
||||
print(" • Native pooling used within each chunk")
|
||||
print(" • MEAN aggregation across all chunks")
|
||||
print(" • Complete semantic coverage for all pooling types")
|
||||
print(" - ✅ Consistent embedding generation")
|
||||
print(" - ✅ Backward compatibility with short text")
|
||||
print("\n📚 For more information, see:")
|
||||
print(
|
||||
" - Documentation: https://docs.vllm.ai/en/latest/models/pooling_models.html"
|
||||
)
|
||||
print(" - Chunked Processing Guide: openai_embedding_long_text.md")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
136
third_party/vllm/examples/pooling/embed/openai_embedding_long_text/service.sh
vendored
Normal file
136
third_party/vllm/examples/pooling/embed/openai_embedding_long_text/service.sh
vendored
Normal file
@@ -0,0 +1,136 @@
|
||||
#!/bin/bash
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# vLLM Embedding Server with Enhanced Chunked Processing
|
||||
# This script starts a vLLM server with chunked processing enabled for long text embedding.
|
||||
# Now supports proper pooling type validation and model-specific configurations.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# Configuration
|
||||
MODEL_NAME=${MODEL_NAME:-"intfloat/multilingual-e5-large"}
|
||||
MODEL_CODE=${MODEL_CODE:-"multilingual-e5-large"}
|
||||
|
||||
PORT=${PORT:-31090}
|
||||
GPU_COUNT=${GPU_COUNT:-1}
|
||||
MAX_EMBED_LEN=${MAX_EMBED_LEN:-3072000}
|
||||
API_KEY=${API_KEY:-"your-api-key"}
|
||||
|
||||
# Enhanced pooling configuration with model-specific defaults
|
||||
POOLING_TYPE=${POOLING_TYPE:-"auto"} # auto, MEAN, CLS, LAST
|
||||
export VLLM_ENABLE_CHUNKED_PROCESSING=true
|
||||
export CUDA_VISIBLE_DEVICES=2,3,4,5
|
||||
|
||||
echo "🚀 Starting vLLM Embedding Server with Enhanced Chunked Processing"
|
||||
echo "=================================================================="
|
||||
|
||||
# Environment variables for optimization
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
# Function to determine optimal pooling type for known models
|
||||
get_optimal_pooling_type() {
|
||||
local model="$1"
|
||||
case "$model" in
|
||||
*"e5-"* | *"multilingual-e5"*)
|
||||
echo "MEAN" # E5 series native pooling
|
||||
;;
|
||||
*"bge-"*)
|
||||
echo "CLS" # BGE series native pooling
|
||||
;;
|
||||
*"gte-"*)
|
||||
echo "LAST" # GTE series native pooling
|
||||
;;
|
||||
*"sentence-t5"* | *"st5"*)
|
||||
echo "MEAN" # Sentence-T5 native pooling
|
||||
;;
|
||||
*"jina-embeddings"*)
|
||||
echo "MEAN" # Jina embeddings native pooling
|
||||
;;
|
||||
*"Qwen"*"Embedding"*)
|
||||
echo "LAST" # Qwen embeddings native pooling
|
||||
;;
|
||||
*)
|
||||
echo "MEAN" # Default native pooling for unknown models
|
||||
;;
|
||||
esac
|
||||
}
|
||||
|
||||
# Auto-detect pooling type if not explicitly set
|
||||
if [ "$POOLING_TYPE" = "auto" ]; then
|
||||
POOLING_TYPE=$(get_optimal_pooling_type "$MODEL_NAME")
|
||||
echo "🔍 Auto-detected pooling type: $POOLING_TYPE for model $MODEL_NAME"
|
||||
fi
|
||||
|
||||
# Display configuration
|
||||
echo "📋 Configuration:"
|
||||
echo " - Model: $MODEL_NAME"
|
||||
echo " - Port: $PORT"
|
||||
echo " - GPU Count: $GPU_COUNT"
|
||||
echo " - Enhanced Chunked Processing: ${VLLM_ENABLE_CHUNKED_PROCESSING}"
|
||||
echo " - Max Embed Length: ${MAX_EMBED_LEN} tokens"
|
||||
echo " - Native Pooling Type: $POOLING_TYPE + Normalization"
|
||||
echo " - Cross-chunk Aggregation: MEAN (automatic)"
|
||||
echo ""
|
||||
|
||||
# Validate GPU availability
|
||||
if command -v nvidia-smi &> /dev/null; then
|
||||
gpu_count=$(nvidia-smi --list-gpus | wc -l)
|
||||
echo "🖥️ Available GPUs: $gpu_count"
|
||||
if [ "$GPU_COUNT" -gt "$gpu_count" ]; then
|
||||
echo "⚠️ Warning: Requested $GPU_COUNT GPUs but only $gpu_count available"
|
||||
echo " Adjusting to use $gpu_count GPUs"
|
||||
GPU_COUNT=$gpu_count
|
||||
fi
|
||||
else
|
||||
echo "⚠️ Warning: nvidia-smi not found. GPU detection skipped."
|
||||
fi
|
||||
|
||||
# Chunked processing uses unified MEAN aggregation
|
||||
echo "ℹ️ Chunked Processing: Using $POOLING_TYPE pooling within chunks, MEAN aggregation across chunks"
|
||||
echo " - All chunks processed for complete semantic coverage"
|
||||
echo " - Weighted averaging based on chunk token counts"
|
||||
|
||||
echo ""
|
||||
echo "🔧 Starting server with enhanced chunked processing configuration..."
|
||||
|
||||
# Build pooler config JSON
|
||||
POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"use_activation\": true, \"enable_chunked_processing\": ${VLLM_ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}"
|
||||
|
||||
# Start vLLM server with enhanced chunked processing
|
||||
vllm serve "$MODEL_NAME" \
|
||||
--tensor-parallel-size "$GPU_COUNT" \
|
||||
--enforce-eager \
|
||||
--pooler-config "$POOLER_CONFIG" \
|
||||
--served-model-name "${MODEL_CODE}" \
|
||||
--api-key "$API_KEY" \
|
||||
--trust-remote-code \
|
||||
--port "$PORT" \
|
||||
--host 0.0.0.0
|
||||
|
||||
echo ""
|
||||
echo "✅ vLLM Embedding Server started successfully!"
|
||||
echo ""
|
||||
echo "📡 Server Information:"
|
||||
echo " - Base URL: http://localhost:$PORT"
|
||||
echo " - Model Code: ${MODEL_CODE}"
|
||||
echo " - API Key: $API_KEY"
|
||||
echo " - Native Pooling: $POOLING_TYPE | Cross-chunk: MEAN"
|
||||
echo ""
|
||||
echo "🧪 Test the server with:"
|
||||
echo " python examples/online_serving/openai_embedding_long_text/client.py"
|
||||
echo ""
|
||||
echo "📚 Enhanced features enabled:"
|
||||
echo " ✅ Intelligent native pooling type detection"
|
||||
echo " ✅ Unified MEAN aggregation for chunked processing"
|
||||
echo " ✅ Model-specific native pooling optimization"
|
||||
echo " ✅ Enhanced max embedding length (${MAX_EMBED_LEN} tokens)"
|
||||
echo " ✅ Complete semantic coverage for all pooling types"
|
||||
echo " ✅ OpenAI-compatible API"
|
||||
echo " ✅ GPU acceleration"
|
||||
echo ""
|
||||
echo "🔧 Advanced usage:"
|
||||
echo " - Set POOLING_TYPE=MEAN|CLS|LAST to override auto-detection"
|
||||
echo " - Set MAX_EMBED_LEN to adjust maximum input length"
|
||||
echo " - All pooling types use MEAN aggregation across chunks"
|
||||
37
third_party/vllm/examples/pooling/embed/openai_embedding_matryoshka_fy_client.py
vendored
Normal file
37
third_party/vllm/examples/pooling/embed/openai_embedding_matryoshka_fy_client.py
vendored
Normal file
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Example Python client for embedding API dimensions using vLLM API server
|
||||
NOTE:
|
||||
start a supported Matryoshka Embeddings model server with `vllm serve`, e.g.
|
||||
vllm serve jinaai/jina-embeddings-v3 --trust-remote-code
|
||||
"""
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||
openai_api_key = "EMPTY"
|
||||
openai_api_base = "http://localhost:8000/v1"
|
||||
|
||||
|
||||
def main():
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
models = client.models.list()
|
||||
model = models.data[0].id
|
||||
|
||||
responses = client.embeddings.create(
|
||||
input=["Follow the white rabbit."],
|
||||
model=model,
|
||||
dimensions=32,
|
||||
)
|
||||
|
||||
for data in responses.data:
|
||||
print(data.embedding) # List of float of len 32
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
7
third_party/vllm/examples/pooling/embed/template/dse_qwen2_vl.jinja
vendored
Normal file
7
third_party/vllm/examples/pooling/embed/template/dse_qwen2_vl.jinja
vendored
Normal file
@@ -0,0 +1,7 @@
|
||||
{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}{% raw %}<|im_start|>system
|
||||
You are a helpful assistant.<|im_end|>
|
||||
{% endraw %}{% endif %}<|im_start|>{{ message['role'] }}{% raw %}
|
||||
{% endraw %}{% if message['content'] is string %}{{ message['content'] }}<|im_end|>{% raw %}
|
||||
{% endraw %}{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>{% raw %}
|
||||
{% endraw %}{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant{% raw %}
|
||||
{% endraw %}{% endif %}<|endoftext|>
|
||||
20
third_party/vllm/examples/pooling/embed/template/nemotron_embed_vl.jinja
vendored
Normal file
20
third_party/vllm/examples/pooling/embed/template/nemotron_embed_vl.jinja
vendored
Normal file
@@ -0,0 +1,20 @@
|
||||
{%- if messages | length > 1 -%}
|
||||
{{ raise_exception('Embedding models should only embed one message at a time') }}
|
||||
{%- endif -%}
|
||||
|
||||
{% set vars = namespace(prefix='', images=[], texts=[]) %}
|
||||
{%- for message in messages -%}
|
||||
{%- if message['role'] == 'query' -%}
|
||||
{%- set vars.prefix = 'query: ' %}
|
||||
{%- elif message['role'] == 'document' -%}
|
||||
{%- set vars.prefix = 'passage: ' %}
|
||||
{%- endif -%}
|
||||
{%- for content in message['content'] -%}
|
||||
{%- if content['type'] == 'text' -%}
|
||||
{%- set vars.texts = vars.texts + [content['text']] %}
|
||||
{%- elif content['type'] == 'image' -%}
|
||||
{%- set vars.images = vars.images + ['<image> '] %}
|
||||
{%- endif -%}
|
||||
{%- endfor -%}
|
||||
{%- endfor -%}
|
||||
{{- bos_token }}{{ vars.prefix }}{{ (vars.images + vars.texts) | join('') }}
|
||||
16
third_party/vllm/examples/pooling/embed/template/vlm2vec_phi3v.jinja
vendored
Normal file
16
third_party/vllm/examples/pooling/embed/template/vlm2vec_phi3v.jinja
vendored
Normal file
@@ -0,0 +1,16 @@
|
||||
{%- if messages | length > 1 -%}
|
||||
{{ raise_exception('Embedding models should only embed one message at a time') }}
|
||||
{%- endif -%}
|
||||
|
||||
{% set vars = namespace(parts=[], next_image_id=1) %}
|
||||
{%- for message in messages -%}
|
||||
{%- for content in message['content'] -%}
|
||||
{%- if content['type'] == 'text' -%}
|
||||
{%- set vars.parts = vars.parts + [content['text']] %}
|
||||
{%- elif content['type'] == 'image' -%}
|
||||
{%- set vars.parts = vars.parts + ['<|image_{i:d}|>'.format(i=vars.next_image_id)] %}
|
||||
{%- set vars.next_image_id = vars.next_image_id + 1 %}
|
||||
{%- endif -%}
|
||||
{%- endfor -%}
|
||||
{%- endfor -%}
|
||||
{{ vars.parts | join(' ') }}
|
||||
15
third_party/vllm/examples/pooling/embed/template/vlm2vec_qwen2vl.jinja
vendored
Normal file
15
third_party/vllm/examples/pooling/embed/template/vlm2vec_qwen2vl.jinja
vendored
Normal file
@@ -0,0 +1,15 @@
|
||||
{%- if messages | length > 1 -%}
|
||||
{{ raise_exception('Embedding models should only embed one message at a time') }}
|
||||
{%- endif -%}
|
||||
|
||||
{% set vars = namespace(parts=[]) %}
|
||||
{%- for message in messages -%}
|
||||
{%- for content in message['content'] -%}
|
||||
{%- if content['type'] == 'text' -%}
|
||||
{%- set vars.parts = vars.parts + [content['text']] %}
|
||||
{%- elif content['type'] == 'image' -%}
|
||||
{%- set vars.parts = vars.parts + ['<|image_pad|>'] %}
|
||||
{%- endif -%}
|
||||
{%- endfor -%}
|
||||
{%- endfor -%}
|
||||
{{ vars.parts | join(' ') }}
|
||||
343
third_party/vllm/examples/pooling/embed/vision_embedding_offline.py
vendored
Normal file
343
third_party/vllm/examples/pooling/embed/vision_embedding_offline.py
vendored
Normal file
@@ -0,0 +1,343 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""
|
||||
This example shows how to use vLLM for running offline inference with
|
||||
the correct prompt format on vision language models for multimodal embedding.
|
||||
|
||||
For most models, the prompt format should follow corresponding examples
|
||||
on HuggingFace model repository.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
|
||||
from PIL.Image import Image
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.multimodal.utils import fetch_image
|
||||
from vllm.utils.print_utils import print_embeddings
|
||||
|
||||
ROOT_DIR = Path(__file__).parent.parent.parent
|
||||
EMBED_TEMPLATE_DIR = ROOT_DIR / "pooling/embed/template/"
|
||||
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
text = "A cat standing in the snow."
|
||||
multi_modal_data = {"image": fetch_image(image_url)}
|
||||
|
||||
|
||||
def run_clip(seed: int):
|
||||
engine_args = EngineArgs(
|
||||
model="openai/clip-vit-base-patch32",
|
||||
runner="pooling",
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
llm = LLM(**asdict(engine_args) | {"seed": seed})
|
||||
|
||||
print("Text embedding output:")
|
||||
outputs = llm.embed(text, use_tqdm=False)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image embedding output:")
|
||||
prompt = "" # For image input, make sure that the prompt text is empty
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
|
||||
def run_e5_v(seed: int):
|
||||
engine_args = EngineArgs(
|
||||
model="royokong/e5-v",
|
||||
runner="pooling",
|
||||
max_model_len=4096,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
llm = LLM(**asdict(engine_args) | {"seed": seed})
|
||||
|
||||
llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n" # noqa: E501
|
||||
|
||||
print("Text embedding output:")
|
||||
prompt_text = llama3_template.format(
|
||||
f"{text}\nSummary above sentence in one word: "
|
||||
)
|
||||
outputs = llm.embed(prompt_text, use_tqdm=False)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image embedding output:")
|
||||
prompt_image = llama3_template.format("<image>\nSummary above image in one word: ")
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": prompt_image,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
|
||||
def run_qwen3_vl(seed: int):
|
||||
try:
|
||||
from qwen_vl_utils import smart_resize
|
||||
except ModuleNotFoundError:
|
||||
print(
|
||||
"WARNING: `qwen-vl-utils` not installed, input images will not "
|
||||
"be automatically resized. This can cause different results "
|
||||
"comparing with HF repo's example. "
|
||||
"You can enable this functionality by `pip install qwen-vl-utils`."
|
||||
)
|
||||
smart_resize = None
|
||||
|
||||
if smart_resize is not None:
|
||||
|
||||
def post_process_image(image: Image) -> Image:
|
||||
width, height = image.size
|
||||
resized_height, resized_width = smart_resize(
|
||||
height,
|
||||
width,
|
||||
factor=32,
|
||||
)
|
||||
return image.resize((resized_width, resized_height))
|
||||
|
||||
multi_modal_data["image"] = post_process_image(multi_modal_data["image"])
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model="Qwen/Qwen3-VL-Embedding-2B",
|
||||
runner="pooling",
|
||||
max_model_len=8192,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
mm_processor_kwargs={"do_resize": False} if smart_resize is not None else None,
|
||||
)
|
||||
default_instruction = "Represent the user's input."
|
||||
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
prompt_text = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_image = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{image_placeholder}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_image_text = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{image_placeholder}{text}<|im_end|>\n<|im_start|>assistant\n"
|
||||
|
||||
llm = LLM(**asdict(engine_args) | {"seed": seed})
|
||||
|
||||
print("Text embedding output:")
|
||||
outputs = llm.embed(prompt_text, use_tqdm=False)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image embedding output:")
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": prompt_image,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image+Text embedding output:")
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": prompt_image_text,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
|
||||
def run_siglip(seed: int):
|
||||
engine_args = EngineArgs(
|
||||
model="google/siglip-base-patch16-224",
|
||||
runner="pooling",
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
llm = LLM(**asdict(engine_args) | {"seed": seed})
|
||||
|
||||
print("Text embedding output:")
|
||||
outputs = llm.embed(text, use_tqdm=False)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image embedding output:")
|
||||
prompt = "" # For image input, make sure that the prompt text is empty
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
|
||||
def run_vlm2vec_phi3v(seed: int):
|
||||
engine_args = EngineArgs(
|
||||
model="TIGER-Lab/VLM2Vec-Full",
|
||||
runner="pooling",
|
||||
max_model_len=4096,
|
||||
trust_remote_code=True,
|
||||
mm_processor_kwargs={"num_crops": 4},
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
llm = LLM(**asdict(engine_args) | {"seed": seed})
|
||||
image_token = "<|image_1|>"
|
||||
|
||||
print("Text embedding output:")
|
||||
prompt_text = f"Find me an everyday image that matches the given caption: {text}"
|
||||
outputs = llm.embed(prompt_text, use_tqdm=False)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image embedding output:")
|
||||
prompt_image = f"{image_token} Find a day-to-day image that looks similar to the provided image." # noqa: E501
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": prompt_image,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image+Text embedding output:")
|
||||
prompt_image_text = (
|
||||
f"{image_token} Represent the given image with the following question: {text}" # noqa: E501
|
||||
)
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": prompt_image_text,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
|
||||
def run_vlm2vec_qwen2vl(seed: int):
|
||||
# vLLM does not support LoRA adapters on multi-modal encoder,
|
||||
# so we merge the weights first
|
||||
from huggingface_hub.constants import HF_HUB_CACHE
|
||||
from peft import PeftConfig, PeftModel
|
||||
from transformers import AutoModelForImageTextToText, AutoProcessor
|
||||
|
||||
from vllm.entrypoints.chat_utils import load_chat_template
|
||||
|
||||
model_id = "TIGER-Lab/VLM2Vec-Qwen2VL-2B"
|
||||
|
||||
base_model = AutoModelForImageTextToText.from_pretrained(model_id)
|
||||
lora_model = PeftModel.from_pretrained(
|
||||
base_model,
|
||||
model_id,
|
||||
config=PeftConfig.from_pretrained(model_id),
|
||||
)
|
||||
model = lora_model.merge_and_unload().to(dtype=base_model.dtype)
|
||||
model._hf_peft_config_loaded = False # Needed to save the merged model
|
||||
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
model_id,
|
||||
# `min_pixels` and `max_pixels` are deprecated for
|
||||
# transformers `preprocessor_config.json`
|
||||
size={"shortest_edge": 3136, "longest_edge": 12845056},
|
||||
)
|
||||
processor.chat_template = load_chat_template(
|
||||
# The original chat template is not correct
|
||||
EMBED_TEMPLATE_DIR / "vlm2vec_qwen2vl.jinja",
|
||||
)
|
||||
|
||||
merged_path = str(
|
||||
Path(HF_HUB_CACHE) / ("models--" + model_id.replace("/", "--") + "-vllm")
|
||||
)
|
||||
print(f"Saving merged model to {merged_path}...")
|
||||
print(
|
||||
"NOTE: This directory is not tracked by `huggingface_hub` "
|
||||
"so you have to delete this manually if you don't want it anymore."
|
||||
)
|
||||
model.save_pretrained(merged_path)
|
||||
processor.save_pretrained(merged_path)
|
||||
print("Done!")
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=merged_path,
|
||||
runner="pooling",
|
||||
max_model_len=4096,
|
||||
mm_processor_kwargs={
|
||||
"min_pixels": 3136,
|
||||
"max_pixels": 12845056,
|
||||
},
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
llm = LLM(**asdict(engine_args) | {"seed": seed})
|
||||
image_token = "<|image_pad|>"
|
||||
|
||||
print("Text embedding output:")
|
||||
prompt_text = f"Find me an everyday image that matches the given caption: {text}"
|
||||
outputs = llm.embed(prompt_text, use_tqdm=False)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image embedding output:")
|
||||
prompt_image = f"{image_token} Find a day-to-day image that looks similar to the provided image." # noqa: E501
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": prompt_image,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
print("Image+Text embedding output:")
|
||||
prompt_image_text = (
|
||||
f"{image_token} Represent the given image with the following question: {text}" # noqa: E501
|
||||
)
|
||||
outputs = llm.embed(
|
||||
{
|
||||
"prompt": prompt_image_text,
|
||||
"multi_modal_data": multi_modal_data,
|
||||
},
|
||||
use_tqdm=False,
|
||||
)
|
||||
print_embeddings(outputs[0].outputs.embedding)
|
||||
|
||||
|
||||
model_example_map = {
|
||||
"clip": run_clip,
|
||||
"e5_v": run_e5_v,
|
||||
"qwen3_vl": run_qwen3_vl,
|
||||
"siglip": run_siglip,
|
||||
"vlm2vec_phi3v": run_vlm2vec_phi3v,
|
||||
"vlm2vec_qwen2vl": run_vlm2vec_qwen2vl,
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
"Script to run a specified VLM through vLLM offline api."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
"-m",
|
||||
type=str,
|
||||
default="vlm2vec_phi3v",
|
||||
choices=model_example_map.keys(),
|
||||
help="The name of the embedding model.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Set the seed when initializing `vllm.LLM`.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
model_example_map[args.model](args.seed)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
415
third_party/vllm/examples/pooling/embed/vision_embedding_online.py
vendored
Normal file
415
third_party/vllm/examples/pooling/embed/vision_embedding_online.py
vendored
Normal file
@@ -0,0 +1,415 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""Example Python client for multimodal embedding API using vLLM API server.
|
||||
|
||||
Refer to each `run_*` function for the command to run the server for that model.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
from typing import Literal
|
||||
|
||||
from openai import OpenAI
|
||||
from openai._types import NOT_GIVEN, NotGiven
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from openai.types.create_embedding_response import CreateEmbeddingResponse
|
||||
from PIL import Image
|
||||
|
||||
from vllm.utils.print_utils import print_embeddings
|
||||
|
||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||
openai_api_key = "EMPTY"
|
||||
openai_api_base = "http://localhost:8000/v1"
|
||||
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
text = "A cat standing in the snow."
|
||||
|
||||
|
||||
def create_chat_embeddings(
|
||||
client: OpenAI,
|
||||
*,
|
||||
messages: list[ChatCompletionMessageParam],
|
||||
model: str,
|
||||
encoding_format: Literal["base64", "float"] | NotGiven = NOT_GIVEN,
|
||||
continue_final_message: bool = False,
|
||||
add_special_tokens: bool = False,
|
||||
) -> CreateEmbeddingResponse:
|
||||
"""
|
||||
Convenience function for accessing vLLM's Chat Embeddings API,
|
||||
which is an extension of OpenAI's existing Embeddings API.
|
||||
"""
|
||||
return client.post(
|
||||
"/embeddings",
|
||||
cast_to=CreateEmbeddingResponse,
|
||||
body={
|
||||
"messages": messages,
|
||||
"model": model,
|
||||
"encoding_format": encoding_format,
|
||||
"continue_final_message": continue_final_message,
|
||||
"add_special_tokens": add_special_tokens,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def run_clip(client: OpenAI, model: str):
|
||||
"""
|
||||
Start the server using:
|
||||
|
||||
vllm serve openai/clip-vit-base-patch32 \
|
||||
--runner pooling
|
||||
"""
|
||||
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Image embedding output:", response.data[0].embedding)
|
||||
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "a photo of a cat"},
|
||||
],
|
||||
}
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Text embedding output:", response.data[0].embedding)
|
||||
|
||||
|
||||
def run_dse_qwen2_vl(client: OpenAI, model: str):
|
||||
"""
|
||||
Start the server using:
|
||||
|
||||
vllm serve MrLight/dse-qwen2-2b-mrl-v1 \
|
||||
--runner pooling \
|
||||
--trust-remote-code \
|
||||
--max-model-len 8192 \
|
||||
--chat-template examples/pooling/embed/template/dse_qwen2_vl.jinja
|
||||
"""
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": image_url,
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "What is shown in this image?"},
|
||||
],
|
||||
}
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Image embedding output:", response.data[0].embedding)
|
||||
|
||||
# MrLight/dse-qwen2-2b-mrl-v1 requires a placeholder image
|
||||
# of the minimum input size
|
||||
buffer = io.BytesIO()
|
||||
image_placeholder = Image.new("RGB", (56, 56))
|
||||
image_placeholder.save(buffer, "png")
|
||||
buffer.seek(0)
|
||||
image_placeholder = base64.b64encode(buffer.read()).decode("utf-8")
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/jpeg;base64,{image_placeholder}",
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "Query: What is the weather like today?"},
|
||||
],
|
||||
}
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Text embedding output:", response.data[0].embedding)
|
||||
|
||||
|
||||
def run_qwen3_vl(client: OpenAI, model: str):
|
||||
"""
|
||||
Start the server using:
|
||||
|
||||
vllm serve Qwen/Qwen3-VL-Embedding-2B \
|
||||
--runner pooling \
|
||||
--max-model-len 8192
|
||||
"""
|
||||
|
||||
default_instruction = "Represent the user's input."
|
||||
|
||||
print("Text embedding output:")
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": default_instruction},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": text},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": ""},
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
continue_final_message=True,
|
||||
add_special_tokens=True,
|
||||
)
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
print("Image embedding output:")
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": default_instruction},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
{"type": "text", "text": ""},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": ""},
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
continue_final_message=True,
|
||||
add_special_tokens=True,
|
||||
)
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
print("Image+Text embedding output:")
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": default_instruction},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{text}",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": ""},
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
continue_final_message=True,
|
||||
add_special_tokens=True,
|
||||
)
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
|
||||
def run_siglip(client: OpenAI, model: str):
|
||||
"""
|
||||
Start the server using:
|
||||
|
||||
vllm serve google/siglip-base-patch16-224 \
|
||||
--runner pooling \
|
||||
--chat-template template_basic.jinja
|
||||
"""
|
||||
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Image embedding output:", response.data[0].embedding)
|
||||
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "a photo of a cat"},
|
||||
],
|
||||
}
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Text embedding output:", response.data[0].embedding)
|
||||
|
||||
|
||||
def run_vlm2vec(client: OpenAI, model: str):
|
||||
"""
|
||||
Start the server using:
|
||||
|
||||
vllm serve TIGER-Lab/VLM2Vec-Full \
|
||||
--runner pooling \
|
||||
--trust-remote-code \
|
||||
--max-model-len 4096 \
|
||||
--chat-template examples/pooling/embed/template/vlm2vec_phi3v.jinja
|
||||
"""
|
||||
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
{"type": "text", "text": "Represent the given image."},
|
||||
],
|
||||
}
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Image embedding output:")
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Represent the given image with the following question: What is in the image.",
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Image+Text embedding output:")
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
response = create_chat_embeddings(
|
||||
client,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "A cat and a dog"},
|
||||
],
|
||||
}
|
||||
],
|
||||
model=model,
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
print("Text embedding output:")
|
||||
print_embeddings(response.data[0].embedding)
|
||||
|
||||
|
||||
model_example_map = {
|
||||
"clip": run_clip,
|
||||
"qwen3_vl": run_qwen3_vl,
|
||||
"dse_qwen2_vl": run_dse_qwen2_vl,
|
||||
"siglip": run_siglip,
|
||||
"vlm2vec": run_vlm2vec,
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
"Script to call a specified VLM through the API. Make sure to serve "
|
||||
"the model with `--runner pooling` before running this."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
choices=model_example_map.keys(),
|
||||
required=True,
|
||||
help="The name of the embedding model.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
models = client.models.list()
|
||||
model_id = models.data[0].id
|
||||
|
||||
model_example_map[args.model](client, model_id)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
60
third_party/vllm/examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py
vendored
Normal file
60
third_party/vllm/examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py
vendored
Normal file
@@ -0,0 +1,60 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import base64
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
from vllm import LLM
|
||||
|
||||
# This example shows how to perform an offline inference that generates
|
||||
# multimodal data. In this specific case this example will take a geotiff
|
||||
# image as input, process it using the multimodal data processor, and
|
||||
# perform inference.
|
||||
# Requirements:
|
||||
# - install TerraTorch v1.1 (or later):
|
||||
# pip install terratorch>=v1.1
|
||||
|
||||
|
||||
def main():
|
||||
torch.set_default_dtype(torch.float16)
|
||||
image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
|
||||
|
||||
img_data = dict(
|
||||
data=image_url,
|
||||
data_format="url",
|
||||
image_format="tiff",
|
||||
out_data_format="b64_json",
|
||||
)
|
||||
|
||||
prompt = dict(data=img_data)
|
||||
|
||||
llm = LLM(
|
||||
model="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
|
||||
skip_tokenizer_init=True,
|
||||
trust_remote_code=True,
|
||||
enforce_eager=True,
|
||||
# Limit the maximum number of parallel requests
|
||||
# to avoid the model going OOM.
|
||||
# The maximum number depends on the available GPU memory
|
||||
max_num_seqs=32,
|
||||
io_processor_plugin="terratorch_segmentation",
|
||||
model_impl="terratorch",
|
||||
enable_mm_embeds=True,
|
||||
)
|
||||
|
||||
pooler_output = llm.encode(prompt, pooling_task="plugin")
|
||||
output = pooler_output[0].outputs
|
||||
|
||||
print(output)
|
||||
decoded_data = base64.b64decode(output.data)
|
||||
|
||||
file_path = os.path.join(os.getcwd(), "offline_prediction.tiff")
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(decoded_data)
|
||||
|
||||
print(f"Output file path: {file_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
421
third_party/vllm/examples/pooling/plugin/prithvi_geospatial_mae_offline.py
vendored
Normal file
421
third_party/vllm/examples/pooling/plugin/prithvi_geospatial_mae_offline.py
vendored
Normal file
@@ -0,0 +1,421 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import argparse
|
||||
import datetime
|
||||
import os
|
||||
|
||||
import albumentations
|
||||
import numpy as np
|
||||
import rasterio
|
||||
import regex as re
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from terratorch.datamodules import Sen1Floods11NonGeoDataModule
|
||||
|
||||
from vllm import LLM
|
||||
|
||||
torch.set_default_dtype(torch.float16)
|
||||
|
||||
NO_DATA = -9999
|
||||
NO_DATA_FLOAT = 0.0001
|
||||
OFFSET = 0
|
||||
PERCENTILE = 99
|
||||
|
||||
datamodule_config = {
|
||||
"bands": ["BLUE", "GREEN", "RED", "NIR_NARROW", "SWIR_1", "SWIR_2"],
|
||||
"batch_size": 16,
|
||||
"constant_scale": 0.0001,
|
||||
"data_root": "/dccstor/geofm-finetuning/datasets/sen1floods11",
|
||||
"drop_last": True,
|
||||
"no_data_replace": 0.0,
|
||||
"no_label_replace": -1,
|
||||
"num_workers": 8,
|
||||
"test_transform": [
|
||||
albumentations.Resize(
|
||||
always_apply=False, height=448, interpolation=1, p=1, width=448
|
||||
),
|
||||
albumentations.pytorch.ToTensorV2(
|
||||
transpose_mask=False, always_apply=True, p=1.0
|
||||
),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class PrithviMAE:
|
||||
def __init__(self, model):
|
||||
self.model = LLM(
|
||||
model=model,
|
||||
skip_tokenizer_init=True,
|
||||
dtype="float16",
|
||||
enforce_eager=True,
|
||||
model_impl="terratorch",
|
||||
enable_mm_embeds=True,
|
||||
)
|
||||
|
||||
def run(self, input_data, location_coords):
|
||||
# merge the inputs into one data structure
|
||||
if input_data is not None and input_data.dtype == torch.float32:
|
||||
input_data = input_data.to(torch.float16)
|
||||
input_data = input_data[0]
|
||||
|
||||
mm_data = {
|
||||
"image": {
|
||||
"pixel_values": input_data,
|
||||
"location_coords": location_coords,
|
||||
}
|
||||
}
|
||||
|
||||
prompt = {"prompt_token_ids": [1], "multi_modal_data": mm_data}
|
||||
outputs = self.model.encode(prompt, pooling_task="plugin", use_tqdm=False)
|
||||
|
||||
return outputs[0].outputs.data
|
||||
|
||||
|
||||
def generate_datamodule():
|
||||
datamodule = Sen1Floods11NonGeoDataModule(
|
||||
data_root=datamodule_config["data_root"],
|
||||
batch_size=datamodule_config["batch_size"],
|
||||
num_workers=datamodule_config["num_workers"],
|
||||
bands=datamodule_config["bands"],
|
||||
drop_last=datamodule_config["drop_last"],
|
||||
test_transform=datamodule_config["test_transform"],
|
||||
)
|
||||
|
||||
return datamodule
|
||||
|
||||
|
||||
def process_channel_group(orig_img, channels):
|
||||
"""
|
||||
Args:
|
||||
orig_img: torch.Tensor representing original image (reference)
|
||||
with shape = (bands, H, W).
|
||||
channels: list of indices representing RGB channels.
|
||||
|
||||
Returns:
|
||||
torch.Tensor with shape (num_channels, height, width)
|
||||
for original image
|
||||
"""
|
||||
|
||||
orig_img = orig_img[channels, ...]
|
||||
valid_mask = torch.ones_like(orig_img, dtype=torch.bool)
|
||||
valid_mask[orig_img == NO_DATA_FLOAT] = False
|
||||
|
||||
# Rescale (enhancing contrast)
|
||||
max_value = max(3000, np.percentile(orig_img[valid_mask], PERCENTILE))
|
||||
min_value = OFFSET
|
||||
|
||||
orig_img = torch.clamp((orig_img - min_value) / (max_value - min_value), 0, 1)
|
||||
|
||||
# No data as zeros
|
||||
orig_img[~valid_mask] = 0
|
||||
|
||||
return orig_img
|
||||
|
||||
|
||||
def read_geotiff(file_path: str):
|
||||
"""Read all bands from *file_path* and return image + meta info.
|
||||
|
||||
Args:
|
||||
file_path: path to image file.
|
||||
|
||||
Returns:
|
||||
np.ndarray with shape (bands, height, width)
|
||||
meta info dict
|
||||
"""
|
||||
|
||||
with rasterio.open(file_path) as src:
|
||||
img = src.read()
|
||||
meta = src.meta
|
||||
try:
|
||||
coords = src.lnglat()
|
||||
except Exception:
|
||||
# Cannot read coords
|
||||
coords = None
|
||||
|
||||
return img, meta, coords
|
||||
|
||||
|
||||
def save_geotiff(image, output_path: str, meta: dict):
|
||||
"""Save multi-band image in Geotiff file.
|
||||
|
||||
Args:
|
||||
image: np.ndarray with shape (bands, height, width)
|
||||
output_path: path where to save the image
|
||||
meta: dict with meta info.
|
||||
"""
|
||||
|
||||
with rasterio.open(output_path, "w", **meta) as dest:
|
||||
for i in range(image.shape[0]):
|
||||
dest.write(image[i, :, :], i + 1)
|
||||
|
||||
return
|
||||
|
||||
|
||||
def _convert_np_uint8(float_image: torch.Tensor):
|
||||
image = float_image.numpy() * 255.0
|
||||
image = image.astype(dtype=np.uint8)
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def load_example(
|
||||
file_paths: list[str],
|
||||
mean: list[float] = None,
|
||||
std: list[float] = None,
|
||||
indices: list[int] | None = None,
|
||||
):
|
||||
"""Build an input example by loading images in *file_paths*.
|
||||
|
||||
Args:
|
||||
file_paths: list of file paths .
|
||||
mean: list containing mean values for each band in the
|
||||
images in *file_paths*.
|
||||
std: list containing std values for each band in the
|
||||
images in *file_paths*.
|
||||
|
||||
Returns:
|
||||
np.array containing created example
|
||||
list of meta info for each image in *file_paths*
|
||||
"""
|
||||
|
||||
imgs = []
|
||||
metas = []
|
||||
temporal_coords = []
|
||||
location_coords = []
|
||||
|
||||
for file in file_paths:
|
||||
img, meta, coords = read_geotiff(file)
|
||||
|
||||
# Rescaling (don't normalize on nodata)
|
||||
img = np.moveaxis(img, 0, -1) # channels last for rescaling
|
||||
if indices is not None:
|
||||
img = img[..., indices]
|
||||
if mean is not None and std is not None:
|
||||
img = np.where(img == NO_DATA, NO_DATA_FLOAT, (img - mean) / std)
|
||||
|
||||
imgs.append(img)
|
||||
metas.append(meta)
|
||||
if coords is not None:
|
||||
location_coords.append(coords)
|
||||
|
||||
try:
|
||||
match = re.search(r"(\d{7,8}T\d{6})", file)
|
||||
if match:
|
||||
year = int(match.group(1)[:4])
|
||||
julian_day = match.group(1).split("T")[0][4:]
|
||||
if len(julian_day) == 3:
|
||||
julian_day = int(julian_day)
|
||||
else:
|
||||
julian_day = (
|
||||
datetime.datetime.strptime(julian_day, "%m%d")
|
||||
.timetuple()
|
||||
.tm_yday
|
||||
)
|
||||
temporal_coords.append([year, julian_day])
|
||||
except Exception as e:
|
||||
print(f"Could not extract timestamp for {file} ({e})")
|
||||
|
||||
imgs = np.stack(imgs, axis=0) # num_frames, H, W, C
|
||||
imgs = np.moveaxis(imgs, -1, 0).astype("float32") # C, num_frames, H, W
|
||||
imgs = np.expand_dims(imgs, axis=0) # add batch di
|
||||
|
||||
return imgs, temporal_coords, location_coords, metas
|
||||
|
||||
|
||||
def run_model(
|
||||
input_data,
|
||||
temporal_coords,
|
||||
location_coords,
|
||||
model,
|
||||
datamodule,
|
||||
img_size,
|
||||
lightning_model=None,
|
||||
):
|
||||
# Reflect pad if not divisible by img_size
|
||||
original_h, original_w = input_data.shape[-2:]
|
||||
pad_h = (img_size - (original_h % img_size)) % img_size
|
||||
pad_w = (img_size - (original_w % img_size)) % img_size
|
||||
input_data = np.pad(
|
||||
input_data, ((0, 0), (0, 0), (0, 0), (0, pad_h), (0, pad_w)), mode="reflect"
|
||||
)
|
||||
|
||||
# Build sliding window
|
||||
|
||||
batch_size = 1
|
||||
# batch = torch.tensor(input_data, device="cpu")
|
||||
batch = torch.tensor(input_data)
|
||||
windows = batch.unfold(3, img_size, img_size).unfold(4, img_size, img_size)
|
||||
h1, w1 = windows.shape[3:5]
|
||||
windows = rearrange(
|
||||
windows, "b c t h1 w1 h w -> (b h1 w1) c t h w", h=img_size, w=img_size
|
||||
)
|
||||
|
||||
# Split into batches if number of windows > batch_size
|
||||
num_batches = windows.shape[0] // batch_size if windows.shape[0] > batch_size else 1
|
||||
windows = torch.tensor_split(windows, num_batches, dim=0)
|
||||
|
||||
if temporal_coords:
|
||||
temporal_coords = torch.tensor(temporal_coords).unsqueeze(0)
|
||||
else:
|
||||
temporal_coords = None
|
||||
if location_coords:
|
||||
location_coords = torch.tensor(location_coords[0]).unsqueeze(0)
|
||||
else:
|
||||
location_coords = None
|
||||
|
||||
# Run Prithvi-EO-V2-300M-TL-Sen1Floods11
|
||||
pred_imgs = []
|
||||
for x in windows:
|
||||
# Apply standardization
|
||||
x = datamodule.test_transform(image=x.squeeze().numpy().transpose(1, 2, 0))
|
||||
x = datamodule.aug(x)["image"]
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model.run(x, location_coords=location_coords)
|
||||
y_hat = pred.argmax(dim=1)
|
||||
|
||||
y_hat = torch.nn.functional.interpolate(
|
||||
y_hat.unsqueeze(1).float(), size=img_size, mode="nearest"
|
||||
)
|
||||
|
||||
pred_imgs.append(y_hat)
|
||||
|
||||
pred_imgs = torch.concat(pred_imgs, dim=0)
|
||||
|
||||
# Build images from patches
|
||||
pred_imgs = rearrange(
|
||||
pred_imgs,
|
||||
"(b h1 w1) c h w -> b c (h1 h) (w1 w)",
|
||||
h=img_size,
|
||||
w=img_size,
|
||||
b=1,
|
||||
c=1,
|
||||
h1=h1,
|
||||
w1=w1,
|
||||
)
|
||||
|
||||
# Cut padded area back to original size
|
||||
pred_imgs = pred_imgs[..., :original_h, :original_w]
|
||||
|
||||
# Squeeze (batch size 1)
|
||||
pred_imgs = pred_imgs[0]
|
||||
|
||||
return pred_imgs
|
||||
|
||||
|
||||
def main(
|
||||
data_file: str,
|
||||
model: str,
|
||||
output_dir: str,
|
||||
rgb_outputs: bool,
|
||||
input_indices: list[int] = None,
|
||||
):
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
model_obj = PrithviMAE(model=model)
|
||||
datamodule = generate_datamodule()
|
||||
img_size = 512 # Size of Sen1Floods11
|
||||
|
||||
input_data, temporal_coords, location_coords, meta_data = load_example(
|
||||
file_paths=[data_file],
|
||||
indices=input_indices,
|
||||
)
|
||||
|
||||
meta_data = meta_data[0] # only one image
|
||||
|
||||
if input_data.mean() > 1:
|
||||
input_data = input_data / 10000 # Convert to range 0-1
|
||||
|
||||
channels = [
|
||||
datamodule_config["bands"].index(b) for b in ["RED", "GREEN", "BLUE"]
|
||||
] # BGR -> RGB
|
||||
|
||||
pred = run_model(
|
||||
input_data, temporal_coords, location_coords, model_obj, datamodule, img_size
|
||||
)
|
||||
# Save pred
|
||||
meta_data.update(count=1, dtype="uint8", compress="lzw", nodata=0)
|
||||
pred_file = os.path.join(
|
||||
output_dir, f"pred_{os.path.splitext(os.path.basename(data_file))[0]}.tiff"
|
||||
)
|
||||
save_geotiff(_convert_np_uint8(pred), pred_file, meta_data)
|
||||
|
||||
# Save image + pred
|
||||
meta_data.update(count=3, dtype="uint8", compress="lzw", nodata=0)
|
||||
|
||||
if input_data.mean() < 1:
|
||||
input_data = input_data * 10000 # Scale to 0-10000
|
||||
|
||||
rgb_orig = process_channel_group(
|
||||
orig_img=torch.Tensor(input_data[0, :, 0, ...]),
|
||||
channels=channels,
|
||||
)
|
||||
rgb_orig = rgb_orig.to(torch.float32)
|
||||
|
||||
pred[pred == 0.0] = np.nan
|
||||
img_pred = rgb_orig * 0.7 + pred * 0.3
|
||||
img_pred[img_pred.isnan()] = rgb_orig[img_pred.isnan()]
|
||||
|
||||
img_pred_file = os.path.join(
|
||||
output_dir, f"rgb_pred_{os.path.splitext(os.path.basename(data_file))[0]}.tiff"
|
||||
)
|
||||
save_geotiff(
|
||||
image=_convert_np_uint8(img_pred),
|
||||
output_path=img_pred_file,
|
||||
meta=meta_data,
|
||||
)
|
||||
|
||||
# Save image rgb
|
||||
if rgb_outputs:
|
||||
name_suffix = os.path.splitext(os.path.basename(data_file))[0]
|
||||
rgb_file = os.path.join(
|
||||
output_dir,
|
||||
f"original_rgb_{name_suffix}.tiff",
|
||||
)
|
||||
save_geotiff(
|
||||
image=_convert_np_uint8(rgb_orig),
|
||||
output_path=rgb_file,
|
||||
meta=meta_data,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser("MAE run inference", add_help=False)
|
||||
|
||||
parser.add_argument(
|
||||
"--data_file",
|
||||
type=str,
|
||||
default="./India_900498_S2Hand.tif",
|
||||
help="Path to the file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
|
||||
help="Path to a checkpoint file to load from.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default="output",
|
||||
help="Path to the directory where to save outputs.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_indices",
|
||||
default=[1, 2, 3, 8, 11, 12],
|
||||
type=int,
|
||||
nargs="+",
|
||||
help="""
|
||||
0-based indices of the six Prithvi channels to be selected from the input.
|
||||
By default selects [1,2,3,8,11,12] for S2L1C data.
|
||||
""",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rgb_outputs",
|
||||
action="store_true",
|
||||
help="If present, output files will only contain RGB channels. "
|
||||
"Otherwise, all bands will be saved.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
main(**vars(args))
|
||||
54
third_party/vllm/examples/pooling/plugin/prithvi_geospatial_mae_online.py
vendored
Normal file
54
third_party/vllm/examples/pooling/plugin/prithvi_geospatial_mae_online.py
vendored
Normal file
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import base64
|
||||
import os
|
||||
|
||||
import requests
|
||||
|
||||
# This example shows how to perform an online inference that generates
|
||||
# multimodal data. In this specific case this example will take a geotiff
|
||||
# image as input, process it using the multimodal data processor, and
|
||||
# perform inference.
|
||||
# Requirements :
|
||||
# - install TerraTorch v1.1 (or later):
|
||||
# pip install terratorch>=v1.1
|
||||
# - start vllm in serving mode with the below args
|
||||
# --model='ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11'
|
||||
# --skip-tokenizer-init --enforce-eager
|
||||
# --io-processor-plugin terratorch_segmentation
|
||||
# --enable-mm-embeds
|
||||
|
||||
|
||||
def main():
|
||||
image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
|
||||
server_endpoint = "http://localhost:8000/pooling"
|
||||
|
||||
request_payload_url = {
|
||||
"data": {
|
||||
"data": image_url,
|
||||
"data_format": "url",
|
||||
"image_format": "tiff",
|
||||
"out_data_format": "b64_json",
|
||||
},
|
||||
"priority": 0,
|
||||
"model": "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
|
||||
}
|
||||
|
||||
ret = requests.post(server_endpoint, json=request_payload_url)
|
||||
|
||||
print(f"response.status_code: {ret.status_code}")
|
||||
print(f"response.reason:{ret.reason}")
|
||||
|
||||
response = ret.json()
|
||||
|
||||
decoded_image = base64.b64decode(response["data"]["data"])
|
||||
|
||||
out_path = os.path.join(os.getcwd(), "online_prediction.tiff")
|
||||
|
||||
with open(out_path, "wb") as f:
|
||||
f.write(decoded_image)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
66
third_party/vllm/examples/pooling/pooling/pooling_online.py
vendored
Normal file
66
third_party/vllm/examples/pooling/pooling/pooling_online.py
vendored
Normal file
@@ -0,0 +1,66 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example online usage of Pooling API.
|
||||
|
||||
Run `vllm serve <model> --runner pooling`
|
||||
to start up the server in vLLM. e.g.
|
||||
|
||||
vllm serve internlm/internlm2-1_8b-reward --trust-remote-code
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
response = requests.post(api_url, headers=headers, json=prompt)
|
||||
return response
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
models_url = base_url + "/v1/models"
|
||||
pooing_url = base_url + "/pooling"
|
||||
|
||||
response = requests.get(models_url)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
# Input like Completions API
|
||||
prompt = {"model": model, "input": "vLLM is great!"}
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
|
||||
print("-" * 50)
|
||||
print("Pooling Response:")
|
||||
pprint.pprint(pooling_response.json())
|
||||
print("-" * 50)
|
||||
|
||||
# Input like Chat API
|
||||
prompt = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "vLLM is great!"}],
|
||||
}
|
||||
],
|
||||
}
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
|
||||
print("Pooling Response:")
|
||||
pprint.pprint(pooling_response.json())
|
||||
print("-" * 50)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
47
third_party/vllm/examples/pooling/score/cohere_rerank_client.py
vendored
Normal file
47
third_party/vllm/examples/pooling/score/cohere_rerank_client.py
vendored
Normal file
@@ -0,0 +1,47 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example of using the OpenAI entrypoint's rerank API which is compatible with
|
||||
the Cohere SDK: https://github.com/cohere-ai/cohere-python
|
||||
Note that `pip install cohere` is needed to run this example.
|
||||
|
||||
run: vllm serve BAAI/bge-reranker-base
|
||||
"""
|
||||
|
||||
import cohere
|
||||
from cohere import Client, ClientV2
|
||||
|
||||
model = "BAAI/bge-reranker-base"
|
||||
|
||||
query = "What is the capital of France?"
|
||||
|
||||
documents = [
|
||||
"The capital of France is Paris",
|
||||
"Reranking is fun!",
|
||||
"vLLM is an open-source framework for fast AI serving",
|
||||
]
|
||||
|
||||
|
||||
def cohere_rerank(
|
||||
client: Client | ClientV2, model: str, query: str, documents: list[str]
|
||||
) -> dict:
|
||||
return client.rerank(model=model, query=query, documents=documents)
|
||||
|
||||
|
||||
def main():
|
||||
# cohere v1 client
|
||||
cohere_v1 = cohere.Client(base_url="http://localhost:8000", api_key="sk-fake-key")
|
||||
rerank_v1_result = cohere_rerank(cohere_v1, model, query, documents)
|
||||
print("-" * 50)
|
||||
print("rerank_v1_result:\n", rerank_v1_result)
|
||||
print("-" * 50)
|
||||
|
||||
# or the v2
|
||||
cohere_v2 = cohere.ClientV2("sk-fake-key", base_url="http://localhost:8000")
|
||||
rerank_v2_result = cohere_rerank(cohere_v2, model, query, documents)
|
||||
print("rerank_v2_result:\n", rerank_v2_result)
|
||||
print("-" * 50)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
92
third_party/vllm/examples/pooling/score/colbert_rerank_online.py
vendored
Normal file
92
third_party/vllm/examples/pooling/score/colbert_rerank_online.py
vendored
Normal file
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example of using ColBERT late interaction models for reranking and scoring.
|
||||
|
||||
ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings
|
||||
and MaxSim scoring for document reranking, providing better accuracy than
|
||||
single-vector models while being more efficient than cross-encoders.
|
||||
|
||||
vLLM supports ColBERT with multiple encoder backbones. Start the server
|
||||
with one of the following:
|
||||
|
||||
# BERT backbone (works out of the box)
|
||||
vllm serve answerdotai/answerai-colbert-small-v1
|
||||
|
||||
# ModernBERT backbone
|
||||
vllm serve lightonai/GTE-ModernColBERT-v1 \
|
||||
--hf-overrides '{"architectures": ["ColBERTModernBertModel"]}'
|
||||
|
||||
# Jina XLM-RoBERTa backbone
|
||||
vllm serve jinaai/jina-colbert-v2 \
|
||||
--hf-overrides '{"architectures": ["ColBERTJinaRobertaModel"]}' \
|
||||
--trust-remote-code
|
||||
|
||||
Then run this script:
|
||||
python colbert_rerank_online.py
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
# Change this to match the model you started the server with
|
||||
MODEL = "answerdotai/answerai-colbert-small-v1"
|
||||
BASE_URL = "http://127.0.0.1:8000"
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
documents = [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"Python is a programming language.",
|
||||
"Deep learning uses neural networks for complex tasks.",
|
||||
"The weather today is sunny.",
|
||||
]
|
||||
|
||||
|
||||
def rerank_example():
|
||||
"""Use the /rerank endpoint to rank documents by query relevance."""
|
||||
print("=== Rerank Example ===")
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": "What is machine learning?",
|
||||
"documents": documents,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
result = response.json()
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
print("\nRanked documents (most relevant first):")
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" Score {score:.4f}: {documents[doc_idx]}")
|
||||
|
||||
|
||||
def score_example():
|
||||
"""Use the /score endpoint for pairwise query-document scoring."""
|
||||
print("\n=== Score Example ===")
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"text_1": "What is machine learning?",
|
||||
"text_2": [
|
||||
"Machine learning is a subset of AI.",
|
||||
"The weather is sunny.",
|
||||
],
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
|
||||
result = response.json()
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
|
||||
def main():
|
||||
rerank_example()
|
||||
score_example()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
166
third_party/vllm/examples/pooling/score/colmodernvbert_rerank_online.py
vendored
Normal file
166
third_party/vllm/examples/pooling/score/colmodernvbert_rerank_online.py
vendored
Normal file
@@ -0,0 +1,166 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example of using ColModernVBERT late interaction model for reranking.
|
||||
|
||||
ColModernVBERT is a multi-modal ColBERT-style model combining a SigLIP
|
||||
vision encoder with a ModernBERT text encoder. It produces per-token
|
||||
embeddings and uses MaxSim scoring for retrieval and reranking.
|
||||
Supports both text and image inputs.
|
||||
|
||||
Start the server with:
|
||||
vllm serve ModernVBERT/colmodernvbert-merged --max-model-len 8192
|
||||
|
||||
Then run this script:
|
||||
python colmodernvbert_rerank_online.py
|
||||
"""
|
||||
|
||||
import requests
|
||||
|
||||
MODEL = "ModernVBERT/colmodernvbert-merged"
|
||||
BASE_URL = "http://127.0.0.1:8000"
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/PNG_transparency_demonstration_1.png/300px-PNG_transparency_demonstration_1.png" # noqa: E501
|
||||
|
||||
|
||||
def rerank_text():
|
||||
"""Text-only reranking via /rerank endpoint."""
|
||||
print("=" * 60)
|
||||
print("1. Text reranking (/rerank)")
|
||||
print("=" * 60)
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": "What is machine learning?",
|
||||
"documents": [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"Python is a programming language.",
|
||||
"Deep learning uses neural networks for complex tasks.",
|
||||
"The weather today is sunny.",
|
||||
],
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print("\n Ranked documents (most relevant first):")
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" [{score:.4f}] {data['documents'][doc_idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def score_text():
|
||||
"""Text-only scoring via /score endpoint."""
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("2. Text scoring (/score)")
|
||||
print("=" * 60)
|
||||
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of France is Paris.",
|
||||
"Berlin is the capital of Germany.",
|
||||
"Python is a programming language.",
|
||||
]
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"text_1": query,
|
||||
"text_2": documents,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print(f"\n Query: {query}\n")
|
||||
for item in result["data"]:
|
||||
idx = item["index"]
|
||||
score = item["score"]
|
||||
print(f" Doc {idx} (score={score:.4f}): {documents[idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def score_text_top_n():
|
||||
"""Text reranking with top_n filtering via /rerank endpoint."""
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("3. Text reranking with top_n=2 (/rerank)")
|
||||
print("=" * 60)
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": "What is the capital of France?",
|
||||
"documents": [
|
||||
"The capital of France is Paris.",
|
||||
"Berlin is the capital of Germany.",
|
||||
"Python is a programming language.",
|
||||
"The Eiffel Tower is in Paris.",
|
||||
],
|
||||
"top_n": 2,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print(f"\n Top {data['top_n']} results:")
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" [{score:.4f}] {data['documents'][doc_idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def rerank_multimodal():
|
||||
"""Multimodal reranking with text and image documents via /rerank."""
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("4. Multimodal reranking: text query vs image document (/rerank)")
|
||||
print("=" * 60)
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": "A colorful logo with transparency",
|
||||
"documents": [
|
||||
{"content": [{"type": "image_url", "image_url": {"url": IMAGE_URL}}]},
|
||||
"Python is a programming language.",
|
||||
"The weather today is sunny.",
|
||||
],
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print("\n Ranked documents (most relevant first):")
|
||||
labels = ["[image]", "Python doc", "Weather doc"]
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" [{score:.4f}] {labels[doc_idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def main():
|
||||
rerank_text()
|
||||
score_text()
|
||||
score_text_top_n()
|
||||
rerank_multimodal()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
258
third_party/vllm/examples/pooling/score/colqwen3_rerank_online.py
vendored
Normal file
258
third_party/vllm/examples/pooling/score/colqwen3_rerank_online.py
vendored
Normal file
@@ -0,0 +1,258 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""
|
||||
Example of using ColQwen3 late interaction model for reranking and scoring.
|
||||
|
||||
ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL.
|
||||
It produces per-token embeddings and uses MaxSim scoring for retrieval
|
||||
and reranking. Supports both text and image inputs.
|
||||
|
||||
Start the server with:
|
||||
vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 50000
|
||||
|
||||
Then run this script:
|
||||
python colqwen3_rerank_online.py
|
||||
"""
|
||||
|
||||
import base64
|
||||
from io import BytesIO
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
MODEL = "TomoroAI/tomoro-colqwen3-embed-4b"
|
||||
BASE_URL = "http://127.0.0.1:8000"
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
# ── Image helpers ──────────────────────────────────────────
|
||||
|
||||
|
||||
def load_image(url: str) -> Image.Image:
|
||||
"""Download an image from URL (handles Wikimedia 403)."""
|
||||
for hdrs in (
|
||||
{},
|
||||
{"User-Agent": "Mozilla/5.0 (compatible; ColQwen3-demo/1.0)"},
|
||||
):
|
||||
resp = requests.get(url, headers=hdrs, timeout=15)
|
||||
if resp.status_code == 403:
|
||||
continue
|
||||
resp.raise_for_status()
|
||||
return Image.open(BytesIO(resp.content)).convert("RGB")
|
||||
raise RuntimeError(f"Could not fetch image from {url}")
|
||||
|
||||
|
||||
def encode_image_base64(image: Image.Image) -> str:
|
||||
"""Encode a PIL image to a base64 data URI."""
|
||||
buf = BytesIO()
|
||||
image.save(buf, format="PNG")
|
||||
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
def make_image_content(image_url: str, text: str = "Describe the image.") -> dict:
|
||||
"""Build a ScoreMultiModalParam dict from an image URL."""
|
||||
image = load_image(image_url)
|
||||
return {
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_base64(image)},
|
||||
},
|
||||
{"type": "text", "text": text},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
# ── Sample image URLs ─────────────────────────────────────
|
||||
|
||||
IMAGE_URLS = {
|
||||
"beijing": "https://upload.wikimedia.org/wikipedia/commons/6/61/Beijing_skyline_at_night.JPG",
|
||||
"london": "https://upload.wikimedia.org/wikipedia/commons/4/49/London_skyline.jpg",
|
||||
"singapore": "https://upload.wikimedia.org/wikipedia/commons/2/27/Singapore_skyline_2022.jpg",
|
||||
}
|
||||
|
||||
# ── Text-only examples ────────────────────────────────────
|
||||
|
||||
|
||||
def rerank_text():
|
||||
"""Text-only reranking via /rerank endpoint."""
|
||||
print("=" * 60)
|
||||
print("1. Text reranking (/rerank)")
|
||||
print("=" * 60)
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": "What is machine learning?",
|
||||
"documents": [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"Python is a programming language.",
|
||||
"Deep learning uses neural networks for complex tasks.",
|
||||
"The weather today is sunny.",
|
||||
],
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print("\n Ranked documents (most relevant first):")
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" [{score:.4f}] {data['documents'][doc_idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def score_text():
|
||||
"""Text-only scoring via /score endpoint."""
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("2. Text scoring (/score)")
|
||||
print("=" * 60)
|
||||
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of France is Paris.",
|
||||
"Berlin is the capital of Germany.",
|
||||
"Python is a programming language.",
|
||||
]
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"text_1": query,
|
||||
"text_2": documents,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print(f"\n Query: {query}\n")
|
||||
for item in result["data"]:
|
||||
idx = item["index"]
|
||||
score = item["score"]
|
||||
print(f" Doc {idx} (score={score:.4f}): {documents[idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def score_text_top_n():
|
||||
"""Text reranking with top_n filtering via /rerank endpoint."""
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("3. Text reranking with top_n=2 (/rerank)")
|
||||
print("=" * 60)
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": "What is the capital of France?",
|
||||
"documents": [
|
||||
"The capital of France is Paris.",
|
||||
"Berlin is the capital of Germany.",
|
||||
"Python is a programming language.",
|
||||
"The Eiffel Tower is in Paris.",
|
||||
],
|
||||
"top_n": 2,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print(f"\n Top {data['top_n']} results:")
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" [{score:.4f}] {data['documents'][doc_idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
# ── Multi-modal examples (text query × image documents) ──
|
||||
|
||||
|
||||
def score_text_vs_images():
|
||||
"""Score a text query against image documents via /score."""
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("4. Multi-modal scoring: text query vs image docs (/score)")
|
||||
print("=" * 60)
|
||||
|
||||
query = "Retrieve the city of Beijing"
|
||||
labels = list(IMAGE_URLS.keys())
|
||||
print(f"\n Loading {len(labels)} images...")
|
||||
image_contents = [make_image_content(IMAGE_URLS[name]) for name in labels]
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"data_1": query,
|
||||
"data_2": image_contents,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print(f'\n Query: "{query}"\n')
|
||||
for item in result["data"]:
|
||||
idx = item["index"]
|
||||
print(f" Doc {idx} [{labels[idx]}] score={item['score']:.4f}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def rerank_text_vs_images():
|
||||
"""Rerank image documents by a text query via /rerank."""
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("5. Multi-modal reranking: text query vs image docs (/rerank)")
|
||||
print("=" * 60)
|
||||
|
||||
query = "Retrieve the city of London"
|
||||
labels = list(IMAGE_URLS.keys())
|
||||
print(f"\n Loading {len(labels)} images...")
|
||||
image_contents = [make_image_content(IMAGE_URLS[name]) for name in labels]
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": query,
|
||||
"documents": image_contents,
|
||||
"top_n": 2,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print(f'\n Query: "{query}"')
|
||||
print(f" Top {data['top_n']} results:\n")
|
||||
for item in result["results"]:
|
||||
idx = item["index"]
|
||||
print(f" [{item['relevance_score']:.4f}] {labels[idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
# ── Main ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def main():
|
||||
# Text-only
|
||||
rerank_text()
|
||||
score_text()
|
||||
score_text_top_n()
|
||||
|
||||
# Multi-modal (text query × image documents)
|
||||
score_text_vs_images()
|
||||
rerank_text_vs_images()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
208
third_party/vllm/examples/pooling/score/convert_model_to_seq_cls.py
vendored
Normal file
208
third_party/vllm/examples/pooling/score/convert_model_to_seq_cls.py
vendored
Normal file
@@ -0,0 +1,208 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
Script to convert Large Language Models (LLMs) to Sequence Classification models.
|
||||
This is particularly useful for converting reranker models that use next-token
|
||||
prediction to a sequence classification format for compatibility with standard
|
||||
classification and rerank pipelines.
|
||||
|
||||
Usage examples:
|
||||
- For BAAI/bge-reranker-v2-gemma:
|
||||
python convert_model_to_seq_cls.py --model_name BAAI/bge-reranker-v2-gemma \
|
||||
--classifier_from_tokens '["Yes"]' --method no_post_processing \
|
||||
--path ./bge-reranker-v2-gemma-seq-cls
|
||||
|
||||
- For mxbai-rerank-v2:
|
||||
python convert_model_to_seq_cls.py --model_name mixedbread-ai/mxbai-rerank-base-v2 \
|
||||
--classifier_from_tokens '["0", "1"]' --method from_2_way_softmax \
|
||||
--path ./mxbai-rerank-base-v2-seq-cls
|
||||
|
||||
- For Qwen3-Reranker:
|
||||
python convert_model_to_seq_cls.py --model_name Qwen/Qwen3-Reranker-0.6B \
|
||||
--classifier_from_tokens '["no", "yes"]' --method from_2_way_softmax \
|
||||
--path ./Qwen3-Reranker-0.6B-seq-cls
|
||||
|
||||
Note: For BAAI/bge-reranker-v2-gemma, "Yes" and "yes" are different tokens.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import torch
|
||||
import transformers
|
||||
|
||||
|
||||
def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device):
|
||||
"""
|
||||
This method extracts the difference between weights for 'true' and 'false' tokens
|
||||
from the language model head to create a single classification weight vector.
|
||||
|
||||
Args:
|
||||
causal_lm: The original causal language model
|
||||
seq_cls_model: The target sequence classification model
|
||||
tokenizer: Model tokenizer
|
||||
tokens: List of two tokens representing [false_token, true_token]
|
||||
device: Target device (cpu/cuda)
|
||||
|
||||
Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
"""
|
||||
assert len(tokens) == 2, (
|
||||
"Method requires exactly two tokens for binary classification"
|
||||
)
|
||||
|
||||
# Get the language model head weights (vocabulary_size x hidden_size)
|
||||
lm_head_weights = causal_lm.lm_head.weight
|
||||
|
||||
# Convert token strings to their corresponding token IDs
|
||||
false_id = tokenizer.convert_tokens_to_ids(tokens[0])
|
||||
true_id = tokenizer.convert_tokens_to_ids(tokens[1])
|
||||
|
||||
# Compute the classification weight as the difference between true and false token weights
|
||||
# This follows the approach in: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
score_weight = lm_head_weights[true_id].to(device).to(
|
||||
torch.float32
|
||||
) - lm_head_weights[false_id].to(device).to(torch.float32)
|
||||
|
||||
# Copy the computed weights to the sequence classification model
|
||||
with torch.no_grad():
|
||||
seq_cls_model.score.weight.copy_(score_weight.unsqueeze(0))
|
||||
if seq_cls_model.score.bias is not None:
|
||||
seq_cls_model.score.bias.zero_()
|
||||
|
||||
|
||||
def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device):
|
||||
"""
|
||||
Directly use token weights from the language model head for classification.
|
||||
|
||||
This method maps each classification label directly to a corresponding token
|
||||
in the vocabulary without additional transformation.
|
||||
|
||||
Args:
|
||||
causal_lm: The original causal language model
|
||||
seq_cls_model: The target sequence classification model
|
||||
tokenizer: Model tokenizer
|
||||
tokens: List of tokens representing class labels
|
||||
device: Target device (cpu/cuda)
|
||||
"""
|
||||
# Get the language model head weights (vocabulary_size x hidden_size)
|
||||
lm_head_weights = causal_lm.lm_head.weight
|
||||
|
||||
# Convert all tokens to their corresponding token IDs
|
||||
token_ids = [tokenizer.convert_tokens_to_ids(t) for t in tokens]
|
||||
|
||||
# Extract weights for the specific tokens (num_tokens x hidden_size)
|
||||
score_weight = lm_head_weights[token_ids].to(device)
|
||||
|
||||
# Copy the weights to the sequence classification model
|
||||
with torch.no_grad():
|
||||
seq_cls_model.score.weight.copy_(score_weight)
|
||||
if seq_cls_model.score.bias is not None:
|
||||
seq_cls_model.score.bias.zero_()
|
||||
|
||||
|
||||
method_map = {
|
||||
function.__name__: function for function in [from_2_way_softmax, no_post_processing]
|
||||
}
|
||||
|
||||
|
||||
def converting(
|
||||
model_name, classifier_from_tokens, path, method, use_sep_token=False, device="cpu"
|
||||
):
|
||||
"""
|
||||
Main conversion function to transform a CausalLM model to SequenceClassification.
|
||||
|
||||
Args:
|
||||
model_name: Name or path of the pretrained model
|
||||
classifier_from_tokens: List of tokens used for classification
|
||||
path: Output path to save the converted model
|
||||
method: Conversion method ('from_2_way_softmax' or 'no_post_processing')
|
||||
use_sep_token: Whether to use separating token in the sequence classification model
|
||||
device: Device to load the model on ('cpu' or 'cuda')
|
||||
"""
|
||||
assert method in method_map, f"Unknown method: {method}"
|
||||
|
||||
# Determine number of labels based on conversion method
|
||||
if method == "from_2_way_softmax":
|
||||
assert len(classifier_from_tokens) == 2
|
||||
num_labels = 1
|
||||
else:
|
||||
num_labels = len(classifier_from_tokens)
|
||||
|
||||
# Load tokenizer and original causal language model
|
||||
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
|
||||
causal_lm = transformers.AutoModelForCausalLM.from_pretrained(
|
||||
model_name, device_map=device
|
||||
)
|
||||
|
||||
# Load an empty sequence classification model with the same architecture
|
||||
seq_cls_model = transformers.AutoModelForSequenceClassification.from_pretrained(
|
||||
model_name,
|
||||
num_labels=num_labels,
|
||||
ignore_mismatched_sizes=True,
|
||||
device_map=device,
|
||||
)
|
||||
|
||||
# Apply the selected conversion method to transfer weights
|
||||
method_map[method](
|
||||
causal_lm, seq_cls_model, tokenizer, classifier_from_tokens, device
|
||||
)
|
||||
|
||||
# Configure separating token settings
|
||||
# Note: `llm as reranker` defaults to not using separating token.
|
||||
seq_cls_model.config.use_sep_token = use_sep_token
|
||||
seq_cls_model.config.sep_token_id = tokenizer.sep_token_id
|
||||
|
||||
# Save the converted model and tokenizer
|
||||
seq_cls_model.save_pretrained(path)
|
||||
tokenizer.save_pretrained(path)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Converting *ForCausalLM models to "
|
||||
"*ForSequenceClassification models."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_name",
|
||||
type=str,
|
||||
default="BAAI/bge-reranker-v2-gemma",
|
||||
help="HuggingFace model name or local path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--classifier_from_tokens",
|
||||
type=str,
|
||||
default='["Yes"]',
|
||||
help="JSON string of tokens used for classification labels",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
type=str,
|
||||
default="no_post_processing",
|
||||
help="Conversion method to use",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-sep-token",
|
||||
action="store_true",
|
||||
help="Enable separating token in the sequence classification model",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--path",
|
||||
type=str,
|
||||
default="./bge-reranker-v2-gemma-seq-cls",
|
||||
help="Output directory to save the converted model",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
converting(
|
||||
model_name=args.model_name,
|
||||
classifier_from_tokens=json.loads(args.classifier_from_tokens),
|
||||
method=args.method,
|
||||
use_sep_token=args.use_sep_token,
|
||||
path=args.path,
|
||||
)
|
||||
104
third_party/vllm/examples/pooling/score/qwen3_reranker_offline.py
vendored
Normal file
104
third_party/vllm/examples/pooling/score/qwen3_reranker_offline.py
vendored
Normal file
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
What is the difference between the official original version and one
|
||||
that has been converted into a sequence classification model?
|
||||
|
||||
Qwen3-Reranker is a language model that doing reranker by using the
|
||||
logits of "no" and "yes" tokens.
|
||||
This requires computing logits for all 151,669 tokens in the vocabulary,
|
||||
making it inefficient and incompatible with vLLM's score() API.
|
||||
|
||||
A conversion method has been proposed to transform the original model into a
|
||||
sequence classification model. This converted model:
|
||||
1. Is significantly more efficient
|
||||
2. Fully supports vLLM's score() API
|
||||
3. Simplifies initialization parameters
|
||||
Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
|
||||
|
||||
For the converted model, initialization would simply be:
|
||||
llm = LLM(model="tomaarsen/Qwen3-Reranker-0.6B-seq-cls", runner="pooling")
|
||||
|
||||
This example demonstrates loading the ORIGINAL model with special overrides
|
||||
to make it compatible with vLLM's score API.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from vllm import LLM
|
||||
|
||||
model_name = "Qwen/Qwen3-Reranker-0.6B"
|
||||
|
||||
|
||||
def get_llm() -> LLM:
|
||||
"""
|
||||
Initializes and returns the LLM model for Qwen3-Reranker.
|
||||
|
||||
Returns:
|
||||
LLM: Configured vLLM instance for reranking tasks.
|
||||
|
||||
Note:
|
||||
This function loads the ORIGINAL Qwen3-Reranker model with specific
|
||||
overrides to make it compatible with vLLM's score API.
|
||||
"""
|
||||
return LLM(
|
||||
# Specify the original model from HuggingFace
|
||||
model=model_name,
|
||||
# Use pooling runner for score task
|
||||
runner="pooling",
|
||||
# HuggingFace model configuration overrides required for compatibility
|
||||
hf_overrides={
|
||||
# Manually route to sequence classification architecture
|
||||
# This tells vLLM to use Qwen3ForSequenceClassification instead of
|
||||
# the default Qwen3ForCausalLM
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
# Specify which token logits to extract from the language model head
|
||||
# The original reranker uses "no" and "yes" token logits for scoring
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
# Enable special handling for original Qwen3-Reranker models
|
||||
# This flag triggers conversion logic that transforms the two token
|
||||
# vectors into a single classification vector
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
# Load the Jinja template for formatting query-document pairs
|
||||
# The template ensures proper formatting for the reranker model
|
||||
template_home = Path(__file__).parent / "template"
|
||||
template_path = "qwen3_reranker.jinja"
|
||||
chat_template = (template_home / template_path).read_text()
|
||||
|
||||
# Sample queries for testing the reranker
|
||||
queries = [
|
||||
"What is the capital of China?",
|
||||
"Explain gravity",
|
||||
]
|
||||
|
||||
# Corresponding documents to be scored against each query
|
||||
documents = [
|
||||
"The capital of China is Beijing.",
|
||||
"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
|
||||
]
|
||||
|
||||
# Initialize the LLM model with the original Qwen3-Reranker configuration
|
||||
llm = get_llm()
|
||||
|
||||
# Compute relevance scores for each query-document pair
|
||||
# The score() method returns a relevance score for each pair
|
||||
# Higher scores indicate better relevance
|
||||
outputs = llm.score(queries, documents, chat_template=chat_template)
|
||||
|
||||
# Extract and print the relevance scores from the outputs
|
||||
# Each output contains a score representing query-document relevance
|
||||
print("-" * 30)
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
print("-" * 30)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
80
third_party/vllm/examples/pooling/score/qwen3_reranker_online.py
vendored
Normal file
80
third_party/vllm/examples/pooling/score/qwen3_reranker_online.py
vendored
Normal file
@@ -0,0 +1,80 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""
|
||||
What is the difference between the official original version and one
|
||||
that has been converted into a sequence classification model?
|
||||
|
||||
Qwen3-Reranker is a language model that doing reranker by using the
|
||||
logits of "no" and "yes" tokens.
|
||||
This requires computing logits for all 151,669 tokens in the vocabulary,
|
||||
making it inefficient and incompatible with vLLM's score() API.
|
||||
|
||||
A conversion method has been proposed to transform the original model into a
|
||||
sequence classification model. This converted model:
|
||||
1. Is significantly more efficient
|
||||
2. Fully supports vLLM's score() API
|
||||
3. Simplifies initialization parameters
|
||||
Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
|
||||
|
||||
For the converted model, initialization would simply be:
|
||||
vllm serve tomaarsen/Qwen3-Reranker-0.6B-seq-cls --runner pooling --chat-template examples/pooling/score/template/qwen3_reranker.jinja
|
||||
|
||||
This example demonstrates loading the ORIGINAL model with special overrides
|
||||
to make it compatible with vLLM's score API.
|
||||
vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
# URL of the vLLM server's score endpoint
|
||||
# Default vLLM server runs on localhost port 8000
|
||||
url = "http://127.0.0.1:8000/score"
|
||||
|
||||
# HTTP headers for the request
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
# Example queries & documents
|
||||
queries = [
|
||||
"What is the capital of China?",
|
||||
"Explain gravity",
|
||||
]
|
||||
documents = [
|
||||
"The capital of China is Beijing.",
|
||||
"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
|
||||
]
|
||||
|
||||
# Request payload for the score API
|
||||
data = {
|
||||
"model": "Qwen/Qwen3-Reranker-0.6B",
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to send a score request to the vLLM server.
|
||||
|
||||
This function sends a POST request to the /score endpoint with
|
||||
the query and documents, then prints the relevance scores.
|
||||
"""
|
||||
# Send POST request to the vLLM server's score endpoint
|
||||
response = requests.post(url, headers=headers, json=data)
|
||||
|
||||
# Check if the request was successful
|
||||
if response.status_code == 200:
|
||||
print("Request successful!")
|
||||
# Pretty print the JSON response containing relevance scores
|
||||
# The response includes scores for each document's relevance to the query
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
else:
|
||||
# Handle request failure
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
42
third_party/vllm/examples/pooling/score/rerank_api_online.py
vendored
Normal file
42
third_party/vllm/examples/pooling/score/rerank_api_online.py
vendored
Normal file
@@ -0,0 +1,42 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example of using the OpenAI entrypoint's rerank API which is compatible with
|
||||
Jina and Cohere https://jina.ai/reranker
|
||||
|
||||
run: vllm serve BAAI/bge-reranker-base
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
url = "http://127.0.0.1:8000/rerank"
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
data = {
|
||||
"model": "BAAI/bge-reranker-base",
|
||||
"query": "What is the capital of France?",
|
||||
"documents": [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
"Horses and cows are both animals",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
response = requests.post(url, headers=headers, json=data)
|
||||
|
||||
# Check the response
|
||||
if response.status_code == 200:
|
||||
print("Request successful!")
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
else:
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
69
third_party/vllm/examples/pooling/score/score_api_online.py
vendored
Normal file
69
third_party/vllm/examples/pooling/score/score_api_online.py
vendored
Normal file
@@ -0,0 +1,69 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example online usage of Score API.
|
||||
|
||||
Run `vllm serve <model> --runner pooling` to start up the server in vLLM.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
response = requests.post(api_url, headers=headers, json=prompt)
|
||||
return response
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model", type=str, default="BAAI/bge-reranker-v2-m3")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
api_url = f"http://{args.host}:{args.port}/score"
|
||||
model_name = args.model
|
||||
|
||||
queries = "What is the capital of Brazil?"
|
||||
documents = "The capital of Brazil is Brasilia."
|
||||
prompt = {"model": model_name, "queries": queries, "documents": documents}
|
||||
score_response = post_http_request(prompt=prompt, api_url=api_url)
|
||||
print("\nPrompt when queries and documents are both strings:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
pprint.pprint(score_response.json())
|
||||
|
||||
queries = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
prompt = {"model": model_name, "queries": queries, "documents": documents}
|
||||
score_response = post_http_request(prompt=prompt, api_url=api_url)
|
||||
print("\nPrompt when queries is string and documents is a list:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
pprint.pprint(score_response.json())
|
||||
|
||||
queries = ["What is the capital of Brazil?", "What is the capital of France?"]
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
prompt = {"model": model_name, "queries": queries, "documents": documents}
|
||||
score_response = post_http_request(prompt=prompt, api_url=api_url)
|
||||
print("\nPrompt when queries and documents are both lists:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
pprint.pprint(score_response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
3
third_party/vllm/examples/pooling/score/template/bge-reranker-v2-gemma.jinja
vendored
Normal file
3
third_party/vllm/examples/pooling/score/template/bge-reranker-v2-gemma.jinja
vendored
Normal file
@@ -0,0 +1,3 @@
|
||||
A: {{ (messages | selectattr("role", "eq", "query") | first).content }}
|
||||
B: {{ (messages | selectattr("role", "eq", "document") | first).content }}
|
||||
Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'.
|
||||
8
third_party/vllm/examples/pooling/score/template/mxbai_rerank_v2.jinja
vendored
Normal file
8
third_party/vllm/examples/pooling/score/template/mxbai_rerank_v2.jinja
vendored
Normal file
@@ -0,0 +1,8 @@
|
||||
<|im_start|>system
|
||||
You are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
|
||||
<|im_start|>user
|
||||
query: {{ (messages | selectattr("role", "eq", "query") | first).content }}
|
||||
document: {{ (messages | selectattr("role", "eq", "document") | first).content }}
|
||||
You are a search relevance expert who evaluates how well documents match search queries. For each query-document pair, carefully analyze the semantic relationship between them, then provide your binary relevance judgment (0 for not relevant, 1 for relevant).
|
||||
Relevance:<|im_end|>
|
||||
<|im_start|>assistant
|
||||
3
third_party/vllm/examples/pooling/score/template/nemotron-rerank.jinja
vendored
Normal file
3
third_party/vllm/examples/pooling/score/template/nemotron-rerank.jinja
vendored
Normal file
@@ -0,0 +1,3 @@
|
||||
question:{{ (messages | selectattr("role", "eq", "query") | first).content }}
|
||||
|
||||
passage:{{ (messages | selectattr("role", "eq", "document") | first).content }}
|
||||
15
third_party/vllm/examples/pooling/score/template/nemotron-vl-rerank.jinja
vendored
Normal file
15
third_party/vllm/examples/pooling/score/template/nemotron-vl-rerank.jinja
vendored
Normal file
@@ -0,0 +1,15 @@
|
||||
{%- set query_msg = (messages | selectattr('role', 'equalto', 'query') | list | first) -%}
|
||||
{%- set doc_msg = (messages | selectattr('role', 'equalto', 'document') | list | first) -%}
|
||||
|
||||
{%- set q = query_msg['content'] -%}
|
||||
{%- set d = doc_msg['content'] -%}
|
||||
|
||||
{# If the doc contains <image> anywhere, hoist a single <image> to the front #}
|
||||
{%- set has_image = ("<image>" in d) -%}
|
||||
{%- set d_clean = d | replace("<image>", "") -%}
|
||||
{%- set q_clean = q | replace("<image>", "") -%}
|
||||
|
||||
{%- if has_image -%}<image>{{ " " }}{%- endif -%}
|
||||
question:{{ q_clean }}{{ " " }}
|
||||
{{ " " }}
|
||||
{{ " " }}passage:{{ d_clean }}
|
||||
11
third_party/vllm/examples/pooling/score/template/qwen3_reranker.jinja
vendored
Normal file
11
third_party/vllm/examples/pooling/score/template/qwen3_reranker.jinja
vendored
Normal file
@@ -0,0 +1,11 @@
|
||||
<|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>: {{ messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a web search query, retrieve relevant passages that answer the query") }}
|
||||
<Query>: {{ messages | selectattr("role", "eq", "query") | map(attribute="content") | first }}
|
||||
<Document>: {{ messages | selectattr("role", "eq", "document") | map(attribute="content") | first }}<|im_end|>
|
||||
<|im_start|>assistant
|
||||
<think>
|
||||
|
||||
</think>
|
||||
|
||||
23
third_party/vllm/examples/pooling/score/template/qwen3_vl_reranker.jinja
vendored
Normal file
23
third_party/vllm/examples/pooling/score/template/qwen3_vl_reranker.jinja
vendored
Normal file
@@ -0,0 +1,23 @@
|
||||
<|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>: {{
|
||||
messages
|
||||
| selectattr("role", "eq", "system")
|
||||
| map(attribute="content")
|
||||
| first
|
||||
| default("Given a search query, retrieve relevant candidates that answer the query.")
|
||||
}}<Query>:{{
|
||||
messages
|
||||
| selectattr("role", "eq", "query")
|
||||
| map(attribute="content")
|
||||
| first
|
||||
}}
|
||||
<Document>:{{
|
||||
messages
|
||||
| selectattr("role", "eq", "document")
|
||||
| map(attribute="content")
|
||||
| first
|
||||
}}<|im_end|>
|
||||
<|im_start|>assistant
|
||||
|
||||
159
third_party/vllm/examples/pooling/score/using_template_offline.py
vendored
Normal file
159
third_party/vllm/examples/pooling/score/using_template_offline.py
vendored
Normal file
@@ -0,0 +1,159 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
def parse_args():
|
||||
"""Parse command line arguments for the reranking example.
|
||||
|
||||
This function sets up the argument parser with default values
|
||||
specific to reranking models, including the model name and
|
||||
runner type.
|
||||
"""
|
||||
parser = FlexibleArgumentParser()
|
||||
# Add all EngineArgs command line arguments to the parser
|
||||
parser = EngineArgs.add_cli_args(parser)
|
||||
|
||||
# Set default values specific to this reranking example
|
||||
# These defaults ensure the script works out-of-the-box for reranking tasks
|
||||
parser.set_defaults(
|
||||
model="nvidia/llama-nemotron-rerank-1b-v2", # Default reranking model
|
||||
runner="pooling", # Required for cross-encoder/reranking models
|
||||
trust_remote_code=True, # Allow loading models with custom code
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def get_chat_template(model: str) -> str:
|
||||
"""Load the appropriate chat template for the specified model.
|
||||
|
||||
Reranking models require specific prompt templates to format
|
||||
query-document pairs correctly. This function maps model names
|
||||
to their corresponding template files.
|
||||
"""
|
||||
# Directory containing all chat template files
|
||||
template_home = Path(__file__).parent / "template"
|
||||
|
||||
# Mapping from model names to their corresponding template files
|
||||
# Each reranking model has its own specific prompt format
|
||||
model_name_to_template_path_map = {
|
||||
"BAAI/bge-reranker-v2-gemma": "bge-reranker-v2-gemma.jinja",
|
||||
"Qwen/Qwen3-Reranker-0.6B": "qwen3_reranker.jinja",
|
||||
"Qwen/Qwen3-Reranker-4B": "qwen3_reranker.jinja",
|
||||
"Qwen/Qwen3-Reranker-8B": "qwen3_reranker.jinja",
|
||||
"tomaarsen/Qwen3-Reranker-0.6B-seq-cls": "qwen3_reranker.jinja",
|
||||
"tomaarsen/Qwen3-Reranker-4B-seq-cls": "qwen3_reranker.jinja",
|
||||
"tomaarsen/Qwen3-Reranker-8B-seq-cls": "qwen3_reranker.jinja",
|
||||
"mixedbread-ai/mxbai-rerank-base-v2": "mxbai_rerank_v2.jinja",
|
||||
"mixedbread-ai/mxbai-rerank-large-v2": "mxbai_rerank_v2.jinja",
|
||||
"nvidia/llama-nemotron-rerank-1b-v2": "nemotron-rerank.jinja",
|
||||
}
|
||||
|
||||
# Get the template filename for the specified model
|
||||
template_path = model_name_to_template_path_map.get(model)
|
||||
|
||||
if template_path is None:
|
||||
raise ValueError(f"This demo does not support model name: {model}.")
|
||||
|
||||
# Read and return the template content
|
||||
return (template_home / template_path).read_text()
|
||||
|
||||
|
||||
def get_hf_overrides(model: str) -> dict[str, Any]:
|
||||
"""Convert Large Language Models (LLMs) to Sequence Classification models.
|
||||
|
||||
note:
|
||||
Some reranking models require special configuration overrides to work
|
||||
correctly with vLLM's score API.
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_offline.py
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
|
||||
"""
|
||||
|
||||
model_name_to_hf_overrides_map = {
|
||||
"BAAI/bge-reranker-v2-gemma": {
|
||||
"architectures": ["GemmaForSequenceClassification"],
|
||||
"classifier_from_token": ["Yes"],
|
||||
"method": "no_post_processing",
|
||||
},
|
||||
"Qwen/Qwen3-Reranker-0.6B": {
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
"Qwen/Qwen3-Reranker-4B": {
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
"Qwen/Qwen3-Reranker-8B": {
|
||||
"architectures": ["Qwen3ForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
"tomaarsen/Qwen3-Reranker-0.6B-seq-cls": {},
|
||||
"tomaarsen/Qwen3-Reranker-4B-seq-cls": {},
|
||||
"tomaarsen/Qwen3-Reranker-8B-seq-cls": {},
|
||||
"mixedbread-ai/mxbai-rerank-base-v2": {
|
||||
"architectures": ["Qwen2ForSequenceClassification"],
|
||||
"classifier_from_token": ["0", "1"],
|
||||
"method": "from_2_way_softmax",
|
||||
},
|
||||
"mixedbread-ai/mxbai-rerank-large-v2": {
|
||||
"architectures": ["Qwen2ForSequenceClassification"],
|
||||
"classifier_from_token": ["0", "1"],
|
||||
"method": "from_2_way_softmax",
|
||||
},
|
||||
"nvidia/llama-nemotron-rerank-1b-v2": {},
|
||||
}
|
||||
|
||||
hf_overrides = model_name_to_hf_overrides_map.get(model)
|
||||
|
||||
if hf_overrides is None:
|
||||
raise ValueError(f"This demo does not support model name: {model}.")
|
||||
|
||||
return hf_overrides
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
"""Main execution function for the reranking example."""
|
||||
|
||||
# Get the overrides for the specified model
|
||||
args.hf_overrides = get_hf_overrides(args.model)
|
||||
|
||||
# Initialize the LLM with all provided arguments
|
||||
llm = LLM(**vars(args))
|
||||
|
||||
# Example query for demonstration
|
||||
query = "how much protein should a female eat?"
|
||||
|
||||
# Example documents to be reranked based on relevance to the query
|
||||
documents = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
|
||||
"Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
|
||||
]
|
||||
|
||||
# Load the appropriate chat template for the selected model
|
||||
# The template formats query-document pairs for the reranking model
|
||||
chat_template = get_chat_template(args.model)
|
||||
|
||||
# Score documents based on relevance to the query
|
||||
# The score method returns relevance scores for each document
|
||||
outputs = llm.score(query, documents, chat_template=chat_template)
|
||||
|
||||
# Display the relevance scores
|
||||
# Higher scores indicate more relevant documents
|
||||
print("-" * 30)
|
||||
print([output.outputs.score for output in outputs])
|
||||
print("-" * 30)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
75
third_party/vllm/examples/pooling/score/using_template_online.py
vendored
Normal file
75
third_party/vllm/examples/pooling/score/using_template_online.py
vendored
Normal file
@@ -0,0 +1,75 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
"""
|
||||
Example of using the rerank API with template.
|
||||
|
||||
This script demonstrates how to interact with a vLLM server running
|
||||
a reranking model via the REST API.
|
||||
Before running this script, start the vLLM server with one of the
|
||||
supported reranking models using the commands below.
|
||||
|
||||
note:
|
||||
Some reranking models require special configuration overrides to work correctly
|
||||
with vLLM's score API.
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_online.py
|
||||
Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
|
||||
|
||||
run:
|
||||
vllm serve BAAI/bge-reranker-v2-gemma --hf_overrides '{"architectures": ["GemmaForSequenceClassification"],"classifier_from_token": ["Yes"],"method": "no_post_processing"}' --chat-template examples/pooling/score/template/bge-reranker-v2-gemma.jinja
|
||||
vllm serve tomaarsen/Qwen3-Reranker-0.6B-seq-cls --chat-template examples/pooling/score/template/qwen3_reranker.jinja
|
||||
vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}' --chat-template examples/pooling/score/template/mxbai_rerank_v2.jinja
|
||||
vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja
|
||||
vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
# URL of the vLLM server's rerank endpoint
|
||||
# Default vLLM server runs on localhost port 8000
|
||||
url = "http://127.0.0.1:8000/rerank"
|
||||
|
||||
# HTTP headers for the request
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
# Example query & documents
|
||||
query = "how much protein should a female eat?"
|
||||
documents = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
|
||||
"Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
|
||||
]
|
||||
|
||||
# Request payload for the rerank API
|
||||
data = {
|
||||
"model": "nvidia/llama-nemotron-rerank-1b-v2", # Model to use for reranking
|
||||
"query": query, # The query to score documents against
|
||||
"documents": documents, # List of documents to be scored
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to send a rerank request to the vLLM server.
|
||||
|
||||
This function sends a POST request to the /rerank endpoint with
|
||||
the query and documents, then prints the relevance scores.
|
||||
"""
|
||||
# Send POST request to the vLLM server's rerank endpoint
|
||||
response = requests.post(url, headers=headers, json=data)
|
||||
|
||||
# Check if the request was successful
|
||||
if response.status_code == 200:
|
||||
print("Request successful!")
|
||||
# Pretty print the JSON response containing relevance scores
|
||||
# The response includes scores for each document's relevance to the query
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
else:
|
||||
# Handle request failure
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
132
third_party/vllm/examples/pooling/score/vision_rerank_api_online.py
vendored
Normal file
132
third_party/vllm/examples/pooling/score/vision_rerank_api_online.py
vendored
Normal file
@@ -0,0 +1,132 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
Example Python client for multimodal rerank API which is compatible with
|
||||
Jina and Cohere https://jina.ai/reranker
|
||||
|
||||
Run `vllm serve <model> --runner pooling` to start up the server in vLLM.
|
||||
e.g.
|
||||
vllm serve jinaai/jina-reranker-m0 --runner pooling
|
||||
|
||||
vllm serve Qwen/Qwen3-VL-Reranker-2B \
|
||||
--runner pooling \
|
||||
--max-model-len 4096 \
|
||||
--hf_overrides '{"architectures": ["Qwen3VLForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' \
|
||||
--chat-template examples/pooling/score/template/qwen3_vl_reranker.jinja
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
query = "A woman playing with her dog on a beach at sunset."
|
||||
document = (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust."
|
||||
)
|
||||
image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
video_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": document,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": video_url},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
models_url = base_url + "/v1/models"
|
||||
rerank_url = base_url + "/rerank"
|
||||
|
||||
response = requests.get(models_url)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
print("Query: string & Document: list of string")
|
||||
prompt = {"model": model, "query": query, "documents": [document]}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: text")
|
||||
prompt = {"model": model, "query": query, "documents": {"content": [documents[0]]}}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image url")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": {"content": [documents[1]]},
|
||||
}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image base64")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": {"content": [documents[2]]},
|
||||
}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: video url")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": {"content": [documents[3]]},
|
||||
}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: text + image url")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": {"content": [documents[0], documents[1]]},
|
||||
}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: list")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": [
|
||||
document,
|
||||
{"content": [documents[0]]},
|
||||
{"content": [documents[1]]},
|
||||
{"content": [documents[0], documents[1]]},
|
||||
],
|
||||
}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
185
third_party/vllm/examples/pooling/score/vision_reranker_offline.py
vendored
Normal file
185
third_party/vllm/examples/pooling/score/vision_reranker_offline.py
vendored
Normal file
@@ -0,0 +1,185 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
This example shows how to use vLLM for running offline inference with
|
||||
vision language reranker models for multimodal scoring tasks.
|
||||
|
||||
Vision language rerankers score the relevance between a text query and
|
||||
multimodal documents (text + images/videos).
|
||||
"""
|
||||
|
||||
from argparse import Namespace
|
||||
from collections.abc import Callable
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
TEMPLATE_HOME = Path(__file__).parent / "template"
|
||||
|
||||
|
||||
query = "A woman playing with her dog on a beach at sunset."
|
||||
document = (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched "
|
||||
"beach at sunset, as the dog offers its paw in a heartwarming display of "
|
||||
"companionship and trust."
|
||||
)
|
||||
image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
video_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": document,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": video_url},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class RerankModelData(NamedTuple):
|
||||
engine_args: EngineArgs
|
||||
chat_template: str | None = None
|
||||
modality: set[str] = {}
|
||||
|
||||
|
||||
def run_jinavl_reranker() -> RerankModelData:
|
||||
engine_args = EngineArgs(
|
||||
model="jinaai/jina-reranker-m0",
|
||||
runner="pooling",
|
||||
max_model_len=32768,
|
||||
trust_remote_code=True,
|
||||
mm_processor_kwargs={
|
||||
"min_pixels": 3136,
|
||||
"max_pixels": 602112,
|
||||
},
|
||||
)
|
||||
return RerankModelData(engine_args=engine_args, modality={"image"})
|
||||
|
||||
|
||||
def run_qwen3_vl_reranker() -> RerankModelData:
|
||||
engine_args = EngineArgs(
|
||||
model="Qwen/Qwen3-VL-Reranker-2B",
|
||||
runner="pooling",
|
||||
max_model_len=16384,
|
||||
# HuggingFace model configuration overrides required for compatibility
|
||||
hf_overrides={
|
||||
# Manually route to sequence classification architecture
|
||||
# This tells vLLM to use Qwen3VLForSequenceClassification instead of
|
||||
# the default Qwen3VLForConditionalGeneration
|
||||
"architectures": ["Qwen3VLForSequenceClassification"],
|
||||
# Specify which token logits to extract from the language model head
|
||||
# The original reranker uses "no" and "yes" token logits for scoring
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
# Enable special handling for original Qwen3-Reranker models
|
||||
# This flag triggers conversion logic that transforms the two token
|
||||
# vectors into a single classification vector
|
||||
"is_original_qwen3_reranker": True,
|
||||
},
|
||||
)
|
||||
chat_template_path = "qwen3_vl_reranker.jinja"
|
||||
chat_template = (TEMPLATE_HOME / chat_template_path).read_text()
|
||||
return RerankModelData(
|
||||
engine_args=engine_args,
|
||||
chat_template=chat_template,
|
||||
modality={"image", "video"},
|
||||
)
|
||||
|
||||
|
||||
model_example_map: dict[str, Callable[[], RerankModelData]] = {
|
||||
"jinavl_reranker": run_jinavl_reranker,
|
||||
"qwen3_vl_reranker": run_qwen3_vl_reranker,
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
description="Demo on using vLLM for offline inference with "
|
||||
"vision language reranker models for multimodal scoring tasks."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-name",
|
||||
"-m",
|
||||
type=str,
|
||||
default="jinavl_reranker",
|
||||
choices=model_example_map.keys(),
|
||||
help="The name of the reranker model.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
# Run the selected reranker model
|
||||
model_request = model_example_map[args.model_name]()
|
||||
engine_args = model_request.engine_args
|
||||
|
||||
llm = LLM(**asdict(engine_args))
|
||||
|
||||
print("Query: string & Document: string")
|
||||
outputs = llm.score(query, document)
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
|
||||
print("Query: string & Document: text")
|
||||
outputs = llm.score(
|
||||
query, {"content": [documents[0]]}, chat_template=model_request.chat_template
|
||||
)
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
|
||||
print("Query: string & Document: image url")
|
||||
outputs = llm.score(
|
||||
query, {"content": [documents[1]]}, chat_template=model_request.chat_template
|
||||
)
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
|
||||
print("Query: string & Document: image base64")
|
||||
outputs = llm.score(
|
||||
query, {"content": [documents[2]]}, chat_template=model_request.chat_template
|
||||
)
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
|
||||
if "video" in model_request.modality:
|
||||
print("Query: string & Document: video url")
|
||||
outputs = llm.score(
|
||||
query,
|
||||
{"content": [documents[3]]},
|
||||
chat_template=model_request.chat_template,
|
||||
)
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
|
||||
print("Query: string & Document: text + image url")
|
||||
outputs = llm.score(
|
||||
query,
|
||||
{"content": [documents[0], documents[1]]},
|
||||
chat_template=model_request.chat_template,
|
||||
)
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
|
||||
print("Query: string & Document: list")
|
||||
outputs = llm.score(
|
||||
query,
|
||||
[
|
||||
document,
|
||||
{"content": [documents[0]]},
|
||||
{"content": [documents[1]]},
|
||||
{"content": [documents[0], documents[1]]},
|
||||
],
|
||||
chat_template=model_request.chat_template,
|
||||
)
|
||||
print("Relevance scores:", [output.outputs.score for output in outputs])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
150
third_party/vllm/examples/pooling/score/vision_score_api_online.py
vendored
Normal file
150
third_party/vllm/examples/pooling/score/vision_score_api_online.py
vendored
Normal file
@@ -0,0 +1,150 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
Example online usage of Score API.
|
||||
|
||||
Run `vllm serve <model> --runner pooling` to start up the server in vLLM.
|
||||
e.g.
|
||||
vllm serve jinaai/jina-reranker-m0 --runner pooling
|
||||
|
||||
vllm serve Qwen/Qwen3-VL-Reranker-2B \
|
||||
--runner pooling \
|
||||
--max-model-len 4096 \
|
||||
--hf_overrides '{"architectures": ["Qwen3VLForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' \
|
||||
--chat-template examples/pooling/score/template/qwen3_vl_reranker.jinja
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
query = "A woman playing with her dog on a beach at sunset."
|
||||
document = (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust."
|
||||
)
|
||||
image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
video_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": document,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": video_url},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
models_url = base_url + "/v1/models"
|
||||
score_url = base_url + "/score"
|
||||
|
||||
response = requests.get(models_url)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
print("Query: string & Document: string")
|
||||
prompt = {"model": model, "queries": query, "documents": document}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: text")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[0]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image url")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[1]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image base64")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[2]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: video url")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[3]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: text + image url")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[0], documents[1]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: list")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": [
|
||||
document,
|
||||
{"content": [documents[0]]},
|
||||
{"content": [documents[1]]},
|
||||
{"content": [documents[0], documents[1]]},
|
||||
],
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: list & Document: list")
|
||||
data = [
|
||||
document,
|
||||
{"content": [documents[0]]},
|
||||
{"content": [documents[1]]},
|
||||
{"content": [documents[0], documents[1]]},
|
||||
]
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": data,
|
||||
"documents": data,
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
54
third_party/vllm/examples/pooling/token_classify/ner_offline.py
vendored
Normal file
54
third_party/vllm/examples/pooling/token_classify/ner_offline.py
vendored
Normal file
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from https://huggingface.co/boltuix/NeuroBERT-NER
|
||||
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser()
|
||||
parser = EngineArgs.add_cli_args(parser)
|
||||
# Set example specific arguments
|
||||
parser.set_defaults(
|
||||
model="boltuix/NeuroBERT-NER",
|
||||
runner="pooling",
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
# Sample prompts.
|
||||
prompts = [
|
||||
"Barack Obama visited Microsoft headquarters in Seattle on January 2025."
|
||||
]
|
||||
|
||||
# Create an LLM.
|
||||
llm = LLM(**vars(args))
|
||||
tokenizer = llm.get_tokenizer()
|
||||
label_map = llm.llm_engine.vllm_config.model_config.hf_config.id2label
|
||||
|
||||
# Run inference
|
||||
outputs = llm.encode(prompts, pooling_task="token_classify")
|
||||
|
||||
for prompt, output in zip(prompts, outputs):
|
||||
logits = output.outputs.data
|
||||
predictions = logits.argmax(dim=-1)
|
||||
|
||||
# Map predictions to labels
|
||||
tokens = tokenizer.convert_ids_to_tokens(output.prompt_token_ids)
|
||||
labels = [label_map[p.item()] for p in predictions]
|
||||
|
||||
# Print results
|
||||
for token, label in zip(tokens, labels):
|
||||
if token not in tokenizer.all_special_tokens:
|
||||
print(f"{token:15} → {label}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
71
third_party/vllm/examples/pooling/token_classify/ner_online.py
vendored
Normal file
71
third_party/vllm/examples/pooling/token_classify/ner_online.py
vendored
Normal file
@@ -0,0 +1,71 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from https://huggingface.co/boltuix/NeuroBERT-NER
|
||||
|
||||
"""
|
||||
Example online usage of Pooling API for Named Entity Recognition (NER).
|
||||
|
||||
Run `vllm serve <model> --runner pooling`
|
||||
to start up the server in vLLM. e.g.
|
||||
|
||||
vllm serve boltuix/NeuroBERT-NER
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
import requests
|
||||
import torch
|
||||
|
||||
|
||||
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
response = requests.post(api_url, headers=headers, json=prompt)
|
||||
return response
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model", type=str, default="boltuix/NeuroBERT-NER")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
from transformers import AutoConfig, AutoTokenizer
|
||||
|
||||
api_url = f"http://{args.host}:{args.port}/pooling"
|
||||
model_name = args.model
|
||||
|
||||
# Load tokenizer and config
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
config = AutoConfig.from_pretrained(model_name)
|
||||
label_map = config.id2label
|
||||
|
||||
# Input text
|
||||
text = "Barack Obama visited Microsoft headquarters in Seattle on January 2025."
|
||||
prompt = {"model": model_name, "input": text}
|
||||
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=api_url)
|
||||
|
||||
# Run inference
|
||||
output = pooling_response.json()["data"][0]
|
||||
logits = torch.tensor(output["data"])
|
||||
predictions = logits.argmax(dim=-1)
|
||||
inputs = tokenizer(text, return_tensors="pt")
|
||||
|
||||
# Map predictions to labels
|
||||
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
|
||||
labels = [label_map[p.item()] for p in predictions]
|
||||
assert len(tokens) == len(predictions)
|
||||
|
||||
# Print results
|
||||
for token, label in zip(tokens, labels):
|
||||
if token not in tokenizer.all_special_tokens:
|
||||
print(f"{token:15} → {label}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
198
third_party/vllm/examples/pooling/token_embed/colqwen3_token_embed_online.py
vendored
Normal file
198
third_party/vllm/examples/pooling/token_embed/colqwen3_token_embed_online.py
vendored
Normal file
@@ -0,0 +1,198 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
Example online usage of Pooling API for ColQwen3 multi-vector retrieval.
|
||||
|
||||
ColQwen3 is a multi-modal late interaction model based on Qwen3-VL that
|
||||
produces per-token embeddings (320-dim, L2-normalized) for both text and
|
||||
image inputs. Similarity is computed via MaxSim scoring.
|
||||
|
||||
This example mirrors the official TomoroAI inference code
|
||||
(https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-4b) but uses the
|
||||
vLLM serving API instead of local HuggingFace model loading.
|
||||
|
||||
Start the server with:
|
||||
vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
|
||||
|
||||
Then run this script:
|
||||
python colqwen3_token_embed_online.py
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
# ── Helpers ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def post_http_request(payload: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
return requests.post(api_url, headers=headers, json=payload)
|
||||
|
||||
|
||||
def load_image(url: str) -> Image.Image:
|
||||
"""Download an image from URL (handles Wikimedia 403)."""
|
||||
for hdrs in ({}, {"User-Agent": "Mozilla/5.0 (compatible; ColQwen3-demo/1.0)"}):
|
||||
resp = requests.get(url, headers=hdrs, timeout=10)
|
||||
if resp.status_code == 403:
|
||||
continue
|
||||
resp.raise_for_status()
|
||||
return Image.open(BytesIO(resp.content)).convert("RGB")
|
||||
raise RuntimeError(f"Could not fetch image from {url}")
|
||||
|
||||
|
||||
def encode_image_base64(image: Image.Image) -> str:
|
||||
"""Encode a PIL image to a base64 data URI."""
|
||||
buf = BytesIO()
|
||||
image.save(buf, format="PNG")
|
||||
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
def compute_maxsim(q_emb: np.ndarray, d_emb: np.ndarray) -> float:
|
||||
"""Compute ColBERT-style MaxSim score between query and document."""
|
||||
sim = q_emb @ d_emb.T
|
||||
return float(sim.max(axis=-1).sum())
|
||||
|
||||
|
||||
# ── Encode functions ────────────────────────────────────────
|
||||
|
||||
|
||||
def encode_queries(texts: list[str], model: str, api_url: str) -> list[np.ndarray]:
|
||||
"""Encode text queries → list of multi-vector embeddings."""
|
||||
resp = post_http_request({"model": model, "input": texts}, api_url)
|
||||
return [np.array(item["data"]) for item in resp.json()["data"]]
|
||||
|
||||
|
||||
def encode_images(image_urls: list[str], model: str, api_url: str) -> list[np.ndarray]:
|
||||
"""Encode image documents → list of multi-vector embeddings.
|
||||
|
||||
Images are sent via the chat-style `messages` field so that the
|
||||
vLLM multimodal processor handles them correctly.
|
||||
"""
|
||||
embeddings = []
|
||||
for url in image_urls:
|
||||
print(f" Loading: {url.split('/')[-1]}...")
|
||||
image = load_image(url)
|
||||
image_uri = encode_image_base64(image)
|
||||
resp = post_http_request(
|
||||
{
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_uri}},
|
||||
{"type": "text", "text": "Describe the image."},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
api_url,
|
||||
)
|
||||
result = resp.json()
|
||||
if resp.status_code != 200 or "data" not in result:
|
||||
print(f" Error ({resp.status_code}): {str(result)[:200]}")
|
||||
continue
|
||||
embeddings.append(np.array(result["data"][0]["data"]))
|
||||
return embeddings
|
||||
|
||||
|
||||
# ── Main ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="TomoroAI/tomoro-colqwen3-embed-4b",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
pooling_url = f"http://{args.host}:{args.port}/pooling"
|
||||
score_url = f"http://{args.host}:{args.port}/score"
|
||||
model = args.model
|
||||
|
||||
# Same sample data as the official TomoroAI example
|
||||
queries = [
|
||||
"Retrieve the city of Singapore",
|
||||
"Retrieve the city of Beijing",
|
||||
"Retrieve the city of London",
|
||||
]
|
||||
image_urls = [
|
||||
"https://upload.wikimedia.org/wikipedia/commons/2/27/Singapore_skyline_2022.jpg",
|
||||
"https://upload.wikimedia.org/wikipedia/commons/6/61/Beijing_skyline_at_night.JPG",
|
||||
"https://upload.wikimedia.org/wikipedia/commons/4/49/London_skyline.jpg",
|
||||
]
|
||||
|
||||
# ── 1) Text query embeddings ────────────────────────────
|
||||
print("=" * 60)
|
||||
print("1. Encode text queries (multi-vector)")
|
||||
print("=" * 60)
|
||||
query_embeddings = encode_queries(queries, model, pooling_url)
|
||||
for i, emb in enumerate(query_embeddings):
|
||||
norm = float(np.linalg.norm(emb[0]))
|
||||
print(f' Query {i}: {emb.shape} (L2 norm: {norm:.4f}) "{queries[i]}"')
|
||||
|
||||
# ── 2) Image document embeddings ────────────────────────
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("2. Encode image documents (multi-vector)")
|
||||
print("=" * 60)
|
||||
doc_embeddings = encode_images(image_urls, model, pooling_url)
|
||||
for i, emb in enumerate(doc_embeddings):
|
||||
print(f" Doc {i}: {emb.shape} {image_urls[i].split('/')[-1]}")
|
||||
|
||||
# ── 3) Cross-modal MaxSim scoring ───────────────────────
|
||||
if doc_embeddings:
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("3. Cross-modal MaxSim scores (text queries × image docs)")
|
||||
print("=" * 60)
|
||||
# Header
|
||||
print(f"{'':>35s}", end="")
|
||||
for j in range(len(doc_embeddings)):
|
||||
print(f" Doc {j:>2d}", end="")
|
||||
print()
|
||||
# Score matrix
|
||||
for i, q_emb in enumerate(query_embeddings):
|
||||
print(f" {queries[i]:<33s}", end="")
|
||||
for j, d_emb in enumerate(doc_embeddings):
|
||||
score = compute_maxsim(q_emb, d_emb)
|
||||
print(f" {score:6.2f}", end="")
|
||||
print()
|
||||
|
||||
# ── 4) Text-only /score endpoint ────────────────────────
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("4. Text-only late interaction scoring (/score endpoint)")
|
||||
print("=" * 60)
|
||||
text_query = "What is the capital of France?"
|
||||
text_docs = [
|
||||
"The capital of France is Paris.",
|
||||
"Berlin is the capital of Germany.",
|
||||
"Python is a programming language.",
|
||||
]
|
||||
resp = post_http_request(
|
||||
{"model": model, "text_1": text_query, "text_2": text_docs},
|
||||
score_url,
|
||||
)
|
||||
print(f' Query: "{text_query}"\n')
|
||||
for item in resp.json()["data"]:
|
||||
idx = item["index"]
|
||||
print(f" Doc {idx} (score={item['score']:.4f}): {text_docs[idx]}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
71
third_party/vllm/examples/pooling/token_embed/jina_embeddings_v4_offline.py
vendored
Normal file
71
third_party/vllm/examples/pooling/token_embed/jina_embeddings_v4_offline.py
vendored
Normal file
@@ -0,0 +1,71 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.inputs.data import TextPrompt
|
||||
from vllm.multimodal.utils import fetch_image
|
||||
|
||||
# Initialize model
|
||||
model = LLM(
|
||||
model="jinaai/jina-embeddings-v4-vllm-text-matching",
|
||||
runner="pooling",
|
||||
max_model_len=1024,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
# Create text prompts
|
||||
text1 = "Ein wunderschöner Sonnenuntergang am Strand"
|
||||
text1_prompt = TextPrompt(prompt=f"Query: {text1}")
|
||||
|
||||
text2 = "浜辺に沈む美しい夕日"
|
||||
text2_prompt = TextPrompt(prompt=f"Query: {text2}")
|
||||
|
||||
# Create image prompt
|
||||
image = fetch_image(
|
||||
"https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
|
||||
)
|
||||
image_prompt = TextPrompt(
|
||||
prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
|
||||
multi_modal_data={"image": image},
|
||||
)
|
||||
|
||||
# Encode all prompts
|
||||
prompts = [text1_prompt, text2_prompt, image_prompt]
|
||||
outputs = model.encode(prompts, pooling_task="token_embed")
|
||||
|
||||
|
||||
def get_embeddings(outputs):
|
||||
VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
|
||||
|
||||
embeddings = []
|
||||
for output in outputs:
|
||||
if VISION_START_TOKEN_ID in output.prompt_token_ids:
|
||||
# Gather only vision tokens
|
||||
img_start_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
|
||||
)[0][0]
|
||||
img_end_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
|
||||
)[0][0]
|
||||
embeddings_tensor = output.outputs.data.detach().clone()[
|
||||
img_start_pos : img_end_pos + 1
|
||||
]
|
||||
else:
|
||||
# Use all tokens for text-only prompts
|
||||
embeddings_tensor = output.outputs.data.detach().clone()
|
||||
|
||||
# Pool and normalize embeddings
|
||||
pooled_output = (
|
||||
embeddings_tensor.sum(dim=0, dtype=torch.float32)
|
||||
/ embeddings_tensor.shape[0]
|
||||
)
|
||||
embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
|
||||
return embeddings
|
||||
|
||||
|
||||
embeddings = get_embeddings(outputs)
|
||||
|
||||
for embedding in embeddings:
|
||||
print(embedding.shape)
|
||||
56
third_party/vllm/examples/pooling/token_embed/multi_vector_retrieval_offline.py
vendored
Normal file
56
third_party/vllm/examples/pooling/token_embed/multi_vector_retrieval_offline.py
vendored
Normal file
@@ -0,0 +1,56 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser()
|
||||
parser = EngineArgs.add_cli_args(parser)
|
||||
# Set example specific arguments
|
||||
parser.set_defaults(
|
||||
model="BAAI/bge-m3",
|
||||
runner="pooling",
|
||||
enforce_eager=True,
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args: Namespace):
|
||||
# Sample prompts.
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
# Create an LLM.
|
||||
# You should pass runner="pooling" for embedding models
|
||||
llm = LLM(**vars(args))
|
||||
|
||||
# Generate embedding. The output is a list of EmbeddingRequestOutputs.
|
||||
outputs = llm.embed(prompts)
|
||||
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for prompt, output in zip(prompts, outputs):
|
||||
embeds = output.outputs.embedding
|
||||
print(len(embeds))
|
||||
|
||||
# Generate embedding for each token. The output is a list of PoolingRequestOutput.
|
||||
outputs = llm.encode(prompts, pooling_task="token_embed")
|
||||
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for prompt, output in zip(prompts, outputs):
|
||||
multi_vector = output.outputs.data
|
||||
print(multi_vector.shape)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
54
third_party/vllm/examples/pooling/token_embed/multi_vector_retrieval_online.py
vendored
Normal file
54
third_party/vllm/examples/pooling/token_embed/multi_vector_retrieval_online.py
vendored
Normal file
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Example online usage of Pooling API for multi vector retrieval.
|
||||
|
||||
Run `vllm serve <model> --runner pooling`
|
||||
to start up the server in vLLM. e.g.
|
||||
|
||||
vllm serve BAAI/bge-m3
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
import requests
|
||||
import torch
|
||||
|
||||
|
||||
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
response = requests.post(api_url, headers=headers, json=prompt)
|
||||
return response
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model", type=str, default="BAAI/bge-m3")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
api_url = f"http://{args.host}:{args.port}/pooling"
|
||||
model_name = args.model
|
||||
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
prompt = {"model": model_name, "input": prompts}
|
||||
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=api_url)
|
||||
for output in pooling_response.json()["data"]:
|
||||
multi_vector = torch.tensor(output["data"])
|
||||
print(multi_vector.shape)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
Reference in New Issue
Block a user