wehub-resource-sync
235 lines
8.1 KiB
Python
235 lines
8.1 KiB
Python
# 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
|