[CI/Build] Reduce more duplicate runner startup in tests (#52570)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
@@ -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():
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user