mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-10 12:35:51 +08:00
99 lines
3.0 KiB
Python
99 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
|
|
import openai # use the official client for correctness check
|
|
import pytest
|
|
import pytest_asyncio
|
|
import torch
|
|
from transformers import AutoConfig
|
|
|
|
from tests.conftest import ImageTestAssets
|
|
from tests.utils import RemoteOpenAIServer
|
|
|
|
# any model with a chat template should work here
|
|
MODEL_NAME = "llava-hf/llava-1.5-7b-hf"
|
|
CONFIG = AutoConfig.from_pretrained(MODEL_NAME)
|
|
MAXIMUM_IMAGES = 2
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def default_image_embeds_server_args() -> list[str]:
|
|
return [
|
|
"--dtype",
|
|
"bfloat16",
|
|
"--max-model-len",
|
|
"2048",
|
|
"--max-num-seqs",
|
|
"4",
|
|
"--enforce-eager",
|
|
"--limit-mm-per-prompt",
|
|
json.dumps({"image": MAXIMUM_IMAGES}),
|
|
]
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def server_with_image_embeds(default_image_embeds_server_args):
|
|
with RemoteOpenAIServer(
|
|
MODEL_NAME, default_image_embeds_server_args, max_wait_seconds=600
|
|
) as remote_server:
|
|
yield remote_server
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def client_with_image_embeds(server_with_image_embeds):
|
|
async with server_with_image_embeds.get_async_client() as async_client:
|
|
yield async_client
|
|
|
|
|
|
def encode_image_embedding_to_base64(image_embedding) -> str:
|
|
"""
|
|
Encode image embedding to base64 string
|
|
"""
|
|
buffer = io.BytesIO()
|
|
torch.save(image_embedding, buffer)
|
|
buffer.seek(0)
|
|
binary_data = buffer.read()
|
|
base64_image_embedding = base64.b64encode(binary_data).decode("utf-8")
|
|
return base64_image_embedding
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
|
@pytest.mark.parametrize("dtype", [torch.half, torch.float16, torch.float32])
|
|
async def test_completions_with_image_embeds(
|
|
client_with_image_embeds: openai.AsyncOpenAI,
|
|
model_name: str,
|
|
image_assets: ImageTestAssets,
|
|
dtype: torch.dtype,
|
|
):
|
|
# Test case: Single image embeds input
|
|
image_embeds = image_assets[0].image_embeds.to(dtype=dtype)
|
|
base64_image_embedding = encode_image_embedding_to_base64(image_embeds)
|
|
chat_completion = await client_with_image_embeds.chat.completions.create(
|
|
messages=[
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": "Describe these images separately. For each image,"
|
|
"reply with a short sentence (no more than 10 words).",
|
|
},
|
|
{
|
|
"type": "image_embeds",
|
|
"image_embeds": base64_image_embedding,
|
|
},
|
|
],
|
|
},
|
|
],
|
|
model=model_name,
|
|
)
|
|
assert chat_completion.choices[0].message.content is not None
|
|
assert isinstance(chat_completion.choices[0].message.content, str)
|
|
assert len(chat_completion.choices[0].message.content) > 0
|