from __future__ import annotations

import importlib
import inspect
from collections.abc import Callable, Sequence
from contextlib import contextmanager
from dataclasses import fields
from typing import TYPE_CHECKING, Any, Literal, TypedDict, get_args, get_type_hints

import torch
from packaging.version import parse as parse_version
from tokenizers.normalizers import Lowercase
from tokenizers.normalizers import Sequence as NormalizerSequence
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForMaskedLM,
    AutoModelForSequenceClassification,
    AutoProcessor,
    BlenderbotConfig,
    BlenderbotSmallConfig,
    FeatureExtractionMixin,
    ImageProcessingMixin,
    LongT5Config,
    M2M100Config,
    MarianConfig,
    MT5Config,
    PegasusConfig,
    PegasusXConfig,
    PretrainedConfig,
    PreTrainedModel,
    PreTrainedTokenizerBase,
    ProcessorMixin,
    ProphetNetConfig,
    SwitchTransformersConfig,
    T5Config,
    UdopConfig,
    UMT5Config,
    WhisperConfig,
)
from transformers import __version__ as transformers_version
from transformers.utils import ModelOutput
from transformers.utils import logging as transformers_logging
from transformers.utils.import_utils import is_peft_available
from transformers.utils.peft_utils import find_adapter_config_file
from typing_extensions import NotRequired

from sentence_transformers.backend import load_onnx_model, load_openvino_model
from sentence_transformers.base.modality import InputFormatter, format_modality, raise_unsupported_modality_error
from sentence_transformers.base.modality_types import (
    MODALITY_TO_PROCESSOR_ARG,
    MessageInput,
    Modality,
    PairInput,
    SingleInput,
)
from sentence_transformers.base.modules.input_module import InputModule
from sentence_transformers.util.decorators import transformer_kwargs_decorator
from sentence_transformers.util.environment import suggest_extra_on_exception

try:
    from typing import Self
except ImportError:
    from typing_extensions import Self

try:
    from transformers import BaseVideoProcessor
except ImportError:

    class BaseVideoProcessor:
        pass


try:
    from transformers import T5Gemma2Config, T5Gemma2TextConfig
except ImportError:

    class T5Gemma2Config:
        pass

    class T5Gemma2TextConfig:
        pass


try:
    from transformers import T5GemmaConfig
except ImportError:

    class T5GemmaConfig:
        pass


try:
    from transformers import MoonshineConfig
except ImportError:

    class MoonshineConfig:
        pass


try:
    from transformers import TimmWrapperConfig
except ImportError:

    class TimmWrapperConfig:
        pass


try:
    from transformers import PaliGemmaProcessor
except ImportError:

    class PaliGemmaProcessor:
        pass


if TYPE_CHECKING and is_peft_available():
    from peft import PeftConfig

logger = transformers_logging.get_logger(__name__)

_TRANSFORMERS_PROCESSOR_SUPPORTS_MODALITY_KWARGS = parse_version(transformers_version) > parse_version("4.56.1")
_TRANSFORMERS_APPLY_CHAT_TEMPLATE_RECOMMENDS_PROCESSOR_KWARGS = parse_version(transformers_version) >= parse_version(
    "5.4.0.dev0"
)
_TRANSFORMERS_SUPPORTS_USE_BIDIRECTIONAL_ATTENTION = parse_version(transformers_version) >= parse_version("4.56.2")
_TRANSFORMERS_SUPPORTS_IS_CAUSAL_FALSE = parse_version(transformers_version) >= parse_version("5.2.0")

# The size/output kwargs of apply_chat_template, which a bare tokenizer wants as top-level kwargs.
_APPLY_CHAT_TEMPLATE_TOP_LEVEL_KWARGS = frozenset({"padding", "truncation", "max_length", "return_tensors"})
_CHAT_TEMPLATE_SUFFIX_CACHE_SIZE = 256

TransformerTask = Literal[
    "feature-extraction", "sequence-classification", "text-generation", "any-to-any", "fill-mask", "retrieval"
]


class _ModalityParamsRequired(TypedDict):
    method: str
    method_output_name: str | None


class ModalityParams(_ModalityParamsRequired, total=False):
    """Parameters for a single modality entry in the modality config.

    The ``format`` key is only used for the ``"message"`` modality and controls how
    message content is structured: ``"structured"`` (list of typed dicts) or ``"flat"``
    (direct string value).
    """

    format: Literal["structured", "flat"]


ModalityConfig = dict[Modality, ModalityParams]


def _resolve_retrieval_model_class(config: PretrainedConfig | PeftConfig) -> type[PreTrainedModel]:
    """Resolve the concrete ``*ForRetrieval`` class (ColPali, ColQwen2, ...) for ``transformer_task="retrieval"``.

    transformers (as of v5.13) has no ``AutoModelForRetrieval``, and ``MODEL_FOR_RETRIEVAL_MAPPING`` is
    incomplete (missing ``colqwen2``), so fall back to the checkpoint's ``architectures`` entry.
    """
    try:
        from transformers.models.auto.modeling_auto import MODEL_FOR_RETRIEVAL_MAPPING

        return MODEL_FOR_RETRIEVAL_MAPPING[type(config)]
    except (ImportError, KeyError):
        pass

    # PEFT adapter checkpoints carry a PeftConfig without `architectures`: resolve from the base
    # model's config instead.
    base_model_name_or_path = getattr(config, "base_model_name_or_path", None)
    if base_model_name_or_path:
        return _resolve_retrieval_model_class(AutoConfig.from_pretrained(base_model_name_or_path))

    import transformers

    for architecture in getattr(config, "architectures", None) or []:
        model_class = getattr(transformers, architecture, None)
        if model_class is not None:
            return model_class

    # Remote-code checkpoints (e.g. ColQwen3) define their class in-repo, so only ``auto_map`` can
    # resolve it. AutoModel still refuses to execute that code without ``trust_remote_code``.
    if "AutoModel" in (getattr(config, "auto_map", None) or {}):
        return AutoModel

    raise ValueError(
        f"Could not resolve a retrieval model class for model_type {getattr(config, 'model_type', None)!r}. "
        'The transformer_task="retrieval" expects a `*ForRetrieval` architecture, e.g. ColPali, '
        "ColQwen2 or ColModernVBert."
    )


TRANSFORMER_TASK_TO_AUTO_MODEL: dict[TransformerTask, Any] = {
    "feature-extraction": AutoModel,  # Used by SentenceTransformer, also covers "image-feature-extraction"
    "sequence-classification": AutoModelForSequenceClassification,  # Used by CrossEncoder
    "text-generation": AutoModelForCausalLM,  # Used by CrossEncoder
    "fill-mask": AutoModelForMaskedLM,  # Used by SparseEncoder
    "retrieval": None,  # Used by MultiVectorEncoder, via _resolve_retrieval_model_class
}

try:
    from transformers import AutoModelForMultimodalLM

    TRANSFORMER_TASK_TO_AUTO_MODEL["any-to-any"] = (
        AutoModelForMultimodalLM  # Used by CrossEncoder, also covers "image-text-to-text"
    )
except ImportError:
    pass

# Default (modality_config, module_output_name) per transformer task.
# Used as the fallback when loading models saved before modality_config was introduced,
# and as the source of default method_output_name / module_output_name during modality inference.
TRANSFORMER_TASK_DEFAULTS: dict[TransformerTask, tuple[ModalityConfig, str]] = {
    "feature-extraction": (
        {"text": {"method": "forward", "method_output_name": "last_hidden_state"}},
        "token_embeddings",
    ),
    "sequence-classification": (
        {"text": {"method": "forward", "method_output_name": "logits"}},
        "scores",
    ),
    "text-generation": (
        {"text": {"method": "forward", "method_output_name": "logits"}},
        "causal_logits",
    ),
    "any-to-any": (
        {"text": {"method": "forward", "method_output_name": "logits"}},
        "causal_logits",
    ),
    "fill-mask": (
        {"text": {"method": "forward", "method_output_name": "logits"}},
        "token_embeddings",
    ),
    # `*ForRetrieval` heads return already-projected, already-normalised token embeddings.
    "retrieval": (
        {
            "text": {"method": "forward", "method_output_name": "embeddings"},
            "image": {"method": "forward", "method_output_name": "embeddings"},
        },
        "token_embeddings",
    ),
}

# Registry of encoder-decoder architectures whose encoder can be loaded standalone.
# Each entry maps a config class to (module_path, encoder_class_name) for lazy importing.
_ENCODER_ONLY_MODELS: list[tuple[type, str, str]] = [
    (T5Config, "transformers", "T5EncoderModel"),
    (MT5Config, "transformers", "MT5EncoderModel"),
    (UMT5Config, "transformers", "UMT5EncoderModel"),
    (UdopConfig, "transformers", "UdopEncoderModel"),
    (LongT5Config, "transformers", "LongT5EncoderModel"),
    (ProphetNetConfig, "transformers", "ProphetNetEncoder"),
    (SwitchTransformersConfig, "transformers", "SwitchTransformersEncoderModel"),
    (BlenderbotConfig, "transformers.models.blenderbot.modeling_blenderbot", "BlenderbotEncoder"),
    (
        BlenderbotSmallConfig,
        "transformers.models.blenderbot_small.modeling_blenderbot_small",
        "BlenderbotSmallEncoder",
    ),
    (M2M100Config, "transformers.models.m2m_100.modeling_m2m_100", "M2M100Encoder"),
    (PegasusConfig, "transformers.models.pegasus.modeling_pegasus", "PegasusEncoder"),
    (PegasusXConfig, "transformers.models.pegasus_x.modeling_pegasus_x", "PegasusXEncoder"),
    (MoonshineConfig, "transformers.models.moonshine.modeling_moonshine", "MoonshineEncoder"),
    (WhisperConfig, "transformers.models.whisper.modeling_whisper", "WhisperEncoder"),
    (MarianConfig, "transformers.models.marian.modeling_marian", "MarianEncoder"),
    # T5Gemma2TextConfig is for loading from an already encoder-only checkpoint.
    # Loading the encoder from a full T5Gemma2Config is handled separately in _load_encoder_only_model.
    (T5Gemma2TextConfig, "transformers.models.t5gemma2.modeling_t5gemma2", "T5Gemma2Encoder"),
]

# Hard-coded modality configs for model types that can't be handled by the general inference path.
# Each entry maps a model_type string to (modality_config, module_output_name, validate_output_names).
# When validate_output_names is True, each modality's method_output_name is validated at runtime
# against the model method's return type via _infer_method_output_name (for transformers v4/v5 compat).
_AUDIO_MODALITY_CONFIG: tuple[ModalityConfig, str, bool] = (
    {
        "audio": {"method": "forward", "method_output_name": "last_hidden_state"},
        ("audio", "text"): {"method": "forward", "method_output_name": "last_hidden_state"},
    },
    "token_embeddings",
    False,
)

_FEATURE_EXTRACTION_EDGE_CASES: dict[str, tuple[ModalityConfig, str, bool]] = {
    # Models with custom get_*_features methods that need output name validation
    "blip": (
        {
            "text": {"method": "get_text_features", "method_output_name": "pooler_output"},
            "image": {"method": "get_image_features", "method_output_name": "pooler_output"},
            ("image", "text"): {"method": "get_multimodal_features", "method_output_name": "pooler_output"},
        },
        "sentence_embedding",
        True,
    ),
    "blip-2": (
        {
            "text": {"method": "get_text_features", "method_output_name": "last_hidden_state"},
            "image": {"method": "get_image_features", "method_output_name": "last_hidden_state"},
        },
        "token_embeddings",
        True,
    ),
    "sam3": (
        {
            "text": {"method": "get_text_features", "method_output_name": "last_hidden_state"},
            "image": {"method": "get_vision_features", "method_output_name": "last_hidden_state"},
        },
        "token_embeddings",
        True,
    ),
    "flava": (
        {
            "text": {"method": "get_text_features", "method_output_name": "pooler_output"},
            "image": {"method": "get_image_features", "method_output_name": "pooler_output"},
        },
        "token_embeddings",
        True,
    ),
    # Models supporting text+image without message format
    "git": (
        {
            "text": {"method": "forward", "method_output_name": "last_hidden_state"},
            ("image", "text"): {"method": "forward", "method_output_name": "last_hidden_state"},
        },
        "token_embeddings",
        False,
    ),
    "visual_bert": (
        {
            "text": {"method": "forward", "method_output_name": "last_hidden_state"},
            ("image", "text"): {"method": "forward", "method_output_name": "last_hidden_state"},
        },
        "token_embeddings",
        False,
    ),
    # Only combined (image, text) input is supported, no text-only nor image-only
    "kosmos-2": (
        {("image", "text"): {"method": "forward", "method_output_name": "last_hidden_state"}},
        "token_embeddings",
        False,
    ),
    "grounding-dino": (
        {("image", "text"): {"method": "forward", "method_output_name": "last_hidden_state"}},
        "token_embeddings",
        False,
    ),
    "paligemma": (
        {("image", "text"): {"method": "forward", "method_output_name": "last_hidden_state"}},
        "token_embeddings",
        False,
    ),
    "vilt": (
        {("image", "text"): {"method": "forward", "method_output_name": "last_hidden_state"}},
        "token_embeddings",
        False,
    ),
    # Image+text supported without message format, plus image-only
    "layoutlmv3": (
        {
            "image": {"method": "forward", "method_output_name": "last_hidden_state"},
            ("image", "text"): {"method": "forward", "method_output_name": "last_hidden_state"},
        },
        "token_embeddings",
        False,
    ),
    # All modalities via forward
    "idefics": (
        {
            "text": {"method": "forward", "method_output_name": "last_hidden_state"},
            "image": {"method": "forward", "method_output_name": "last_hidden_state"},
            ("image", "text"): {"method": "forward", "method_output_name": "last_hidden_state"},
        },
        "token_embeddings",
        False,
    ),
    # Audio encoder models: we load only the encoder, so text decoding is not available
    "hubert": _AUDIO_MODALITY_CONFIG,
    "moonshine": _AUDIO_MODALITY_CONFIG,
    "sew": _AUDIO_MODALITY_CONFIG,
    "sew-d": _AUDIO_MODALITY_CONFIG,
    "unispeech-sat": _AUDIO_MODALITY_CONFIG,
    "unispeech": _AUDIO_MODALITY_CONFIG,
    "wav2vec2": _AUDIO_MODALITY_CONFIG,
    "wav2vec2-conformer": _AUDIO_MODALITY_CONFIG,
    "wavlm": _AUDIO_MODALITY_CONFIG,
    "whisper": _AUDIO_MODALITY_CONFIG,
    "voxtral_realtime": _AUDIO_MODALITY_CONFIG,
}

_FILL_MASK_EDGE_CASES: dict[str, tuple[ModalityConfig, str, bool]] = {
    # wav2vec2's forward outputs 'logits' rather than 'last_hidden_state' for fill-mask,
    # but accepts audio input rather than text
    "wav2vec2": (
        {
            "audio": {"method": "forward", "method_output_name": "logits"},
        },
        "token_embeddings",
        False,
    ),
}

_TEXT_GENERATION_EDGE_CASES = {
    # Models supporting text+image without message format, but no image-only
    "git": (
        {
            "text": {"method": "forward", "method_output_name": "logits"},
            ("image", "text"): {"method": "forward", "method_output_name": "logits"},
        },
        "causal_logits",
        False,
    ),
    # The Whisper decoder is text only
    "whisper": (
        {
            "text": {"method": "forward", "method_output_name": "logits"},
        },
        "causal_logits",
        False,
    ),
}

_ANY_TO_ANY_EDGE_CASES = {
    # Models supporting text+image without message format, but no image-only
    "git": (
        {
            "text": {"method": "forward", "method_output_name": "logits"},
            ("image", "text"): {"method": "forward", "method_output_name": "logits"},
        },
        "causal_logits",
        False,
    ),
    # Only combined (image, text) input is supported, no text-only nor image-only
    "blip": (
        {("image", "text"): {"method": "forward", "method_output_name": "logits"}},
        "causal_logits",
        False,
    ),
    "blip-2": (
        {("image", "text"): {"method": "forward", "method_output_name": "logits"}},
        "causal_logits",
        False,
    ),
    "kosmos-2": (
        {("image", "text"): {"method": "forward", "method_output_name": "logits"}},
        "causal_logits",
        False,
    ),
    "paligemma": (
        {("image", "text"): {"method": "forward", "method_output_name": "logits"}},
        "causal_logits",
        False,
    ),
    # Models supporting text+audio, but no text-only
    "voxtral_realtime": (
        {
            "audio": {"method": "forward", "method_output_name": "logits"},
            ("audio", "text"): {"method": "forward", "method_output_name": "logits"},
        },
        "causal_logits",
        False,
    ),
}

_EDGE_CASE_MODALITY_CONFIGS: dict[str, dict[str, tuple[ModalityConfig, str, bool]]] = {
    "feature-extraction": _FEATURE_EXTRACTION_EDGE_CASES,
    "text-generation": _TEXT_GENERATION_EDGE_CASES,
    "any-to-any": _ANY_TO_ANY_EDGE_CASES,
    "fill-mask": _FILL_MASK_EDGE_CASES,
}


def _has_lowercase(normalizer) -> bool:
    """Check whether a tokenizers normalizer (or sequence of normalizers) includes Lowercase."""
    if normalizer is None:
        return False
    if isinstance(normalizer, Lowercase):
        return True
    if isinstance(normalizer, NormalizerSequence):
        return any(isinstance(n, Lowercase) for n in normalizer)
    return False


@contextmanager
def set_temporary_class_attrs(cls, **overrides):
    originals = {name: getattr(cls, name, None) for name in overrides}
    try:
        for name, value in overrides.items():
            setattr(cls, name, value)
        yield
    finally:
        for name, value in originals.items():
            setattr(cls, name, value)


class ProcessingKwargs(TypedDict, total=False):
    """Keyword arguments applied when *calling* the processor during preprocessing.

    Valid keys: ``"common"``, ``"text"``, ``"audio"``, ``"image"``, ``"video"``, ``"chat_template"``.
    Modality and ``"common"`` kwargs override built-in defaults. ``"chat_template"`` kwargs are
    forwarded to ``apply_chat_template``, except the Sentence Transformers ``restore_suffix`` flag.
    """

    common: dict[str, Any]
    text: dict[str, Any]
    audio: dict[str, Any]
    image: dict[str, Any]
    video: dict[str, Any]
    chat_template: dict[str, Any]


class QueryExpansionConfig(TypedDict):
    """ColBERT-style query expansion config. When set, :meth:`Transformer.preprocess` expands every
    ``task="query"`` batch, regardless of the downstream modules: only meaningful with a
    :class:`MultiVectorMask` downstream (which scores the expansion positions), so leave it ``None``
    for dense / sparse / cross-encoder pipelines. Keys:

    * ``strategy``:

      - ``"fixed"`` (classic ColBERT): every query leaves preprocessing at exactly ``length``
        tokens. Shorter queries are padded up to ``length`` and the pads swapped for ``token``
        (defaults to the tokenizer's ``mask_token``), longer queries are truncated down to it.
        ``length`` is both the floor and the ceiling: for query batches it overrides
        ``Transformer.query_length``.
      - ``"min"``: ``length`` is only the floor. Shorter queries are expanded exactly like
        ``"fixed"``, longer queries pass through untruncated instead of silently losing content.
        ``Transformer.query_length`` (or a trainer ``max_length`` override) still applies as the
        ceiling. On corpora where queries stay under ``length``, ``"min"`` and ``"fixed"`` produce
        identical batches.

    * ``attend``: whether the encoder attends to the expansion positions. ``False`` (default,
      classic ColBERT) leaves them at ``attention_mask=0``: the encoder skips them in the forward
      pass, MaxSim still scores them, and they soak up query context via the asymmetric attention
      mask. ``True`` forces ``attention_mask=1`` so the encoder also attends to them, matching
      PyLate's ``attend_to_expansion_tokens``. Used by some ColBERT variants.

    * ``token``: defaults to the tokenizer's ``mask_token``.

    * ``length``: pad target (total tokens after expansion, the width floor). **Required**
      (classic ColBERT uses 32).

    Query expansion exists for classic mask-token checkpoints (PyLate / Stanford-NLP ColBERT):
    their attention-mask surgery and fixed-length pad target cannot be expressed in a chat
    template. Modern chat-template backbones (the colpali-engine ColQwen / ColGemma / ColIdefics
    families) should instead own their query augmentation in the checkpoint's chat template,
    which receives the ``task`` kwarg and can e.g. append the family's suffix tokens for
    ``task == "query"`` renders.
    """

    strategy: Literal["fixed", "min"]
    attend: NotRequired[bool]
    token: NotRequired[str | None]
    length: NotRequired[int]


_VALID_QUERY_EXPANSION_KEYS: set[str] = {"strategy", "attend", "token", "length"}
_VALID_QUERY_EXPANSION_STRATEGIES: tuple[str, ...] = ("fixed", "min")


def _normalize_query_expansion(
    query_expansion: QueryExpansionConfig | None,
) -> QueryExpansionConfig | None:
    """Validate the dict, fill defaults, and return a fully-populated config (or ``None`` for the
    no-expansion case).
    """
    if query_expansion is None:
        return None
    unknown = set(query_expansion) - _VALID_QUERY_EXPANSION_KEYS
    if unknown:
        raise ValueError(
            f"query_expansion has unknown keys {sorted(unknown)!r}. "
            f"Valid keys: {sorted(_VALID_QUERY_EXPANSION_KEYS)!r}."
        )
    strategy = query_expansion.get("strategy")
    if strategy not in _VALID_QUERY_EXPANSION_STRATEGIES:
        raise ValueError(
            f"query_expansion['strategy'] must be one of {_VALID_QUERY_EXPANSION_STRATEGIES!r}, got {strategy!r}."
        )
    attend = query_expansion.get("attend", False)
    if not isinstance(attend, bool):
        raise ValueError(f"query_expansion['attend'] must be a bool, got {attend!r}.")
    token = query_expansion.get("token")
    normalized: QueryExpansionConfig = {"strategy": strategy, "attend": attend, "token": token}
    length = query_expansion.get("length")
    if not isinstance(length, int) or isinstance(length, bool) or length < 1:
        raise ValueError(
            f"query_expansion strategy={strategy!r} requires 'length': a positive int specifying "
            "the pad target (total tokens after expansion). Classic ColBERT uses 32."
        )
    normalized["length"] = length
    return normalized


def _count_media_per_sample(messages: list[list[dict[str, Any]]]) -> tuple[list[int], list[int]]:
    """Count images and videos per sample from the message structure.

    Some VLM processors flatten per-sample visual tokens into single tensors (e.g.
    ``pixel_values`` shape ``(total_visual_tokens, hidden_dim)``), losing the per-sample
    association. By counting before the processor call, we get reliable per-sample counts
    that downstream minibatching can use directly.
    """
    num_images: list[int] = []
    num_videos: list[int] = []
    for sample_messages in messages:
        img_count = 0
        vid_count = 0
        for msg in sample_messages:
            content = msg.get("content", [])
            if not isinstance(content, list):
                continue
            for item in content:
                if isinstance(item, dict):
                    item_type = item.get("type", "")
                    if item_type == "image":
                        img_count += 1
                    elif item_type == "video":
                        vid_count += 1
        num_images.append(img_count)
        num_videos.append(vid_count)
    return num_images, num_videos


# Forward-kwarg filter constants: ST/tokenizer bookkeeping to never forward, plus common kwargs to always allow.
_NON_MODEL_FEATURE_KEYS = frozenset(
    {
        # Sentence Transformers' own bookkeeping.
        "modality",
        "num_images_per_sample",
        "num_videos_per_sample",
        "prompt_length",
        "query_expansion_positions",
        "task",
        # Tokenizer extras from return_* flags (return_offsets_mapping, etc.), not forward inputs.
        "offset_mapping",
        "overflow_to_sample_mapping",
        "special_tokens_mask",
    }
)
_FORWARD_SAFETY_NET_KEYS = frozenset({"input_ids", "attention_mask", "token_type_ids", "inputs_embeds", "return_dict"})


class Transformer(InputModule):
    """Hugging Face AutoModel wrapper that handles loading, preprocessing, and inference.

    Loads the appropriate model class (e.g. BERT, RoBERTa, CLIP, Whisper) based on the model configuration
    and the specified ``transformer_task``. Supports text, image, audio, and video modalities depending on
    the underlying model. This module is typically the first module in a
    :class:`~sentence_transformers.sentence_transformer.model.SentenceTransformer`, :class:`~sentence_transformers.sparse_encoder.model.SparseEncoder`,
    or :class:`~sentence_transformers.cross_encoder.model.CrossEncoder` pipeline.

    Args:
        model_name_or_path (str): Hugging Face model name or path to a local model directory.
        transformer_task (str, optional): The task determining which ``AutoModel``-like class to load.
            Supported values:

            - ``"feature-extraction"`` (default): :class:`~transformers.AutoModel`, e.g. used by
              :class:`~sentence_transformers.sentence_transformer.model.SentenceTransformer`.
            - ``"sequence-classification"``: :class:`~transformers.AutoModelForSequenceClassification`,
              e.g. used by :class:`~sentence_transformers.cross_encoder.model.CrossEncoder`.
            - ``"text-generation"``: :class:`~transformers.AutoModelForCausalLM`, e.g. used by generative
              :class:`~sentence_transformers.cross_encoder.model.CrossEncoder` models. Sets the ``tokenizer`` padding_side to "left".
            - ``"any-to-any"``: :class:`~transformers.AutoModelForMultimodalLM`, e.g. used by multimodal generative
              :class:`~sentence_transformers.cross_encoder.model.CrossEncoder` models (requires transformers v5+). Sets the
              ``tokenizer`` padding_side to "left".
            - ``"fill-mask"``: :class:`~transformers.AutoModelForMaskedLM`, e.g. used by
              :class:`~sentence_transformers.sparse_encoder.model.SparseEncoder`.

            Defaults to ``"feature-extraction"``.
        model_kwargs (dict[str, Any], optional): Keyword arguments forwarded to
            ``AutoModel.from_pretrained`` when loading the model. Particularly useful options include:

            - ``torch_dtype``: Override the default ``torch.dtype`` and load the model under a specific
              dtype. Can be ``torch.float16``, ``torch.bfloat16``, ``torch.float32``, or ``"auto"`` to
              use the dtype from the model's ``config.json``.
            - ``attn_implementation``: The attention implementation to use. For example ``"eager"``,
              ``"sdpa"``, or ``"flash_attention_2"``. If you ``pip install kernels``, then
              ``"flash_attention_2"`` should work without having to install ``flash_attn``. It is
              frequently the fastest option. Defaults to ``"sdpa"`` when available (torch>=2.1.1).
            - ``device_map``: Controls how the model is placed across devices, e.g. ``"auto"`` for
              model parallelism, or a single device like ``"cuda:1"`` to load the backbone directly onto one GPU.
            - ``provider``: For ``backend="onnx"``, the ONNX execution provider
              (e.g. ``"CUDAExecutionProvider"``).
            - ``file_name``: For ``backend="onnx"`` or ``"openvino"``, the filename to load
              (e.g. for optimized or quantized models).
            - ``export``: For ``backend="onnx"`` or ``"openvino"``, whether to export the model to the
              backend format. Also set automatically if the exported file doesn't exist.

            See the `PreTrainedModel.from_pretrained
            <https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.from_pretrained>`_
            documentation for more details. Defaults to None.
        processor_kwargs (dict[str, Any], optional): Keyword arguments forwarded to
            ``AutoProcessor.from_pretrained`` when loading the processor/tokenizer. See the
            `AutoTokenizer.from_pretrained
            <https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoTokenizer.from_pretrained>`_
            documentation for more details. Defaults to None.
        config_kwargs (dict[str, Any], optional): Keyword arguments forwarded to
            ``AutoConfig.from_pretrained`` when loading the config. See the `AutoConfig.from_pretrained
            <https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoConfig.from_pretrained>`_
            documentation for more details. Defaults to None.
        processing_kwargs (dict[str, dict[str, Any]], optional): Keyword arguments applied when *calling*
            the processor during preprocessing. This is a nested dict whose keys are modality names
            (``"text"``, ``"audio"``, ``"image"``, ``"video"``), ``"common"`` for kwargs shared across all
            modalities, or ``"chat_template"`` for kwargs forwarded to ``apply_chat_template`` (e.g.
            ``{"add_generation_prompt": True}``). Modality and common kwargs override the built-in defaults.
            The ``"chat_template"`` dict also accepts ``"restore_suffix"`` (default ``True``), a Sentence
            Transformers flag: fixed tokens a chat template appends after the content (e.g. an assistant
            prefill) are restored when truncation drops them, so models that read the final token position
            keep working. Saved to and loaded from the model configuration file. Defaults to None.
        backend (str, optional): Backend used for model inference. Can be ``"torch"`` (default), ``"onnx"``,
            or ``"openvino"``. Defaults to ``"torch"``.
        modality_config (dict, optional): Custom modality configuration mapping modality names to method and
            output name dicts. When provided, ``module_output_name`` must also be set. The ``"message"``
            modality entry may include a ``"format"`` key (``"structured"``, ``"flat"``, or ``"auto"``)
            to control how chat-template inputs are formatted. Defaults to None.
        module_output_name (str, optional): The name of the output feature this module creates (e.g.
            ``"token_embeddings"``, ``"scores"``). Required when ``modality_config`` is provided.
            Defaults to None.
        unpad_inputs (bool, optional): Controls whether text-only inputs are concatenated without
            padding for faster inference using flash attention's variable-length functions. Non-text
            inputs (images, audio, video) are always padded normally. If ``None`` (default), unpadding
            is enabled automatically when all prerequisites are met (flash attention with variable-length
            support, ``"torch"`` backend, and the ``"feature-extraction"`` or ``"fill-mask"`` task).
            Set to ``False`` to force
            padding, which is needed for architectures that don't support unpadded inputs (e.g.
            ``qwen2_vl``). Set to ``True`` to request unpadding explicitly. A warning is logged if the
            prerequisites are not met. Defaults to None.
        query_length (int, optional): Per-task text max-length applied when ``task="query"`` reaches
            :meth:`preprocess` (e.g. via :meth:`encode_query`). Content-cap only (truncation). ``None``
            (default) falls back to the tokenizer's ``model_max_length``. Used by ColBERT-style
            retrieval to cap queries shorter than documents. Not the pad target: query expansion
            specifies its pad target via ``query_expansion["length"]``, which overrides this (a
            smaller ``query_length`` raises when the expansion is assigned).
        document_length (int, optional): Per-task text max-length applied when ``task="document"``
            reaches :meth:`preprocess`. ``None`` (default) falls back to the tokenizer's
            ``model_max_length``.
        query_expansion (QueryExpansionConfig, optional): ColBERT-style query expansion config. See
            :class:`QueryExpansionConfig`. ``None`` (default) means no expansion. When set, every
            ``task="query"`` batch is expanded during preprocessing, so only set this with a
            :class:`MultiVectorMask` downstream (leave it ``None`` for dense / sparse /
            cross-encoder pipelines).
        max_seq_length (int, optional): Truncate any inputs longer than this value. Prefer setting
            ``model_max_length`` via ``processor_kwargs`` instead. Defaults to None.
        do_lower_case (bool, optional): If true, lowercases the input (independent of whether the model
            is cased or not). Rarely needed. Defaults to False.
        tokenizer_name_or_path (str, optional): Name or path of the tokenizer. When None,
            ``model_name_or_path`` is used. Deprecated. Defaults to None.
    """

    config_file_name: str = "sentence_bert_config.json"
    config_keys: list[str] = [
        "transformer_task",
        "modality_config",
        "module_output_name",
        "processing_kwargs",
        "unpad_inputs",
        "query_length",
        "document_length",
        "query_expansion",
    ]
    # Config keys held at their default for dense / sparse / cross-encoder use are omitted from the
    # saved ``sentence_bert_config.json`` (see :meth:`get_config_dict`). The multi-vector machinery is
    # paid-for only when a model actually opts into it.
    _DEFAULT_CONFIG_VALUES: dict[str, Any] = {
        "query_length": None,
        "document_length": None,
        "query_expansion": None,
    }
    save_in_root: bool = True
    _VALID_PROCESSING_KWARGS_KEYS: set[str] = {"common", "text", "audio", "image", "video", "chat_template"}

    @transformer_kwargs_decorator
    def __init__(
        self,
        model_name_or_path: str,
        *,
        transformer_task: TransformerTask = "feature-extraction",
        model_kwargs: dict[str, Any] | None = None,
        processor_kwargs: dict[str, Any] | None = None,
        config_kwargs: dict[str, Any] | None = None,
        processing_kwargs: ProcessingKwargs | None = None,
        backend: Literal["torch", "onnx", "openvino"] = "torch",
        modality_config: ModalityConfig | None = None,
        module_output_name: str | None = None,
        unpad_inputs: bool | None = None,
        query_length: int | None = None,
        document_length: int | None = None,
        query_expansion: QueryExpansionConfig | None = None,
        max_seq_length: int | None = None,
        do_lower_case: bool = False,
        tokenizer_name_or_path: str | None = None,
    ) -> None:
        super().__init__()
        if transformer_task not in TRANSFORMER_TASK_TO_AUTO_MODEL:
            if transformer_task == "any-to-any":
                raise ImportError(
                    "The 'any-to-any' transformer task requires transformers v5+. "
                    "Please upgrade transformers with `pip install transformers>=5.0.0`."
                )
            raise ValueError(
                f"Unsupported transformer_task '{transformer_task}'. Supported tasks are: {list(TRANSFORMER_TASK_TO_AUTO_MODEL.keys())}"
            )
        self.transformer_task: TransformerTask = transformer_task
        if model_kwargs is None:
            model_kwargs = {}
        if processor_kwargs is None:
            processor_kwargs = {}
        if config_kwargs is None:
            config_kwargs = {}
        self.processing_kwargs: ProcessingKwargs = processing_kwargs or {}
        unknown_keys = set(self.processing_kwargs) - self._VALID_PROCESSING_KWARGS_KEYS
        if unknown_keys:
            logger.warning(
                f"Unknown keys in `processing_kwargs`: {unknown_keys}. "
                f"Valid keys are: {sorted(self._VALID_PROCESSING_KWARGS_KEYS)}. "
                "Unknown keys will be ignored. Did you mean to nest them under "
                "'common' or a modality key ('text', 'audio', 'image', 'video')?"
            )
        self.backend = backend
        self.do_lower_case = do_lower_case
        self.track_media_counts = False
        self._prompt_length_mapping = {}
        self._method_signature_cache: dict[str, set[str]] = {}
        self._chat_template_suffix_cache: dict[Any, list[int]] = {}
        # Per-task max-length overrides applied to the text modality. When ``task`` reaches
        # :meth:`preprocess` as ``"query"`` / ``"document"`` we pick the matching attribute. ``None``
        # leaves the tokenizer's default behaviour intact.
        self.query_length = query_length
        self.document_length = document_length

        config, is_peft_model = self._load_config(model_name_or_path, backend, config_kwargs)
        self._warn_on_unsupported_attention_config(config)

        if (
            transformer_task == "sequence-classification"
            and "num_labels" not in config_kwargs
            and (
                config.architectures is None
                or not any(arch.endswith("ForSequenceClassification") for arch in config.architectures)
            )
        ):
            # If we're loading a model for sequence-classification, but the base architecture is not for sequence-classification,
            # and num_labels is not specified, we default to 1 label for CrossEncoder-like behavior
            config.num_labels = 1

        self.model = self._load_model(
            model_name_or_path, transformer_task, config, backend, is_peft_model, **model_kwargs
        )

        # Unwrap PEFT first: ``PeftModel.forward`` hides the base model's multimodal params (e.g. ``pixel_values``).
        forward_model = self.model
        if is_peft_available():
            from peft import PeftModel

            if isinstance(forward_model, PeftModel):
                forward_model = forward_model.get_base_model()
        # ``None`` = forward accepts ``**kwargs`` (pass all but undeclared ST bookkeeping).
        # Set = declared params plus safety net.
        forward_signature = inspect.signature(forward_model.forward)
        self._declared_forward_params = set(forward_signature.parameters)
        self.model_forward_params: set[str] | None = None
        if not any(p.kind is inspect.Parameter.VAR_KEYWORD for p in forward_signature.parameters.values()):
            self.model_forward_params = self._declared_forward_params | _FORWARD_SAFETY_NET_KEYS

        if max_seq_length is not None and "model_max_length" not in processor_kwargs:
            processor_kwargs["model_max_length"] = max_seq_length
        with suggest_extra_on_exception():
            self.processor = AutoProcessor.from_pretrained(
                tokenizer_name_or_path if tokenizer_name_or_path is not None else model_name_or_path,
                **processor_kwargs,
            )

        # Cap the tokenizer model_max_length at the model's max_position_embeddings
        if self.tokenizer is not None:
            # NOTE: xlnet uses a hardcoded config.max_position_embeddings != -1 to denote no max_length
            if (
                "model_max_length" not in processor_kwargs
                and hasattr(self.config, "max_position_embeddings")
                and self.config.max_position_embeddings != -1
            ):
                self.tokenizer.model_max_length = min(
                    self.tokenizer.model_max_length, self.config.max_position_embeddings
                )

            if do_lower_case:
                # NOTE: All Transformers v5 tokenizers are fast tokenizers, but we keep the v4 branch for compatibility
                if self.tokenizer.is_fast:
                    normalizer = self.tokenizer.backend_tokenizer.normalizer
                    if not _has_lowercase(normalizer):
                        new_normalizers = [Lowercase()]
                        if isinstance(normalizer, NormalizerSequence):
                            new_normalizers += list(normalizer)
                        elif normalizer is not None:
                            new_normalizers.append(normalizer)
                        self.tokenizer.backend_tokenizer.normalizer = NormalizerSequence(new_normalizers)
                else:
                    # Some v4 Tokenizers have do_lower_case as property without a setter, and those often
                    # have a basic_tokenizer on which do_lower_case can be set.
                    try:
                        self.tokenizer.do_lower_case = do_lower_case
                    except AttributeError:
                        self.tokenizer.basic_tokenizer.do_lower_case = do_lower_case

        # Causal models require left padding so the last position is always a real token,
        # which is needed for logits_to_keep=1 and LogitScore.
        if self.transformer_task in ("text-generation", "any-to-any"):
            # A multimodal processor pads via its tokenizer and has no padding_side of its own that
            # transformers reads. Setting one anyway plants an attribute that looks authoritative but isn't.
            if hasattr(self.processor, "padding_side"):
                self.processor.padding_side = "left"
            if hasattr(self.processor, "tokenizer"):
                self.processor.tokenizer.padding_side = "left"

        # Extract message format from modality_config if provided, otherwise let InputFormatter infer it
        if modality_config is not None and "message" in modality_config:
            message_format = modality_config["message"].get("format", "auto")
        else:
            message_format = "auto"
        self.input_formatter = InputFormatter(
            model_type=self.config.model_type, message_format=message_format, processor=self.processor
        )

        if modality_config is not None:
            self.modality_config = modality_config
            if module_output_name is None:
                raise ValueError(
                    "Loading the Transformer module with a custom modality_config requires also providing "
                    "module_output_name with the name of the output feature that this module should create, "
                    'for example "token_embeddings" or "sentence_embedding".'
                )
            self.module_output_name = module_output_name
            for modality_key, params in modality_config.items():
                if not isinstance(params, dict) or "method" not in params or "method_output_name" not in params:
                    raise ValueError(
                        f"Invalid modality_config entry for {modality_key!r}: each entry must be a dict with "
                        f"'method' and 'method_output_name' keys, but got {params!r}"
                    )
        else:
            self.modality_config, self.module_output_name = self.infer_modalities(self.model, self.processor)
        logger.debug(f"Active modality config: {self.modality_config}")
        self.input_formatter.supported_modalities = list(self.modality_config.keys())

        if tokenizer_name_or_path is not None:
            logger.warning(
                "The `tokenizer_name_or_path` argument is deprecated and will be removed in a future version. "
                "Please use the same path for the model and processor."
            )
            self.model.config.tokenizer_class = self.processor.__class__.__name__

        # Evaluate whether we can skip padding
        self.unpad_inputs = unpad_inputs

        # ColBERT-style query expansion. The setter normalizes the dict, fills defaults, and
        # runs the FA2 compatibility check for non-attend expansion.
        self.query_expansion = query_expansion

    @property
    def unpad_inputs(self) -> bool | None:
        """Whether text-only inputs are concatenated without padding for faster inference.

        Non-text inputs (images, audio, video) are always padded normally.
        ``None`` auto-detects, ``False`` forces padding, ``True`` requests unpadding.
        Re-evaluates on every assignment, so it can be changed after loading::

            model = SentenceTransformer("my-model", model_kwargs={"attn_implementation": "flash_attention_2"})
            model[0].unpad_inputs = False  # Force padding for models that need it
        """
        return self._unpad_inputs

    @unpad_inputs.setter
    def unpad_inputs(self, value: bool | None) -> None:
        self._unpad_inputs = value
        if value is False:
            self.can_flatten_inputs = False
        else:
            self.can_flatten_inputs = self._can_flatten_inputs()
            if value is True and not self.can_flatten_inputs:
                logger.warning(
                    "unpad_inputs=True was set, but the prerequisites for skipping padding are not met. "
                    "Falling back to padded inputs."
                )

    @property
    def query_expansion(self) -> QueryExpansionConfig | None:
        """ColBERT-style query expansion config. See :class:`QueryExpansionConfig`. Validated and
        normalized on assignment, so re-setting after construction is safe::

            model[0].query_expansion = {"strategy": "fixed", "attend": True, "length": 32}
            model[0].query_expansion = None  # turn off
        """
        return self._query_expansion

    @query_expansion.setter
    def query_expansion(self, value: QueryExpansionConfig | None) -> None:
        expansion = _normalize_query_expansion(value)
        # ``model`` may not exist yet during very early init. Skip tokenizer + FA2 checks until it does.
        if expansion is not None and getattr(self, "model", None) is not None:
            self._validate_query_expansion_token(expansion)
            if not expansion["attend"] and self._is_flash_attention_requested():
                raise ValueError(
                    "FlashAttention-2 is incompatible with query_expansion attend=False. "
                    "FA2 strips attention_mask=0 positions, so the [MASK] expansion tokens used by MaxSim "
                    "never receive an attention update. Pass attn_implementation='sdpa' (preserves semantics) "
                    "or set attend=True (changes semantics)."
                )
        # Expansion tokenizes with max_length=length, so a smaller query_length content cap is
        # inexpressible: raise instead of silently ignoring it.
        query_length = getattr(self, "query_length", None)
        if expansion is not None and query_length is not None and query_length < expansion["length"]:
            raise ValueError(
                f"query_length={query_length} is smaller than query_expansion['length']={expansion['length']}. "
                "The expansion length is the query width floor, so a smaller query_length cap is "
                "contradictory. Drop query_length, raise it to at least the expansion length, or lower "
                "the expansion length."
            )
        self._query_expansion = expansion

    def _validate_query_expansion_token(self, expansion: QueryExpansionConfig) -> None:
        """Tokenizer-aware checks. The structural ``_normalize_query_expansion`` runs without the
        tokenizer, so this method catches the two silent-wrong-behavior cases that require it:

        * ``token=None`` falls back to ``mask_token_id``, then ``eos_token_id``. If the tokenizer
          has neither, the preprocess swap silently no-ops and the encoder sees raw pad tokens.
        * An explicit ``token`` that isn't in the tokenizer's vocabulary resolves to
          ``unk_token_id`` (or ``None``), silently inserting unk tokens at expansion positions.
        """
        if self.tokenizer is None:
            return
        strategy = expansion["strategy"]
        token = expansion["token"]
        if token is None:
            # token=None falls back to mask_token, then eos_token (PyLate's chain:
            # [MASK] for encoder models, EOS for decoder models without one).
            if self.tokenizer.mask_token_id is None and self.tokenizer.eos_token_id is None:
                raise ValueError(
                    f"query_expansion strategy={strategy!r} with token=None falls back to the "
                    "tokenizer's mask_token, then its eos_token, but this tokenizer has neither. "
                    "Set ``query_expansion={..., 'token': '<your_token>'}`` explicitly."
                )
            return
        token_id = self.tokenizer.convert_tokens_to_ids(token)
        unk_id = self.tokenizer.unk_token_id
        if token_id is None or (token_id == unk_id and token != self.tokenizer.unk_token):
            raise ValueError(
                f"query_expansion token={token!r} isn't in this tokenizer's vocabulary (resolves "
                f"to {'None' if token_id is None else 'unk_token_id'}). Check your model family's "
                "expected token: PaliGemma uses '<pad>', Qwen2/2.5/3/Omni use '<|endoftext|>', "
                "Gemma3 uses '<eos>', Idefics3 / ModernVBERT use '<end_of_utterance>'."
            )

    def _can_flatten_inputs(self) -> bool:
        """Determine whether text-only inputs can be flattened (concatenated without padding) for more efficient inference.

        When enabled, text-only inputs are concatenated into a single sequence and processed using flash
        attention's variable-length functions, eliminating padding overhead and significantly speeding up
        inference. Non-text inputs (images, audio, video) are always padded normally, even when this
        returns True.

        This requires:
        1. The ``"feature-extraction"`` or ``"fill-mask"`` task. Position-reading model heads
           (e.g. ``AutoModelForSequenceClassification``, which pools specific positions) are
           incompatible with flattened inputs, but the masked language modelling head is strictly
           position-wise, so its logits stay correct on a flattened sequence.
        2. The ``"text"`` modality must be supported by the model.
        3. All modality call methods must be ``"forward"``. ``get_..._features`` methods apply heads
           that are incompatible with flattened inputs.
        4. The ``"torch"`` backend with an attention-interface-compatible model.
        5. Flash attention with variable-length function support.

        Note: Some architectures don't work with unpadded inputs (e.g. ``qwen2_vl``). Use
        ``unpad_inputs=False`` to disable this optimization for such models.

        Returns:
            bool: True if text-only inputs can be flattened for efficient inference.
        """
        if (
            self.transformer_task not in ("feature-extraction", "fill-mask")
            or "text" not in self.modality_config
            or self.backend != "torch"
            or not getattr(self.model, "is_backend_compatible", lambda: False)()
            or any(params["method"] != "forward" for params in self.modality_config.values())
        ):
            return False

        try:
            from transformers import DataCollatorWithFlattening
            from transformers.modeling_flash_attention_utils import lazy_import_flash_attention
        except ImportError:
            logger.debug(
                "Consider upgrading to transformers >= 5.0.0 to skip padding for text-only inputs, "
                "which can significantly speed up processing."
            )
            return False

        if not self._is_flash_attention_requested():
            return False

        (_, flash_varlen_fn, *_), _ = lazy_import_flash_attention(self.config._attn_implementation)
        if flash_varlen_fn is None:
            return False

        logger.debug(
            "Using flattened inputs with flash attention variable-length functions to avoid padding overhead for text-only inputs."
        )
        self.data_collator = DataCollatorWithFlattening(
            return_seq_idx=True,  # Not always necessary, perhaps only Mamba/Bamba?
            return_flash_attn_kwargs=True,
            return_position_ids=True,  # Crucial for performance
        )
        # Ensure the flash attention keys survive forward-kwargs filtering: they are not named
        # parameters in the model's forward signature, but the model passes them through to its
        # attention layers via **kwargs. A None model_forward_params already passes everything.
        if self.model_forward_params is not None:
            self.model_forward_params |= {
                "cu_seq_lens_q",
                "cu_seq_lens_k",
                "max_length_q",
                "max_length_k",
                "seq_idx",
            }
        self._flatten_position_offset = self._infer_flatten_position_offset()
        return True

    def _infer_flatten_position_offset(self) -> int:
        # DataCollatorWithFlattening emits 0-based position_ids, but RoBERTa-family embeddings
        # (XLM-R, MPNet, ...) expect positions from padding_idx + 1. No 0-based or rotary module
        # pairs an int padding_idx with a learned position_embeddings table.
        for module in self.model.modules():
            padding_idx = getattr(module, "padding_idx", None)
            if isinstance(padding_idx, int) and isinstance(
                getattr(module, "position_embeddings", None), torch.nn.Embedding
            ):
                return padding_idx + 1
        return 0

    @property
    def max_seq_length(self) -> int | None:
        """The maximum input sequence length. Reads from the tokenizer if available, otherwise
        falls back to ``max_position_embeddings`` from the model config."""
        if self.tokenizer is not None:
            return self.tokenizer.model_max_length

        # Get text config for multi-modal models that don't have a tokenizer
        if hasattr(self.model.config, "get_text_config"):
            text_config = self.model.config.get_text_config()
        else:
            text_config = self.model.config

        if hasattr(text_config, "max_position_embeddings"):
            return text_config.max_position_embeddings
        return None

    @max_seq_length.setter
    def max_seq_length(self, value: int | None) -> None:
        """Set the maximum input sequence length. Only effective when a tokenizer is available."""
        if self.tokenizer is not None:
            self.tokenizer.model_max_length = value

    @property
    def auto_model(self) -> PreTrainedModel:
        """The underlying transformer model."""
        return self.model

    @property
    def config(self) -> PretrainedConfig:
        """The underlying model configuration."""
        return self.model.config

    @property
    def modalities(self) -> list[Modality]:
        """The list of supported input modalities (e.g. ``"text"``, ``"image"``, ``("image", "text")``)."""
        return list(self.modality_config.keys())

    @property
    def tokenizer(self) -> PreTrainedTokenizerBase | None:
        """The tokenizer, extracted from the processor. Returns ``None`` for non-text processors."""
        if isinstance(self.processor, PreTrainedTokenizerBase):
            return self.processor
        return getattr(self.processor, "tokenizer", None)

    def _chat_template_declares_task(self) -> bool:
        """Whether the active chat template (the processor's, else the tokenizer's) uses a ``task``
        variable, e.g. to append query augmentation tokens for ``task == "query"`` renders.
        Reuses the template introspection transformers' own kwarg routing is based on."""
        try:
            from transformers.utils.chat_template_utils import _get_template_variables
        except ImportError:
            # Without the introspection-based kwarg routing, an extra template kwarg is harmless.
            return True
        template = getattr(self.processor, "chat_template", None)
        if template is None and self.tokenizer is not None:
            template = getattr(self.tokenizer, "chat_template", None)
        templates = template.values() if isinstance(template, dict) else [template]
        return any(isinstance(value, str) and "task" in _get_template_variables(value) for value in templates)

    def _should_flatten_inputs(
        self,
        modality: Modality,
        processor_inputs: dict[str, Any],
        **kwargs: Any,
    ) -> bool:
        """Whether to pack variable-length text inputs into a single flat sequence (FA2 unpadding).

        Only safe for text-only inputs, since :class:`DataCollatorWithFlattening` only handles
        ``input_ids`` / ``labels``. Subclasses can override to opt out for specific tasks.
        """
        # FA2 unpadding would drop the pad positions before `preprocess` can swap them for the expansion
        # token, so fall back to padded for expansion queries.
        if kwargs.get("task") == "query" and self.query_expansion is not None:
            return False
        return self.can_flatten_inputs and (
            modality == "text"
            or (modality == "message" and self.input_formatter.is_text_only_messages(processor_inputs["message"]))
        )

    def _is_flash_attention_requested(self) -> bool:
        """Whether the loaded model was configured for FlashAttention-2 (via ``attn_implementation``)."""
        try:
            from transformers.utils.generic import is_flash_attention_requested
        except ImportError:
            return False
        attn_impl = getattr(self.config, "_attn_implementation", None)
        if attn_impl is None:
            # E.g. ONNX/OpenVINO-loaded configs never materialize _attn_implementation.
            return False
        return is_flash_attention_requested(requested_attention_implementation=attn_impl)

    def preprocess(
        self,
        inputs: Sequence[SingleInput | PairInput],
        prompt: str | None = None,
        processing_kwargs: ProcessingKwargs | None = None,
        **kwargs,
    ) -> dict[str, Any]:
        """Preprocess inputs into model-ready features.

        Args:
            inputs: List of inputs. Can contain strings, dicts with modality keys, PIL images,
                or numpy/torch arrays for audio/video.
            prompt: Optional prompt to prepend to text inputs or inject as a system message.
            processing_kwargs: Per-call overrides for the processor kwargs configured on the module
                via the ``processing_kwargs`` constructor argument. Same nested structure (modality
                keys plus ``"common"`` and ``"chat_template"``). Per-call values are merged on top of
                the instance-level kwargs with shallow per-modality merge, so individual settings
                (e.g. only ``max_length``) can be overridden without replacing the entire modality
                dict. Defaults to None.
            **kwargs: Additional keyword arguments forwarded to prompt length computation
                (e.g. ``task``). Only used when ``prompt`` is provided for text inputs.

        Returns:
            Dictionary containing preprocessed tensors with a ``modality`` key indicating the
            input type and optionally a ``prompt_length`` key for prompt-aware pooling.
        """
        if not inputs:
            return {}

        common_kwargs = {"return_tensors": "pt"}
        modality_kwargs = {
            "text": {"padding": True, "truncation": "longest_first"},
            "audio": {"padding": True},
            "image": {},
            "video": {},
        }
        if self.config.model_type == "whisper":
            # Whisper requires inputs to be exactly 30 seconds long. WhisperFeatureExtractor pads
            # to that length by default, but ST's general "audio" modality kwargs above set
            # padding=True ("longest") which would override that default. Restore "max_length".
            modality_kwargs["audio"]["padding"] = "max_length"

        # Apply the per-task max-length override for text inputs (e.g. ColBERT-style
        # query_length/document_length). A caller-supplied max_length (e.g. the trainer's
        # args.max_length via the data collator) wins over the module lengths, and caller-supplied
        # processing_kwargs still wins below.
        task = kwargs.get("task")
        max_length_override = kwargs.get("max_length")
        if max_length_override is not None:
            task_max_length = max_length_override
        elif task == "query":
            task_max_length = self.query_length
        elif task == "document":
            task_max_length = self.document_length
        else:
            task_max_length = None
        # Expansion queries manage their own max_length below (they pad to the config length).
        # Skip the generic path for them.
        expansion = self.query_expansion if task == "query" else None
        if (
            expansion is not None
            and expansion["strategy"] == "fixed"
            and max_length_override is not None
            and max_length_override != expansion["length"]
        ):
            logger.warning_once(
                f"query_expansion fixes queries to {expansion['length']} tokens, so the "
                f"max_length={max_length_override} override is ignored for query batches."
            )
        if task_max_length is not None and expansion is None:
            modality_kwargs["text"]["max_length"] = task_max_length

        effective_processing_kwargs = self._merge_processing_kwargs(processing_kwargs)
        if "common" in effective_processing_kwargs:
            common_kwargs.update(effective_processing_kwargs["common"])
        for modality_key in modality_kwargs:
            if overrides := effective_processing_kwargs.get(modality_key):  # type: ignore[arg-type]
                modality_kwargs[modality_key].update(overrides)

        # strategy='fixed': pad to the expansion length so the post-tokenization swap below has
        # pad positions to replace. Applied after the merge so processing_kwargs can't break the
        # fixed width. strategy='min' only applies the ceiling here, the width floor is enforced
        # post-tokenization.
        if expansion is not None:
            text_overrides = effective_processing_kwargs.get("text") or {}
            if expansion["strategy"] == "fixed":
                if "padding" in text_overrides or "max_length" in text_overrides:
                    logger.warning_once(
                        "processing_kwargs overrides the text padding or max_length, but query "
                        "expansion requires a fixed width: re-applying padding='max_length' with the "
                        "expansion length for query batches."
                    )
                modality_kwargs["text"]["max_length"] = expansion["length"]
                modality_kwargs["text"]["padding"] = "max_length"
            elif task_max_length is not None and "max_length" not in text_overrides:
                if task_max_length < expansion["length"]:
                    logger.warning_once(
                        f"query_expansion guarantees at least {expansion['length']} query tokens, so the "
                        f"max_length={task_max_length} override is raised to it for query batches."
                    )
                modality_kwargs["text"]["max_length"] = max(task_max_length, expansion["length"])

        chat_template_kwargs = effective_processing_kwargs.get("chat_template", {})
        # Task-aware templates can branch on the task, e.g. to append query augmentation tokens.
        # Only forwarded when declared: transformers treats unknown kwargs as processor kwargs.
        if task is not None and self._chat_template_declares_task():
            chat_template_kwargs = {"task": task, **chat_template_kwargs}

        modality, processor_inputs, extra_modality_kwargs = self.input_formatter.parse_inputs(inputs)

        # `*ForRetrieval` processors have no document treatment for text (always query-rendered):
        # warn instead of silently scoring query-formatted "documents".
        if self.transformer_task == "retrieval" and modality == "text" and task is not None and task != "query":
            logger.warning_once(
                f"This retrieval model's processor always renders text as a query, but task={task!r} was "
                "requested: these inputs get the query prefix and augmentation tokens. *ForRetrieval "
                "models expect image documents."
            )

        for modality_key, extra_kwargs in extra_modality_kwargs.items():
            modality_kwargs[modality_key].update(extra_kwargs)

        # Flatten inputs to avoid padding overhead when using flash attention variable-length functions.
        # Subclasses can override `_should_flatten_inputs` to opt out of flattening for specific tasks
        # (e.g. ColBERT-style query expansion, which repurposes padding positions as [MASK] tokens and
        # therefore can't have them dropped by FA2 unpadding).
        should_flatten = self._should_flatten_inputs(modality, processor_inputs, **kwargs)
        if should_flatten:
            del common_kwargs["return_tensors"]
            modality_kwargs["text"].pop("padding", None)
            modality_kwargs["text"]["return_attention_mask"] = False

        # Always convert to the message format if it's supported, since it's most flexible with e.g. defaults
        if "message" in self.modality_config and modality != "message":
            if expansion is not None and modality == "text":
                # The expansion length would be applied as max_length to the chat-rendered
                # sequence, right-truncating it: at small lengths every query collapses to the
                # identical template preamble.
                raise ValueError(
                    f"query_expansion strategy={expansion['strategy']!r} cannot be combined with a "
                    "chat-template (message) backbone: queries are rendered through the chat template "
                    "and then padded / truncated to the expansion length, which destroys the query "
                    "content. Chat-template backbones should own their query augmentation in the "
                    "template itself (it receives task='query'), or pin modality_config to the plain "
                    "text modality."
                )
            modality, processor_inputs = self.input_formatter.batch_to_message(modality, processor_inputs)
        elif modality not in self.modality_config:
            raise_unsupported_modality_error(inputs, modality, list(self.modality_config.keys()), "Transformer module")

        # Checked after the conversion so it covers pairs that parse_inputs already turned into messages
        if modality == "message":
            self._verify_pair_roles_supported(processor_inputs["message"], chat_template_kwargs)

        # Incorporate prompt into inputs if applicable
        prompt_length = None
        if prompt and modality == "message":
            processor_inputs["message"] = self.input_formatter.prepend_prompt_to_messages(
                processor_inputs["message"], prompt
            )
            # Models relying on the message format don't support excluding prompt tokens in mean pooling,
            # so we don't track prompt length.
        elif prompt and modality == "text":
            processor_inputs["text"] = self.input_formatter.prepend_prompt_to_texts(processor_inputs["text"], prompt)
            prompt_length = self._get_prompt_length(prompt, **kwargs)

        # Track per-sample image/video counts before the processor flattens them into single tensors.
        # Losses that minibatch VLM inputs (e.g. CachedMNRL) use these counts to slice visual tensors.
        # Only used if the Trainer updated track_media_counts to True. Not gated on self.training:
        # trainer.evaluate() computes the eval loss under model.eval() and needs the counts too.
        num_images_per_sample = None
        num_videos_per_sample = None
        if self.track_media_counts and modality == "message":
            num_images_per_sample, num_videos_per_sample = _count_media_per_sample(processor_inputs["message"])

        with suggest_extra_on_exception():
            processor_output = self._call_processor(
                modality,
                processor_inputs,
                modality_kwargs,
                common_kwargs,
                chat_template_kwargs=chat_template_kwargs,
            )

        if num_images_per_sample is not None and "image_grid_thw" in processor_output:
            processor_output["num_images_per_sample"] = torch.tensor(num_images_per_sample, dtype=torch.long)
        if num_videos_per_sample is not None and "video_grid_thw" in processor_output:
            processor_output["num_videos_per_sample"] = torch.tensor(num_videos_per_sample, dtype=torch.long)

        if should_flatten:
            # DataCollatorWithFlattening expects list[dict], but the processor returns dict[str, list].
            per_sample = [dict(zip(processor_output, values)) for values in zip(*processor_output.values())]
            processor_output = self.data_collator(per_sample)
            processor_output.pop("labels", None)
            if self._flatten_position_offset:
                processor_output["position_ids"] = processor_output["position_ids"] + self._flatten_position_offset

        processor_output["modality"] = modality
        if prompt_length is not None:
            processor_output["prompt_length"] = prompt_length

        if self.transformer_task in ("text-generation", "any-to-any"):
            self._verify_left_padding(processor_output, modality_kwargs, common_kwargs)
            processor_output["logits_to_keep"] = 1

        # ColBERT-style query expansion: swap pad positions within the width floor for the expansion
        # token id, and mark them via ``query_expansion_positions`` so downstream modules score them.
        expansion = self.query_expansion if task == "query" else None
        if expansion is not None and "attention_mask" in processor_output and "input_ids" in processor_output:
            floor = expansion["length"]
            input_ids = processor_output["input_ids"]
            attention_mask = processor_output["attention_mask"]
            # 'min' tokenizes with batch-longest padding, which can sit below the floor: widen up to
            # it. The 0-filler is overwritten by the expansion swap below.
            if input_ids.shape[1] < floor:
                missing = floor - input_ids.shape[1]
                input_ids = torch.nn.functional.pad(input_ids, (0, missing), value=0)
                attention_mask = torch.nn.functional.pad(attention_mask, (0, missing), value=0)
                processor_output["input_ids"] = input_ids
                processor_output["attention_mask"] = attention_mask
                if "token_type_ids" in processor_output:
                    processor_output["token_type_ids"] = torch.nn.functional.pad(
                        processor_output["token_type_ids"], (0, missing), value=0
                    )
            # Per position rather than per batch: 'min' rows can carry real pads beyond the floor
            # that must stay excluded from scoring. For 'fixed', width == floor: every pad counts.
            expansion_positions = (attention_mask == 0) & (
                torch.arange(input_ids.shape[1], device=attention_mask.device) < floor
            )
            # A 'fixed' row with no expansion positions was truncated to the width (silent content
            # loss). Under 'min' such rows are simply long queries.
            if expansion["strategy"] == "fixed" and not expansion_positions.any(dim=-1).all():
                logger.warning_once(
                    f"One or more queries filled the entire query_expansion length of {floor} tokens: any "
                    "content beyond it was truncated, as this model always processes queries at exactly "
                    "this width."
                )
            if self.tokenizer is not None:
                if expansion["token"] is not None:
                    expansion_id = self.tokenizer.convert_tokens_to_ids(expansion["token"])
                else:
                    # PyLate's fallback chain: [MASK] for encoder models, EOS for decoder models.
                    expansion_id = self.tokenizer.mask_token_id
                    if expansion_id is None:
                        expansion_id = self.tokenizer.eos_token_id
                if expansion_id is not None:
                    processor_output["input_ids"] = input_ids.masked_fill(expansion_positions, expansion_id)
            processor_output["query_expansion_positions"] = expansion_positions
            # Scoped to the expansion positions, not ones_like: 'min' rows can carry real trailing
            # pads. attend=False keeps the classic ColBERT "skip during forward" trick.
            if expansion["attend"]:
                processor_output["attention_mask"] = attention_mask.masked_fill(expansion_positions, 1)

        return processor_output

    def _verify_pair_roles_supported(
        self, messages_batch: list[list[dict[str, Any]]], chat_template_kwargs: dict[str, Any]
    ) -> None:
        """Raise if the chat template cannot carry the ``query``/``document`` roles a pair is mapped to.

        A template that only branches on ``system``/``user``/``assistant`` renders a pair to an empty
        prompt, so the model would score an input holding neither the query nor the document. Published
        rerankers ship a template that handles both roles, so this is reached when a backbone is used,
        or fine-tuned, without one.
        """
        # Checked before the probe, which costs a render, so batches without a pair never pay for it
        if not self.input_formatter.has_pair_roles(messages_batch):
            return
        # Drop what never reaches the Jinja render: size kwargs (tokenization only) and the ST-only restore_suffix
        probe_kwargs = {
            key: value
            for key, value in chat_template_kwargs.items()
            if key not in _APPLY_CHAT_TEMPLATE_TOP_LEVEL_KWARGS and key != "restore_suffix"
        }
        failure = self.input_formatter.pair_roles_failure(probe_kwargs)
        if failure is None:
            return
        model_name = getattr(self.config, "name_or_path", None) or "this model"
        if self.input_formatter.message_format == "structured":
            example = (
                '    <Query>: {{ (messages | selectattr("role", "eq", "query") | first).content[0].text }}\n'
                '    <Document>: {{ (messages | selectattr("role", "eq", "document") | first).content[0].text }}\n'
            )
        else:
            example = (
                '    <Query>: {{ messages | selectattr("role", "eq", "query") | map(attribute="content") | first }}\n'
                '    <Document>: {{ messages | selectattr("role", "eq", "document") | map(attribute="content") | first }}\n'
            )
        raise ValueError(
            f"The chat template of {model_name} cannot carry a 'query'/'document' pair ({failure}). Pair inputs "
            "are mapped to one message per role, so set a chat template that renders both roles, e.g.:\n"
            f"{example}"
            'via `processor_kwargs={"chat_template": ...}` when loading, `model.processor.chat_template = ...` '
            "afterwards, or a `chat_template.jinja` file alongside the model."
        )

    def _verify_left_padding(
        self,
        processor_output: dict[str, Any],
        modality_kwargs: dict[str, dict[str, Any]],
        common_kwargs: dict[str, Any],
    ) -> None:
        """Verify that causal inputs are left-padded, i.e. that every sample ends in a real token.

        Checks the produced ``attention_mask`` rather than ``padding_side``, because the side that
        actually applies can come from the tokenizer attribute or from either the ``"text"`` or
        ``"common"`` processing_kwargs bucket, and which of those wins differs per processor type and
        per call path. The output is the only signal that covers every route.
        """
        attention_mask = processor_output.get("attention_mask")
        if not isinstance(attention_mask, torch.Tensor) or attention_mask.ndim != 2:
            return
        if attention_mask[:, -1].all():
            return

        # Only reached when padding is wrong, so it's worth working out which route caused it. Either
        # bucket can be the one that wins, depending on the processor type, so blame whichever is at fault.
        kwargs_side = next(
            (
                side
                for side in (modality_kwargs["text"].get("padding_side"), common_kwargs.get("padding_side"))
                if side is not None and side != "left"
            ),
            None,
        )
        if kwargs_side is not None:
            fix = f'Your ``processing_kwargs`` set the padding side to {kwargs_side!r}. Remove it or set it to "left".'
        else:
            padder_path = "processor.tokenizer" if hasattr(self.processor, "tokenizer") else "processor"
            fix = (
                f'The module sets ``{padder_path}.padding_side`` to "left" on load, so it was changed afterwards. '
                f'Restore it with ``{padder_path}.padding_side = "left"``.'
            )
        raise ValueError(
            "At least one sample ends in a padding token, but causal models require left padding so that "
            "the last token position is always a real token. "
            "This is needed for efficient logit computation (logits_to_keep=1) and for LogitScore. " + fix
        )

    def _merge_processing_kwargs(self, per_call: ProcessingKwargs | None) -> ProcessingKwargs:
        """Merge per-call ``processing_kwargs`` on top of the instance-level ``self.processing_kwargs``.

        The merge is shallow per top-level key: for each modality (or ``"common"`` /
        ``"chat_template"``), per-call entries override individual instance entries, leaving
        non-overridden settings intact. Unknown top-level keys in ``per_call`` are warned about
        and ignored.
        """
        if not per_call:
            return self.processing_kwargs

        unknown_keys = set(per_call) - self._VALID_PROCESSING_KWARGS_KEYS
        if unknown_keys:
            logger.warning_once(
                f"Unknown keys in per-call `processing_kwargs`: {unknown_keys}. "
                f"Valid keys are: {sorted(self._VALID_PROCESSING_KWARGS_KEYS)}. Unknown keys will be ignored."
            )

        merged: ProcessingKwargs = {}
        for key in self._VALID_PROCESSING_KWARGS_KEYS:
            instance_value = self.processing_kwargs.get(key, {})
            per_call_value = per_call.get(key, {})
            if instance_value or per_call_value:
                merged[key] = {**instance_value, **per_call_value}
        return merged

    def forward(self, features: dict[str, Any], **kwargs) -> dict[str, Any]:
        """Forward pass through the transformer model.

        Dispatches to the appropriate model method based on the ``modality`` key in ``features``
        and writes the result into ``features[self.module_output_name]``.

        Args:
            features: Input features dictionary produced by :meth:`preprocess`. Must contain the
                keys expected by the underlying model (e.g. ``input_ids``, ``pixel_values``, etc.).
                A ``modality`` key selects the modality config to use. Defaults to ``"text"``
                when absent.
            **kwargs: Additional keyword arguments forwarded to the model method (override features).

        Returns:
            The updated ``features`` dict with the model output stored under ``self.module_output_name``
            (e.g. ``token_embeddings``, ``sentence_embedding``, ``scores``, or ``causal_logits``).
            May also include ``all_layer_embeddings`` if ``output_hidden_states`` is enabled.
        """

        modality_name: Modality = features.get("modality", "text")
        modality_params = self.modality_config[modality_name]
        method_name = modality_params["method"]
        method_output_name = modality_params["method_output_name"]
        if isinstance(method_output_name, str):
            method_output_name = (method_output_name,)
        elif isinstance(method_output_name, list):
            method_output_name = tuple(method_output_name)

        # kwargs override features
        all_kwargs = {**features, **kwargs, "return_dict": True}
        model_method = getattr(self.model, method_name, None)
        if model_method is None:
            raise ValueError(f"Model does not have the requested '{method_name}' method")

        if method_name == "forward":
            if self.model_forward_params is None:
                # forward accepts **kwargs: drop ST's own bookkeeping, unless the forward names it and
                # so wants to condition on it (e.g. a custom model with task adapters keyed on ``task``).
                filtered_kwargs = {
                    key: value
                    for key, value in all_kwargs.items()
                    if key not in _NON_MODEL_FEATURE_KEYS or key in self._declared_forward_params
                }
            else:
                filtered_kwargs = {key: value for key, value in all_kwargs.items() if key in self.model_forward_params}
        else:
            method_params = self._method_signature_cache.get(method_name)
            if method_params is None:
                method_params = set(inspect.signature(model_method).parameters)
                self._method_signature_cache[method_name] = method_params
            filtered_kwargs = {key: value for key, value in all_kwargs.items() if key in method_params}

        # Auto-enable output_hidden_states when the output path traverses hidden_states
        if method_output_name is not None and "hidden_states" in method_output_name:
            filtered_kwargs["output_hidden_states"] = True

        model_output = model_method(**filtered_kwargs)

        embedding = model_output
        if method_output_name is not None:
            for output_key in method_output_name:
                try:
                    embedding = embedding[output_key]
                except (KeyError, TypeError):
                    # Some models (e.g. chinese_clip) only expose output fields via attribute access,
                    # not dictionary-style indexing. See https://github.com/huggingface/transformers/issues/44079
                    try:
                        embedding = getattr(embedding, output_key)
                    except AttributeError:
                        raise AttributeError(
                            f"Could not access output key {output_key!r} via indexing or attribute access "
                            f"on {type(embedding).__name__}."
                        )

        if embedding.ndim == 4:
            # Some image models return (batch_size, num_channels, height, width) instead of (batch_size, seq_len, hidden_size)
            # We flatten the height and width dimensions and transpose to get (batch_size, height*width, num_channels)
            # which a subsequent Pooling layer can handle to remove the height*width dimension
            embedding = embedding.flatten(2).transpose(1, 2)

        features[self.module_output_name] = embedding

        # If the AutoModel is wrapped with a PeftModel(ForFeatureExtraction), then it may have added virtual tokens
        # We need to extend the attention mask to include these virtual tokens, or the pooling will fail
        if "input_ids" in features and "attention_mask" in features and is_peft_available():
            from peft import PeftModel

            if isinstance(self.model, PeftModel) and self.model.active_peft_config.is_prompt_learning:
                batch_size = features["input_ids"].shape[0]
                attention_mask = features["attention_mask"]
                prefix_attention_mask = torch.ones(
                    batch_size, self.model.active_peft_config.num_virtual_tokens, device=attention_mask.device
                )
                features["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1)

        if (
            hasattr(self.model.config, "output_hidden_states")
            and self.model.config.output_hidden_states
            and "hidden_states" in model_output
        ):
            features["all_layer_embeddings"] = model_output["hidden_states"]

        return features

    def get_embedding_dimension(self) -> int:
        """Get the output embedding dimension from the transformer model.

        Returns:
            int: The hidden dimension size of the model's embeddings.

        Raises:
            ValueError: If the embedding dimension cannot be determined from the model config.
        """
        # Edge case for timm models
        if isinstance(self.model.config, TimmWrapperConfig):
            return self.model.config.num_features

        # `*ForRetrieval` heads project down to `config.embedding_dim`, not the backbone hidden size.
        if self.transformer_task == "retrieval" and getattr(self.model.config, "embedding_dim", None) is not None:
            return self.model.config.embedding_dim

        def get_hidden_size_from_config(config):
            # If we're directly outputting sentence embeddings from the transformer (e.g., using the pooler output),
            # then we should check for projection_dim first, as that's likely the dimension of the sentence embeddings
            # after a projection layer
            if hasattr(config, "projection_dim") and self.module_output_name == "sentence_embedding":
                return config.projection_dim

            if hasattr(config, "hidden_size"):
                return config.hidden_size
            for attr_name in ("neck_hidden_sizes", "hidden_sizes", "embed_dims"):
                if hasattr(config, attr_name):
                    value = getattr(config, attr_name)
                    if isinstance(value, list):
                        if value:
                            return value[-1]
                    else:
                        return value
            if hasattr(config, "hidden_dim"):
                return config.hidden_dim
            return None

        if (hidden_size := get_hidden_size_from_config(self.model.config)) is not None:
            return hidden_size

        # Text config hidden size has priority
        if hasattr(self.model.config, "text_config"):
            if (hidden_size := get_hidden_size_from_config(self.model.config.text_config)) is not None:
                return hidden_size

        # Afterwards we check all sub-configs
        if hasattr(self.model.config, "sub_configs"):
            for sub_config_name in self.model.config.sub_configs.keys():
                sub_config = getattr(self.model.config, sub_config_name)
                if (hidden_size := get_hidden_size_from_config(sub_config)) is not None:
                    return hidden_size

        raise ValueError(
            f"Could not determine embedding dimension from model config. Config type: {type(self.model.config).__name__}."
        )

    def _call_processor(
        self,
        modality: Modality,
        processor_inputs: dict[str, list],
        modality_kwargs: dict[str, dict[str, Any]],
        common_kwargs: dict[str, Any],
        chat_template_kwargs: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Call the appropriate processor with the correct arguments.

        Dispatches based on the processor type and modality:

        1. **Message modality**: delegates to :meth:`_process_chat_messages`.
        2. **Multi-modal processor** (:class:`ProcessorMixin`): delegates to
           :meth:`_call_multimodal_processor`, which handles both legacy (flat kwargs) and
           transformers v5 (per-modality kwargs) calling conventions.
        3. **Single-modality processor** (tokenizer, feature extractor, image/video processor):
           delegates to :meth:`_call_single_modality_processor`, which matches the processor
           type and passes the primary input as a positional argument when available.

        Args:
            modality: The modality or tuple of modalities being processed.
            processor_inputs: Dictionary of processor argument names to lists of values.
            modality_kwargs: Per-modality configuration kwargs (keys: ``"text"``, ``"image"``,
                ``"audio"``, ``"video"``).
            common_kwargs: Common kwargs passed to all processor calls (e.g. ``padding``,
                ``return_tensors``).
            chat_template_kwargs: Kwargs forwarded to ``apply_chat_template`` when the modality
                is ``"message"``. Ignored for non-message modalities.

        Returns:
            Processor output dictionary.
        """
        if modality == "message":
            return self._process_chat_messages(
                processor_inputs["message"],
                modality_kwargs,
                common_kwargs,
                chat_template_kwargs=chat_template_kwargs,
            )

        if isinstance(self.processor, ProcessorMixin):
            return self._call_multimodal_processor(modality, processor_inputs, modality_kwargs, common_kwargs)

        return self._call_single_modality_processor(modality, processor_inputs, modality_kwargs, common_kwargs)

    def _call_multimodal_processor(
        self,
        modality: Modality,
        processor_inputs: dict[str, list],
        modality_kwargs: dict[str, dict[str, Any]],
        common_kwargs: dict[str, Any],
    ) -> dict[str, Any]:
        """Call a :class:`ProcessorMixin` processor, handling both legacy and v5 calling conventions."""
        # `*ForRetrieval` processors reject raw image strings, and the retrieval path skips the chat
        # template that resolves them elsewhere: load URLs / paths / base64 here (also converts to RGB).
        if self.transformer_task == "retrieval" and "image" in processor_inputs:
            from transformers.image_utils import load_image

            processor_inputs = {
                **processor_inputs,
                "image": [
                    load_image(image) if isinstance(image, str) else image for image in processor_inputs["image"]
                ],
            }

        # Convert modality keys to processor argument names (e.g., "image" -> "images")
        processor_inputs = {MODALITY_TO_PROCESSOR_ARG.get(key, key): value for key, value in processor_inputs.items()}

        # Some transformers processors are still outdated, and don't accept common_kwargs, etc.
        if (
            self.config.model_type in {"clipseg", "whisper", "sam3"}
            or not _TRANSFORMERS_PROCESSOR_SUPPORTS_MODALITY_KWARGS
        ):
            # Check against the only valid multimodal modality for these architectures
            if modality == ("audio", "text"):
                # Audio must have priority for whisper, to correctly set padding to max_length
                kwargs = {**modality_kwargs["text"], **modality_kwargs["audio"]}
            else:
                kwargs = modality_kwargs[modality]
            return self.processor(**processor_inputs, **kwargs, **common_kwargs)

        # This is the much cleaner transformers v5 approach
        return self.processor(
            **processor_inputs,
            text_kwargs=modality_kwargs["text"],
            images_kwargs=modality_kwargs["image"],
            audio_kwargs=modality_kwargs["audio"],
            videos_kwargs=modality_kwargs["video"],
            common_kwargs=common_kwargs,
        )

    def _call_single_modality_processor(
        self,
        modality: Modality,
        processor_inputs: dict[str, list],
        modality_kwargs: dict[str, dict[str, Any]],
        common_kwargs: dict[str, Any],
    ) -> dict[str, Any]:
        """Call a single-modality processor (tokenizer, feature extractor, image/video processor)."""
        # Check in order: text, audio, video, image (video before image due to inheritance)
        processor_type_checks = [
            ("text", PreTrainedTokenizerBase, modality_kwargs["text"]),
            ("audio", FeatureExtractionMixin, modality_kwargs["audio"]),
            ("video", BaseVideoProcessor, modality_kwargs["video"]),
            ("image", ImageProcessingMixin, modality_kwargs["image"]),
        ]

        for modality_type, processor_class, type_kwargs in processor_type_checks:
            if not isinstance(self.processor, processor_class):
                continue

            call_kwargs = {**type_kwargs, **common_kwargs}

            # Tokenizers and feature extractors expect the primary input as the first positional arg
            if modality_type in processor_inputs:
                primary_input = processor_inputs.pop(modality_type)
                return self.processor(primary_input, **processor_inputs, **call_kwargs)
            return self.processor(**processor_inputs, **call_kwargs)

        raise RuntimeError(
            f"Could not determine how to call processor of type {type(self.processor).__name__} "
            f"for modality '{format_modality(modality)}'"
        )

    def _process_chat_messages(
        self,
        messages: list[list[MessageInput]],
        modality_kwargs: dict[str, dict[str, Any]],
        common_kwargs: dict[str, Any],
        chat_template_kwargs: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Process chat messages using the processor's chat template.

        ``chat_template_kwargs`` is forwarded to ``apply_chat_template``, except the Sentence Transformers
        ``restore_suffix`` flag, which only toggles the truncation suffix restore.
        """
        if "message" not in self.modality_config:
            raise ValueError(
                f"The model does not support 'message' modality, but the input looks like a chat message. "
                f"Supported modalities: {list(self.modality_config.keys())}"
            )

        chat_template_kwargs = dict(chat_template_kwargs or {})

        # `restore_suffix` is a Sentence Transformers flag, not an apply_chat_template arg, so pop it (default on).
        restore_suffix = chat_template_kwargs.pop("restore_suffix", True)

        # Truncation/max_length may ride in any kwarg bucket (chat-template wins). Tokenizers also truncate
        # when max_length is set while truncation was left unset, so only an explicit opt-out disables the
        # restore.
        text_kwargs = modality_kwargs["text"]
        truncation = chat_template_kwargs.get(
            "truncation", text_kwargs.get("truncation", common_kwargs.get("truncation"))
        )
        max_length = chat_template_kwargs.get(
            "max_length", text_kwargs.get("max_length", common_kwargs.get("max_length"))
        )
        truncation_active = bool(truncation) or (truncation is None and max_length is not None)
        restore_chat_template_suffix = restore_suffix and truncation_active and truncation != "do_not_truncate"

        features = self._apply_chat_template(messages, modality_kwargs, common_kwargs, chat_template_kwargs)

        if restore_chat_template_suffix:
            if max_length is None:
                max_length = self.max_seq_length
            self._restore_chat_template_suffix(
                features, messages, chat_template_kwargs, text_kwargs, common_kwargs, max_length
            )
        return features

    def _apply_chat_template(
        self,
        messages: list[list[MessageInput]],
        modality_kwargs: dict[str, dict[str, Any]],
        common_kwargs: dict[str, Any],
        chat_template_kwargs: dict[str, Any],
    ) -> dict[str, Any]:
        """Call ``apply_chat_template`` with the kwarg routing the processor type expects: nested kwargs for
        a ProcessorMixin, flat with the size kwargs hoisted to the top level for a bare tokenizer. Both the
        real batch and the suffix skeleton renders go through this (the given dicts are never mutated).
        """
        # Fold size kwargs from the chat-template bucket into the text kwargs (chat wins), else
        # apply_chat_template gets them twice (bare tokenizer) or misroutes them (ProcessorMixin).
        if size_keys := _APPLY_CHAT_TEMPLATE_TOP_LEVEL_KWARGS & chat_template_kwargs.keys():
            chat_template_kwargs = dict(chat_template_kwargs)
            folded = {key: chat_template_kwargs.pop(key) for key in size_keys}
            modality_kwargs = {**modality_kwargs, "text": {**modality_kwargs["text"], **folded}}

        # PaliGemmaProcessor refuses text-only inputs. Fall back to the tokenizer's apply_chat_template
        # (same token ids). Extend the isinstance check if other processors share the constraint.
        if (
            isinstance(self.processor, PaliGemmaProcessor)
            and self.tokenizer is not None
            and self.input_formatter.is_text_only_messages(messages)
        ):
            top_level_kwarg_names = {"padding", "truncation", "max_length", "return_tensors"}
            top_level_kwargs = {key: common_kwargs.pop(key) for key in top_level_kwarg_names & common_kwargs.keys()}
            top_level_kwargs |= {
                key: modality_kwargs["text"].pop(key) for key in top_level_kwarg_names & modality_kwargs["text"].keys()
            }
            return self.tokenizer.apply_chat_template(
                messages,
                tokenize=True,
                return_dict=True,
                **top_level_kwargs,
                **modality_kwargs["text"],
                **chat_template_kwargs,
            )

        if isinstance(self.processor, ProcessorMixin):
            # Transformers v5.4.0 prefers us to pass processor_kwargs as a single dict, but there's still some top level
            # kwargs that need to be hoisted out for backwards compatibility.
            if _TRANSFORMERS_APPLY_CHAT_TEMPLATE_RECOMMENDS_PROCESSOR_KWARGS:
                return self.processor.apply_chat_template(
                    messages,
                    tokenize=True,
                    return_dict=True,
                    return_tensors=common_kwargs.get("return_tensors"),
                    load_audio_from_video=modality_kwargs["video"].get("load_audio_from_video", False),
                    processor_kwargs={
                        "text_kwargs": modality_kwargs["text"],
                        "images_kwargs": modality_kwargs["image"],
                        "audio_kwargs": modality_kwargs["audio"],
                        "videos_kwargs": modality_kwargs["video"],
                        "common_kwargs": common_kwargs,
                    },
                    **chat_template_kwargs,
                )
            return self.processor.apply_chat_template(
                messages,
                tokenize=True,
                return_dict=True,
                text_kwargs=modality_kwargs["text"],
                images_kwargs=modality_kwargs["image"],
                audio_kwargs=modality_kwargs["audio"],
                videos_kwargs=modality_kwargs["video"],
                common_kwargs=common_kwargs,
                **chat_template_kwargs,
            )

        # A bare tokenizer's apply_chat_template wants padding/truncation/max_length/return_tensors as
        # top-level kwargs, not nested in tokenizer_kwargs or common_kwargs, so we hoist them out of copies.
        text_kwargs = dict(modality_kwargs["text"])
        common_kwargs = dict(common_kwargs)
        top_level_kwargs = {
            key: common_kwargs.pop(key) for key in _APPLY_CHAT_TEMPLATE_TOP_LEVEL_KWARGS & common_kwargs.keys()
        }
        top_level_kwargs |= {
            key: text_kwargs.pop(key) for key in _APPLY_CHAT_TEMPLATE_TOP_LEVEL_KWARGS & text_kwargs.keys()
        }
        return self.processor.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            tokenizer_kwargs=text_kwargs,
            common_kwargs=common_kwargs,
            **top_level_kwargs,
            **chat_template_kwargs,
        )

    def _restore_chat_template_suffix(
        self,
        features: dict[str, Any],
        messages: list[list[MessageInput]],
        chat_template_kwargs: dict[str, Any],
        tokenizer_kwargs: dict[str, Any],
        common_kwargs: dict[str, Any] | None = None,
        max_length: int | None = None,
    ) -> None:
        """Restore the trailing chat-template suffix that truncation may have removed (in place).

        Chat templates append fixed tokens after the content (e.g. an assistant prefill or EOS) that
        right-truncation cuts off first. That is harmless for mean/CLS pooling but breaks models that read
        the final position (causal-LM rerankers, last-token pooling), so the suffix is written back over the
        final real tokens of every truncated row. This must happen post-hoc: the tokenizer truncates the
        rendered string blind to template structure, with no upstream option to trim only the content.

        Only rows whose real length reached ``max_length`` can have been truncated, so shorter rows are
        skipped. The suffix is derived (and cached) per row, handling batches that mix role layouts.
        """
        input_ids = features.get("input_ids")
        # Anything but torch tensors and the flattening path's token lists (e.g. numpy) is left untouched.
        if not isinstance(input_ids, (torch.Tensor, list)):
            return

        # An unknown max_length (None) becomes 0, so the gates below never skip a row.
        max_length = max_length or 0

        is_tensor = isinstance(input_ids, torch.Tensor)
        if is_tensor:
            # unsqueeze returns a view, so the in-place writes below propagate to features["input_ids"]
            ids = input_ids if input_ids.dim() == 2 else input_ids.unsqueeze(0)
            width = ids.shape[-1]
            # No row can have reached max_length when even the padded width is below it.
            if width < max_length:
                return
            # The mask gives each row's real length and padding side, so the suffix lands on real tokens.
            # Without one, real tokens can't be told apart from padding and nothing can be restored safely.
            attention_mask = features.get("attention_mask")
            if not isinstance(attention_mask, torch.Tensor) or attention_mask.numel() == 0:
                return
            am = attention_mask if attention_mask.dim() == 2 else attention_mask.unsqueeze(0)
            real_lengths = am.long().sum(dim=1).tolist()
            left_padded = (am[:, 0] == 0).tolist()
            n_rows = ids.shape[0]
        else:
            # The flattening path (see preprocess) drops return_tensors, giving ragged unpadded token lists.
            rows = input_ids
            real_lengths = [len(row) for row in rows]
            n_rows = len(rows)

        for index in range(min(n_rows, len(messages))):
            # Truncation cuts to exactly max_length, so shorter rows kept their tail. Real lengths stay
            # correct when pad_to_multiple_of pads a truncated row past max_length.
            if real_lengths[index] < max_length:
                continue
            suffix = self._chat_template_suffix_ids(
                messages[index], chat_template_kwargs, tokenizer_kwargs, common_kwargs
            )
            keep = min(len(suffix), real_lengths[index])
            if keep == 0:
                continue
            if is_tensor:
                end = width if left_padded[index] else real_lengths[index]
                ids[index, end - keep : end] = torch.as_tensor(suffix[-keep:], dtype=ids.dtype, device=ids.device)
            else:
                rows[index][-keep:] = suffix[-keep:]

    def _chat_template_suffix_ids(
        self,
        conversation: list[MessageInput],
        chat_template_kwargs: dict[str, Any],
        tokenizer_kwargs: dict[str, Any],
        common_kwargs: dict[str, Any] | None = None,
    ) -> list[int]:
        """Derive one conversation's trailing chat-template suffix (cached).

        The conversation is rendered twice with different fillers: the longest common token suffix of the
        renders is the template tail, so no per-model token count is hard-coded. ``system`` messages are
        kept verbatim so a prompt the template renders in the tail is captured too.
        """
        # Media rows are skipped: cutting placeholder tokens would desync them from their media features.
        # Message shapes the check can't inspect are skipped too rather than failing tokenization.
        try:
            if not conversation or not self.input_formatter.is_text_only_messages([conversation]):
                return []
        except Exception as exc:
            logger.debug(f"Could not derive the chat-template suffix: {exc}")
            return []
        # Size/output kwargs in any bucket must not reach the untruncated skeleton renders or the cache key.
        chat_template_kwargs, tokenizer_kwargs, common_kwargs = (
            {key: value for key, value in kwargs.items() if key not in _APPLY_CHAT_TEMPLATE_TOP_LEVEL_KWARGS}
            for kwargs in (chat_template_kwargs, tokenizer_kwargs, common_kwargs or {})
        )
        # The skeleton key covers everything the renders see except the fillered content (roles, system
        # prompts, fields like ``name``), so conversations that render differently can't share an entry.
        # repr() keeps the key hashable even when a value is a list (e.g. a ``tools`` arg).
        cache = self._chat_template_suffix_cache
        cache_key = (
            repr(sorted(chat_template_kwargs.items())),
            repr(sorted(tokenizer_kwargs.items())),
            repr(sorted(common_kwargs.items())),
            repr(self._message_skeleton(conversation, "")),
        )
        if cache_key in cache:
            return cache[cache_key]

        suffix: list[int] = []
        try:
            # The fillers tokenize differently, so the common token suffix is template tail, never content.
            first = self._render_message_skeleton(
                conversation, "0", chat_template_kwargs, tokenizer_kwargs, common_kwargs
            )
            second = self._render_message_skeleton(
                conversation, "1 2 3 4", chat_template_kwargs, tokenizer_kwargs, common_kwargs
            )
            count = 0
            for left, right in zip(reversed(first), reversed(second)):
                if left != right:
                    break
                count += 1
            # count == min(len) means the renders never diverged (no variable content): don't trust the run.
            if 0 < count < min(len(first), len(second)):
                suffix = first[len(first) - count :]
        except Exception as exc:  # never let suffix derivation break tokenization
            logger.debug(f"Could not derive the chat-template suffix: {exc}")

        # Cache the result (even an empty one) and evict the oldest entry past the cap. Concurrent encode()
        # calls can race on the shared dict, so the eviction is best effort.
        cache[cache_key] = suffix
        if len(cache) > _CHAT_TEMPLATE_SUFFIX_CACHE_SIZE:
            try:
                cache.pop(next(iter(cache)))
            except (KeyError, RuntimeError):
                pass
        return suffix

    def _message_skeleton(self, conversation: list[MessageInput], filler: str) -> list[MessageInput]:
        """Replace each non-system message content with ``filler``, keeping everything else verbatim.

        Keeping system messages whole lets the derivation capture a prompt rendered in the tail. Fillering
        the rest keeps the renders tiny regardless of input length.
        """
        skeleton = []
        for message in conversation:
            if not isinstance(message, dict) or message.get("role") == "system":
                skeleton.append(message)
                continue
            new_content = [{"type": "text", "text": filler}] if isinstance(message.get("content"), list) else filler
            skeleton.append({**message, "content": new_content})
        return skeleton

    def _render_message_skeleton(
        self,
        conversation: list[MessageInput],
        filler: str,
        chat_template_kwargs: dict[str, Any],
        tokenizer_kwargs: dict[str, Any],
        common_kwargs: dict[str, Any],
    ) -> list[int]:
        """Tokenize one conversation's skeleton (non-system content replaced by ``filler``), untruncated.

        The render shares ``_apply_chat_template`` with the real batch, so the suffix tokenizes exactly like
        the rows it is restored onto, just without padding, truncation, or tensors.
        """
        skeleton = self._message_skeleton(conversation, filler)
        output = self._apply_chat_template(
            [skeleton],
            modality_kwargs={
                "text": {**tokenizer_kwargs, "padding": False, "truncation": False},
                "audio": {},
                "image": {},
                "video": {},
            },
            common_kwargs=common_kwargs,
            chat_template_kwargs=chat_template_kwargs,
        )
        # Coerce to a plain list and unwrap the single-conversation batch (some processors return tensors).
        ids = output["input_ids"]
        if isinstance(ids, torch.Tensor):
            ids = ids.tolist()
        if ids and isinstance(ids[0], list):
            ids = ids[0]
        return ids

    def _get_prompt_length(self, prompt: str, **kwargs) -> int | None:
        """Return the length of the prompt in tokens, excluding any trailing special token.

        Returns None if the processor does not produce ``input_ids``.
        """
        # TODO: Perhaps mirror the _chat_template_suffix_ids implementation with 2 forwards and then checking the common prefix
        # The prompt is measured under plain tokenization: with the task kept, query expansion would pad
        # the lone prompt to the expansion length and report that width instead of the prompt's own size.
        # Task is also excluded from the cache key so the same prompt is measured once across tasks.
        kwargs = {key: value for key, value in kwargs.items() if key != "task"}
        cache_key = (prompt, *sorted(kwargs.items()))
        if cache_key in self._prompt_length_mapping:
            return self._prompt_length_mapping[cache_key]

        tokenized_prompt = self.preprocess([prompt], **kwargs)
        if "input_ids" not in tokenized_prompt:
            self._prompt_length_mapping[cache_key] = None
            return None
        prompt_length = tokenized_prompt["input_ids"].shape[-1]
        # If the tokenizer adds a trailing special token (EOS, SEP, etc.), exclude it from the prompt length
        tokenizer = self.tokenizer
        last_token = tokenized_prompt["input_ids"][..., -1].item()
        if tokenizer is not None and hasattr(tokenizer, "all_special_ids") and last_token in tokenizer.all_special_ids:
            prompt_length -= 1
        self._prompt_length_mapping[cache_key] = prompt_length
        return prompt_length

    def _load_config(
        self, model_name_or_path: str, backend: str, config_kwargs: dict[str, Any]
    ) -> tuple[PeftConfig | PretrainedConfig, bool]:
        """Loads the transformers or PEFT configuration

        Args:
            model_name_or_path (str): The model name on Hugging Face (e.g. 'sentence-transformers/all-MiniLM-L6-v2')
                or the path to a local model directory.
            backend (str): The backend used for model inference. Can be `torch`, `onnx`, or `openvino`.
            config_kwargs (dict[str, Any]): Keyword arguments passed to the Hugging Face Transformers config.

        Returns:
            tuple[PeftConfig | PretrainedConfig, bool]: The model configuration and a boolean indicating whether the model is a PEFT model.
        """
        adapter_config_file = find_adapter_config_file(
            model_name_or_path,
            cache_dir=config_kwargs.get("cache_dir"),
            token=config_kwargs.get("token"),
            revision=config_kwargs.get("revision"),
            subfolder=config_kwargs.get("subfolder", ""),
            local_files_only=config_kwargs.get("local_files_only", False),
        )
        if adapter_config_file is not None:
            if backend != "torch":
                # TODO: Consider following these steps automatically so we can load PEFT models with other backends
                raise ValueError(
                    "PEFT models can currently only be loaded with the `torch` backend. "
                    'To use other backends, load the model with `backend="torch"`, call `model.transformers_model.merge_and_unload()`, '
                    "save that model with `model.save_pretrained()` and then load the model with the desired backend."
                )
            if not is_peft_available():
                raise ImportError(
                    "Loading a PEFT model requires installing the `peft` package. You can install it via `pip install peft`."
                )
            from peft import PeftConfig

            return PeftConfig.from_pretrained(model_name_or_path, **config_kwargs), True

        return AutoConfig.from_pretrained(model_name_or_path, **config_kwargs), False

    @staticmethod
    def _warn_on_unsupported_attention_config(config: PeftConfig | PretrainedConfig) -> None:
        """Warn if the config requests bidirectional attention settings not supported by the installed transformers version."""
        if not isinstance(config, PretrainedConfig):
            return

        configs_to_check: list[PretrainedConfig] = [config]
        if hasattr(config, "sub_configs"):
            for sub_config_name in config.sub_configs.keys():
                sub_config = getattr(config, sub_config_name, None)
                if isinstance(sub_config, PretrainedConfig):
                    configs_to_check.append(sub_config)

        def warn_if_unsupported(param: str, expected_value: bool, is_supported: bool, min_version: str) -> None:
            if is_supported or not any(getattr(cfg, param, None) is expected_value for cfg in configs_to_check):
                return
            logger.warning(
                f"The model config specifies `{param}={expected_value}`, but the installed "
                f"`transformers` version ({transformers_version}) may not support this parameter, "
                "in which case it may be silently ignored and the model will fall back to its default attention "
                "behavior (often causal), likely producing degraded model outputs. Consider upgrading with "
                f'`pip install -U "transformers>={min_version}"`.'
            )

        warn_if_unsupported(
            "use_bidirectional_attention", True, _TRANSFORMERS_SUPPORTS_USE_BIDIRECTIONAL_ATTENTION, "4.56.2"
        )
        warn_if_unsupported("is_causal", False, _TRANSFORMERS_SUPPORTS_IS_CAUSAL_FALSE, "5.2.0")

    def _load_model(
        self,
        model_name_or_path: str,
        transformer_task: TransformerTask,
        config: PeftConfig | PretrainedConfig,
        backend: str,
        is_peft_model: bool,
        **model_kwargs,
    ) -> PreTrainedModel:
        """Loads the transformers or PEFT model into the `auto_model` attribute

        Args:
            model_name_or_path (str): The model name on Hugging Face (e.g. 'sentence-transformers/all-MiniLM-L6-v2')
                or the path to a local model directory.
            config ("PeftConfig" | PretrainedConfig): The model configuration.
            backend (str): The backend used for model inference. Can be `torch`, `onnx`, or `openvino`.
            is_peft_model (bool): Whether the model is a PEFT model.
            model_kwargs (dict[str, Any]): Keyword arguments passed to the Hugging Face Transformers model.
        """
        if backend == "torch":
            # When loading a PEFT model, we load the base model first. The revision
            # (e.g. "main") refers to the adapter checkpoint, not the base model, so
            # we must not pass it to the base model's from_pretrained.
            if is_peft_model:
                model_kwargs.pop("revision", None)

            if transformer_task == "feature-extraction":
                model = self._load_encoder_only_model(model_name_or_path, config, **model_kwargs)
                if model is not None:
                    return model

            if transformer_task == "retrieval":
                # No AutoModelForRetrieval exists: resolve the concrete `*ForRetrieval` class per checkpoint.
                model_cls = _resolve_retrieval_model_class(config)
            else:
                model_cls = TRANSFORMER_TASK_TO_AUTO_MODEL[transformer_task]
            return model_cls.from_pretrained(model_name_or_path, config=config, **model_kwargs)
        elif backend == "onnx":
            return load_onnx_model(
                model_name_or_path=model_name_or_path,
                config=config,
                task_name=transformer_task,
                **model_kwargs,
            )
        elif backend == "openvino":
            return load_openvino_model(
                model_name_or_path=model_name_or_path,
                config=config,
                task_name=transformer_task,
                **model_kwargs,
            )
        else:
            raise ValueError(f"Unsupported backend '{backend}'. `backend` should be `torch`, `onnx`, or `openvino`.")

    def _load_encoder_only_model(
        self,
        model_name_or_path: str,
        config: PretrainedConfig,
        **model_kwargs,
    ) -> PreTrainedModel | None:
        """Load encoder-only variants for encoder-decoder architectures.

        Checks :data:`_ENCODER_ONLY_MODELS` for standard mappings and handles a few special cases
        (T5Gemma, T5Gemma2) that require extra configuration before loading.

        Returns the loaded model, or None if the config doesn't match any encoder-only architecture.
        """

        def _load_encoder(model_cls, load_config=None, **extra_class_attrs):
            with set_temporary_class_attrs(
                model_cls, _keys_to_ignore_on_load_unexpected=["decoder.*"], **extra_class_attrs
            ):
                return model_cls.from_pretrained(model_name_or_path, config=load_config or config, **model_kwargs)

        # Special cases that need extra handling before/during loading
        if isinstance(config, T5GemmaConfig):
            from transformers import T5GemmaEncoderModel

            config.is_encoder_decoder = False
            return _load_encoder(T5GemmaEncoderModel)

        if isinstance(config, T5Gemma2Config):
            from transformers.models.t5gemma2.modeling_t5gemma2 import T5Gemma2Encoder

            # T5Gemma2Encoder expects the encoder sub-config, not the full composite config
            return _load_encoder(T5Gemma2Encoder, load_config=config.encoder, base_model_prefix="model.encoder")

        # Standard encoder-only models from the registry
        for config_cls, module_path, class_name in _ENCODER_ONLY_MODELS:
            if isinstance(config, config_cls):
                encoder_cls = getattr(importlib.import_module(module_path), class_name)
                return _load_encoder(encoder_cls)

        return None

    def infer_modalities(
        self,
        model: PreTrainedModel,
        processor: ProcessorMixin
        | PreTrainedTokenizerBase
        | FeatureExtractionMixin
        | BaseVideoProcessor
        | ImageProcessingMixin,
    ) -> tuple[ModalityConfig, str]:
        """Infer the modality configuration and module output name from the model and processor.

        First checks :meth:`infer_modalities_edge_cases` for hard-coded overrides, then falls back
        to general inference based on the processor type and model forward signature.
        """
        default_modality_config, default_module_output_name = TRANSFORMER_TASK_DEFAULTS[self.transformer_task]
        default_method_output_name = default_modality_config["text"]["method_output_name"]

        if (result := self.infer_modalities_edge_cases(model, processor)) is not None:
            modality_config, module_output_name = result
            # Edge-case models may also support the message format via chat templates
            if hasattr(processor, "chat_template") and processor.chat_template is not None:
                if "message" not in modality_config:
                    modality_config["message"] = {
                        **modality_config.get(
                            "text", {"method": "forward", "method_output_name": default_method_output_name}
                        ),
                        "format": self.input_formatter.message_format,
                    }
            return modality_config, module_output_name

        modalities = self.infer_modalities_from_processor(processor)
        if hasattr(processor, "chat_template") and processor.chat_template is not None:
            modalities.append("message")

        # Inspect forward to see if it can be used for all modalities, or if we need modality-specific methods.
        # If we can't inspect the method return type, we assume it has the default output name.
        output_fields = self._get_method_output_fields(model.forward)
        if output_fields is None or default_method_output_name in output_fields:
            modality_config: ModalityConfig = {}
            for modality in modalities:
                entry = ModalityParams(method="forward", method_output_name=default_method_output_name)
                if modality == "message":
                    entry["format"] = self.input_formatter.message_format
                modality_config[modality] = entry
            return modality_config, default_module_output_name

        # For feature-extraction, if there's no 'last_hidden_state', we can check for modality-specific methods like get_..._features
        if self.transformer_task == "feature-extraction":
            modality_config: ModalityConfig = {}
            for modality in modalities:
                if modality == "message":
                    continue

                method_name = f"get_{modality}_features"
                if hasattr(model, method_name):
                    method = getattr(model, method_name)
                    method_output_fields = self._get_method_output_fields(method)
                    if method_output_fields and "pooler_output" in method_output_fields:
                        modality_config[modality] = {"method": method_name, "method_output_name": "pooler_output"}
                    else:
                        modality_config[modality] = {"method": method_name, "method_output_name": None}

            return modality_config, "sentence_embedding"

        return {
            modality: {"method": "forward", "method_output_name": default_method_output_name}
            for modality in modalities
        }, default_module_output_name

    def infer_modalities_edge_cases(
        self,
        model: PreTrainedModel,
        processor: ProcessorMixin | PreTrainedTokenizerBase | FeatureExtractionMixin | ImageProcessingMixin,
    ) -> tuple[ModalityConfig, str] | None:
        """Return a ``(modality_config, module_output_name)`` for model types that cannot be handled
        by the general :meth:`infer_modalities` inference path, or ``None`` to fall through.

        Looks up the model type in the task-specific edge case configs from
        :data:`_EDGE_CASE_MODALITY_CONFIGS`. For entries that require output name validation
        (transformers v4/v5 compat), resolves each modality's ``method_output_name`` against the
        actual model method via :meth:`_infer_method_output_name`.
        """
        task_edge_cases = _EDGE_CASE_MODALITY_CONFIGS.get(self.transformer_task)
        if task_edge_cases is None:
            return None

        entry = task_edge_cases.get(model.config.model_type)
        if entry is None:
            return None

        raw_config, module_output_name, validate_output_names = entry
        if not validate_output_names:
            return raw_config, module_output_name

        modality_config: ModalityConfig = {}
        for modality, params in raw_config.items():
            if not hasattr(model, params["method"]):
                logger.warning_once(
                    f"Model does not have method {params['method']!r} for modality {modality!r}. Skipping."
                )
                continue
            method = getattr(model, params["method"])
            modality_config[modality] = {
                "method": params["method"],
                "method_output_name": self._infer_method_output_name(params["method_output_name"], method),
            }
        return modality_config, module_output_name

    def infer_modalities_from_processor(
        self,
        processor: ProcessorMixin | PreTrainedTokenizerBase | FeatureExtractionMixin | ImageProcessingMixin,
    ) -> list[Modality]:
        """Determine which modalities the processor supports by inspecting its attributes or type."""
        processor_attribute_mapping: dict[str, Modality] = {
            "tokenizer": "text",
            "image_processor": "image",
            "feature_extractor": "audio",
            "video_processor": "video",
        }
        if isinstance(processor, ProcessorMixin):
            processor_attributes = self._get_processor_attributes() or {}
            return [
                modality_name
                for processor_attribute, modality_name in processor_attribute_mapping.items()
                if processor_attribute in processor_attributes
            ]

        modality_checks: dict[Modality, type] = {
            "text": PreTrainedTokenizerBase,
            "audio": FeatureExtractionMixin,
            "video": BaseVideoProcessor,
            "image": ImageProcessingMixin,
        }
        for modality_name, processor_class in modality_checks.items():
            if isinstance(processor, processor_class):
                return [modality_name]

        logger.warning(
            f"Could not determine modalities from processor of type {type(processor).__name__}. "
            "Returning an empty modality list."
        )
        return []

    def _get_processor_attributes(self) -> list[str] | None:
        """Get the attributes of the processor if available. Will be removed in the future as transformers v5
        becomes the minimum requirement.

        Returns:
            list[str] | None: The processor attribute names, or None if not available.
        """
        if hasattr(self.processor, "get_attributes"):  # Transformers v5+
            return self.processor.get_attributes()
        elif hasattr(self.processor, "attributes"):  # Transformers v4
            return self.processor.attributes
        return None

    @staticmethod
    def _get_method_output_fields(method: Callable) -> list[str] | None:
        """Extract the output field names from a method's return type annotation.

        Args:
            method (Callable): The method to inspect.

        Returns:
            list[str] | None: List of output field names, or None if not found.
        """

        def find_model_output_class(type_annotation):
            if isinstance(type_annotation, type) and issubclass(type_annotation, ModelOutput):
                return type_annotation
            for sub_annotation in get_args(type_annotation):
                if (result := find_model_output_class(sub_annotation)) is not None:
                    return result
            return None

        try:
            return_annotation = get_type_hints(method).get("return", None)
        except Exception:
            return None
        output_class = find_model_output_class(return_annotation)
        if output_class is None:
            return None
        return [field.name for field in fields(output_class)]

    @staticmethod
    def _infer_method_output_name(method_output_name: str, method: Callable) -> str | None:
        """Validate that ``method_output_name`` is present in the method's return type annotation.

        Returns the name if found, or ``None`` if the method's output type does not include it.
        Primarily needed for transformers v4 compatibility: v5 often allows ``pooler_output``
        from ``get_..._features`` methods, but v4 didn't use ``BaseModelOutputWithPooling`` yet.
        """
        output_fields = Transformer._get_method_output_fields(method) or []
        if method_output_name in output_fields:
            return method_output_name
        return None

    def save(self, output_path: str, *args, safe_serialization: bool = True, **kwargs) -> None:
        """Save the model, processor, and module config to ``output_path``."""
        self.model.save_pretrained(output_path, safe_serialization=safe_serialization)
        self.processor.save_pretrained(output_path)
        self.save_config(output_path)

    @classmethod
    def load(
        cls,
        model_name_or_path: str,
        # Loading arguments
        subfolder: str = "",
        token: bool | str | None = None,
        cache_folder: str | None = None,
        revision: str | None = None,
        local_files_only: bool = False,
        # Module-specific arguments
        trust_remote_code: bool = False,
        model_kwargs: dict[str, Any] | None = None,
        processor_kwargs: dict[str, Any] | None = None,
        config_kwargs: dict[str, Any] | None = None,
        backend: str = "torch",
        init_defaults: dict[str, Any] | None = None,
        **kwargs,
    ) -> Self:
        """Load a Transformer module from a pretrained model directory or Hugging Face model name."""
        init_kwargs = cls._load_init_kwargs(
            model_name_or_path=model_name_or_path,
            subfolder=subfolder,
            token=token,
            cache_folder=cache_folder,
            revision=revision,
            local_files_only=local_files_only,
            trust_remote_code=trust_remote_code,
            model_kwargs=model_kwargs,
            processor_kwargs=processor_kwargs,
            config_kwargs=config_kwargs,
            backend=backend,
        )
        # init_kwargs has priority over init_defaults
        for key, value in (init_defaults or {}).items():
            init_kwargs.setdefault(key, value)
        return cls(model_name_or_path=model_name_or_path, **init_kwargs)

    @classmethod
    def _load_init_kwargs(
        cls,
        model_name_or_path: str,
        # Loading arguments
        subfolder: str = "",
        token: bool | str | None = None,
        cache_folder: str | None = None,
        revision: str | None = None,
        local_files_only: bool = False,
        # Module-specific arguments
        trust_remote_code: bool = False,
        model_kwargs: dict[str, Any] | None = None,
        processor_kwargs: dict[str, Any] | None = None,
        config_kwargs: dict[str, Any] | None = None,
        backend: str = "torch",
        **kwargs,
    ) -> dict[str, Any]:
        """Build the kwargs dict for ``__init__`` by merging config file, hub kwargs, and caller overrides.

        Priority (highest to lowest): caller kwargs > hub kwargs > config file values.
        """
        config = cls.load_config(
            model_name_or_path=model_name_or_path,
            subfolder=subfolder,
            token=token,
            cache_folder=cache_folder,
            revision=revision,
            local_files_only=local_files_only,
        )

        hub_kwargs = {
            "subfolder": subfolder,
            "token": token,
            "cache_dir": cache_folder,  # Transformers uses `cache_dir` instead of `cache_folder`
            "revision": revision,
            "local_files_only": local_files_only,
            "trust_remote_code": trust_remote_code,
        }

        # 3rd priority: config file
        # Config files may use the old key names (model_args, tokenizer_args, config_args) for backwards compat.
        # We normalize them to the new names (model_kwargs, processor_kwargs, config_kwargs) here so we don't
        # trigger deprecation warnings from the transformer_kwargs_decorator on Transformer.__init__.
        _OLD_TO_NEW = {
            "model_args": "model_kwargs",
            "tokenizer_args": "processor_kwargs",
            "config_args": "config_kwargs",
        }
        for old_name, new_name in _OLD_TO_NEW.items():
            if old_name in config:
                config[new_name] = config.pop(old_name)
            config.setdefault(new_name, {})

        # 2nd priority: hub_kwargs
        config["model_kwargs"].update(hub_kwargs)
        config["processor_kwargs"].update(hub_kwargs)
        config["config_kwargs"].update(hub_kwargs)

        # 1st priority: kwargs passed to SentenceTransformer
        if model_kwargs:
            config["model_kwargs"].update(model_kwargs)
        if processor_kwargs:
            config["processor_kwargs"].update(processor_kwargs)
        if config_kwargs:
            config["config_kwargs"].update(config_kwargs)

        return {**config, "backend": backend}

    @classmethod
    def load_config(
        cls,
        model_name_or_path: str,
        subfolder: str = "",
        config_filename: str | None = None,
        token: bool | str | None = None,
        cache_folder: str | None = None,
        revision: str | None = None,
        local_files_only: bool = False,
    ) -> dict[str, Any]:
        """Load the module config, trying several legacy config filenames for backward compatibility.

        Handles deserialization of ``modality_config`` tuple keys (stored as comma-separated strings
        in JSON) and strips ``trust_remote_code`` from all sub-dicts for security.
        """
        config_filenames = (
            [config_filename]
            if config_filename
            else [
                "sentence_bert_config.json",
                "sentence_roberta_config.json",
                "sentence_distilbert_config.json",
                "sentence_camembert_config.json",
                "sentence_albert_config.json",
                "sentence_xlm-roberta_config.json",
                "sentence_xlnet_config.json",
            ]
        )
        for config_filename in config_filenames:
            config = super().load_config(
                model_name_or_path=model_name_or_path,
                subfolder=subfolder,
                config_filename=config_filename,
                token=token,
                cache_folder=cache_folder,
                revision=revision,
                local_files_only=local_files_only,
            )
            if config:
                break

        # Don't allow configs to set trust_remote_code
        for key in (
            "model_args",
            "model_kwargs",
            "tokenizer_args",
            "processor_kwargs",
            "config_args",
            "config_kwargs",
        ):
            if key in config and "trust_remote_code" in config[key]:
                config[key].pop("trust_remote_code")

        if "modality_config" in config:
            # Deserialize modality_config keys: "+" separates tuple modalities (e.g. "image+text")
            valid_single_modalities = {"text", "image", "audio", "video", "message"}
            deserialized_modality_config = {}
            for modality_key, params in config["modality_config"].items():
                if "+" in modality_key:
                    parts = tuple(modality_key.split("+"))
                    invalid = [p for p in parts if p not in valid_single_modalities]
                    if invalid:
                        logger.warning(
                            f"Ignoring unknown modality components {invalid!r} in modality_config key {modality_key!r}."
                        )
                        continue
                    deserialized_modality_config[parts] = params
                else:
                    if modality_key not in valid_single_modalities:
                        logger.warning(f"Ignoring unknown modality key {modality_key!r} in modality_config.")
                        continue
                    deserialized_modality_config[modality_key] = params
            config["modality_config"] = deserialized_modality_config

        else:
            # This method is only called if this model has a modules.json, i.e. it's already been saved
            # with Sentence Transformers. So, if modality_config is not in the config, we can assume it
            # was saved with an older version where Transformer was text-only, and so we can set the
            # modality_config accordingly for backward compatibility. Otherwise, we might infer and use
            # the 'message' format and get different results than what previously worked.
            config["modality_config"], config["module_output_name"] = cls._get_default_modality_config(config)

        return config

    @staticmethod
    def _get_default_modality_config(config: dict[str, Any]) -> tuple[ModalityConfig, str]:
        """Get the default modality configuration for the current transformer task.

        Returns:
            tuple[ModalityConfig, str]: A tuple of (modality_config, module_output_name).
                The modality_config maps modality keys to dicts with 'method' and 'method_output_name'.
                The module_output_name is the name of the output feature this module creates.
        """
        return TRANSFORMER_TASK_DEFAULTS[config.get("transformer_task", "feature-extraction")]

    def get_config_dict(self) -> dict[str, Any]:
        """Return the config dict for serialization, with tuple modality keys joined as plus-separated strings."""
        config_dict = super().get_config_dict()
        config_dict["modality_config"] = {
            format_modality(modality): params for modality, params in self.modality_config.items()
        }
        if not self.processing_kwargs:
            config_dict.pop("processing_kwargs", None)
        if self.unpad_inputs is None:
            config_dict.pop("unpad_inputs", None)
        # Drop multi-vector knobs held at their dense/sparse/cross-encoder defaults so
        # ``sentence_bert_config.json`` for non-multi-vector models stays uncluttered.
        for key, default in self._DEFAULT_CONFIG_VALUES.items():
            if config_dict.get(key) == default:
                config_dict.pop(key, None)
        return config_dict

    def __repr__(self) -> str:
        return f"Transformer({dict(self.get_config_dict(), architecture=self.model.__class__.__name__)})"