e32014ac1d
### What this PR does / why we need it? Support pooling models (like `bge-reranker-v2-m3`) in vllm-ascend, this pr covered the three model types of embed (cls_token, mean_token, lasttoken). After this [commit](17373dcd93), vllm has provided support for adapting pooling models on the v1 engine. This PR includes corresponding adaptations on the vllm-ascend side. Fixes #1960 - vLLM version: v0.12.0 - vLLM main:ad32e3e19c--------- Signed-off-by: lianyibo <lianyibo1@kunlunit.com> Signed-off-by: MengqingCao <cmq0113@163.com> Co-authored-by: MengqingCao <cmq0113@163.com>
35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
import torch
|
|
from modelscope import snapshot_download # type: ignore[import-untyped]
|
|
from transformers import AutoModelForSequenceClassification
|
|
|
|
from tests.e2e.conftest import HfRunner, VllmRunner
|
|
|
|
|
|
def test_classify_correctness() -> None:
|
|
|
|
model_name = snapshot_download("Howeee/Qwen2.5-1.5B-apeach")
|
|
|
|
prompts = [
|
|
"Hello, my name is",
|
|
"The president of the United States is",
|
|
"The capital of France is",
|
|
"The future of AI is what",
|
|
]
|
|
with VllmRunner(
|
|
model_name,
|
|
runner="pooling",
|
|
max_model_len=None,
|
|
cudagraph_capture_sizes=[4],
|
|
) as vllm_runner:
|
|
vllm_outputs = vllm_runner.classify(prompts)
|
|
|
|
with HfRunner(model_name,
|
|
dtype="float32",
|
|
auto_cls=AutoModelForSequenceClassification) as hf_runner:
|
|
hf_outputs = hf_runner.classify(prompts)
|
|
|
|
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
|
|
hf_output = torch.tensor(hf_output)
|
|
vllm_output = torch.tensor(vllm_output)
|
|
assert torch.allclose(hf_output, vllm_output, 1e-2)
|