modeling_vela_reranker.py

"""Vela's 20-exit reranker, loaded through Transformers AutoModel."""

import json
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import ClassVar

import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file, save_file
from torch import nn
from transformers import ModernBertConfig, ModernBertModel, PreTrainedModel
from transformers.utils import ModelOutput

from .modernbert_sdpa_layout import install_rocm_sdpa_layout_guard

MIN_HEAD_DIMENSION = 2
TOKEN_INPUT_RANK = 2


class VelaRerankerConfig(ModernBertConfig):
    model_type = "modernbert"

    def __init__(self, layer_indices=None, dim_indices=None, **kwargs):
        super().__init__(**kwargs)
        self.layer_indices = [3, 6, 11, 22] if layer_indices is None else layer_indices
        self.dim_indices = (
            [768, 512, 256, 128, 64] if dim_indices is None else dim_indices
        )


@dataclass
class VelaRerankerOutput(ModelOutput):
    logits: torch.Tensor | None = None
    all_scores: dict[str, torch.Tensor] | None = None


class VelaReranker(PreTrainedModel):
    config_class = VelaRerankerConfig
    base_model_prefix = "encoder"
    _supports_sdpa = True
    _no_split_modules: ClassVar[list[str]] = ["ModernBertEncoderLayer"]
    _keep_in_fp32_modules: ClassVar[list[str]] = ["layer_heads"]

    def __init__(self, config, encoder=None):
        super().__init__(config)
        expected = {
            "version": 1,
            "intermediate_normalization": "final_norm",
            "final_normalization": "final_norm",
            "pooling": "cls",
            "head_dtype": "float32",
        }
        if getattr(config, "representation_contract", None) != expected:
            raise ValueError("Unsupported reranker representation contract")
        self.layer_indices = list(config.layer_indices)
        self.dim_indices = list(config.dim_indices)
        if (
            self.layer_indices != sorted(set(self.layer_indices))
            or not self.layer_indices
            or self.layer_indices[-1] != config.num_hidden_layers
            or any(type(x) is not int or x < 1 for x in self.layer_indices)
        ):
            raise ValueError("Invalid reranker layer indices")
        if (
            self.dim_indices != sorted(set(self.dim_indices), reverse=True)
            or not self.dim_indices
            or self.dim_indices[0] > config.hidden_size
            or any(
                type(x) is not int or x < MIN_HEAD_DIMENSION for x in self.dim_indices
            )
        ):
            raise ValueError("Invalid reranker dimensions")
        native = config.to_dict()
        for key in ("auto_map", "architectures", "layer_indices", "dim_indices"):
            native.pop(key, None)
        native["model_type"] = "modernbert"
        encoder_config = ModernBertConfig.from_dict(native)
        if hasattr(encoder_config, "reference_compile"):
            encoder_config.reference_compile = False
        encoder_config._attn_implementation = config._attn_implementation or "sdpa"
        self.encoder = install_rocm_sdpa_layout_guard(
            ModernBertModel(encoder_config) if encoder is None else encoder
        )
        self.layer_heads = nn.ModuleDict(
            {
                str(layer): nn.ModuleDict(
                    {
                        str(dim): nn.Sequential(
                            nn.Linear(dim, dim // 2),
                            nn.GELU(),
                            nn.Dropout(0.1),
                            nn.Linear(dim // 2, 1),
                        )
                        for dim in self.dim_indices
                    }
                )
                for layer in self.layer_indices
            }
        )
        self.post_init()

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path,
        *model_args,
        config=None,
        revision=None,
        cache_dir=None,
        token=None,
        local_files_only=False,
        force_download=False,
        torch_dtype=None,
        dtype=None,
        attn_implementation="sdpa",
        **kwargs,
    ):
        """Load the encoder and every trained head from one immutable snapshot."""
        for key in (
            "_from_auto",
            "_commit_hash",
            "trust_remote_code",
            "adapter_kwargs",
        ):
            kwargs.pop(key, None)
        if model_args or kwargs:
            raise TypeError(f"Unsupported loading arguments: {sorted(kwargs)}")
        path = Path(pretrained_model_name_or_path)
        if not path.is_dir():
            pinned = getattr(config, "_commit_hash", None) or revision
            path = Path(
                snapshot_download(
                    pretrained_model_name_or_path,
                    revision=pinned,
                    cache_dir=cache_dir,
                    token=token,
                    local_files_only=local_files_only,
                    force_download=force_download,
                    allow_patterns=[
                        "config.json",
                        "model.safetensors",
                        "model-*.safetensors",
                        "model.safetensors.index.json",
                        "classification_heads.safetensors",
                        "matryoshka_config.json",
                    ],
                )
            )
        if config is None:
            config = VelaRerankerConfig.from_pretrained(path, local_files_only=True)
        saved = json.loads((path / "matryoshka_config.json").read_text())
        if (
            saved["layer_indices"] != config.layer_indices
            or saved["dim_indices"] != config.dim_indices
            or saved.get("pooling_strategy") != "cls"
            or saved.get("representation_contract") != config.representation_contract
        ):
            raise ValueError("Head metadata does not match the model configuration")
        native = config.to_dict()
        for key in ("auto_map", "architectures", "layer_indices", "dim_indices"):
            native.pop(key, None)
        encoder_config = ModernBertConfig.from_dict(native)
        if hasattr(encoder_config, "reference_compile"):
            encoder_config.reference_compile = False
        precision = dtype if dtype is not None else torch_dtype
        encoder, loading = ModernBertModel.from_pretrained(
            path,
            config=encoder_config,
            local_files_only=True,
            torch_dtype=precision or torch.float32,
            attn_implementation=attn_implementation or "sdpa",
            output_loading_info=True,
        )
        if any(
            loading.get(key)
            for key in (
                "missing_keys",
                "unexpected_keys",
                "mismatched_keys",
                "error_msgs",
            )
        ):
            raise ValueError("Encoder checkpoint is incomplete or incompatible")
        model = cls(config, encoder=encoder)
        heads = load_file(path / "classification_heads.safetensors", device="cpu")
        if any(tensor.dtype != torch.float32 for tensor in heads.values()):
            raise ValueError("Saved reranker heads must be float32")
        model.layer_heads.load_state_dict(heads, strict=True)
        return model.eval()

    def save_pretrained(self, save_directory, **kwargs):
        """Retain the native encoder and separate-head artifact layout."""
        path = Path(save_directory)
        path.mkdir(parents=True, exist_ok=True)
        self.encoder.save_pretrained(path, **kwargs)
        state = {
            key: value.detach().cpu().contiguous()
            for key, value in self.layer_heads.state_dict().items()
        }
        if any(value.dtype != torch.float32 for value in state.values()):
            raise ValueError("Reranker heads must retain float32 precision")
        save_file(state, path / "classification_heads.safetensors")
        metadata = {
            "layer_indices": self.layer_indices,
            "dim_indices": self.dim_indices,
            "hidden_size": self.config.hidden_size,
            "num_layers": self.config.num_hidden_layers,
            "pooling_strategy": "cls",
            "representation_contract": self.config.representation_contract,
        }
        (path / "matryoshka_config.json").write_text(
            json.dumps(metadata, indent=2) + "\n"
        )
        self.config.auto_map = {
            "AutoConfig": "modeling_vela_reranker.VelaRerankerConfig",
            "AutoModel": "modeling_vela_reranker.VelaReranker",
        }
        self.config.architectures = ["ModernBertModel"]
        self.config.save_pretrained(path)
        target = path / "modeling_vela_reranker.py"
        if Path(__file__).resolve() != target.resolve():
            shutil.copyfile(__file__, target)
        guard = Path(__file__).with_name("modernbert_sdpa_layout.py")
        guard_target = path / guard.name
        if guard.resolve() != guard_target.resolve():
            shutil.copyfile(guard, guard_target)

    def forward(
        self,
        input_ids,
        attention_mask=None,
        position_ids=None,
        layer_idx=None,
        dim_idx=None,
        return_all_scores=False,
    ):
        if (
            input_ids.ndim != TOKEN_INPUT_RANK
            or not 1 <= input_ids.shape[1] <= self.config.max_position_embeddings
        ):
            raise ValueError("Input must fit the configured token capacity")
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids)
        if attention_mask.shape != input_ids.shape:
            raise ValueError("Attention mask shape must match input_ids")
        layers = self.layer_indices if layer_idx is None else [layer_idx]
        dims = self.dim_indices if dim_idx is None else [dim_idx]
        if any(layer not in self.layer_indices for layer in layers) or any(
            dim not in self.dim_indices for dim in dims
        ):
            raise ValueError("Requested exit is not present in this checkpoint")
        if any(
            parameter.dtype != torch.float32
            for parameter in self.layer_heads.parameters()
        ):
            raise ValueError("Reranker heads must retain float32 precision")
        scores = {}
        with torch.autocast(device_type=input_ids.device.type, enabled=False):
            outputs = self.encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                output_hidden_states=True,
                return_dict=True,
            )
            if len(outputs.hidden_states) != self.config.num_hidden_layers + 1:
                raise ValueError("Unexpected encoder hidden-state layout")
            for layer in layers:
                hidden = (
                    outputs.last_hidden_state
                    if layer == self.config.num_hidden_layers
                    else self.encoder.final_norm(outputs.hidden_states[layer])
                )
                pooled = hidden[:, 0]
                for dim in dims:
                    scores[f"layer_{layer}_dim_{dim}"] = self.layer_heads[str(layer)][
                        str(dim)
                    ](pooled[:, :dim].float()).squeeze(-1)
        primary = f"layer_{layers[-1]}_dim_{dims[0]}"
        return VelaRerankerOutput(
            logits=scores[primary], all_scores=scores if return_all_scores else None
        )