modeling_llava_eurobert_audio.py

"""
LlavaEuroBertAudioForEmbedding: Qwen3VL vision + Qwen2.5-Omni audio + EuroBERT text.

Architecture:
  - Vision: Qwen3VLVisionModel (with RoPE, 3D Conv3d patch embed, all layers)
  - Merger: PretrainedMerger (top-level, NOT inside vision_tower)
  - Audio: Qwen2_5OmniAudioEncoder (Qwen2.5-Omni) + Linear projector
  - Text: LlamaModel (EuroBERT, bidirectional)
  - LM head: Identity (embedding model, no vocab projection)

Modality loading:
  model = AutoModel.from_pretrained(path, trust_remote_code=True, modality="omni")   # all components (default)
  model = AutoModel.from_pretrained(path, trust_remote_code=True, modality="vision") # no audio tower/projector
  model = AutoModel.from_pretrained(path, trust_remote_code=True, modality="audio")  # no vision tower/merger
"""

from typing import List, Optional

import torch
import torch.nn as nn
from transformers import LlamaConfig, PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.llama.modeling_llama import LlamaModel
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLVisionConfig
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel
from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniAudioEncoderConfig
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder


_VALID_MODALITIES = ("omni", "vision", "audio", "text")


class PretrainedMerger(nn.Module):
    def __init__(self, hidden_size, out_hidden_size, spatial_merge_size=2):
        super().__init__()
        self.hidden_size = hidden_size * (spatial_merge_size**2)
        self.norm = nn.LayerNorm(hidden_size, eps=1e-6)
        self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
        self.act = nn.GELU()
        self.linear_fc2 = nn.Linear(self.hidden_size, out_hidden_size)

    def forward(self, x):
        x = self.norm(x)
        x = x.view(-1, self.hidden_size)
        x = self.linear_fc2(self.act(self.linear_fc1(x)))
        return x


class LlavaEuroBertAudioConfig(PretrainedConfig):
    model_type = "llava_eurobert_audio"

    def __init__(
        self,
        vision_config=None,
        text_config=None,
        audio_config=None,
        image_token_index=None,
        audio_token_id=None,
        audio_start_token_id=None,
        audio_end_token_id=None,
        projector_hidden_act="gelu",
        tie_word_embeddings=False,
        modality="omni",
        **kwargs,
    ):
        if isinstance(vision_config, dict):
            vision_config = PretrainedConfig(**vision_config)
        self.vision_config = vision_config or PretrainedConfig()
        if isinstance(text_config, dict):
            text_config = PretrainedConfig(**text_config)
        self.text_config = text_config or PretrainedConfig()
        if isinstance(audio_config, dict):
            audio_config = PretrainedConfig(**audio_config)
        self.audio_config = audio_config or PretrainedConfig()
        self.image_token_index = image_token_index
        self.audio_token_id = audio_token_id
        self.audio_start_token_id = audio_start_token_id
        self.audio_end_token_id = audio_end_token_id
        self.projector_hidden_act = projector_hidden_act
        if modality not in _VALID_MODALITIES:
            raise ValueError(f"modality must be one of {_VALID_MODALITIES}, got '{modality}'")
        self.modality = modality
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

    def get_text_config(self, **kwargs):
        return self.text_config


class LlavaEuroBertAudioForEmbedding(PreTrainedModel):
    config_class = LlavaEuroBertAudioConfig
    supports_gradient_checkpointing = True
    _supports_sdpa = True
    _supports_flash_attn_2 = True
    _supports_flash_attn = True
    _supports_attention_backend = True
    _tied_weights_keys = []
    _keys_to_ignore_on_load_missing = ["lm_head.weight"]
    _keys_to_ignore_on_load_unexpected = []

    @classmethod
    def is_backend_compatible(cls) -> bool:
        return True

    def __init__(self, config: LlavaEuroBertAudioConfig):
        super().__init__(config)

        modality = getattr(config, "modality", "omni")
        if modality not in _VALID_MODALITIES:
            raise ValueError(f"modality must be one of {_VALID_MODALITIES}, got '{modality}'")
        self._modality = modality
        # propagate attn_implementation into the inner towers. The EuroBERT
        # text encoder is bidirectional/non-causal; flash-attn varlen does not
        # support it, so the text tower is kept on sdpa.
        attn_impl = self.config._attn_implementation
        text_attn_impl = "sdpa" if (attn_impl and "flash" in str(attn_impl)) else attn_impl

        vision_cfg = config.vision_config
        if not isinstance(vision_cfg, Qwen3VLVisionConfig):
            if hasattr(vision_cfg, "to_dict"):
                d = vision_cfg.to_dict()
            else:
                d = dict(vision_cfg)
            d.pop("model_type", None)
            d.pop("transformers_version", None)
            vision_cfg = Qwen3VLVisionConfig(**d)

        vision_cfg.deepstack_visual_indexes = []
        vision_cfg._attn_implementation = attn_impl
        spatial_merge_size = getattr(vision_cfg, "spatial_merge_size", 2)

        text_cfg = config.text_config
        if not isinstance(text_cfg, LlamaConfig):
            txt_dict = text_cfg.to_dict() if hasattr(text_cfg, 'to_dict') else dict(text_cfg)
            _saved_attn_impl = getattr(text_cfg, "_attn_implementation", None)
            text_cfg = LlamaConfig(**txt_dict)
            # Propagate attn_implementation — vLLM's transformers backend sets
            # this to "vllm" on the parent text_config to route attention through
            # its packed-sequence-aware kernels. If we re-instantiate LlamaConfig
            # without carrying this over, HF falls back to sdpa/flash_attention_2
            # which treats the packed batch as a single long sequence and leaks
            # attention across requests.
            if _saved_attn_impl is not None:
                text_cfg._attn_implementation = _saved_attn_impl
        text_hidden = text_cfg.hidden_size

        self._spatial_merge_size = spatial_merge_size

        if modality not in ("audio", "text"):
            self.vision_tower = Qwen3VLVisionModel(vision_cfg)
            self.vision_tower.merger = nn.Identity()
            self.vision_tower.deepstack_merger_list = nn.ModuleList()
            self.vision_tower.deepstack_visual_indexes = []
            self.merger = PretrainedMerger(
                vision_cfg.hidden_size, text_hidden, spatial_merge_size
            )

        self.multi_modal_projector = nn.Identity()
        text_cfg._attn_implementation = text_attn_impl
        self.language_model = LlamaModel(text_cfg)
        self.lm_head = nn.Identity()

        for layer in self.language_model.layers:
            layer.self_attn.is_causal = False

        if modality not in ("vision", "text"):
            aud_cfg = config.audio_config
            aud_dict = aud_cfg.to_dict() if hasattr(aud_cfg, 'to_dict') else aud_cfg
            audio_encoder_config = Qwen2_5OmniAudioEncoderConfig(**aud_dict)
            audio_encoder_config._attn_implementation = attn_impl
            self.audio_tower = Qwen2_5OmniAudioEncoder(audio_encoder_config)
            self.audio_tower.proj = nn.Identity()  # fused into audio_projector(s)
            output_dim = aud_dict.get('d_model', 1280)  # fused: audio_projector(s) now take d_model
            self.audio_projector = nn.Linear(output_dim, text_hidden)

        ignore = []
        if modality in ("audio", "text"):
            ignore.extend([r"^vision_tower\.", r"^merger\."])
        if modality in ("vision", "text"):
            ignore.extend([r"^audio_tower\.", r"^audio_projector\."])
        if ignore:
            self._keys_to_ignore_on_load_unexpected = ignore

        self.post_init()

    @property
    def modality(self) -> str:
        return self._modality

    def get_input_embeddings(self):
        return self.language_model.embed_tokens

    def set_input_embeddings(self, value):
        self.language_model.embed_tokens = value

    def get_output_embeddings(self):
        return None

    def get_image_features(
        self,
        pixel_values: torch.FloatTensor,
        image_grid_thw: torch.LongTensor,
        num_image_tokens: Optional[int] = None,
    ) -> List[torch.Tensor]:
        if self._modality in ("audio", "text"):
            raise ValueError(
                f"Vision inputs are not available in {self._modality}-only mode. "
                "Load with modality='omni' or modality='vision'."
            )

        vision_output = self.vision_tower(
            hidden_states=pixel_values, grid_thw=image_grid_thw
        )
        if isinstance(vision_output, tuple):
            raw_hidden = vision_output[0]
        elif hasattr(vision_output, "pooler_output") and vision_output.pooler_output is not None:
            raw_hidden = vision_output.pooler_output
        else:
            raw_hidden = vision_output[0]

        image_features = self.merger(raw_hidden)

        merge_sq = self._spatial_merge_size ** 2
        split_sizes = (image_grid_thw.prod(-1) // merge_sq).tolist()
        return list(torch.split(image_features, split_sizes))

    def get_audio_features(
        self,
        input_features: torch.FloatTensor,
        feature_attention_mask: Optional[torch.LongTensor] = None,
    ) -> torch.Tensor:
        if self._modality in ("vision", "text"):
            raise ValueError(
                f"Audio inputs are not available in {self._modality}-only mode. "
                "Load with modality='omni' or modality='audio'."
            )

        batch_size = input_features.shape[0]
        if batch_size > 1:
            # Serialize per-sample so the packed-frames GEMM shape stays invariant
            # across batch sizes. Makes batched audio bit-exact to B=1 in bf16,
            # and is substantially faster for B>=16 because B=1 hits a
            # well-optimized kernel while the packed-B=N path thrashes on a
            # (total_frames)^2 sdpa matrix.
            outs = [
                self.get_audio_features(
                    input_features[i : i + 1],
                    feature_attention_mask[i : i + 1] if feature_attention_mask is not None else None,
                )
                for i in range(batch_size)
            ]
            return torch.cat(outs, dim=0)
        if feature_attention_mask is not None:
            feature_lens = feature_attention_mask.sum(-1).long()
            packed = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
        else:
            feature_lens = torch.full(
                (batch_size,), input_features.shape[2],
                device=input_features.device, dtype=torch.long,
            )
            packed = input_features.transpose(1, 2).reshape(-1, input_features.shape[1]).T
        aftercnn_lens, _ = self.audio_tower._get_feat_extract_output_lengths(feature_lens)
        audio_output = self.audio_tower(
            packed, feature_lens=feature_lens, aftercnn_lens=aftercnn_lens,
        )
        return self.audio_projector(audio_output.last_hidden_state)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        input_features: Optional[torch.FloatTensor] = None,
        feature_attention_mask: Optional[torch.LongTensor] = None,
        cache_position: Optional[torch.LongTensor] = None,
        output_hidden_states: Optional[bool] = None,
        **kwargs,
    ):
        image_grid_thw = kwargs.pop("image_grid_thw", None)
        num_image_tokens = kwargs.pop("num_image_tokens", None)
        pixel_values_videos = kwargs.pop("pixel_values_videos", None)
        video_grid_thw = kwargs.pop("video_grid_thw", None)
        num_video_tokens = kwargs.pop("num_video_tokens", None)
        kwargs.pop("spatial_shapes", None)
        kwargs.pop("pixel_attention_mask", None)

        if pixel_values is not None and self._modality in ("audio", "text"):
            raise ValueError(
                f"Vision inputs are not available in {self._modality}-only mode. "
                "Load with modality='omni' or modality='vision'."
            )
        if input_features is not None and self._modality in ("vision", "text"):
            raise ValueError(
                f"Audio inputs are not available in {self._modality}-only mode. "
                "Load with modality='omni' or modality='audio'."
            )

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError(
                "You must specify exactly one of input_ids or inputs_embeds"
            )

        if inputs_embeds is None:
            inputs_embeds = self.get_input_embeddings()(input_ids)

        # Image and video both use config.image_token_index (the processor
        # remaps <|video_pad|> to <image>). Multipart inputs with both
        # modalities need a single combined scatter; otherwise the first
        # masked_scatter sees True positions for image+video but only enough
        # source rows for image, hitting an assertion. Concatenate features
        # in part order (image then video — matches input_ids ordering when
        # custom_st emits image parts before video parts).
        all_features = []
        if pixel_values is not None and image_grid_thw is not None:
            all_features.extend(self.get_image_features(
                pixel_values=pixel_values,
                image_grid_thw=image_grid_thw,
                num_image_tokens=num_image_tokens,
            ))
        if pixel_values_videos is not None and video_grid_thw is not None:
            all_features.extend(self.get_image_features(
                pixel_values=pixel_values_videos,
                image_grid_thw=video_grid_thw,
                num_image_tokens=num_video_tokens,
            ))
        if all_features:
            features = torch.cat(all_features, dim=0).to(
                inputs_embeds.device, inputs_embeds.dtype
            )
            mask = (
                (input_ids == self.config.image_token_index)
                .unsqueeze(-1)
                .expand_as(inputs_embeds)
            )
            inputs_embeds = inputs_embeds.masked_scatter(mask, features)

        if input_features is not None:
            audio_embeds = self.get_audio_features(
                input_features, feature_attention_mask
            )
            audio_embeds_flat = audio_embeds.reshape(
                -1, audio_embeds.shape[-1]
            ).to(inputs_embeds.device, inputs_embeds.dtype)
            audio_mask = (
                (input_ids == self.config.audio_token_id)
                .unsqueeze(-1)
                .expand_as(inputs_embeds)
            )
            inputs_embeds = inputs_embeds.masked_scatter(
                audio_mask, audio_embeds_flat
            )

        if attention_mask is not None and attention_mask.dim() == 2:
            dtype = inputs_embeds.dtype
            seq_len = inputs_embeds.shape[1]
            bidi_mask = attention_mask[:, None, None, :].to(dtype=dtype)
            bidi_mask = (1.0 - bidi_mask) * torch.finfo(dtype).min
            attention_mask = bidi_mask.expand(-1, -1, seq_len, -1)

        if position_ids is not None and position_ids.dim() == 3:
            position_ids = position_ids[0]

        # vLLM's transformers backend passes `return_dict=False` via kwargs; we
        # always use dict-style output internally so the attribute access below
        # works, and BaseModelOutputWithPast is indexable like a tuple so vLLM's
        # `output[0]` contract still holds.
        kwargs.pop("return_dict", None)
        outputs = self.language_model(
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            **kwargs,  # carries `attention_instances` when called from vLLM
        )

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)

        return BaseModelOutputWithPast(
            last_hidden_state=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def embed(self, truncate_dim=None, **inputs):
        """Encode processor outputs into L2-normalized last-token embeddings.

        Matryoshka: pass `truncate_dim=N` to get an N-dim unit-norm vector
        (truncation is applied before L2-normalization).
        """
        attention_mask = inputs.get("attention_mask", None)
        self.eval()
        with torch.no_grad():
            out = self(**inputs)
        hidden = out.last_hidden_state
        if attention_mask is not None and attention_mask.dim() == 2:
            idx = attention_mask.sum(dim=1) - 1
        else:
            idx = torch.full(
                (hidden.shape[0],), hidden.shape[1] - 1,
                device=hidden.device, dtype=torch.long,
            )
        pooled = hidden[torch.arange(hidden.shape[0], device=hidden.device), idx]
        if truncate_dim is not None:
            pooled = pooled[:, :truncate_dim]
        return torch.nn.functional.normalize(pooled, dim=-1)


# ---------------------------------------------------------------------------
# vLLM registration (side-effect on module import)
#
# Triggered via config.json "auto_map.AutoConfig" → this module → ModelRegistry.
# Pure-HF / sentence-transformers usage is unaffected: if vLLM is not installed
# or not importable, the try block is skipped silently.
# ---------------------------------------------------------------------------

def _register_vllm() -> None:
    # Register LlavaEuroBertAudioForVLLMEmbedding so vLLM's resolve_model_cls
    # picks it up instead of falling back to TransformersMultiModalForCausalLM.
    #
    # Flow: _try_resolve_transformers imports our file (is_backend_compatible
    # returns False → returns None), then the arch loop finds our registered
    # class. Also patches KVCacheManager to auto-disable caching for encoder-
    # only models (no KV cache groups). Pure-HF usage is unaffected.
    import importlib.util as _iu
    if _iu.find_spec("vllm") is None:
        return
    try:
        # ----- bootstrap sibling vLLM port file -----
        # When loaded via transformers' `trust_remote_code=True`, only the
        # modeling_*.py referenced in auto_map is fetched into the
        # transformers_modules cache — `vllm_llava_eurobert_audio.py` is
        # NOT. Pull it from HF Hub before registering and put it on
        # sys.path / PYTHONPATH so vLLM's inspect_model_cls subprocess
        # can `import vllm_llava_eurobert_audio` (it spawns a fresh
        # interpreter that doesn't inherit our sys.modules).
        import os, sys, importlib, shutil
        pkg = __package__ or ""
        current_dir = os.path.dirname(os.path.abspath(__file__))
        sibling_name = "vllm_llava_eurobert_audio"
        sibling_path = os.path.join(current_dir, sibling_name + ".py")
        if not os.path.exists(sibling_path):
            # Two cache layouts to handle:
            #  - 4-part: transformers_modules.<ns>.<safe-repo>.<sha>  (HF cache load,
            #    e.g. AutoModel.from_pretrained("jinaai/jina-embeddings-v5-omni-...")).
            #  - 2-part: transformers_modules.<safe-repo>             (flat local_dir
            #    load, e.g. from_pretrained("/path/to/snapshot") — namespace + sha
            #    are dropped by transformers' dynamic_module_utils).
            # In the 2-part case we hardcode the namespace because every model that
            # ships this file lives under jinaai/.
            parts = pkg.split(".") if pkg else []
            repo_id, revision = None, None
            if len(parts) >= 4 and parts[0] == "transformers_modules":
                repo_name = parts[2].replace("_hyphen_", "-").replace("_dot_", ".")
                repo_id = f"{parts[1]}/{repo_name}"
                revision = parts[3]
            elif len(parts) >= 2 and parts[0] == "transformers_modules":
                repo_name = parts[1].replace("_hyphen_", "-").replace("_dot_", ".")
                repo_id = f"jinaai/{repo_name}"
            if repo_id:
                from huggingface_hub import hf_hub_download
                downloaded = hf_hub_download(
                    repo_id=repo_id,
                    filename=sibling_name + ".py",
                    revision=revision,
                )
                shutil.copy(downloaded, sibling_path)
        if current_dir not in sys.path:
            sys.path.insert(0, current_dir)
        existing = os.environ.get("PYTHONPATH", "")
        if current_dir not in existing.split(os.pathsep):
            os.environ["PYTHONPATH"] = (
                current_dir if not existing else current_dir + os.pathsep + existing
            )
        # Try top-level import first so flat local-dir loads (from_pretrained with a
        # local path, Docker image pre-fetches, etc.) work even when __package__ has
        # fewer than 4 dot-separated parts and the HF download path is skipped.
        try:
            _lla = importlib.import_module(sibling_name)
        except (ImportError, ModuleNotFoundError):
            if pkg:
                _lla = importlib.import_module("." + sibling_name, package=pkg)
            else:
                raise
        _ = _lla.LlavaEuroBertAudioForVLLMEmbedding  # keep reference

        # vLLM 0.19.1+ no longer auto-registers via @MULTIMODAL_REGISTRY;
        # explicit registration is required for architecture resolution.
        #
        # Use lazy string registration ("module:class") instead of direct class
        # reference. When the class is registered directly, vLLM pickles the
        # _RegisteredModel (containing the class object) and sends it to the
        # EngineCore spawn-subprocess. Unpickling requires importing
        # "transformers_modules.jinaai...hash.vllm_llava_eurobert_audio", but
        # the transformers_modules root is not yet in sys.path at that point in
        # the fresh subprocess. The lazy string form stores only two strings —
        # no class reference — so it deserializes cleanly. When load_model_cls()
        # is later called, it does importlib.import_module("vllm_llava_eurobert_audio")
        # which finds the file in current_dir via the PYTHONPATH set above.
        from vllm import ModelRegistry
        ModelRegistry.register_model(
            "LlavaEuroBertAudioForEmbedding",
            "vllm_llava_eurobert_audio:LlavaEuroBertAudioForVLLMEmbedding",
        )

        # With runner="pooling", vLLM sets convert_type="embed" which skips the
        # early _try_resolve_transformers path (requires convert_type="none").
        # Our module is only imported inside the LAST _try_resolve_transformers
        # call. By then the for-loop has already run without finding us.
        # _try_resolve_transformers then checks is_backend_compatible() (True),
        # calls _get_transformers_backend_cls() → by default "TransformersMultiModal-
        # ForCausalLM" which can't handle video. Patch it to return our arch name
        # so _try_load_model_cls finds LlavaEuroBertAudioForVLLMEmbedding instead.
        try:
            try:
                from vllm.config.model import ModelConfig as _ModelConfig
            except ImportError:
                from vllm.config import ModelConfig as _ModelConfig
            if not getattr(_ModelConfig, "_jinaai_llava_eurobert_get_backend_patch", False):
                _orig_get_backend = _ModelConfig._get_transformers_backend_cls
                def _patched_get_backend(self):
                    try:
                        hf_archs = getattr(self.hf_config, "architectures", []) or []
                    except Exception:
                        hf_archs = []
                    if "LlavaEuroBertAudioForEmbedding" in hf_archs:
                        return "LlavaEuroBertAudioForEmbedding"
                    return _orig_get_backend(self)
                _ModelConfig._get_transformers_backend_cls = _patched_get_backend
                _ModelConfig._jinaai_llava_eurobert_get_backend_patch = True
        except Exception as _e:
            import warnings
            warnings.warn(
                f"jina-embeddings-v5-omni nano: _get_transformers_backend_cls patch "
                f"failed ({type(_e).__name__}: {_e}); vLLM will fall back to "
                f"TransformersMultiModalForCausalLM which cannot handle video inputs.",
                stacklevel=2,
            )

        try:
            from vllm.model_executor.models.transformers.base import Base
        except ImportError:
            from vllm.model_executor.models.transformers import TransformersBase as Base
        if getattr(Base, "_jinaai_v5_bidirectional_patch", False):
            return

        try:
            from vllm.attention.backends.abstract import AttentionType
        except ImportError:
            from vllm.v1.attention.backend import AttentionType
        try:
            from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
        except ImportError:
            from vllm.model_executor.layers.attention.encoder_only_attention import EncoderOnlyAttention
        from vllm.distributed.utils import get_pp_indices

        def _create_attention_instances(self):
            # Build encoder-only attention instances when the HF model has at
            # least one bidirectional layer (is_causal=False). This handles
            # the EuroBERT-style text tower inside otherwise multimodal configs,
            # where vLLM's default heuristic would force causal attention and
            # collapse multi-request batches to near-identical embeddings.
            #
            # The stock path is kept for fully-causal text towers (e.g. Qwen3VL
            # in small).
            is_encoder_layer = lambda m: not getattr(m, "is_causal", True)
            if not any(is_encoder_layer(m) for m in self.model.modules()):
                return _orig_create(self)

            text_config = self.text_config
            num_heads = self.model_config.get_num_attention_heads(
                self.parallel_config
            )
            head_size = self.model_config.get_head_size()
            num_kv_heads = self.model_config.get_num_kv_heads(
                self.parallel_config
            )
            logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None)

            pp_rank = self.pp_group.rank_in_group
            pp_size = self.pp_group.world_size
            start, end = get_pp_indices(
                text_config.num_hidden_layers, pp_rank, pp_size
            )

            attention_instances = {}
            for i in range(start, end):
                per_layer_sliding_window = None
                if (
                    hasattr(self.config, "layer_types")
                    and self.config.layer_types[i] == "sliding_attention"
                ):
                    per_layer_sliding_window = self.config.sliding_window

                attention_instances[i] = EncoderOnlyAttention(
                    num_heads=num_heads,
                    head_size=head_size,
                    scale=head_size**-0.5,
                    num_kv_heads=num_kv_heads,
                    cache_config=self.cache_config,
                    quant_config=self.quant_config,
                    logits_soft_cap=logits_soft_cap,
                    per_layer_sliding_window=per_layer_sliding_window,
                    prefix=f"{i}.attn",
                    attn_type=AttentionType.ENCODER_ONLY,
                )
            return attention_instances

        _orig_create = Base.create_attention_instances
        Base.create_attention_instances = _create_attention_instances
        Base._jinaai_v5_bidirectional_patch = True

        # Encoder-only attention has zero KV cache groups; vLLM's KVCacheManager
        # asserts "Only one block size is supported" with `enable_caching=True`
        # (the default). Without this patch, users of nano repos must explicitly
        # pass `enable_prefix_caching=False` to `LLM(...)`. Auto-disable caching
        # for this configuration so the defaults just work.
        from vllm.v1.core import kv_cache_manager as _kvm
        _orig_kvm_init = _kvm.KVCacheManager.__init__

        def _kvm_init(self, kv_cache_config, *args, **kwargs):
            if not kv_cache_config.kv_cache_groups:
                kwargs["enable_caching"] = False
            return _orig_kvm_init(self, kv_cache_config, *args, **kwargs)

        _kvm.KVCacheManager.__init__ = _kvm_init
    except Exception as e:
        import warnings
        warnings.warn(
            f"jina-embeddings-v5-omni nano: bidirectional attention patch "
            f"failed ({type(e).__name__}: {e}); embeddings will be incorrect "
            f"under vLLM. This is a model-loader bug, not a vLLM bug.",
            stacklevel=2,
        )


_register_vllm()