Load Vela Reranker directly with Transformers

README.md

@@ -35,4 +35,21 @@ Vela Reranker ranks passages by their relevance to a query.
 
 Higher scores indicate greater relevance. The default uses the 22-layer, 768-dimensional exit.
 
+## Quick start
+
+Install `torch`, `transformers`, and `safetensors`.
+
+```python
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+model_id = "llm-semantic-router/Vela-1.0-Encoder-307M-Reranker"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModel.from_pretrained(model_id, trust_remote_code=True).eval()
+inputs = tokenizer(["When does the library open?"], ["The library opens in the morning."], padding=True, truncation=False, return_tensors="pt")
+with torch.inference_mode():
+    scores = model(**inputs).logits
+print(scores)
+```
+
 [Explore the Vela model collection](https://huggingface.co/collections/llm-semantic-router/vela-10-router-models-6aa555ba70cc6997d6d67798)

config.json

@@ -57,5 +57,22 @@
   "sparse_pred_ignore_index": -100,
   "sparse_prediction": false,
   "transformers_version": "4.57.6",
-  "vocab_size": 256000
+  "vocab_size": 256000,
+  "layer_indices": [
+    3,
+    6,
+    11,
+    22
+  ],
+  "dim_indices": [
+    768,
+    512,
+    256,
+    128,
+    64
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_vela_reranker.VelaRerankerConfig",
+    "AutoModel": "modeling_vela_reranker.VelaReranker"
+  }
 }

modeling_vela_reranker.py

@@ -0,0 +1,271 @@
+"""Vela's 20-exit reranker, loaded through Transformers AutoModel."""
+
+from dataclasses import dataclass
+import json
+from pathlib import Path
+import shutil
+
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file, save_file
+
+import torch
+from torch import nn
+from transformers import ModernBertConfig, ModernBertModel, PreTrainedModel
+from transformers.utils import ModelOutput
+
+
+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 = ["ModernBertEncoderLayer"]
+    _keep_in_fp32_modules = ["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 < 2 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 = 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)
+
+    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 != 2
+            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
+        )