mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-16 16:36:02 +08:00
86ae693f20
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
64 lines
1.5 KiB
Python
64 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.utils import set_default_torch_num_threads
|
|
|
|
from ....conftest import VllmRunner
|
|
|
|
|
|
def generate_test_mm_data():
|
|
mm_data = {
|
|
"pixel_values": torch.full((6, 512, 512), 1.0, dtype=torch.float16),
|
|
"location_coords": torch.full((1, 2), 1.0, dtype=torch.float16),
|
|
}
|
|
return mm_data
|
|
|
|
|
|
def _run_test(
|
|
vllm_runner: type[VllmRunner],
|
|
model: str,
|
|
) -> None:
|
|
|
|
prompt = [
|
|
{
|
|
# This model deals with no text input
|
|
"prompt_token_ids": [1],
|
|
"multi_modal_data": generate_test_mm_data(),
|
|
} for _ in range(10)
|
|
]
|
|
|
|
with (
|
|
set_default_torch_num_threads(1),
|
|
vllm_runner(
|
|
model,
|
|
runner="pooling",
|
|
dtype=torch.float16,
|
|
enforce_eager=True,
|
|
skip_tokenizer_init=True,
|
|
# Limit the maximum number of sequences to avoid the
|
|
# test going OOM during the warmup run
|
|
max_num_seqs=32,
|
|
) as vllm_model,
|
|
):
|
|
vllm_model.encode(prompt)
|
|
|
|
|
|
MODELS = ["christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM"]
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.parametrize("model", MODELS)
|
|
def test_models_image(
|
|
hf_runner,
|
|
vllm_runner,
|
|
image_assets,
|
|
model: str,
|
|
) -> None:
|
|
_run_test(
|
|
vllm_runner,
|
|
model,
|
|
)
|