[Model] Support pooling models (#3122)
### 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>
This commit is contained in:
34
tests/e2e/singlecard/pooling/test_classification.py
Normal file
34
tests/e2e/singlecard/pooling/test_classification.py
Normal file
@@ -0,0 +1,34 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user