| 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 | |
| 8 | @pytest.mark.parametrize( |
| 9 | "model", |
| 10 | ["Rami/multi-label-class-classification-on-github-issues"], |
| 11 | ) |
| 12 | @pytest.mark.parametrize("dtype", ["half"]) |
| 13 | def test_classify_models( |
| 14 | hf_runner, |
| 15 | vllm_runner, |
| 16 | example_prompts, |
| 17 | model: str, |
| 18 | dtype: str, |
| 19 | ) -> None: |
| 20 | with vllm_runner(model, max_model_len=512, dtype=dtype) as vllm_model: |
| 21 | vllm_outputs = vllm_model.classify(example_prompts) |
| 22 | |
| 23 | with hf_runner( |
| 24 | model, dtype=dtype, auto_cls=AutoModelForSequenceClassification |
| 25 | ) as hf_model: |
| 26 | hf_outputs = hf_model.classify(example_prompts) |
| 27 | |
| 28 | for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): |
| 29 | hf_output = torch.tensor(hf_output) |
| 30 | vllm_output = torch.tensor(vllm_output) |
| 31 | |
| 32 | assert torch.allclose( |
| 33 | hf_output, vllm_output, 1e-3 if dtype == "float" else 1e-2 |
| 34 | ) |