From 908733aca703f90476f91f400a77cbe31fc81d3f Mon Sep 17 00:00:00 2001 From: 22quinn <33176974+22quinn@users.noreply.github.com> Date: Sun, 18 May 2025 07:00:09 -0700 Subject: [PATCH] [Model] Use sigmoid for single-label classification (#18313) Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com> --- vllm/model_executor/layers/pooler.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/pooler.py b/vllm/model_executor/layers/pooler.py index c5bf603a14df3..6abbc90819a82 100644 --- a/vllm/model_executor/layers/pooler.py +++ b/vllm/model_executor/layers/pooler.py @@ -242,9 +242,16 @@ class PoolerHead(nn.Module): if self.softmax: if isinstance(pooled_data, list): - pooled_data = [F.softmax(data, dim=-1) for data in pooled_data] + pooled_data = [ + F.softmax(data, dim=-1) + if data.shape[-1] >= 2 else F.sigmoid(data) + for data in pooled_data + ] else: - pooled_data = F.softmax(pooled_data, dim=-1) + if pooled_data.shape[-1] >= 2: + pooled_data = F.softmax(pooled_data, dim=-1) + else: + pooled_data = F.sigmoid(pooled_data) return pooled_data