| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | from transformers import AutoModelForSequenceClassification |
| 6 | |
| 7 | from tests.models.language.pooling.embed_utils import run_embedding_correctness_test |
| 8 | |
| 9 | |
| 10 | @pytest.mark.parametrize( |
| 11 | "model", |
| 12 | ["jason9693/Qwen2.5-1.5B-apeach"], |
| 13 | ) |
| 14 | @pytest.mark.parametrize("dtype", ["half"]) |
| 15 | def test_classify_models( |
| 16 | hf_runner, |
| 17 | vllm_runner, |
| 18 | example_prompts, |
| 19 | model: str, |
| 20 | dtype: str, |
| 21 | ) -> None: |
| 22 | # example_prompts is too short for testing prefix_caching |
| 23 | example_prompts = [s * 10 for s in example_prompts] |
| 24 | |
| 25 | with vllm_runner( |
| 26 | model, max_model_len=512, dtype=dtype, enable_prefix_caching=True |
| 27 | ) as vllm_model: |
| 28 | vllm_config = vllm_model.llm.llm_engine.vllm_config |
| 29 | cache_config = vllm_config.cache_config |
| 30 | assert cache_config.enable_prefix_caching |
| 31 | |
| 32 | # First Run |
| 33 | vllm_model.classify(example_prompts) |
| 34 | |
| 35 | # assert prefix_caching works |
| 36 | pooling_outputs = vllm_model.llm.encode( |
| 37 | example_prompts, pooling_task="classify" |
| 38 | ) |
| 39 | for output in pooling_outputs: |
| 40 | assert output.num_cached_tokens > 0 |
| 41 | vllm_outputs = [req_output.outputs.data for req_output in pooling_outputs] |
| 42 | |
| 43 | with hf_runner( |
| 44 | model, dtype=dtype, auto_cls=AutoModelForSequenceClassification |
| 45 | ) as hf_model: |
| 46 | hf_outputs = hf_model.classify(example_prompts) |
| 47 | |
| 48 | for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): |
| 49 | hf_output = torch.tensor(hf_output) |
| 50 | vllm_output = torch.tensor(vllm_output) |
| 51 | |
| 52 | assert torch.allclose( |
| 53 | hf_output, vllm_output, 1e-3 if dtype == "float" else 1e-2 |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | @pytest.mark.parametrize( |
| 58 | "model", |
| 59 | ["Qwen/Qwen3-Embedding-0.6B"], |
| 60 | ) |
| 61 | @pytest.mark.parametrize("dtype", ["half"]) |
| 62 | def test_embed_models( |
| 63 | hf_runner, |
| 64 | vllm_runner, |
| 65 | example_prompts, |
| 66 | model: str, |
| 67 | dtype: str, |
| 68 | ): |
| 69 | # example_prompts is too short for testing prefix_caching |
| 70 | example_prompts = [str(s).strip() * 10 for s in example_prompts] |
| 71 | |
| 72 | with vllm_runner( |
| 73 | model, |
| 74 | runner="pooling", |
| 75 | max_model_len=None, |
| 76 | enable_prefix_caching=True, |
| 77 | ) as vllm_model: |
| 78 | vllm_config = vllm_model.llm.llm_engine.vllm_config |
| 79 | cache_config = vllm_config.cache_config |
| 80 | assert cache_config.enable_prefix_caching |
| 81 | |
| 82 | # First Run |
| 83 | vllm_model.embed(example_prompts) |
| 84 | |
| 85 | # assert prefix_caching works |
| 86 | pooling_outputs = vllm_model.llm.encode(example_prompts, pooling_task="embed") |
| 87 | for output in pooling_outputs: |
| 88 | assert output.num_cached_tokens > 0 |
| 89 | vllm_outputs = [req_output.outputs.data for req_output in pooling_outputs] |
| 90 | |
| 91 | with hf_runner( |
| 92 | model, |
| 93 | is_sentence_transformer=True, |
| 94 | ) as hf_model: |
| 95 | run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs) |
| 96 | |
| 97 | |
| 98 | @pytest.mark.parametrize( |
| 99 | "model", |
| 100 | [ |
| 101 | "intfloat/e5-small", |
| 102 | "Alibaba-NLP/gte-Qwen2-1.5B-instruct", # is_causal == False |
| 103 | "papluca/xlm-roberta-base-language-detection", |
| 104 | ], |
| 105 | ) |
| 106 | @pytest.mark.parametrize("dtype", ["half"]) |
| 107 | def test_non_causal_models( |
| 108 | hf_runner, vllm_runner, example_prompts, model: str, dtype: str |
| 109 | ) -> None: |
| 110 | with vllm_runner(model, max_model_len=512, dtype=dtype) as vllm_model: |
| 111 | vllm_config = vllm_model.llm.llm_engine.vllm_config |
| 112 | cache_config = vllm_config.cache_config |
| 113 | assert not cache_config.enable_prefix_caching |