A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Model] Add multi_label_classification support (#23173)
Signed-off-by: wang.yuqi <noooop@126.com>
wang.yuqi committed
Aug 19, 2025 at 20:54 UTC
f856c33ce934721595d06c25b58ab0a467392d60
3 files changed
+57
-1
tests/conftest.py
+9
-1
@@ -456,7 +456,15 @@ class HfRunner:
456
outputs = []
457
for inputs in all_inputs:
458
output = self.model(**self.wrap_device(inputs))
459
- logits = output.logits.softmax(dim=-1)[0].tolist()
459
+
460
+ problem_type = getattr(self.config, "problem_type", "")
461
+
462
+ if problem_type == "regression":
463
+ logits = output.logits[0].tolist()
464
+ elif problem_type == "multi_label_classification":
465
+ logits = output.logits.sigmoid()[0].tolist()
466
+ else:
467
+ logits = output.logits.softmax(dim=-1)[0].tolist()
468
outputs.append(logits)
469
470
return outputs
tests/models/language/pooling/test_multilabel_classification_support.py
new
+33
@@ -0,0 +1,33 @@
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(model,
24
+ dtype=dtype,
25
+ auto_cls=AutoModelForSequenceClassification) 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(hf_output, vllm_output,
33
+ 1e-3 if dtype == "float" else 1e-2)
vllm/model_executor/layers/pooler.py
+15
@@ -172,6 +172,15 @@ def get_tasks(pooling_metadata: PoolingMetadata) -> list[PoolingTask]:
172
173
174
def get_classification_activation_function(config: PretrainedConfig):
175
+ # Implement alignment with transformers ForSequenceClassificationLoss
176
+ # https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92
177
+ problem_type = getattr(config, "problem_type", "")
178
+ if problem_type == "regression":
179
+ return PoolerIdentity()
180
+ if problem_type == "single_label_classification":
181
+ return PoolerClassify()
182
+ if problem_type == "multi_label_classification":
183
+ return PoolerMultiLabelClassify()
184
return PoolerClassify()
185
186
@@ -409,6 +418,12 @@ class PoolerNormalize(PoolerActivation):
418
return x.to(pooled_data.dtype)
419
420
421
+class PoolerMultiLabelClassify(PoolerActivation):
422
+
423
+ def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
424
+ return F.sigmoid(pooled_data.float()).to(pooled_data.dtype)
425
+
426
+
427
class PoolerClassify(PoolerActivation):
428
429
def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor: