| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for vllm.model_executor.layers.pooler.activations.""" |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | |
| 11 | from vllm.model_executor.layers.pooler.activations import ( |
| 12 | LambdaPoolerActivation, |
| 13 | PoolerClassify, |
| 14 | PoolerIdentity, |
| 15 | PoolerMultiLabelClassify, |
| 16 | PoolerNormalize, |
| 17 | get_act_fn, |
| 18 | resolve_classifier_act_fn, |
| 19 | ) |
| 20 | |
| 21 | |
| 22 | # --------------------------------------------------------------------------- |
| 23 | # PoolerIdentity |
| 24 | # --------------------------------------------------------------------------- |
| 25 | class TestPoolerIdentity: |
| 26 | def test_returns_input_unchanged(self): |
| 27 | pooler = PoolerIdentity() |
| 28 | x = torch.randn(4, 128) |
| 29 | out = pooler(x) |
| 30 | assert torch.equal(out, x) |
| 31 | |
| 32 | def test_forward_list(self): |
| 33 | pooler = PoolerIdentity() |
| 34 | tensors = [torch.randn(128), torch.randn(256)] |
| 35 | out = pooler(tensors) |
| 36 | assert len(out) == 2 |
| 37 | for orig, result in zip(tensors, out): |
| 38 | assert torch.equal(orig, result) |
| 39 | |
| 40 | |
| 41 | # --------------------------------------------------------------------------- |
| 42 | # PoolerNormalize |
| 43 | # --------------------------------------------------------------------------- |
| 44 | class TestPoolerNormalize: |
| 45 | def test_output_has_unit_norm(self): |
| 46 | pooler = PoolerNormalize() |
| 47 | x = torch.randn(4, 128) |
| 48 | out = pooler(x) |
| 49 | norms = torch.linalg.norm(out, dim=-1) |
| 50 | assert torch.allclose(norms, torch.ones(4), atol=1e-5) |
| 51 | |
| 52 | def test_single_vector(self): |
| 53 | pooler = PoolerNormalize() |
| 54 | x = torch.randn(1, 64) |
| 55 | out = pooler(x) |
| 56 | norm = torch.linalg.norm(out, dim=-1) |
| 57 | assert torch.allclose(norm, torch.ones(1), atol=1e-5) |
| 58 | |
| 59 | def test_forward_list(self): |
| 60 | pooler = PoolerNormalize() |
| 61 | tensors = [torch.randn(1, 64), torch.randn(1, 128)] |
| 62 | out = pooler(tensors) |
| 63 | for t in out: |
| 64 | norm = torch.linalg.norm(t, dim=-1) |
| 65 | assert torch.allclose(norm, torch.ones(1), atol=1e-5) |
| 66 | |
| 67 | |
| 68 | # --------------------------------------------------------------------------- |
| 69 | # PoolerMultiLabelClassify |
| 70 | # --------------------------------------------------------------------------- |
| 71 | class TestPoolerMultiLabelClassify: |
| 72 | def test_output_in_zero_one(self): |
| 73 | pooler = PoolerMultiLabelClassify() |
| 74 | x = torch.randn(4, 10) |
| 75 | out = pooler(x) |
| 76 | assert (out >= 0).all() and (out <= 1).all() |
| 77 | |
| 78 | def test_large_positive_maps_near_one(self): |
| 79 | pooler = PoolerMultiLabelClassify() |
| 80 | x = torch.full((1, 3), 100.0) |
| 81 | out = pooler(x) |
| 82 | assert torch.allclose(out, torch.ones(1, 3), atol=1e-4) |
| 83 | |
| 84 | def test_large_negative_maps_near_zero(self): |
| 85 | pooler = PoolerMultiLabelClassify() |
| 86 | x = torch.full((1, 3), -100.0) |
| 87 | out = pooler(x) |
| 88 | assert torch.allclose(out, torch.zeros(1, 3), atol=1e-4) |
| 89 | |
| 90 | |
| 91 | # --------------------------------------------------------------------------- |
| 92 | # PoolerClassify |
| 93 | # --------------------------------------------------------------------------- |
| 94 | class TestPoolerClassify: |
| 95 | def test_infers_from_shape_when_num_labels_none(self): |
| 96 | pooler = PoolerClassify(num_labels=None) |
| 97 | assert pooler.num_labels is None |
| 98 | x = torch.randn(2, 5) |
| 99 | out = pooler(x) |
| 100 | sums = out.sum(dim=-1) |
| 101 | assert torch.allclose(sums, torch.ones(2), atol=1e-5) |
| 102 | |
| 103 | def test_sigmoid_when_num_labels_lt_2(self): |
| 104 | pooler = PoolerClassify(num_labels=1) |
| 105 | x = torch.zeros(1, 1) |
| 106 | out = pooler(x) |
| 107 | assert torch.allclose(out, torch.tensor([[0.5]]), atol=1e-5) |
| 108 | |
| 109 | def test_num_labels_zero_uses_sigmoid(self): |
| 110 | pooler = PoolerClassify(num_labels=0) |
| 111 | assert pooler.num_labels == 0 |
| 112 | x = torch.zeros(1, 3) |
| 113 | out = pooler(x) |
| 114 | assert torch.allclose(out, torch.full((1, 3), 0.5), atol=1e-5) |
| 115 | |
| 116 | def test_num_labels_ge_2_uses_softmax(self): |
| 117 | pooler = PoolerClassify(num_labels=4) |
| 118 | assert pooler.num_labels == 4 |
| 119 | x = torch.randn(2, 4) |
| 120 | out = pooler(x) |
| 121 | sums = out.sum(dim=-1) |
| 122 | assert torch.allclose(sums, torch.ones(2), atol=1e-5) |
| 123 | |
| 124 | def test_default_num_labels_is_none(self): |
| 125 | pooler = PoolerClassify() |
| 126 | assert pooler.num_labels is None |
| 127 | |
| 128 | |
| 129 | # --------------------------------------------------------------------------- |
| 130 | # LambdaPoolerActivation |
| 131 | # --------------------------------------------------------------------------- |
| 132 | class TestLambdaPoolerActivation: |
| 133 | def test_applies_custom_fn(self): |
| 134 | pooler = LambdaPoolerActivation(nn.ReLU()) |
| 135 | x = torch.tensor([[-1.0, 2.0, -3.0]]) |
| 136 | out = pooler(x) |
| 137 | expected = torch.tensor([[0.0, 2.0, 0.0]]) |
| 138 | assert torch.equal(out, expected) |
| 139 | |
| 140 | def test_forward_list(self): |
| 141 | pooler = LambdaPoolerActivation(nn.ReLU()) |
| 142 | tensors = [torch.tensor([-1.0, 2.0]), torch.tensor([3.0, -4.0])] |
| 143 | out = pooler(tensors) |
| 144 | assert torch.equal(out[0], torch.tensor([0.0, 2.0])) |
| 145 | assert torch.equal(out[1], torch.tensor([3.0, 0.0])) |
| 146 | |
| 147 | |
| 148 | # --------------------------------------------------------------------------- |
| 149 | # get_act_fn factory |
| 150 | # --------------------------------------------------------------------------- |
| 151 | class TestGetActFn: |
| 152 | @staticmethod |
| 153 | def _make_config(**kwargs): |
| 154 | return SimpleNamespace(**kwargs) |
| 155 | |
| 156 | def test_regression(self): |
| 157 | cfg = self._make_config(problem_type="regression") |
| 158 | result = get_act_fn(cfg) |
| 159 | assert isinstance(result, PoolerIdentity) |
| 160 | |
| 161 | def test_single_label_classification(self): |
| 162 | cfg = self._make_config( |
| 163 | problem_type="single_label_classification", num_labels=3 |
| 164 | ) |
| 165 | result = get_act_fn(cfg) |
| 166 | assert isinstance(result, PoolerClassify) |
| 167 | assert result.num_labels == 3 |
| 168 | |
| 169 | def test_multi_label_classification(self): |
| 170 | cfg = self._make_config(problem_type="multi_label_classification") |
| 171 | result = get_act_fn(cfg) |
| 172 | assert isinstance(result, PoolerMultiLabelClassify) |
| 173 | |
| 174 | def test_sentence_transformers_activation(self): |
| 175 | cfg = self._make_config( |
| 176 | problem_type="", |
| 177 | sentence_transformers={ |
| 178 | "activation_fn": "torch.nn.modules.activation.Sigmoid" |
| 179 | }, |
| 180 | ) |
| 181 | result = get_act_fn(cfg) |
| 182 | assert isinstance(result, PoolerClassify) |
| 183 | |
| 184 | def test_sbert_activation(self): |
| 185 | cfg = self._make_config( |
| 186 | problem_type="", |
| 187 | sbert_ce_default_activation_function=( |
| 188 | "torch.nn.modules.activation.Sigmoid" |
| 189 | ), |
| 190 | ) |
| 191 | result = get_act_fn(cfg) |
| 192 | assert isinstance(result, PoolerClassify) |
| 193 | |
| 194 | def test_default_fallback(self): |
| 195 | cfg = self._make_config(problem_type="") |
| 196 | result = get_act_fn(cfg) |
| 197 | assert isinstance(result, PoolerClassify) |
| 198 | |
| 199 | def test_sentence_transformers_takes_priority(self): |
| 200 | cfg = self._make_config( |
| 201 | problem_type="", |
| 202 | sentence_transformers={"activation_fn": "torch.nn.modules.linear.Identity"}, |
| 203 | sbert_ce_default_activation_function=( |
| 204 | "torch.nn.modules.activation.Sigmoid" |
| 205 | ), |
| 206 | ) |
| 207 | result = get_act_fn(cfg) |
| 208 | assert isinstance(result, PoolerIdentity) |
| 209 | |
| 210 | def test_rejects_non_torch_activation(self): |
| 211 | cfg = self._make_config( |
| 212 | problem_type="", |
| 213 | sentence_transformers={"activation_fn": "os.system"}, |
| 214 | ) |
| 215 | with pytest.raises(ValueError, match="restricted"): |
| 216 | get_act_fn(cfg) |
| 217 | |
| 218 | |
| 219 | # --------------------------------------------------------------------------- |
| 220 | # resolve_classifier_act_fn |
| 221 | # --------------------------------------------------------------------------- |
| 222 | class TestResolveClassifierActFn: |
| 223 | def test_delegates_to_get_act_fn_when_none(self): |
| 224 | model_config = SimpleNamespace( |
| 225 | hf_config=SimpleNamespace(num_labels=3, problem_type="") |
| 226 | ) |
| 227 | result = resolve_classifier_act_fn(model_config, act_fn=None) |
| 228 | assert isinstance(result, PoolerClassify) |
| 229 | assert result.num_labels == 3 |
| 230 | |
| 231 | def test_passes_through_provided_act_fn(self): |
| 232 | custom = PoolerIdentity() |
| 233 | result = resolve_classifier_act_fn(None, act_fn=custom) |
| 234 | assert result is custom |