@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a388252ac4

Fork 0

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

Add explicit pooling classes for the Transformers backend (#25322)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Harry Mellor committed Sep 30, 2025 at 23:07 UTC a388252ac49ec642bec5a8a63b226d3058e64a54

7 files changed +296 -138

tests/models/registry.py

+2 -1

@@ -657,7 +657,8 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

657 }

658

659 _TRANSFORMERS_BACKEND_MODELS = {

660 - "TransformersModel": _HfExamplesInfo("Qwen/Qwen3-Embedding-0.6B"),

660 + "TransformersEmbeddingModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5", min_transformers_version="4.57.0.dev0"), # noqa: E501

661 + "TransformersForSequenceClassification": _HfExamplesInfo("papluca/xlm-roberta-base-language-detection", min_transformers_version="4.57.0.dev0"), # noqa: E501

662 "TransformersForCausalLM": _HfExamplesInfo("hmellor/Ilama-3.2-1B", trust_remote_code=True), # noqa: E501

663 "TransformersForMultimodalLM": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),

664 }

tests/models/test_transformers.py

+41 -58

@@ -9,9 +9,16 @@ from vllm.platforms import current_platform

10 from ..conftest import HfRunner, VllmRunner

11 from ..utils import multi_gpu_test, prep_prompts

12 +from .registry import HF_EXAMPLE_MODELS

13 from .utils import check_embeddings_close, check_logprobs_close

16 +def get_model(arch: str) -> str:

17 + model_info = HF_EXAMPLE_MODELS.get_hf_info(arch)

18 + model_info.check_transformers_version(on_fail="skip")

19 + return model_info.default

20 +

21 +

22 def check_implementation(

23 runner_ref: type[Union[HfRunner, VllmRunner]],

24 runner_test: type[VllmRunner],

@@ -170,71 +177,47 @@ def test_embed_loading(vllm_runner, model):

177

178

179 @pytest.mark.parametrize(

173 - "model",

174 - [

175 - # Encoder model

176 - "BAAI/bge-base-en-v1.5",

177 - ])

178 -def test_embed_correctness(hf_runner, vllm_runner, example_prompts, model):

179 - import transformers

180 - from packaging.version import Version

181 - installed = Version(transformers.__version__)

182 - required = Version("4.57.0.dev0")

183 - if installed < required:

184 - pytest.skip("Encoder models with the Transformers backend require "

185 - f"transformers>={required}, but got {installed}")

186 -

187 - with vllm_runner(model, max_model_len=512,

188 - model_impl="transformers") as vllm_model:

180 + "arch",

181 + ["TransformersEmbeddingModel", "TransformersForSequenceClassification"])

182 +def test_pooling(hf_runner, vllm_runner, example_prompts, arch):

183 + model = get_model(arch)

184 +

185 + vllm_kwargs = dict(

186 + max_model_len=None,

187 + model_impl="transformers",

188 + compilation_config=dict(cudagraph_capture_sizes=[8]),

189 + )

190 +

191 + hf_kwargs = dict()

192 + if arch == "TransformersEmbeddingModel":

193 + hf_kwargs["is_sentence_transformer"] = True

194 + elif arch == "TransformersForSequenceClassification":

195 + from transformers import AutoModelForSequenceClassification

196 + hf_kwargs["auto_cls"] = AutoModelForSequenceClassification

197 +

198 + # The example_prompts has ending "\n", for example:

199 + # "Write a short story about a robot that dreams for the first time.\n"

200 + # sentence_transformers will strip the input texts, see:

201 + # https://github.com/UKPLab/sentence-transformers/blob/v3.1.1/sentence_transformers/models/Transformer.py#L159

202 + # This makes the input_ids different between hf_model and vllm_model.

203 + # So we need to strip the input texts to avoid test failing.

204 + example_prompts = [str(s).strip() for s in example_prompts]

205 +

206 + with (vllm_runner(model, **vllm_kwargs) as

207 + vllm_model, hf_runner(model, **hf_kwargs) as hf_model):

208 model_config = vllm_model.llm.llm_engine.model_config

209 assert model_config.using_transformers_backend()

210

192 - vllm_outputs = vllm_model.embed(example_prompts)

193 -

194 - with hf_runner(model, is_sentence_transformer=True) as hf_model:

195 - hf_outputs = hf_model.encode(example_prompts)

211 + if arch == "TransformersEmbeddingModel":

212 + vllm_outputs = vllm_model.embed(example_prompts)

213 + hf_outputs = hf_model.encode(example_prompts)

214 + elif arch == "TransformersForSequenceClassification":

215 + vllm_outputs = vllm_model.classify(example_prompts)

216 + hf_outputs = hf_model.classify(example_prompts)

217

218 check_embeddings_close(

219 embeddings_0_lst=hf_outputs,

220 embeddings_1_lst=vllm_outputs,

221 name_0="hf",

222 name_1="vllm",

202 - tol=1e-2,

223 )

204 -

205 -

206 -@pytest.mark.parametrize(

207 - "model",

208 - ["jason9693/Qwen2.5-1.5B-apeach"],

209 -)

210 -@pytest.mark.parametrize("dtype", ["float"])

211 -def test_classify(

212 - hf_runner,

213 - vllm_runner,

214 - example_prompts,

215 - model: str,

216 - dtype: str,

217 -) -> None:

218 - import torch

219 - from transformers import AutoModelForSequenceClassification

220 -

221 - with vllm_runner(model,

222 - max_model_len=512,

223 - dtype=dtype,

224 - model_impl="transformers") as vllm_model:

225 - model_config = vllm_model.llm.llm_engine.model_config

226 - assert model_config.using_transformers_backend()

227 -

228 - vllm_outputs = vllm_model.classify(example_prompts)

229 -

230 - with hf_runner(model,

231 - dtype=dtype,

232 - auto_cls=AutoModelForSequenceClassification) as hf_model:

233 - hf_outputs = hf_model.classify(example_prompts)

234 -

235 - for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):

236 - hf_output = torch.tensor(hf_output)

237 - vllm_output = torch.tensor(vllm_output)

238 -

239 - assert torch.allclose(hf_output, vllm_output,

240 - 1e-3 if dtype == "float" else 1e-2)

vllm/config/model.py

+27 -6

@@ -19,6 +19,7 @@ import vllm.envs as envs

19 from vllm.config.multimodal import (MMCacheType, MMEncoderTPMode,

20 MultiModalConfig)

21 from vllm.config.pooler import PoolerConfig

22 +from vllm.config.scheduler import RunnerType

23 from vllm.config.utils import assert_hashable, config

24 from vllm.logger import init_logger

25 from vllm.platforms import current_platform

@@ -40,7 +41,6 @@ if TYPE_CHECKING:

41 import vllm.model_executor.models as me_models

42 from vllm.config.load import LoadConfig

43 from vllm.config.parallel import ParallelConfig

43 - from vllm.config.scheduler import RunnerType

44 from vllm.model_executor.layers.quantization import QuantizationMethods

45 from vllm.v1.sample.logits_processor import LogitsProcessor

46 else:

@@ -52,13 +52,12 @@ else:

52 "vllm.model_executor.models")

53 LoadConfig = Any

54 ParallelConfig = Any

55 - RunnerType = Any

55 QuantizationMethods = Any

56 LogitsProcessor = Any

58 logger = init_logger(__name__)

61 -RunnerOption = Literal["auto", "generate", "pooling", "draft"]

60 +RunnerOption = Literal["auto", RunnerType]

61 ConvertType = Literal["none", "embed", "classify", "reward"]

62 ConvertOption = Literal["auto", ConvertType]

63 TaskOption = Literal["auto", "generate", "embedding", "embed", "classify",

@@ -668,8 +667,28 @@ class ModelConfig:

667 def _get_transformers_backend_cls(self) -> str:

668 """Determine which Transformers backend class will be used if

669 `model_impl` is set to `transformers` or `auto`."""

671 - if getattr(self, "runner_type", self.runner) == "pooling":

672 - return "TransformersModel"

670 + # Check if the architecture we're wrapping has defaults

671 + runner = None

672 + convert = None

673 + if defaults := try_match_architecture_defaults(self.architectures[0]):

674 + _, (runner, convert) = defaults

675 + # Overwrite with user-specified values

676 + if self.runner != "auto":

677 + runner = self.runner

678 + if self.convert not in {"auto", "none"}:

679 + convert = self.convert

680 + # Fall back to default values if still not set

681 + if runner is None:

682 + runner = "generate"

683 + if convert in {None, "none"}:

684 + convert = "embed"

685 + # Resolve Transformers backend pooling classes

686 + if runner == "pooling":

687 + if convert == "embed":

688 + return "TransformersEmbeddingModel"

689 + if convert == "classify":

690 + return "TransformersForSequenceClassification"

691 + # Resolve Transformers backend generate classes

692 if self.hf_config != self.hf_text_config:

693 # If 'hf_text_config' is the same as 'hf_config'. If not, it is

694 # probably a composite config, i.e. multimodal

@@ -678,7 +697,9 @@ class ModelConfig:

697

698 def using_transformers_backend(self) -> bool:

699 """Check if the model is using the Transformers backend class."""

681 - return self.architecture == self._get_transformers_backend_cls()

700 + used_cls = self._model_info.architecture

701 + transformers_backend_cls = self._get_transformers_backend_cls()

702 + return used_cls == transformers_backend_cls

703

704 @property

705 def registry(self):

vllm/config/utils.py

+13

@@ -4,6 +4,7 @@

4 import ast

5 import inspect

6 import textwrap

7 +from collections.abc import Iterable

8 from dataclasses import MISSING, Field, field, fields, is_dataclass, replace

9 from typing import TYPE_CHECKING, Any, Protocol, TypeVar

@@ -52,6 +53,18 @@ def get_field(cls: ConfigType, name: str) -> Field:

53 f"{cls.__name__}.{name} must have a default value or default factory.")

56 +def getattr_iter(object: object, names: Iterable[str], default: Any) -> Any:

57 + """

58 + A helper function that retrieves an attribute from an object which may

59 + have multiple possible names. This is useful when fetching attributes from

60 + arbitrary `transformers.PretrainedConfig` instances.

61 + """

62 + for name in names:

63 + if hasattr(object, name):

64 + return getattr(object, name)

65 + return default

66 +

67 +

68 def contains_object_print(text: str) -> bool:

69 """

70 Check if the text looks like a printed Python object, e.g.

vllm/model_executor/models/registry.py

+3 -2

@@ -307,9 +307,10 @@ _TRANSFORMERS_SUPPORTED_MODELS = {

307 }

308

309 _TRANSFORMERS_BACKEND_MODELS = {

310 - "TransformersModel": ("transformers", "TransformersModel"),

310 + "TransformersEmbeddingModel": ("transformers_pooling", "TransformersEmbeddingModel"), # noqa: E501

311 + "TransformersForSequenceClassification": ("transformers_pooling", "TransformersForSequenceClassification"), # noqa: E501

312 "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),

312 - "TransformersForMultimodalLM": ("transformers", "TransformersForMultimodalLM"), # noqa: E501

313 + "TransformersForMultimodalLM": ("transformers", "TransformersForMultimodalLM"), # noqa: E501

314 }

315 # yapf: enable

316

vllm/model_executor/models/transformers.py

+10 -71

@@ -31,6 +31,7 @@ from vllm.attention import Attention, AttentionType

31 from vllm.compilation.decorators import support_torch_compile

32 from vllm.config import (CacheConfig, DeviceConfig, ModelConfig,

33 ParallelConfig, VllmConfig)

34 +from vllm.config.utils import getattr_iter

35 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

36 from vllm.distributed.utils import get_pp_indices

37 from vllm.logger import init_logger

@@ -486,10 +487,13 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

487

488 # Input embeddings

489 if not isinstance(self.model.get_input_embeddings(), PPMissingLayer):

490 + names = ("embedding_size", "hidden_size")

491 + embedding_dim = getattr_iter(self.text_config, names, None)

492 + assert embedding_dim is not None

493 self.model.set_input_embeddings(

494 VocabParallelEmbedding(

495 self.text_config.vocab_size,

492 - self.text_config.hidden_size,

496 + embedding_dim=embedding_dim,

497 org_num_embeddings=self.text_config.vocab_size,

498 quant_config=self.quant_config,

499 ))

@@ -645,7 +649,9 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

649 attn_type=attn_type)

650 return attention_instances

651

648 - def init_parameters(self, module: nn.Module):

652 + def init_parameters(self,

653 + module: nn.Module,

654 + dtype: Optional[torch.dtype] = None):

655 """

656 If a `parameter` is on the `meta` device, then its parent

657 `module` is the original module created by:

@@ -659,11 +665,11 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

665 if param.device == torch.device("meta"):

666 new_param = nn.Parameter(

667 torch.empty_like(param.data,

662 - dtype=self.model_config.dtype,

668 + dtype=dtype or self.model_config.dtype,

669 device=self.device_config.device))

670 setattr(module, name, new_param)

671 for child in module.children():

666 - self.init_parameters(child)

672 + self.init_parameters(child, dtype)

673

674 def forward(

675 self,

@@ -712,73 +718,6 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):

718 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

719

720

715 -@support_torch_compile(enable_if=can_enable_torch_compile)

716 -class TransformersModel(TransformersBase):

717 - hf_to_vllm_mapper = WeightsMapper(

718 - orig_to_new_prefix={

719 - # Handle BERT-like models

720 - "bert": "model",

721 - # Add `model.` prefix for base model checkpoints

722 - "": "model.",

723 - # Remove `model.` prefix if it was already there

724 - "model.model.": "model.",

725 - # Pooling adapters will be adjacent to `model`

726 - "model.pooler": "pooler",

727 - "model.score": "score",

728 - # Classifier adapter's classifier layer is renamed to score

729 - "model.classifier": "score",

730 - },

731 - orig_to_new_suffix={

732 - # Replace legacy suffixes used for norms

733 - ".gamma": ".weight",

734 - ".beta": ".bias",

735 - })

736 -

737 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

738 - super().__init__(vllm_config=vllm_config, prefix=prefix)

739 -

740 - # After creating a pooling model, `pooler` will be duplicated.

741 - # The one inside `model` comes from the Transformers modelling code.

742 - # The one after `model` is an adapter from vLLM.

743 - # We want to use the adapter so we nullify the original pooler.

744 - if getattr(self.model, "pooler", None) is not None:

745 - self.skip_prefixes.append("pooler.")

746 - self.model.pooler = torch.nn.Identity()

747 -

748 - # Some encoder models have the position_ids buffer in the checkpoint.

749 - # vLLM will always pass position_ids as an argument, so we skip loading

750 - # the buffer if it exists

751 - self.skip_substrs.append("position_ids")

752 -

753 - # Some encoder models have the bias of the final classifier layer

754 - # in the checkpoint. vLLM does not use this bias, so we skip loading

755 - # it if it exists

756 - self.skip_substrs.append("score.bias")

757 -

758 - def create_attention_instances(

759 - self, attn_type: AttentionType = AttentionType.DECODER):

760 - # TODO(hmellor): Better way to detect encoder models

761 - # In encoder models, the attention layers will have `is_causal=False`

762 - is_encoder = lambda m: not getattr(m, "is_causal", True)

763 - # vLLM does not support encoder-decoder models, so if any encoder layer

764 - # is found, we assume the whole model is an encoder model

765 - if any(is_encoder(m) for m in self.model.modules()):

766 - attn_type = AttentionType.ENCODER_ONLY

767 -

768 - # Check minimum transformers version for encoder models support

769 - if attn_type == AttentionType.ENCODER_ONLY:

770 - import transformers

771 - from packaging.version import Version

772 - installed = Version(transformers.__version__)

773 - required = Version("4.57.0.dev0")

774 - if installed < required:

775 - raise ValueError(

776 - "Encoder models with the Transformers backend require "

777 - f"transformers>={required}, but got {installed}")

778 -

779 - return super().create_attention_instances(attn_type)

780 -

781 -

721 @support_torch_compile(enable_if=can_enable_torch_compile)

722 class TransformersForCausalLM(TransformersBase):

723

vllm/model_executor/models/transformers_pooling.py new

+200

@@ -0,0 +1,200 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

5 +#

6 +# Licensed under the Apache License, Version 2.0 (the "License");

7 +# you may not use this file except in compliance with the License.

8 +# You may obtain a copy of the License at

9 +#

10 +# http://www.apache.org/licenses/LICENSE-2.0

11 +#

12 +# Unless required by applicable law or agreed to in writing, software

13 +# distributed under the License is distributed on an "AS IS" BASIS,

14 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

15 +# See the License for the specific language governing permissions and

16 +# limitations under the License.

17 +"""Wrapper around `transformers` models for pooling tasks."""

18 +from typing import Optional, Union

19 +

20 +import torch

21 +from transformers import AutoModelForSequenceClassification

22 +

23 +from vllm.attention import AttentionType

24 +from vllm.compilation.decorators import support_torch_compile

25 +from vllm.config import VllmConfig

26 +from vllm.model_executor.layers.pooler import (ClassifierPooler, CLSPool,

27 + DispatchPooler, Pooler)

28 +from vllm.sequence import IntermediateTensors

29 +

30 +from .interfaces_base import VllmModelForPooling

31 +from .transformers import TransformersBase, can_enable_torch_compile

32 +from .utils import WeightsMapper

33 +

34 +

35 +class TransformersPoolingBase(TransformersBase, VllmModelForPooling):

36 + hf_to_vllm_mapper = WeightsMapper(

37 + # These are applied in order, so the order matters!

38 + orig_to_new_prefix={

39 + # Handle BERT-like models

40 + "roberta": "model",

41 + "bert": "model",

42 + # Add `model.` prefix for base model checkpoints

43 + "": "model.",

44 + # Remove `model.` prefix if it was already there

45 + "model.model.": "model.",

46 + # Classifier/scoring heads will be adjacent to `model`

47 + "model.score": "classifier",

48 + "model.classifier": "classifier",

49 + },

50 + orig_to_new_suffix={

51 + # Replace legacy suffixes used for norms

52 + ".gamma": ".weight",

53 + ".beta": ".bias",

54 + })

55 +

56 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

57 + super().__init__(vllm_config=vllm_config, prefix=prefix)

58 +

59 + # Skip unsupported/unwanted output embeddings layers

60 + self.skip_prefixes.extend([

61 + "model.lm_head.", "model.predictions.", "model.qa_outputs.",

62 + "model.embeddings_project.", "model.discriminator_predictions."

63 + ])

64 +

65 + # Some encoder models have the position_ids buffer in the checkpoint.

66 + # vLLM will always pass position_ids as an argument, so we skip loading

67 + # the buffer if it exists

68 + self.skip_substrs.append("position_ids")

69 +

70 + # Some encoder models have the bias of the final classifier layer

71 + # in the checkpoint. vLLM does not use this bias, so we skip loading

72 + # it if it exists

73 + self.skip_substrs.append("score.bias")

74 +

75 + # roberta-like models an extra padding in positions.

76 + # FIXME(Isotr0py): This is quite hacky for roberta edge case,

77 + # we should find a better way to handle this.

78 + self.is_roberta = "roberta" in self.text_config.model_type

79 + self.padding_idx = self.text_config.pad_token_id

80 +

81 + def create_attention_instances(

82 + self, attn_type: AttentionType = AttentionType.DECODER):

83 + # TODO(hmellor): Better way to detect encoder models

84 + # In encoder models, the attention layers will have `is_causal=False`

85 + is_encoder = lambda m: not getattr(m, "is_causal", True)

86 + # vLLM does not support encoder-decoder models, so if any encoder layer

87 + # is found, we assume the whole model is an encoder model

88 + if any(is_encoder(m) for m in self.model.modules()):

89 + attn_type = AttentionType.ENCODER_ONLY

90 +

91 + # Check minimum transformers version for encoder models support

92 + if attn_type == AttentionType.ENCODER_ONLY:

93 + import transformers

94 + from packaging.version import Version

95 + installed = Version(transformers.__version__)

96 + required = Version("4.57.0.dev0")

97 + if installed < required:

98 + raise ValueError(

99 + "Encoder models with the Transformers backend require "

100 + f"transformers>={required}, but got {installed}")

101 +

102 + return super().create_attention_instances(attn_type)

103 +

104 + def forward(

105 + self,

106 + input_ids: Optional[torch.Tensor],

107 + positions: torch.Tensor,

108 + intermediate_tensors: Optional[IntermediateTensors] = None,

109 + inputs_embeds: Optional[torch.Tensor] = None,

110 + ) -> Union[torch.Tensor, IntermediateTensors]:

111 + if self.is_roberta:

112 + # RoBERTa-specific positions padding

113 + positions += self.padding_idx + 1

114 + return super().forward(input_ids=input_ids,

115 + positions=positions,

116 + intermediate_tensors=intermediate_tensors,

117 + inputs_embeds=inputs_embeds)

118 +

119 +

120 +@support_torch_compile(enable_if=can_enable_torch_compile)

121 +class TransformersEmbeddingModel(TransformersPoolingBase):

122 + default_pooling_type = "CLS"

123 +

124 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

125 + super().__init__(vllm_config=vllm_config, prefix=prefix)

126 +

127 + pooler_config = vllm_config.model_config.pooler_config

128 + assert pooler_config is not None

129 +

130 + self.pooler = DispatchPooler({

131 + "encode": Pooler.for_encode(pooler_config),

132 + "embed": Pooler.for_embed(pooler_config),

133 + })

134 +

135 +

136 +@support_torch_compile(enable_if=can_enable_torch_compile)

137 +class TransformersForSequenceClassification(TransformersPoolingBase):

138 + default_pooling_type = "CLS"

139 +

140 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

141 + super().__init__(vllm_config=vllm_config, prefix=prefix)

142 +

143 + pooler_config = vllm_config.model_config.pooler_config

144 + assert pooler_config is not None

145 +

146 + # Certain information about the the model and classifier can only be

147 + # inferred from the `ForSequenceClassification` class. Therefore, we

148 + # instantiate it on the "meta" device to avoid allocating GPU memory.

149 + with torch.device("meta"):

150 + seq_cls_model = AutoModelForSequenceClassification.from_config(

151 + self.config,

152 + torch_dtype=self.model_config.dtype,

153 + trust_remote_code=self.model_config.trust_remote_code,

154 + )

155 +

156 + # When used for sequence classification, some models have their

157 + # pooling layers removed. Make sure this is reflected in vLLM.

158 + for module in seq_cls_model.modules():

159 + if hasattr(module, "pooler") and module.pooler is None:

160 + self.model.pooler = None

161 + break

162 + if self.model.pooler is not None:

163 + raise ValueError(

164 + "Sequence classification models with pooling layers are not "

165 + "supported yet in the Transformers backend.")

166 +

167 + # Unlike `lm_head`, `classifier` is not always `nn.Linear`.

168 + self.classifier = seq_cls_model.classifier

169 + self.init_parameters(self.classifier,

170 + dtype=self.model_config.head_dtype)

171 +

172 + class ClassifierWithReshape(self.classifier.__class__):

173 + """CLSPool has already been applied in `pooling`.

174 + Add dim to match expected input shape of `classifier.forward`."""

175 +

176 + def forward(self, *args, **kwargs):

177 + if len(args) > 0:

178 + args = (args[0].unsqueeze(1), *args[1:])

179 + return super().forward(*args, **kwargs)

180 +

181 + self.classifier.__class__ = ClassifierWithReshape

182 +

183 + self.pooler = DispatchPooler({

184 + "encode":

185 + Pooler.for_encode(pooler_config),

186 + "classify":

187 + ClassifierPooler(

188 + pooling=CLSPool(),

189 + classifier=self.classifier,

190 + act_fn=ClassifierPooler.act_fn_for_seq_cls(

191 + vllm_config.model_config),

192 + ),

193 + "score":

194 + ClassifierPooler(

195 + pooling=CLSPool(),

196 + classifier=self.classifier,

197 + act_fn=ClassifierPooler.act_fn_for_cross_encoder(

198 + vllm_config.model_config),

199 + ),

200 + })