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)
|