| 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 AutoModel |
| 6 | |
| 7 | from tests.models.utils import check_embeddings_close |
| 8 | from vllm.config import PoolerConfig |
| 9 | |
| 10 | |
| 11 | @pytest.mark.parametrize( |
| 12 | "model", |
| 13 | ["BAAI/bge-m3"], |
| 14 | ) |
| 15 | @pytest.mark.parametrize("dtype", ["float"]) |
| 16 | @torch.inference_mode |
| 17 | def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype: str): |
| 18 | with vllm_runner( |
| 19 | model, |
| 20 | runner="pooling", |
| 21 | pooler_config=PoolerConfig(task="token_embed"), |
| 22 | max_model_len=None, |
| 23 | dtype=dtype, |
| 24 | ) as vllm_model: |
| 25 | vllm_outputs = vllm_model.token_embed(example_prompts) |
| 26 | |
| 27 | with hf_runner(model, auto_cls=AutoModel, dtype=dtype) as hf_model: |
| 28 | tokenizer = hf_model.tokenizer |
| 29 | hf_outputs = [] |
| 30 | for prompt in example_prompts: |
| 31 | inputs = tokenizer([prompt], return_tensors="pt") |
| 32 | inputs = hf_model.wrap_device(inputs) |
| 33 | output = hf_model.model(**inputs) |
| 34 | embedding = output.last_hidden_state[0].float() |
| 35 | hf_outputs.append(embedding.cpu()) |
| 36 | |
| 37 | for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): |
| 38 | check_embeddings_close( |
| 39 | embeddings_0_lst=hf_output, |
| 40 | embeddings_1_lst=vllm_output, |
| 41 | name_0="hf", |
| 42 | name_1="vllm", |
| 43 | tol=1e-2, |
| 44 | ) |