# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for vllm.model_executor.layers.pooler.activations.""" from types import SimpleNamespace import pytest import torch import torch.nn as nn from vllm.model_executor.layers.pooler.activations import ( LambdaPoolerActivation, PoolerClassify, PoolerIdentity, PoolerMultiLabelClassify, PoolerNormalize, get_act_fn, resolve_classifier_act_fn, ) # --------------------------------------------------------------------------- # PoolerIdentity # --------------------------------------------------------------------------- class TestPoolerIdentity: def test_returns_input_unchanged(self): pooler = PoolerIdentity() x = torch.randn(4, 128) out = pooler(x) assert torch.equal(out, x) def test_forward_list(self): pooler = PoolerIdentity() tensors = [torch.randn(128), torch.randn(256)] out = pooler(tensors) assert len(out) == 2 for orig, result in zip(tensors, out): assert torch.equal(orig, result) # --------------------------------------------------------------------------- # PoolerNormalize # --------------------------------------------------------------------------- class TestPoolerNormalize: def test_output_has_unit_norm(self): pooler = PoolerNormalize() x = torch.randn(4, 128) out = pooler(x) norms = torch.linalg.norm(out, dim=-1) assert torch.allclose(norms, torch.ones(4), atol=1e-5) def test_single_vector(self): pooler = PoolerNormalize() x = torch.randn(1, 64) out = pooler(x) norm = torch.linalg.norm(out, dim=-1) assert torch.allclose(norm, torch.ones(1), atol=1e-5) def test_forward_list(self): pooler = PoolerNormalize() tensors = [torch.randn(1, 64), torch.randn(1, 128)] out = pooler(tensors) for t in out: norm = torch.linalg.norm(t, dim=-1) assert torch.allclose(norm, torch.ones(1), atol=1e-5) # --------------------------------------------------------------------------- # PoolerMultiLabelClassify # --------------------------------------------------------------------------- class TestPoolerMultiLabelClassify: def test_output_in_zero_one(self): pooler = PoolerMultiLabelClassify() x = torch.randn(4, 10) out = pooler(x) assert (out >= 0).all() and (out <= 1).all() def test_large_positive_maps_near_one(self): pooler = PoolerMultiLabelClassify() x = torch.full((1, 3), 100.0) out = pooler(x) assert torch.allclose(out, torch.ones(1, 3), atol=1e-4) def test_large_negative_maps_near_zero(self): pooler = PoolerMultiLabelClassify() x = torch.full((1, 3), -100.0) out = pooler(x) assert torch.allclose(out, torch.zeros(1, 3), atol=1e-4) # --------------------------------------------------------------------------- # PoolerClassify # --------------------------------------------------------------------------- class TestPoolerClassify: def test_infers_from_shape_when_num_labels_none(self): pooler = PoolerClassify(num_labels=None) assert pooler.num_labels is None x = torch.randn(2, 5) out = pooler(x) sums = out.sum(dim=-1) assert torch.allclose(sums, torch.ones(2), atol=1e-5) def test_sigmoid_when_num_labels_lt_2(self): pooler = PoolerClassify(num_labels=1) x = torch.zeros(1, 1) out = pooler(x) assert torch.allclose(out, torch.tensor([[0.5]]), atol=1e-5) def test_num_labels_zero_uses_sigmoid(self): pooler = PoolerClassify(num_labels=0) assert pooler.num_labels == 0 x = torch.zeros(1, 3) out = pooler(x) assert torch.allclose(out, torch.full((1, 3), 0.5), atol=1e-5) def test_num_labels_ge_2_uses_softmax(self): pooler = PoolerClassify(num_labels=4) assert pooler.num_labels == 4 x = torch.randn(2, 4) out = pooler(x) sums = out.sum(dim=-1) assert torch.allclose(sums, torch.ones(2), atol=1e-5) def test_default_num_labels_is_none(self): pooler = PoolerClassify() assert pooler.num_labels is None # --------------------------------------------------------------------------- # LambdaPoolerActivation # --------------------------------------------------------------------------- class TestLambdaPoolerActivation: def test_applies_custom_fn(self): pooler = LambdaPoolerActivation(nn.ReLU()) x = torch.tensor([[-1.0, 2.0, -3.0]]) out = pooler(x) expected = torch.tensor([[0.0, 2.0, 0.0]]) assert torch.equal(out, expected) def test_forward_list(self): pooler = LambdaPoolerActivation(nn.ReLU()) tensors = [torch.tensor([-1.0, 2.0]), torch.tensor([3.0, -4.0])] out = pooler(tensors) assert torch.equal(out[0], torch.tensor([0.0, 2.0])) assert torch.equal(out[1], torch.tensor([3.0, 0.0])) # --------------------------------------------------------------------------- # get_act_fn factory # --------------------------------------------------------------------------- class TestGetActFn: @staticmethod def _make_config(**kwargs): return SimpleNamespace(**kwargs) def test_regression(self): cfg = self._make_config(problem_type="regression") result = get_act_fn(cfg) assert isinstance(result, PoolerIdentity) def test_single_label_classification(self): cfg = self._make_config( problem_type="single_label_classification", num_labels=3 ) result = get_act_fn(cfg) assert isinstance(result, PoolerClassify) assert result.num_labels == 3 def test_multi_label_classification(self): cfg = self._make_config(problem_type="multi_label_classification") result = get_act_fn(cfg) assert isinstance(result, PoolerMultiLabelClassify) def test_sentence_transformers_activation(self): cfg = self._make_config( problem_type="", sentence_transformers={ "activation_fn": "torch.nn.modules.activation.Sigmoid" }, ) result = get_act_fn(cfg) assert isinstance(result, PoolerClassify) def test_sbert_activation(self): cfg = self._make_config( problem_type="", sbert_ce_default_activation_function=( "torch.nn.modules.activation.Sigmoid" ), ) result = get_act_fn(cfg) assert isinstance(result, PoolerClassify) def test_default_fallback(self): cfg = self._make_config(problem_type="") result = get_act_fn(cfg) assert isinstance(result, PoolerClassify) def test_sentence_transformers_takes_priority(self): cfg = self._make_config( problem_type="", sentence_transformers={"activation_fn": "torch.nn.modules.linear.Identity"}, sbert_ce_default_activation_function=( "torch.nn.modules.activation.Sigmoid" ), ) result = get_act_fn(cfg) assert isinstance(result, PoolerIdentity) def test_rejects_non_torch_activation(self): cfg = self._make_config( problem_type="", sentence_transformers={"activation_fn": "os.system"}, ) with pytest.raises(ValueError, match="restricted"): get_act_fn(cfg) # --------------------------------------------------------------------------- # resolve_classifier_act_fn # --------------------------------------------------------------------------- class TestResolveClassifierActFn: def test_delegates_to_get_act_fn_when_none(self): model_config = SimpleNamespace( hf_config=SimpleNamespace(num_labels=3, problem_type="") ) result = resolve_classifier_act_fn(model_config, act_fn=None) assert isinstance(result, PoolerClassify) assert result.num_labels == 3 def test_passes_through_provided_act_fn(self): custom = PoolerIdentity() result = resolve_classifier_act_fn(None, act_fn=custom) assert result is custom