File size: 9,008 Bytes
9bbba62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 | """
Test suite for Transformer module with sequence-classification task.
Tests various tiny model architectures that support sequence classification,
which is the default task used by CrossEncoder.
"""
from __future__ import annotations
from contextlib import nullcontext
import pytest
import torch
from packaging.version import Version
from transformers import __version__ as transformers_version
from transformers.models.auto.modeling_auto import (
MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES,
MODEL_MAPPING_NAMES,
)
from sentence_transformers.base.modules import Transformer
from sentence_transformers.util.tensor import batch_to_device
from .conftest import (
EXPECT_FORWARD_FAIL,
EXPECT_IMAGE_ONLY_FAILURE,
EXPECT_IMAGE_VIDEO_FAILURE,
EXPECT_MULTIMODAL_FAILURE,
EXPECT_MULTIMODAL_SUCCESS,
FAULTY_CHECKPOINTS,
REQUIRES_CUDA,
TINY_MODEL_MAPPING,
TRANSFORMERS_V4_XFAIL_ARCHITECTURES,
XFAIL_ARCHITECTURES,
create_modality_pair_samples,
create_modality_samples,
load_transformer,
modify_processor_for_pairs,
)
# Architectures that fail specifically for sequence-classification
# (beyond the general XFAIL_ARCHITECTURES)
XFAIL_SEQUENCE_CLASSIFICATION = [
"luke", # The LUKE tokenize doesn't work conveniently with text pairs
]
def _get_seq_cls_archs() -> list[str]:
"""Get architectures that support sequence-classification."""
return [
key
for key, value in TINY_MODEL_MAPPING.items()
if value is not None
and key not in XFAIL_ARCHITECTURES
and key not in XFAIL_SEQUENCE_CLASSIFICATION
and (key not in TRANSFORMERS_V4_XFAIL_ARCHITECTURES or Version(transformers_version) >= Version("5.0.0"))
and key not in FAULTY_CHECKPOINTS
and key in MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
and key in MODEL_MAPPING_NAMES # Ensure we also have a base model
]
@pytest.fixture(params=_get_seq_cls_archs(), scope="class")
def arch(request):
"""Get the model architecture name for sequence-classification task."""
return request.param
@pytest.fixture(scope="class")
def arch_model(arch):
model = load_transformer(arch, transformer_task="sequence-classification")
return arch, model
@pytest.fixture
def arch_model_modalities(arch_model):
"""Create a Transformer instance and return it with its supported modalities."""
try:
arch, model = arch_model
modalities = model.modalities
return arch, model, modalities
except Exception as e:
pytest.fail(f"Failed to get modalities: {e}")
class TestSequenceClassificationArchitectures:
"""Test suite for Transformer module with sequence-classification task."""
def test_get_embedding_dimension(self, arch_model):
"""Test that embedding dimension can be retrieved."""
arch, model = arch_model
dim = model.get_embedding_dimension()
assert isinstance(dim, int)
assert dim > 0
def test_save_load(self, arch_model, tmp_path):
"""Test saving and loading the model."""
arch, model = arch_model
save_path = tmp_path / "model"
save_path.mkdir(exist_ok=True)
model.save(str(save_path))
loaded_model = Transformer(str(save_path))
assert loaded_model.get_embedding_dimension() == model.get_embedding_dimension()
def test_module_output_name(self, arch_model):
"""Test that the module output name is 'scores' for sequence-classification."""
arch, model = arch_model
assert model.module_output_name == "scores"
def test_inference_with_supported_modalities(self, arch_model_modalities, subtests):
"""Test inference with each supported modality (single and multi-modal)."""
arch, model, modalities = arch_model_modalities
if arch in REQUIRES_CUDA:
if not torch.cuda.is_available():
pytest.skip(f"{arch} requires CUDA for inference, but CUDA is not available.")
else:
model = model.to("cuda")
# Create all valid test samples for the model's supported modalities
test_samples = create_modality_samples(
model, modalities, n=2, message_format=model.input_formatter.message_format
)
for modality_desc, inputs in test_samples.items():
with subtests.test(msg=f"Testing {modality_desc}"):
context = nullcontext()
expected_fail = EXPECT_FORWARD_FAIL.get(arch)
if expected_fail is None and arch in EXPECT_FORWARD_FAIL:
context = pytest.raises(Exception)
elif expected_fail is not None and modality_desc in expected_fail:
context = pytest.raises(Exception)
elif arch in EXPECT_IMAGE_VIDEO_FAILURE and "image" in modality_desc and "video" in modality_desc:
context = pytest.raises((ValueError, TypeError, AttributeError))
elif arch in EXPECT_IMAGE_ONLY_FAILURE and "image" in modality_desc and "+" not in modality_desc:
context = pytest.raises((ValueError, TypeError, AttributeError))
elif (
model.module_output_name == "sentence_embedding"
and "+" in modality_desc
and arch not in EXPECT_MULTIMODAL_SUCCESS
) or ("+" in modality_desc and arch in EXPECT_MULTIMODAL_FAILURE):
context = pytest.raises(ValueError)
with context:
try:
features = model.preprocess(inputs)
except ValueError as exc:
if (
"Could not make a flat list of images from" in str(exc)
and "image" in modality_desc
and ("url" in modality_desc or "path" in modality_desc)
):
pytest.skip(
f"The {arch!r} architecture with an older transformers version doesn't support image URLs, skipping this modality format."
)
raise
except KeyError as exc:
if "'height'" in str(exc) and "image" in modality_desc:
pytest.skip(
f"The {arch!r} architecture with an older transformers version doesn't yet nicely extract the size of image inputs, skipping this modality format."
)
raise
if arch in REQUIRES_CUDA:
features = batch_to_device(features, torch.device("cuda"))
with torch.no_grad():
output = model.forward(features)
assert model.module_output_name in output, (
f"Expected '{model.module_output_name}' in output for {modality_desc}"
)
output_tensor = output[model.module_output_name]
assert output_tensor.shape[0] == len(inputs), f"Batch size mismatch for {modality_desc}"
assert output_tensor.shape[-1] == 1, f"Output dimension mismatch for {modality_desc}"
def test_inference_with_supported_modality_pairs(self, arch_model_modalities, subtests):
"""Test inference with pair inputs for each supported modality combination."""
arch, model, modalities = arch_model_modalities
if arch in REQUIRES_CUDA:
if not torch.cuda.is_available():
pytest.skip(f"{arch} requires CUDA for inference, but CUDA is not available.")
else:
model = model.to("cuda")
if "message" in modalities:
modify_processor_for_pairs(model)
test_pairs = create_modality_pair_samples(model, modalities, n=2)
for pair_desc, pairs in test_pairs.items():
with subtests.test(msg=f"Testing {pair_desc}"):
context = nullcontext()
if arch in EXPECT_FORWARD_FAIL and EXPECT_FORWARD_FAIL[arch] is None:
context = pytest.raises(Exception)
with context:
features = model.preprocess(pairs)
if arch in REQUIRES_CUDA:
features = batch_to_device(features, torch.device("cuda"))
with torch.no_grad():
output = model.forward(features)
assert model.module_output_name in output, (
f"Expected '{model.module_output_name}' in output for {pair_desc}"
)
output_tensor = output[model.module_output_name]
assert output_tensor.shape[0] == len(pairs), f"Batch size mismatch for {pair_desc}"
assert output_tensor.shape[-1] == 1, f"Output dimension mismatch for {pair_desc}"
|