| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import ABC, abstractmethod |
| 4 | from collections.abc import Callable |
| 5 | from typing import TypeVar |
| 6 | |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | import torch.nn.functional as F |
| 10 | from transformers import PreTrainedConfig |
| 11 | |
| 12 | from vllm.config import ModelConfig |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | def get_act_fn( |
| 20 | config: PreTrainedConfig, |
| 21 | static_num_labels: bool = True, |
| 22 | ) -> "PoolerActivation": |
| 23 | # get classification act_fn |
| 24 | # Implement alignment with transformers ForSequenceClassificationLoss |
| 25 | # https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92 |
| 26 | num_labels: int | None = None |
| 27 | if static_num_labels: |
| 28 | num_labels = getattr(config, "num_labels", 0) |
| 29 | |
| 30 | problem_type = getattr(config, "problem_type", "") |
| 31 | if problem_type == "regression": |
| 32 | return PoolerIdentity() |
| 33 | if problem_type == "single_label_classification": |
| 34 | return PoolerClassify(num_labels=num_labels) |
| 35 | if problem_type == "multi_label_classification": |
| 36 | return PoolerMultiLabelClassify() |
| 37 | |
| 38 | # get cross_encoder act_fn |
| 39 | function_name: str | None = None |
| 40 | if ( |
| 41 | hasattr(config, "sentence_transformers") |
| 42 | and "activation_fn" in config.sentence_transformers |
| 43 | ): |
| 44 | function_name = config.sentence_transformers["activation_fn"] |
| 45 | elif ( |
| 46 | hasattr(config, "sbert_ce_default_activation_function") |
| 47 | and config.sbert_ce_default_activation_function is not None |
| 48 | ): |
| 49 | function_name = config.sbert_ce_default_activation_function |
| 50 | |
| 51 | if function_name is not None: |
| 52 | if not function_name.startswith("torch.nn.modules."): |
| 53 | raise ValueError( |
| 54 | "Loading of activation functions is restricted to " |
| 55 | "torch.nn.modules for security reasons" |
| 56 | ) |
| 57 | fn = resolve_obj_by_qualname(function_name)() |
| 58 | return PoolerActivation.wraps(fn) |
| 59 | |
| 60 | return PoolerClassify(num_labels=num_labels) |
| 61 | |
| 62 | |
| 63 | def resolve_classifier_act_fn( |
| 64 | model_config: ModelConfig, |
| 65 | static_num_labels: bool = True, |
| 66 | act_fn: "PoolerActivation | None" = None, |
| 67 | ) -> "PoolerActivation": |
| 68 | if act_fn is None: |
| 69 | return get_act_fn(model_config.hf_config, static_num_labels) |
| 70 | |
| 71 | if not callable(act_fn): |
| 72 | raise TypeError(f"Expected a callable activation function, got {type(act_fn)}") |
| 73 | return act_fn |
| 74 | |
| 75 | |
| 76 | _T = TypeVar("_T", torch.Tensor, list[torch.Tensor]) |
| 77 | |
| 78 | |
| 79 | class PoolerActivation(nn.Module, ABC): |
| 80 | @staticmethod |
| 81 | def wraps(module: nn.Module) -> "PoolerActivation": |
| 82 | if isinstance(module, nn.Identity): |
| 83 | return PoolerIdentity() |
| 84 | if isinstance(module, (nn.Sigmoid, nn.Softmax)): |
| 85 | return PoolerClassify() |
| 86 | |
| 87 | return LambdaPoolerActivation(module) |
| 88 | |
| 89 | @abstractmethod |
| 90 | def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor: |
| 91 | raise NotImplementedError |
| 92 | |
| 93 | def forward(self, pooled_data: _T) -> _T: |
| 94 | # shape: |
| 95 | # classify -> (batch_size, num_classes) |
| 96 | # embed -> (batch_size, embedding_size) or list(embedding_size) |
| 97 | if isinstance(pooled_data, list): |
| 98 | return [self.forward_chunk(data) for data in pooled_data] |
| 99 | |
| 100 | return self.forward_chunk(pooled_data) |
| 101 | |
| 102 | |
| 103 | class PoolerIdentity(PoolerActivation): |
| 104 | def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor: |
| 105 | return pooled_data |
| 106 | |
| 107 | |
| 108 | class PoolerNormalize(PoolerActivation): |
| 109 | def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor: |
| 110 | return F.normalize(pooled_data, p=2, dim=-1) |
| 111 | |
| 112 | |
| 113 | class PoolerMultiLabelClassify(PoolerActivation): |
| 114 | def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor: |
| 115 | return F.sigmoid(pooled_data) |
| 116 | |
| 117 | |
| 118 | class PoolerClassify(PoolerActivation): |
| 119 | def __init__(self, *, num_labels: int | None = None) -> None: |
| 120 | super().__init__() |
| 121 | |
| 122 | if num_labels == 0: |
| 123 | logger.warning( |
| 124 | "num_labels should be > 0 for classification " |
| 125 | "models, falling back to sigmoid. " |
| 126 | "Please check if the configuration is correct." |
| 127 | ) |
| 128 | |
| 129 | self.num_labels = num_labels |
| 130 | |
| 131 | def extra_repr(self) -> str: |
| 132 | return f"num_labels={self.num_labels}" |
| 133 | |
| 134 | def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor: |
| 135 | num_labels = self.num_labels |
| 136 | if num_labels is None: |
| 137 | num_labels = pooled_data.shape[-1] |
| 138 | |
| 139 | if num_labels < 2: |
| 140 | return F.sigmoid(pooled_data) |
| 141 | |
| 142 | return F.softmax(pooled_data, dim=-1) |
| 143 | |
| 144 | |
| 145 | class LambdaPoolerActivation(PoolerActivation): |
| 146 | def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> None: |
| 147 | super().__init__() |
| 148 | |
| 149 | self.fn = fn |
| 150 | |
| 151 | def extra_repr(self) -> str: |
| 152 | name = getattr(self.fn, "__name__", None) |
| 153 | if name is None: |
| 154 | name = self.fn.__class__.__name__ |
| 155 | return f"fn={name}" |
| 156 | |
| 157 | def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor: |
| 158 | return self.fn(pooled_data) |