[CI/Build] Reduce more duplicate runner startup in tests (#52570)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
Isotr0py
2026-08-17 05:19:07 +00:00
committed by GitHub
parent 311b3513af
commit 53e211d292
10 changed files with 647 additions and 916 deletions
+58 -108
View File
@@ -212,14 +212,13 @@ def colbert_extra_kwargs(colbert_spec):
return colbert_spec["extra_kwargs"]
def test_colbert_token_embed(
@pytest.fixture(scope="module")
def colbert_model(
vllm_runner,
colbert_model_name,
colbert_dim,
colbert_max_model_len,
colbert_extra_kwargs,
):
"""Test that ColBERT model produces token embeddings."""
with vllm_runner(
colbert_model_name,
runner="pooling",
@@ -228,112 +227,84 @@ def test_colbert_token_embed(
enforce_eager=True,
**colbert_extra_kwargs,
) as vllm_model:
outputs = vllm_model.token_embed([TEXTS_1[0]])
yield vllm_model
assert len(outputs) == 1
emb = torch.as_tensor(outputs[0])
assert emb.dim() == 2
assert emb.shape[1] == colbert_dim
assert emb.shape[0] > 1
def test_colbert_token_embed(
colbert_model,
colbert_dim,
):
"""Test that ColBERT model produces token embeddings."""
outputs = colbert_model.token_embed([TEXTS_1[0]])
assert len(outputs) == 1
emb = torch.as_tensor(outputs[0])
assert emb.dim() == 2
assert emb.shape[1] == colbert_dim
assert emb.shape[0] > 1
def test_colbert_late_interaction_1_to_1(
vllm_runner,
colbert_model_name,
colbert_max_model_len,
colbert_extra_kwargs,
colbert_model,
):
"""Test ColBERT late interaction scoring with 1:1 query-document pair."""
with vllm_runner(
colbert_model_name,
runner="pooling",
dtype=DTYPE,
max_model_len=colbert_max_model_len,
enforce_eager=True,
**colbert_extra_kwargs,
) as vllm_model:
q_outputs = vllm_model.token_embed([TEXTS_1[0]])
d_outputs = vllm_model.token_embed([TEXTS_2[0]])
q_outputs = colbert_model.token_embed([TEXTS_1[0]])
d_outputs = colbert_model.token_embed([TEXTS_2[0]])
q_emb = torch.as_tensor(q_outputs[0])
d_emb = torch.as_tensor(d_outputs[0])
q_emb = torch.as_tensor(q_outputs[0])
d_emb = torch.as_tensor(d_outputs[0])
manual_score = compute_maxsim_score(q_emb, d_emb).item()
manual_score = compute_maxsim_score(q_emb, d_emb).item()
vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2[0])
vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2[0])
assert len(vllm_scores) == 1
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
assert len(vllm_scores) == 1
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
def test_colbert_late_interaction_1_to_N(
vllm_runner,
colbert_model_name,
colbert_max_model_len,
colbert_extra_kwargs,
colbert_model,
):
"""Test ColBERT late interaction scoring with 1:N query-documents."""
with vllm_runner(
colbert_model_name,
runner="pooling",
dtype=DTYPE,
max_model_len=colbert_max_model_len,
enforce_eager=True,
**colbert_extra_kwargs,
) as vllm_model:
q_outputs = vllm_model.token_embed([TEXTS_1[0]])
d_outputs = vllm_model.token_embed(TEXTS_2)
q_outputs = colbert_model.token_embed([TEXTS_1[0]])
d_outputs = colbert_model.token_embed(TEXTS_2)
q_emb = torch.as_tensor(q_outputs[0])
q_emb = torch.as_tensor(q_outputs[0])
manual_scores = []
for d_out in d_outputs:
d_emb = torch.as_tensor(d_out)
manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
manual_scores = []
for d_out in d_outputs:
d_emb = torch.as_tensor(d_out)
manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2)
vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2)
assert len(vllm_scores) == 2
for i in range(2):
assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
assert len(vllm_scores) == 2
for i in range(2):
assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
def test_colbert_late_interaction_N_to_N(
vllm_runner,
colbert_model_name,
colbert_max_model_len,
colbert_extra_kwargs,
colbert_model,
):
"""Test ColBERT late interaction scoring with N:N query-documents."""
with vllm_runner(
colbert_model_name,
runner="pooling",
dtype=DTYPE,
max_model_len=colbert_max_model_len,
enforce_eager=True,
**colbert_extra_kwargs,
) as vllm_model:
q_outputs = vllm_model.token_embed(TEXTS_1)
d_outputs = vllm_model.token_embed(TEXTS_2)
q_outputs = colbert_model.token_embed(TEXTS_1)
d_outputs = colbert_model.token_embed(TEXTS_2)
manual_scores = []
for q_out, d_out in zip(q_outputs, d_outputs):
q_emb = torch.as_tensor(q_out)
d_emb = torch.as_tensor(d_out)
manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
manual_scores = []
for q_out, d_out in zip(q_outputs, d_outputs):
q_emb = torch.as_tensor(q_out)
d_emb = torch.as_tensor(d_out)
manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2)
vllm_scores = colbert_model.score(TEXTS_1, TEXTS_2)
assert len(vllm_scores) == 2
for i in range(2):
assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
assert len(vllm_scores) == 2
for i in range(2):
assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
def test_colbert_relevance_ordering(
vllm_runner,
colbert_model_name,
colbert_max_model_len,
colbert_extra_kwargs,
colbert_model,
):
"""Test that ColBERT scores relevant documents higher than irrelevant."""
query = "What is machine learning?"
@@ -343,40 +314,19 @@ def test_colbert_relevance_ordering(
"Deep learning uses neural networks.",
]
with vllm_runner(
colbert_model_name,
runner="pooling",
dtype=DTYPE,
max_model_len=colbert_max_model_len,
enforce_eager=True,
**colbert_extra_kwargs,
) as vllm_model:
scores = vllm_model.score(query, documents)
scores = colbert_model.score(query, documents)
assert len(scores) == 3
assert scores[0] > scores[1], "ML doc should score higher than Python doc"
assert scores[2] > scores[1], "DL doc should score higher than Python doc"
assert len(scores) == 3
assert scores[0] > scores[1], "ML doc should score higher than Python doc"
assert scores[2] > scores[1], "DL doc should score higher than Python doc"
def test_colbert_embed_not_supported(
vllm_runner,
colbert_model_name,
colbert_max_model_len,
colbert_extra_kwargs,
colbert_model,
):
"""Test that ColBERT model does not support 'embed' task."""
with (
vllm_runner(
colbert_model_name,
runner="pooling",
dtype=DTYPE,
max_model_len=colbert_max_model_len,
enforce_eager=True,
**colbert_extra_kwargs,
) as vllm_model,
pytest.raises(ValueError, match="Embedding API is not supported"),
):
vllm_model.embed([TEXTS_1[0]])
"""Test that ColBERT model does not support the embed task."""
with pytest.raises(ValueError, match="Embedding API is not supported"):
colbert_model.embed([TEXTS_1[0]])
@pytest.mark.parametrize(
@@ -22,60 +22,47 @@ calculus, each contributing unique perspectives that would shape this new
field."""
@pytest.fixture(scope="module")
def vllm_model(vllm_runner):
with vllm_runner(
MODEL_NAME, runner="pooling", max_model_len=max_model_len
) as model:
yield model
def test_smaller_truncation_size(
vllm_runner, model_name=MODEL_NAME, input_str=input_str
vllm_model,
):
truncate_prompt_tokens = 10
with vllm_runner(
model_name, runner="pooling", max_model_len=max_model_len
) as vllm_model:
vllm_output = vllm_model.llm.embed(
input_str,
tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
)
vllm_output = vllm_model.llm.embed(
input_str,
tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
)
prompt_tokens = vllm_output[0].prompt_token_ids
assert len(prompt_tokens) == truncate_prompt_tokens
def test_max_truncation_size(vllm_runner, model_name=MODEL_NAME, input_str=input_str):
def test_max_truncation_size(vllm_model):
truncate_prompt_tokens = -1
with vllm_runner(
model_name, runner="pooling", max_model_len=max_model_len
) as vllm_model:
vllm_output = vllm_model.llm.embed(
input_str,
tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
)
vllm_output = vllm_model.llm.embed(
input_str,
tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
)
prompt_tokens = vllm_output[0].prompt_token_ids
assert len(prompt_tokens) == max_model_len
def test_bigger_truncation_size(
vllm_runner, model_name=MODEL_NAME, input_str=input_str
):
def test_bigger_truncation_size(vllm_model):
truncate_prompt_tokens = max_model_len + 1
with (
pytest.raises(VLLMValidationError),
vllm_runner(
model_name, runner="pooling", max_model_len=max_model_len
) as vllm_model,
):
llm_output = vllm_model.llm.embed(
with pytest.raises(VLLMValidationError):
vllm_model.llm.embed(
input_str,
tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
)
assert (
llm_output
== f"""truncate_prompt_tokens value
({truncate_prompt_tokens}) is greater than
max_model_len ({max_model_len}). Please, select
a smaller truncation size."""
)
@@ -20,6 +20,7 @@ VLLM_PROMPT = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
HF_PROMPT = ""
# Whisper expects 16kHz audio
WHISPER_SAMPLE_RATE = 16000
BEAM_WIDTHS = (1, 2)
@pytest.fixture(autouse=True)
@@ -127,13 +128,11 @@ def check_model_available(model: str) -> None:
@pytest.mark.parametrize("dtype", ["half"])
@pytest.mark.parametrize("max_tokens", [64])
@pytest.mark.parametrize("beam_width", [1, 2])
def test_beam_search_encoder_decoder(
hf_runner,
vllm_runner,
dtype: str,
max_tokens: int,
beam_width: int,
resampled_assets,
) -> None:
"""Test beam search with encoder-decoder models (Whisper)."""
@@ -146,12 +145,15 @@ def test_beam_search_encoder_decoder(
]
with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSpeechSeq2Seq) as hf_model:
hf_outputs = hf_model.generate_beam_search(
hf_prompts,
beam_width=beam_width,
max_tokens=max_tokens,
audios=resampled_assets,
)
hf_outputs_by_beam_width = [
hf_model.generate_beam_search(
hf_prompts,
beam_width=beam_width,
max_tokens=max_tokens,
audios=resampled_assets,
)
for beam_width in BEAM_WIDTHS
]
# Test both explicit encoder/decoder prompts
vllm_prompts = [
@@ -179,38 +181,46 @@ def test_beam_search_encoder_decoder(
limit_mm_per_prompt={"audio": 2},
enforce_eager=True,
) as vllm_model:
vllm_outputs = vllm_model.generate_beam_search(
vllm_prompts,
beam_width=beam_width,
max_tokens=max_tokens,
)
for i in range(len(vllm_prompts)):
hf_output_ids, hf_output_texts = hf_outputs[i]
vllm_output_ids, vllm_output_texts = vllm_outputs[i]
for j, (hf_text, vllm_text) in enumerate(
zip(hf_output_texts, vllm_output_texts)
):
print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:")
print(hf_text)
print(f">>>{j}-th vllm output:")
print(vllm_text)
# Check that we got the same number of beams
assert len(hf_output_ids) == len(vllm_output_ids)
# For encoder-decoder models, we primarily want to verify that:
# 1. Beam search completes without errors
# 2. We get the expected number of beams
# 3. Outputs are reasonable (non-empty, diverse beams)
for j in range(len(vllm_output_ids)):
# Check that outputs are not empty
assert len(vllm_output_ids[j]) > 0, f"Prompt {i}, beam {j}: empty output"
# Check that decoded text is not empty
assert len(vllm_output_texts[j].strip()) > 0, (
f"Prompt {i}, beam {j}: empty text output"
vllm_outputs_by_beam_width = [
vllm_model.generate_beam_search(
vllm_prompts,
beam_width=beam_width,
max_tokens=max_tokens,
)
for beam_width in BEAM_WIDTHS
]
for beam_width, hf_outputs, vllm_outputs in zip(
BEAM_WIDTHS, hf_outputs_by_beam_width, vllm_outputs_by_beam_width
):
for i in range(len(vllm_prompts)):
hf_output_ids, hf_output_texts = hf_outputs[i]
vllm_output_ids, vllm_output_texts = vllm_outputs[i]
for j, (hf_text, vllm_text) in enumerate(
zip(hf_output_texts, vllm_output_texts)
):
print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:")
print(hf_text)
print(f">>>{j}-th vllm output:")
print(vllm_text)
# Check that we got the same number of beams
assert len(hf_output_ids) == len(vllm_output_ids) == beam_width
# For encoder-decoder models, we primarily want to verify that:
# 1. Beam search completes without errors
# 2. We get the expected number of beams
# 3. Outputs are reasonable (non-empty, diverse beams)
for j in range(len(vllm_output_ids)):
# Check that outputs are not empty
assert len(vllm_output_ids[j]) > 0, (
f"Prompt {i}, beam {j}: empty output"
)
# Check that decoded text is not empty
assert len(vllm_output_texts[j].strip()) > 0, (
f"Prompt {i}, beam {j}: empty text output"
)
def test_parse_language_detection_output():
+66 -97
View File
@@ -26,8 +26,7 @@ MODELS = ["openai/clip-vit-base-patch32"]
def _run_test(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
input_texts: list[str],
input_images: PromptImageInput,
input_cases: list[tuple[list[str], PromptImageInput]],
model: str,
*,
dtype: str,
@@ -39,105 +38,75 @@ def _run_test(
with vllm_runner(
model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77
) as vllm_model:
vllm_outputs = vllm_model.embed(input_texts, images=input_images)
vllm_outputs_per_case = [
vllm_model.embed(input_texts, images=input_images)
for input_texts, input_images in input_cases
]
with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model:
all_inputs = hf_model.get_inputs(input_texts, images=input_images)
all_outputs = []
for inputs in all_inputs:
inputs = hf_model.wrap_device(inputs)
if "pixel_values" in inputs:
pooled_output = hf_model.model.get_image_features(
pixel_values=inputs.pixel_values,
)
else:
pooled_output = hf_model.model.get_text_features(
input_ids=inputs.input_ids,
attention_mask=inputs.attention_mask,
)
if not isinstance(pooled_output, torch.Tensor):
pooled_output = pooled_output.pooler_output
pooled_output = pooled_output.squeeze(0)
all_outputs.append(pooled_output.tolist())
hf_outputs = all_outputs
check_embeddings_close(
embeddings_0_lst=hf_outputs,
embeddings_1_lst=vllm_outputs,
name_0="hf",
name_1="vllm",
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_models_text(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images, # type: ignore
model,
dtype=dtype,
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_models_image(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
input_texts_images = [
(text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images,
model,
dtype=dtype,
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_models_text_image_no_crash(
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
texts = [HF_TEXT_PROMPTS[0]]
images = [image_assets[0].pil_image]
with vllm_runner(
model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77
) as vllm_model:
texts = [HF_TEXT_PROMPTS[0]]
images = [input_cases[1][1][0]]
with pytest.raises(ValueError, match="not both"):
vllm_model.embed(texts, images=images)
# Should still be able to run subsequent requests
vllm_model.embed(texts)
vllm_model.embed([""], images=images)
with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model:
hf_outputs_per_case = []
for input_texts, input_images in input_cases:
all_inputs = hf_model.get_inputs(input_texts, images=input_images)
hf_outputs = []
for inputs in all_inputs:
inputs = hf_model.wrap_device(inputs)
if "pixel_values" in inputs:
pooled_output = hf_model.model.get_image_features(
pixel_values=inputs.pixel_values,
)
else:
pooled_output = hf_model.model.get_text_features(
input_ids=inputs.input_ids,
attention_mask=inputs.attention_mask,
)
if not isinstance(pooled_output, torch.Tensor):
pooled_output = pooled_output.pooler_output
pooled_output = pooled_output.squeeze(0)
hf_outputs.append(pooled_output.tolist())
hf_outputs_per_case.append(hf_outputs)
for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
check_embeddings_close(
embeddings_0_lst=hf_outputs,
embeddings_1_lst=vllm_outputs,
name_0="hf",
name_1="vllm",
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_models(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
text_images = [None] * len(HF_TEXT_PROMPTS)
images = [asset.pil_image for asset in image_assets]
input_cases = [
(HF_TEXT_PROMPTS, text_images),
(HF_IMAGE_PROMPTS, images),
]
_run_test(
hf_runner,
vllm_runner,
input_cases, # type: ignore[arg-type]
model,
dtype=dtype,
)
@@ -17,29 +17,34 @@ COLBERT_DIM = 128
DTYPE = "half"
# -----------------------------------------------------------------------
# Text-only tests
# -----------------------------------------------------------------------
def test_colmodernvbert_text_token_embed(vllm_runner):
"""Text query produces per-token embeddings with shape (seq_len, 128)."""
@pytest.fixture(scope="module")
def colmodernvbert_model(vllm_runner):
with vllm_runner(
MODEL_NAME,
runner="pooling",
dtype=DTYPE,
enforce_eager=True,
) as vllm_model:
outputs = vllm_model.token_embed(["What is machine learning?"])
assert len(outputs) == 1
emb = torch.tensor(outputs[0])
assert emb.dim() == 2
assert emb.shape[1] == COLBERT_DIM
assert emb.shape[0] > 1
yield vllm_model
def test_colmodernvbert_text_relevance_ordering(vllm_runner):
# -----------------------------------------------------------------------
# Text-only tests
# -----------------------------------------------------------------------
def test_colmodernvbert_text_token_embed(colmodernvbert_model):
"""Text query produces per-token embeddings with shape (seq_len, 128)."""
outputs = colmodernvbert_model.token_embed(["What is machine learning?"])
assert len(outputs) == 1
emb = torch.tensor(outputs[0])
assert emb.dim() == 2
assert emb.shape[1] == COLBERT_DIM
assert emb.shape[0] > 1
def test_colmodernvbert_text_relevance_ordering(colmodernvbert_model):
"""Relevant documents score higher than irrelevant ones."""
query = "What is machine learning?"
documents = [
@@ -47,40 +52,28 @@ def test_colmodernvbert_text_relevance_ordering(vllm_runner):
"The weather in Paris is mild in spring.",
]
with vllm_runner(
MODEL_NAME,
runner="pooling",
dtype=DTYPE,
enforce_eager=True,
) as vllm_model:
scores = vllm_model.score(query, documents)
scores = colmodernvbert_model.score(query, documents)
assert len(scores) == 2
assert scores[0] > scores[1], "ML doc should score higher than weather doc"
assert len(scores) == 2
assert scores[0] > scores[1], "ML doc should score higher than weather doc"
def test_colmodernvbert_text_late_interaction(vllm_runner):
def test_colmodernvbert_text_late_interaction(colmodernvbert_model):
"""MaxSim scoring via vLLM matches manual computation."""
query = "What is the capital of France?"
doc = "The capital of France is Paris."
with vllm_runner(
MODEL_NAME,
runner="pooling",
dtype=DTYPE,
enforce_eager=True,
) as vllm_model:
q_out = vllm_model.token_embed([query])
d_out = vllm_model.token_embed([doc])
q_out = colmodernvbert_model.token_embed([query])
d_out = colmodernvbert_model.token_embed([doc])
q_emb = torch.tensor(q_out[0])
d_emb = torch.tensor(d_out[0])
manual_score = compute_maxsim_score(q_emb, d_emb).item()
q_emb = torch.tensor(q_out[0])
d_emb = torch.tensor(d_out[0])
manual_score = compute_maxsim_score(q_emb, d_emb).item()
vllm_scores = vllm_model.score(query, doc)
vllm_scores = colmodernvbert_model.score(query, doc)
assert len(vllm_scores) == 1
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
assert len(vllm_scores) == 1
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
# -----------------------------------------------------------------------
@@ -88,28 +81,22 @@ def test_colmodernvbert_text_late_interaction(vllm_runner):
# -----------------------------------------------------------------------
def test_colmodernvbert_image_token_embed(vllm_runner, image_assets):
def test_colmodernvbert_image_token_embed(colmodernvbert_model, image_assets):
"""Image input produces per-token embeddings including vision tokens."""
with vllm_runner(
MODEL_NAME,
runner="pooling",
dtype=DTYPE,
enforce_eager=True,
) as vllm_model:
image = image_assets[0].pil_image
inputs = vllm_model.get_inputs(
[""],
images=[image],
)
req_outputs = vllm_model.llm.encode(
inputs,
pooling_task="token_embed",
)
outputs = [req_output.outputs.data for req_output in req_outputs]
image = image_assets[0].pil_image
inputs = colmodernvbert_model.get_inputs(
[""],
images=[image],
)
req_outputs = colmodernvbert_model.llm.encode(
inputs,
pooling_task="token_embed",
)
outputs = [req_output.outputs.data for req_output in req_outputs]
assert len(outputs) == 1
emb = torch.tensor(outputs[0])
assert emb.dim() == 2
assert emb.shape[1] == COLBERT_DIM
# Should have at least the image tokens (64 after pixel shuffle)
assert emb.shape[0] >= 64
assert len(outputs) == 1
emb = torch.tensor(outputs[0])
assert emb.dim() == 2
assert emb.shape[1] == COLBERT_DIM
# Should have at least the image tokens (64 after pixel shuffle)
assert emb.shape[0] >= 64
+110 -195
View File
@@ -75,75 +75,51 @@ def _make_image_mm_param(
def _run_token_embed_test(
vllm_runner: type[VllmRunner],
vllm_model: VllmRunner,
model: str,
*,
dtype: str,
) -> None:
"""Verify per-token embedding shape and L2 normalization."""
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=4096,
enforce_eager=True,
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
) as vllm_model:
outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
assert len(outputs) == 1
emb = torch.tensor(outputs[0])
# Token embeddings should be 2D: [num_tokens, embed_dim]
assert emb.dim() == 2
assert emb.shape[1] == EMBED_DIMS[model]
assert emb.shape[0] > 1
assert len(outputs) == 1
emb = torch.tensor(outputs[0])
# Token embeddings should be 2D: [num_tokens, embed_dim]
assert emb.dim() == 2
assert emb.shape[1] == EMBED_DIMS[model]
assert emb.shape[0] > 1
# Verify L2 normalization
norms = torch.norm(emb, p=2, dim=-1)
torch.testing.assert_close(
norms,
torch.ones_like(norms),
rtol=1e-2,
atol=1e-2,
)
# Verify L2 normalization
norms = torch.norm(emb, p=2, dim=-1)
torch.testing.assert_close(
norms,
torch.ones_like(norms),
rtol=1e-2,
atol=1e-2,
)
def _run_late_interaction_test(
vllm_runner: type[VllmRunner],
model: str,
*,
dtype: str,
vllm_model: VllmRunner,
) -> None:
"""Verify MaxSim scoring matches manual computation."""
from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=4096,
enforce_eager=True,
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
) as vllm_model:
q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
q_emb = torch.tensor(q_outputs[0])
d_emb = torch.tensor(d_outputs[0])
q_emb = torch.tensor(q_outputs[0])
d_emb = torch.tensor(d_outputs[0])
manual_score = compute_maxsim_score(q_emb, d_emb).item()
manual_score = compute_maxsim_score(q_emb, d_emb).item()
vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
assert len(vllm_scores) == 1
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
assert len(vllm_scores) == 1
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
def _run_relevance_test(
vllm_runner: type[VllmRunner],
model: str,
*,
dtype: str,
vllm_model: VllmRunner,
) -> None:
"""Verify that relevant documents score higher than irrelevant ones."""
query = "What is machine learning?"
@@ -153,59 +129,18 @@ def _run_relevance_test(
"Deep learning uses neural networks for complex tasks.",
]
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=4096,
enforce_eager=True,
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
) as vllm_model:
scores = vllm_model.score(query, documents)
scores = vllm_model.score(query, documents)
assert len(scores) == 3
assert scores[0] > scores[1], "ML doc should score higher than weather doc"
assert scores[2] > scores[1], "DL doc should score higher than weather doc"
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colpali_token_embed(
vllm_runner,
model: str,
dtype: str,
) -> None:
_run_token_embed_test(vllm_runner, model, dtype=dtype)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colpali_late_interaction_scoring(
vllm_runner,
model: str,
dtype: str,
) -> None:
_run_late_interaction_test(vllm_runner, model, dtype=dtype)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colpali_relevance_ordering(
vllm_runner,
model: str,
dtype: str,
) -> None:
_run_relevance_test(vllm_runner, model, dtype=dtype)
assert len(scores) == 3
assert scores[0] > scores[1], "ML doc should score higher than weather doc"
assert scores[2] > scores[1], "DL doc should score higher than weather doc"
# ── Multimodal scoring tests ────────────────────────────────
def _run_multimodal_text_query_image_docs_test(
vllm_runner: type[VllmRunner],
model: str,
*,
dtype: str,
vllm_model: VllmRunner,
) -> None:
"""Score a text query against image documents via the multimodal path."""
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
@@ -218,6 +153,84 @@ def _run_multimodal_text_query_image_docs_test(
]
attention_backend = "FLASH_ATTN" if current_platform.is_cuda() else None
scores = vllm_model.llm.score(query, image_docs)
assert len(scores) == 2
for s in scores:
assert isinstance(s.outputs.score, float)
def _run_multimodal_mixed_docs_test(
vllm_model: VllmRunner,
) -> None:
"""Score a text query against a mix of text and image documents."""
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
query = "What is the capital of France?"
documents: list = [
"The capital of France is Paris.",
_make_image_mm_param(red_image),
]
scores = vllm_model.llm.score(query, documents)
assert len(scores) == 2
for s in scores:
assert isinstance(s.outputs.score, float)
# Text document about France should score higher than a random image
assert scores[0].outputs.score > scores[1].outputs.score
def _run_multimodal_image_query_text_docs_test(
vllm_model: VllmRunner,
) -> None:
"""Score an image query against text documents."""
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
image_query = _make_image_mm_param(red_image, text="red color")
documents = [
"A bright red sports car.",
"The weather forecast shows rain tomorrow.",
]
scores = vllm_model.llm.score(image_query, documents)
assert len(scores) == 2
for s in scores:
assert isinstance(s.outputs.score, float)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colpali_default_runner(
vllm_runner,
model: str,
dtype: str,
) -> None:
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=4096,
enforce_eager=True,
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
) as vllm_model:
_run_token_embed_test(vllm_model, model)
_run_late_interaction_test(vllm_model)
_run_relevance_test(vllm_model)
_run_multimodal_mixed_docs_test(vllm_model)
_run_multimodal_image_query_text_docs_test(vllm_model)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colpali_v2_multimodal_text_query_image_docs(
vllm_runner,
monkeypatch: pytest.MonkeyPatch,
model: str,
dtype: str,
) -> None:
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
with vllm_runner(
model,
runner="pooling",
@@ -229,102 +242,4 @@ def _run_multimodal_text_query_image_docs_test(
kernel_config={"enable_flashinfer_autotune": False},
) as vllm_model:
assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
scores = vllm_model.llm.score(query, image_docs)
assert len(scores) == 2
for s in scores:
assert isinstance(s.outputs.score, float)
def _run_multimodal_mixed_docs_test(
vllm_runner: type[VllmRunner],
model: str,
*,
dtype: str,
) -> None:
"""Score a text query against a mix of text and image documents."""
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
query = "What is the capital of France?"
documents: list = [
"The capital of France is Paris.",
_make_image_mm_param(red_image),
]
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=4096,
enforce_eager=True,
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
) as vllm_model:
scores = vllm_model.llm.score(query, documents)
assert len(scores) == 2
for s in scores:
assert isinstance(s.outputs.score, float)
# Text document about France should score higher than a random image
assert scores[0].outputs.score > scores[1].outputs.score
def _run_multimodal_image_query_text_docs_test(
vllm_runner: type[VllmRunner],
model: str,
*,
dtype: str,
) -> None:
"""Score an image query against text documents."""
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
image_query = _make_image_mm_param(red_image, text="red color")
documents = [
"A bright red sports car.",
"The weather forecast shows rain tomorrow.",
]
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=4096,
enforce_eager=True,
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
) as vllm_model:
scores = vllm_model.llm.score(image_query, documents)
assert len(scores) == 2
for s in scores:
assert isinstance(s.outputs.score, float)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colpali_multimodal_text_query_image_docs(
vllm_runner,
monkeypatch: pytest.MonkeyPatch,
model: str,
dtype: str,
) -> None:
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
_run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colpali_multimodal_mixed_docs(
vllm_runner,
model: str,
dtype: str,
) -> None:
_run_multimodal_mixed_docs_test(vllm_runner, model, dtype=dtype)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colpali_multimodal_image_query_text_docs(
vllm_runner,
model: str,
dtype: str,
) -> None:
_run_multimodal_image_query_text_docs_test(vllm_runner, model, dtype=dtype)
_run_multimodal_text_query_image_docs_test(vllm_model)
@@ -35,74 +35,65 @@ TEXT_DOCUMENTS = [
DTYPE = "half"
def _run_token_embed_test(
vllm_runner: type[VllmRunner],
model: str,
*,
dtype: str,
) -> None:
"""Verify per-token embedding shape and L2 normalization."""
@pytest.fixture(scope="module", params=MODELS)
def colqwen3_5_model(request, vllm_runner):
model = request.param
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
dtype=DTYPE,
max_model_len=4096,
enforce_eager=True,
) as vllm_model:
outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
yield model, vllm_model
assert len(outputs) == 1
emb = torch.tensor(outputs[0])
# Token embeddings should be 2D: [num_tokens, embed_dim]
assert emb.dim() == 2
assert emb.shape[1] == EMBED_DIMS[model]
assert emb.shape[0] > 1
# Verify L2 normalization
norms = torch.norm(emb, p=2, dim=-1)
torch.testing.assert_close(
norms,
torch.ones_like(norms),
rtol=1e-2,
atol=1e-2,
)
def _run_token_embed_test(
vllm_model: VllmRunner,
model: str,
) -> None:
"""Verify per-token embedding shape and L2 normalization."""
outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
assert len(outputs) == 1
emb = torch.tensor(outputs[0])
# Token embeddings should be 2D: [num_tokens, embed_dim]
assert emb.dim() == 2
assert emb.shape[1] == EMBED_DIMS[model]
assert emb.shape[0] > 1
# Verify L2 normalization
norms = torch.norm(emb, p=2, dim=-1)
torch.testing.assert_close(
norms,
torch.ones_like(norms),
rtol=1e-2,
atol=1e-2,
)
def _run_late_interaction_test(
vllm_runner: type[VllmRunner],
model: str,
*,
dtype: str,
vllm_model: VllmRunner,
) -> None:
"""Verify MaxSim scoring matches manual computation."""
from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=4096,
enforce_eager=True,
) as vllm_model:
q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
q_emb = torch.tensor(q_outputs[0])
d_emb = torch.tensor(d_outputs[0])
q_emb = torch.tensor(q_outputs[0])
d_emb = torch.tensor(d_outputs[0])
manual_score = compute_maxsim_score(q_emb, d_emb).item()
manual_score = compute_maxsim_score(q_emb, d_emb).item()
vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
assert len(vllm_scores) == 1
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
assert len(vllm_scores) == 1
assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
def _run_relevance_test(
vllm_runner: type[VllmRunner],
model: str,
*,
dtype: str,
vllm_model: VllmRunner,
) -> None:
"""Verify that relevant documents score higher than irrelevant ones."""
query = "What is machine learning?"
@@ -112,48 +103,26 @@ def _run_relevance_test(
"Deep learning uses neural networks for complex tasks.",
]
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=4096,
enforce_eager=True,
) as vllm_model:
scores = vllm_model.score(query, documents)
scores = vllm_model.score(query, documents)
assert len(scores) == 3
assert scores[0] > scores[1], "ML doc should score higher than weather doc"
assert scores[2] > scores[1], "DL doc should score higher than weather doc"
assert len(scores) == 3
assert scores[0] > scores[1], "ML doc should score higher than weather doc"
assert scores[2] > scores[1], "DL doc should score higher than weather doc"
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colqwen3_5_token_embed(
vllm_runner,
model: str,
dtype: str,
) -> None:
_run_token_embed_test(vllm_runner, model, dtype=dtype)
def test_colqwen3_5_token_embed(colqwen3_5_model) -> None:
model, vllm_model = colqwen3_5_model
_run_token_embed_test(vllm_model, model)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colqwen3_5_late_interaction_scoring(
vllm_runner,
model: str,
dtype: str,
) -> None:
_run_late_interaction_test(vllm_runner, model, dtype=dtype)
def test_colqwen3_5_late_interaction_scoring(colqwen3_5_model) -> None:
_, vllm_model = colqwen3_5_model
_run_late_interaction_test(vllm_model)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", [DTYPE])
def test_colqwen3_5_relevance_ordering(
vllm_runner,
model: str,
dtype: str,
) -> None:
_run_relevance_test(vllm_runner, model, dtype=dtype)
def test_colqwen3_5_relevance_ordering(colqwen3_5_model) -> None:
_, vllm_model = colqwen3_5_model
_run_relevance_test(vllm_model)
@pytest.mark.parametrize(
@@ -9,12 +9,14 @@ Tests for the LlamaNemotronVL model family:
Both variants share a SigLIP vision encoder with a bidirectional LLaMA backbone.
"""
from collections.abc import Sequence
from io import BytesIO
from pathlib import Path
import pybase64 as base64
import pytest
import torch
from PIL import Image
from transformers import AutoModel, AutoModelForSequenceClassification, AutoProcessor
from vllm.entrypoints.chat_utils import (
@@ -54,17 +56,15 @@ MODELS = ["nvidia/llama-nemotron-embed-vl-1b-v2"]
def _run_test(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
input_texts: list[str],
input_images: PromptImageInput,
input_cases: list[tuple[list[str], PromptImageInput]],
model: str,
*,
dtype: str,
) -> None:
"""Run embedding comparison test between HF and vLLM.
"""Compare HF and vLLM embeddings for all input cases.
NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
"""
# Run vLLM inference first
with vllm_runner(
model,
runner="pooling",
@@ -74,91 +74,70 @@ def _run_test(
trust_remote_code=True,
**ROCM_ENGINE_KWARGS,
) as vllm_model:
vllm_outputs = vllm_model.embed(input_texts, images=input_images)
vllm_outputs_per_case = [
vllm_model.embed(input_texts, images=input_images)
for input_texts, input_images in input_cases
]
# Run HF inference using the model's encode_queries/encode_documents API
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
hf_outputs = []
for text, image in zip(input_texts, input_images):
with torch.inference_mode():
if text.startswith(QUERY_PREFIX):
# Strip prefix and use encode_queries for query texts
query_text = text[len(QUERY_PREFIX) :]
embedding = hf_model.model.encode_queries([query_text])
elif text.startswith(PASSAGE_PREFIX):
# Strip prefix and use encode_documents for passages/images
passage_text = text[len(PASSAGE_PREFIX) :]
if image is not None:
# Image document - pass image to encode_documents
embedding = hf_model.model.encode_documents(
images=[image],
texts=[passage_text],
)
hf_outputs_per_case = []
for input_texts, input_images in input_cases:
hf_outputs = []
for text, image in zip(input_texts, input_images):
with torch.inference_mode():
if text.startswith(QUERY_PREFIX):
query_text = text[len(QUERY_PREFIX) :]
embedding = hf_model.model.encode_queries([query_text])
elif text.startswith(PASSAGE_PREFIX):
passage_text = text[len(PASSAGE_PREFIX) :]
if image is not None:
embedding = hf_model.model.encode_documents(
images=[image],
texts=[passage_text],
)
else:
embedding = hf_model.model.encode_documents(
texts=[passage_text]
)
else:
# Text-only document
embedding = hf_model.model.encode_documents(
texts=[passage_text]
raise ValueError(
f"Text must start with {QUERY_PREFIX!r} "
f"or {PASSAGE_PREFIX!r}"
)
else:
raise ValueError(
f"Text must start with '{QUERY_PREFIX}' or '{PASSAGE_PREFIX}'"
)
hf_outputs.append(embedding[0].tolist())
hf_outputs.append(embedding[0].tolist())
hf_outputs_per_case.append(hf_outputs)
check_embeddings_close(
embeddings_0_lst=hf_outputs,
embeddings_1_lst=vllm_outputs,
name_0="hf",
name_1="vllm",
)
for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
check_embeddings_close(
embeddings_0_lst=hf_outputs,
embeddings_1_lst=vllm_outputs,
name_0="hf",
name_1="vllm",
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_models_text(
def test_models(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
"""Test text-only embedding."""
input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images, # type: ignore
model,
dtype=dtype,
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_models_image(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
"""Test image embedding."""
input_texts_images = [
(text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
"""Test text and image embedding."""
text_images = [None] * len(HF_TEXT_PROMPTS)
images = [asset.pil_image for asset in image_assets]
input_cases = [
(HF_TEXT_PROMPTS, text_images),
(HF_IMAGE_PROMPTS, images),
]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images,
input_cases, # type: ignore[arg-type]
model,
dtype=dtype,
)
@@ -188,8 +167,11 @@ RERANKER_TEXT_DOCS = [
RERANKER_IMAGE_QUERY = "photo of a red stop sign on a street"
RerankerDocument = tuple[str | None, Image.Image | None]
RerankerCase = tuple[str, Sequence[RerankerDocument]]
def _pil_to_data_uri(image) -> str:
def _pil_to_data_uri(image: Image.Image) -> str:
buf = BytesIO()
image.save(buf, format="PNG")
b64 = base64.b64encode(buf.getvalue()).decode()
@@ -200,10 +182,9 @@ def _run_hf_reranker(
hf_runner: type[HfRunner],
model: str,
dtype: str,
query: str,
docs: list,
) -> list[float]:
"""Run HF reranker inference; docs is a list of (doc_text, doc_image|None)."""
input_cases: Sequence[RerankerCase],
) -> list[list[float]]:
"""Run all HF reranker cases in one model lifecycle."""
with hf_runner(
model,
dtype=dtype,
@@ -217,35 +198,37 @@ def _run_hf_reranker(
use_thumbnail=True,
rerank_max_length=2048,
)
examples = [
{
"question": query,
"doc_text": doc_text if doc_text is not None else "",
"doc_image": doc_image if doc_image is not None else "",
scores_per_case = []
for query, docs in input_cases:
examples = [
{
"question": query,
"doc_text": doc_text if doc_text is not None else "",
"doc_image": doc_image if doc_image is not None else "",
}
for doc_text, doc_image in docs
]
batch_dict = processor.process_queries_documents_crossencoder(examples)
batch_dict = {
k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v
for k, v in batch_dict.items()
}
for doc_text, doc_image in docs
]
batch_dict = processor.process_queries_documents_crossencoder(examples)
batch_dict = {
k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v
for k, v in batch_dict.items()
}
with torch.inference_mode():
logits = hf_model.model(**batch_dict, return_dict=True).logits
# vLLM applies sigmoid activation to the raw logits before returning
# scores; apply the same here so both sides are comparable.
scores = torch.sigmoid(logits.squeeze(-1).float())
return scores.detach().cpu().tolist()
with torch.inference_mode():
logits = hf_model.model(**batch_dict, return_dict=True).logits
# vLLM applies sigmoid activation to raw logits before returning scores.
scores = torch.sigmoid(logits.squeeze(-1).float())
scores_per_case.append(scores.detach().cpu().tolist())
return scores_per_case
def _run_vllm_reranker(
vllm_runner: type[VllmRunner],
model: str,
dtype: str,
query: str,
docs: list,
) -> list[float]:
"""Run vLLM reranker inference; docs is a list of (doc_text, doc_image|None)."""
input_cases: Sequence[RerankerCase],
) -> list[list[float]]:
"""Run all vLLM reranker cases in one model lifecycle."""
with vllm_runner(
model,
runner="pooling",
@@ -255,57 +238,59 @@ def _run_vllm_reranker(
trust_remote_code=True,
**ROCM_ENGINE_KWARGS,
) as vllm_model:
has_images = any(img is not None for _, img in docs)
scores_per_case = []
for query, docs in input_cases:
has_images = any(img is not None for _, img in docs)
if not has_images:
# Text-only path: use the simple string score API.
queries = [query] * len(docs)
doc_texts = [doc_text for doc_text, _ in docs]
outputs = vllm_model.score(
queries,
doc_texts,
chat_template=_RERANKER_SCORE_TEMPLATE,
)
else:
# Multimodal path: build ScoreMultiModalParam for each pair.
query_params = [
ScoreMultiModalParam(
content=[
ChatCompletionContentPartTextParam(
type="text",
text=query,
)
]
if not has_images:
queries = [query] * len(docs)
doc_texts = [doc_text for doc_text, _ in docs]
outputs = vllm_model.score(
queries,
doc_texts,
chat_template=_RERANKER_SCORE_TEMPLATE,
)
] * len(docs)
doc_params = []
for doc_text, doc_image in docs:
content: list = []
if doc_image is not None:
content.append(
ChatCompletionContentPartImageParam(
type="image_url",
image_url={"url": _pil_to_data_uri(doc_image)},
)
else:
query_params = [
ScoreMultiModalParam(
content=[
ChatCompletionContentPartTextParam(
type="text",
text=query,
)
]
)
if doc_text:
content.append(
ChatCompletionContentPartTextParam(
type="text",
text=doc_text,
] * len(docs)
doc_params = []
for doc_text, doc_image in docs:
content: list = []
if doc_image is not None:
content.append(
ChatCompletionContentPartImageParam(
type="image_url",
image_url={"url": _pil_to_data_uri(doc_image)},
)
)
)
doc_params.append(ScoreMultiModalParam(content=content))
if doc_text:
content.append(
ChatCompletionContentPartTextParam(
type="text",
text=doc_text,
)
)
doc_params.append(ScoreMultiModalParam(content=content))
raw_outputs = vllm_model.llm.score(
query_params,
doc_params,
chat_template=_RERANKER_SCORE_TEMPLATE,
)
outputs = [o.outputs.score for o in raw_outputs]
raw_outputs = vllm_model.llm.score(
query_params,
doc_params,
chat_template=_RERANKER_SCORE_TEMPLATE,
)
outputs = [output.outputs.score for output in raw_outputs]
return outputs
scores_per_case.append(outputs)
return scores_per_case
def _run_reranker_test(
@@ -313,50 +298,44 @@ def _run_reranker_test(
vllm_runner: type[VllmRunner],
model: str,
dtype: str,
query: str,
docs: list,
input_cases: Sequence[RerankerCase],
) -> None:
"""Compare HF and vLLM reranker scores.
"""Compare HF and vLLM reranker scores for all input cases.
NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
"""
vllm_scores = _run_vllm_reranker(vllm_runner, model, dtype, query, docs)
hf_scores = _run_hf_reranker(hf_runner, model, dtype, query, docs)
vllm_scores_per_case = _run_vllm_reranker(vllm_runner, model, dtype, input_cases)
hf_scores_per_case = _run_hf_reranker(hf_runner, model, dtype, input_cases)
assert len(hf_scores) == len(vllm_scores), (
f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}"
)
# NOTE: ROCm shows slightly higher numerical variance dues to different attention
# backend between vLLM and HF; use a marginally looser tolerance
# ROCm has slightly higher variance because vLLM and HF use different
# attention backends.
rel_tol = 0.022 if current_platform.is_rocm() else 0.02
for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)):
assert hf_score == pytest.approx(vllm_score, rel=rel_tol), (
f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}"
for hf_scores, vllm_scores in zip(hf_scores_per_case, vllm_scores_per_case):
assert len(hf_scores) == len(vllm_scores), (
f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}"
)
for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)):
assert hf_score == pytest.approx(vllm_score, rel=rel_tol), (
f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}"
)
@pytest.mark.parametrize("model", RERANKER_MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_reranker_text(
hf_runner,
vllm_runner,
model: str,
dtype: str,
) -> None:
"""Test reranking with text-only query and text documents."""
docs = [(text, None) for text in RERANKER_TEXT_DOCS]
_run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_TEXT_QUERY, docs)
@pytest.mark.parametrize("model", RERANKER_MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_reranker_image_doc(
def test_reranker(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
"""Test reranking with text query against image documents."""
docs = [(None, asset.pil_image) for asset in image_assets]
_run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_IMAGE_QUERY, docs)
"""Test reranking with text and image documents."""
text_docs: list[RerankerDocument] = [(text, None) for text in RERANKER_TEXT_DOCS]
image_docs: list[RerankerDocument] = [
(None, asset.pil_image) for asset in image_assets
]
input_cases: list[RerankerCase] = [
(RERANKER_TEXT_QUERY, text_docs),
(RERANKER_IMAGE_QUERY, image_docs),
]
_run_reranker_test(hf_runner, vllm_runner, model, dtype, input_cases)
+80 -113
View File
@@ -33,16 +33,11 @@ MODELS = [
def _run_test(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
input_texts: list[str],
input_images: PromptImageInput,
input_cases: list[tuple[list[str], PromptImageInput, dict[str, Any]]],
model: str,
*,
dtype: str,
tokenization_kwargs: dict[str, Any] | None = None,
) -> None:
if tokenization_kwargs is None:
tokenization_kwargs = {}
with vllm_runner(
model,
runner="pooling",
@@ -51,116 +46,88 @@ def _run_test(
max_model_len=64,
gpu_memory_utilization=0.7,
) as vllm_model:
vllm_outputs = vllm_model.embed(
input_texts, images=input_images, tokenization_kwargs=tokenization_kwargs
)
vllm_outputs_per_case = [
vllm_model.embed(
input_texts,
images=input_images,
tokenization_kwargs=tokenization_kwargs,
)
for input_texts, input_images, tokenization_kwargs in input_cases
]
with hf_runner(model, dtype=dtype, auto_cls=SiglipModel) as hf_model:
all_inputs = hf_model.get_inputs(
input_texts, images=input_images, tokenization_kwargs=tokenization_kwargs
)
all_outputs = []
for inputs in all_inputs:
inputs = hf_model.wrap_device(inputs)
if "pixel_values" in inputs:
pooled_output = hf_model.model.get_image_features(
pixel_values=inputs.pixel_values,
)
else:
pooled_output = hf_model.model.get_text_features(
input_ids=inputs.input_ids,
)
if not isinstance(pooled_output, torch.Tensor):
pooled_output = pooled_output.pooler_output
pooled_output = pooled_output.squeeze(0)
all_outputs.append(pooled_output.tolist())
hf_outputs = all_outputs
check_embeddings_close(
embeddings_0_lst=hf_outputs,
embeddings_1_lst=vllm_outputs,
name_0="hf",
name_1="vllm",
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_models_text(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images, # type: ignore
model,
dtype=dtype,
tokenization_kwargs={
"padding": "max_length",
"max_length": 64,
}, # siglip2 was trained with this padding setting.
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_models_image(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
input_texts_images = [
(text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images,
model,
dtype=dtype,
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_models_text_image_no_crash(
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
texts = [HF_TEXT_PROMPTS[0]]
images = [image_assets[0].pil_image]
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
enforce_eager=True,
max_model_len=64,
gpu_memory_utilization=0.7,
) as vllm_model:
texts = [HF_TEXT_PROMPTS[0]]
images = [input_cases[1][1][0]]
with pytest.raises(ValueError, match="not both"):
vllm_model.embed(texts, images=images)
vllm_model.embed(texts)
vllm_model.embed([""], images=images)
with hf_runner(model, dtype=dtype, auto_cls=SiglipModel) as hf_model:
hf_outputs_per_case = []
for input_texts, input_images, tokenization_kwargs in input_cases:
all_inputs = hf_model.get_inputs(
input_texts,
images=input_images,
tokenization_kwargs=tokenization_kwargs,
)
hf_outputs = []
for inputs in all_inputs:
inputs = hf_model.wrap_device(inputs)
if "pixel_values" in inputs:
pooled_output = hf_model.model.get_image_features(
pixel_values=inputs.pixel_values,
)
else:
pooled_output = hf_model.model.get_text_features(
input_ids=inputs.input_ids,
)
if not isinstance(pooled_output, torch.Tensor):
pooled_output = pooled_output.pooler_output
pooled_output = pooled_output.squeeze(0)
hf_outputs.append(pooled_output.tolist())
hf_outputs_per_case.append(hf_outputs)
for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
check_embeddings_close(
embeddings_0_lst=hf_outputs,
embeddings_1_lst=vllm_outputs,
name_0="hf",
name_1="vllm",
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_models(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
text_images = [None] * len(HF_TEXT_PROMPTS)
images = [asset.pil_image for asset in image_assets]
input_cases = [
(
HF_TEXT_PROMPTS,
text_images,
{
"padding": "max_length",
"max_length": 64,
},
),
(HF_IMAGE_PROMPTS, images, {}),
]
_run_test(
hf_runner,
vllm_runner,
input_cases, # type: ignore[arg-type]
model,
dtype=dtype,
)
+15 -17
View File
@@ -17,6 +17,15 @@ HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
}
)
IMAGE_SIZE_FACTOR_GROUPS = (
# Single-scale
(1.0,),
# Single-scale, batched
(1.0, 1.0, 1.0),
# Multi-scale
(0.25, 0.5, 1.0),
)
def run_awq_test(
vllm_runner: type[VllmRunner],
@@ -24,7 +33,7 @@ def run_awq_test(
source_model: str,
quant_model: str,
*,
size_factors: list[float],
size_factor_groups: tuple[tuple[float, ...], ...],
dtype: str,
max_tokens: int,
num_logprobs: int,
@@ -33,11 +42,12 @@ def run_awq_test(
):
images = [asset.pil_image for asset in image_assets]
inputs_per_image = [
inputs_per_image_and_size_group = [
(
[prompt for _ in size_factors],
[rescale_image_size(image, factor) for factor in size_factors],
)
for size_factors in size_factor_groups
for image, prompt in zip(images, HF_IMAGE_PROMPTS)
]
@@ -60,7 +70,7 @@ def run_awq_test(
vllm_model.generate_greedy_logprobs(
prompts, max_tokens, num_logprobs=num_logprobs, images=images
)
for prompts, images in inputs_per_image
for prompts, images in inputs_per_image_and_size_group
]
with vllm_runner(
@@ -77,7 +87,7 @@ def run_awq_test(
vllm_model.generate_greedy_logprobs(
prompts, max_tokens, num_logprobs=num_logprobs, images=images
)
for prompts, images in inputs_per_image
for prompts, images in inputs_per_image_and_size_group
]
for source_outputs, quant_outputs in zip(
@@ -128,17 +138,6 @@ def test_awq_load(
("source_model", "quant_model"),
[("OpenGVLab/InternVL2-2B", "OpenGVLab/InternVL2-2B-AWQ")],
)
@pytest.mark.parametrize(
"size_factors",
[
# Single-scale
[1.0],
# Single-scale, batched
[1.0, 1.0, 1.0],
# Multi-scale
[0.25, 0.5, 1.0],
],
)
@pytest.mark.parametrize("dtype", ["half"])
@pytest.mark.parametrize("max_tokens", [128])
@pytest.mark.parametrize("num_logprobs", [5])
@@ -148,7 +147,6 @@ def test_awq_models(
image_assets,
source_model,
quant_model,
size_factors,
dtype,
max_tokens,
num_logprobs,
@@ -158,7 +156,7 @@ def test_awq_models(
image_assets,
source_model,
quant_model,
size_factors=size_factors,
size_factor_groups=IMAGE_SIZE_FACTOR_GROUPS,
dtype=dtype,
max_tokens=max_tokens,
num_logprobs=num_logprobs,