hf model and vllm integration files

config.json

@@ -0,0 +1,77 @@
+{
+  "architectures": [
+    "EmbedderModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_embedder.EmbedderConfig",
+    "AutoModel": "modeling_embedder.EmbedderModel"
+  },
+  "base_model_name": "nomic-ai/nomic-embed-text-v1.5",
+  "dropout": 0.0,
+  "dtype": "float32",
+  "encoder_config": {
+    "_name_or_path": "nomic-ai/nomic-embed-text-v1.5",
+    "activation_function": "swiglu",
+    "architectures": [
+      "NomicBertModel"
+    ],
+    "attn_pdrop": 0.0,
+    "auto_map": {
+      "AutoConfig": "nomic-ai/nomic-bert-2048--configuration_hf_nomic_bert.NomicBertConfig",
+      "AutoModel": "nomic-ai/nomic-bert-2048--modeling_hf_nomic_bert.NomicBertModel",
+      "AutoModelForMaskedLM": "nomic-ai/nomic-bert-2048--modeling_hf_nomic_bert.NomicBertForPreTraining",
+      "AutoModelForMultipleChoice": "nomic-ai/nomic-bert-2048--modeling_hf_nomic_bert.NomicBertForMultipleChoice",
+      "AutoModelForQuestionAnswering": "nomic-ai/nomic-bert-2048--modeling_hf_nomic_bert.NomicBertForQuestionAnswering",
+      "AutoModelForSequenceClassification": "nomic-ai/nomic-bert-2048--modeling_hf_nomic_bert.NomicBertForSequenceClassification",
+      "AutoModelForTokenClassification": "nomic-ai/nomic-bert-2048--modeling_hf_nomic_bert.NomicBertForTokenClassification"
+    },
+    "bos_token_id": null,
+    "causal": false,
+    "dense_seq_output": true,
+    "dtype": "float32",
+    "embd_pdrop": 0.0,
+    "eos_token_id": null,
+    "fused_bias_fc": true,
+    "fused_dropout_add_ln": true,
+    "initializer_range": 0.02,
+    "layer_norm_epsilon": 1e-12,
+    "max_trained_positions": 2048,
+    "mlp_fc1_bias": false,
+    "mlp_fc2_bias": false,
+    "model_type": "nomic_bert",
+    "n_embd": 768,
+    "n_head": 12,
+    "n_inner": 3072,
+    "n_layer": 12,
+    "n_positions": 8192,
+    "pad_vocab_size_multiple": 64,
+    "parallel_block": false,
+    "parallel_block_tied_norm": false,
+    "prenorm": false,
+    "qkv_proj_bias": false,
+    "reorder_and_upcast_attn": false,
+    "resid_pdrop": 0.0,
+    "rotary_emb_base": 1000,
+    "rotary_emb_fraction": 1.0,
+    "rotary_emb_interleaved": false,
+    "rotary_emb_scale_base": null,
+    "rotary_scaling_factor": null,
+    "scale_attn_by_inverse_layer_idx": false,
+    "scale_attn_weights": true,
+    "summary_activation": null,
+    "summary_first_dropout": 0.0,
+    "summary_proj_to_labels": true,
+    "summary_type": "cls_index",
+    "summary_use_proj": true,
+    "type_vocab_size": 2,
+    "use_cache": true,
+    "use_flash_attn": true,
+    "use_rms_norm": false,
+    "use_xentropy": true,
+    "vocab_size": 30528
+  },
+  "encoder_only": false,
+  "model_type": "embedder",
+  "num_blocks": 2,
+  "transformers_version": "4.57.3"
+}

configuration_embedder.py

@@ -0,0 +1,37 @@
+from typing import Any
+
+from transformers import AutoConfig, PretrainedConfig
+
+
+class EmbedderConfig(PretrainedConfig):
+    model_type = "embedder"
+
+    def __init__(
+        self,
+        base_model_name: str = "nomic-ai/nomic-embed-text-v1.5",
+        num_blocks: int = 2,
+        dropout: float = 0.0,
+        encoder_only: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.base_model_name = base_model_name
+        self.num_blocks = num_blocks
+        self.dropout = dropout
+        self.encoder_only = encoder_only
+        self.encoder_config = AutoConfig.from_pretrained(
+            base_model_name,
+            trust_remote_code=True,
+        )
+        self.auto_map = {
+            "AutoConfig": "configuration_embedder.EmbedderConfig",
+            "AutoModel": "modeling_embedder.EmbedderModel",
+        }
+
+    def __getattr__(self, name: str) -> Any:
+        if name != "encoder_config" and hasattr(self.encoder_config, name):
+            return getattr(self.encoder_config, name)
+        raise AttributeError(name)
+
+
+EmbedderConfig.register_for_auto_class()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e8b233a6180cf2e969f267c82df53de264a7d8433b723d666f0c5c9e85bd6a8
+size 561127080

modeling_embedder.py

@@ -0,0 +1,78 @@
+from typing import Any, Optional, cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import nn
+from transformers import AutoModel, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from .configuration_embedder import EmbedderConfig
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, dropout: float):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.ReLU(),
+        )
+        self.norm = nn.LayerNorm(dim)
+        self.dropout = nn.Dropout(dropout)
+        self.proj = nn.Linear(hidden_dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.net(x)
+        x = self.dropout(x)
+        x = self.proj(x)
+        return cast(torch.Tensor, self.norm(x + residual))
+
+
+class Head(nn.Module):
+    def __init__(self, dim: int, num_blocks: int = 1, dropout: float = 0):
+        super().__init__()
+        self.blocks = nn.Sequential(
+            *[EncoderBlock(dim=dim, hidden_dim=dim, dropout=dropout) for _ in range(num_blocks)]
+        )
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.blocks(x)
+        x = self.proj(x)
+        return x
+
+
+class EmbedderModel(PreTrainedModel):
+    config_class = EmbedderConfig  # type: ignore[assignment]
+    base_model_prefix = "model"
+    _supports_attention_backend = True
+
+    def __init__(self, config: EmbedderConfig):
+        super().__init__(config)
+        self.encoder = AutoModel.from_config(
+            config.encoder_config,
+            trust_remote_code=True,
+        )
+        self._init_requires_grad(self.encoder)
+        self.head = Head(
+            dim=self.encoder.embeddings.word_embeddings.embedding_dim,
+            num_blocks=config.num_blocks,
+            dropout=config.dropout,
+        )
+
+    def _init_requires_grad(self, module: nn.Module) -> None:
+        for p in module.parameters():
+            p.requires_grad = False
+
+    def forward(
+        self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs: Any
+    ) -> BaseModelOutput:
+        hidden_states = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+        if not self.config.encoder_only:
+            emb = self.head(hidden_states)  # B, T, D
+        emb = F.normalize(emb, dim=-1)
+        return BaseModelOutput(last_hidden_state=emb)  # type: ignore[arg-type]
+
+
+EmbedderModel.register_for_auto_class()

vllm_configuration_embedder.py

@@ -0,0 +1,92 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+from vllm.model_executor.models.config import VerifyAndUpdateConfig
+
+
+class EmbedderModelConfig(VerifyAndUpdateConfig):
+
+    @staticmethod
+    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
+        from copy import deepcopy
+
+        from vllm.transformers_utils.config import set_default_rope_theta
+
+        config = vllm_config.model_config.hf_config
+        assert config.__class__.__name__ == "EmbedderConfig"
+        assert config.activation_function in ["swiglu", "gelu"]
+        config.position_embedding_type = getattr(config, "position_embedding_type", "rope")
+
+        if config.activation_function == "swiglu":
+            config.hidden_act = "silu"
+        else:
+            config.hidden_act = config.activation_function
+
+        assert config.mlp_fc1_bias == config.mlp_fc2_bias == config.qkv_proj_bias
+        config.bias = config.qkv_proj_bias
+
+        assert config.rotary_emb_scale_base is None
+        assert not config.rotary_emb_interleaved
+
+        config.layer_norm_eps = config.layer_norm_epsilon
+        config.intermediate_size = config.n_inner
+        config.hidden_size = config.n_embd
+        config.num_hidden_layers = config.n_layer
+
+        head_dim = config.hidden_size // config.num_attention_heads
+        rotary_emb_dim = int(head_dim * config.rotary_emb_fraction)
+        max_trained_positions = getattr(config, "max_trained_positions", 2048)
+
+        set_default_rope_theta(config, default_theta=config.rotary_emb_base)
+
+        config.rotary_kwargs = {
+            "head_size": head_dim,
+            "rotary_dim": rotary_emb_dim,
+            "max_position": max_trained_positions,
+            "rope_parameters": config.rope_parameters,
+        }
+
+        # we ignore config.rotary_scaling_factor so that for datasets shorter
+        # than max_trained_positions 2048, the results are consistent
+        # with SentenceTransformer.
+        # The context extension uses vllm style rope_theta and rope_parameters.
+        # See #17785 #18755
+        if not vllm_config.model_config.hf_overrides and vllm_config.model_config.original_max_model_len is None:
+            # Default
+            # Reset max_model_len to max_trained_positions.
+            # nomic-embed-text-v2-moe the length is set to 512
+            # by sentence_bert_config.json.
+            max_model_len = min(vllm_config.model_config.max_model_len, max_trained_positions)  # type: ignore[unreachable]
+
+            vllm_config.recalculate_max_model_len(max_model_len)
+
+        else:
+            # We need to re-verify max_model_len to avoid lengths
+            # greater than position_embedding.
+            model_config = vllm_config.model_config
+            hf_text_config = model_config.hf_text_config
+
+            if isinstance(model_config.hf_overrides, dict):
+                # hf_overrides_kw
+                max_model_len = model_config.hf_overrides.get("max_model_len", vllm_config.model_config.max_model_len)
+            else:
+                # hf_overrides_fn
+                # This might be overridden by sentence_bert_config.json.
+                max_model_len = vllm_config.model_config.max_model_len
+
+            # reset hf_text_config for recalculate_max_model_len.
+            if hasattr(hf_text_config, "max_model_len"):
+                delattr(hf_text_config, "max_model_len")
+            hf_text_config.max_position_embeddings = max_trained_positions
+            hf_text_config.rope_parameters = config.rotary_kwargs["rope_parameters"]
+
+            # The priority of sentence_bert_config.json is higher
+            # than max_position_embeddings
+            encoder_config = deepcopy(model_config.encoder_config)
+            if encoder_config:
+                encoder_config.pop("max_seq_length", None)
+            model_config.encoder_config = encoder_config
+
+            vllm_config.recalculate_max_model_len(max_model_len)

vllm_modeling_embedder.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Optional, cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.model_executor.models.bert_with_rope import NomicBertModel
+from vllm.model_executor.models.interfaces_base import default_pooling_type
+from vllm.model_executor.models.utils import WeightsMapper
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, dropout: float):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.ReLU(),
+        )
+        self.norm = nn.LayerNorm(dim)
+        self.dropout = nn.Dropout(dropout)
+        self.proj = nn.Linear(hidden_dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.net(x)
+        x = self.dropout(x)
+        x = self.proj(x)
+        return cast(torch.Tensor, self.norm(x + residual))
+
+
+class Head(nn.Module):
+    def __init__(self, dim: int, num_blocks: int = 1, dropout: float = 0):
+        super().__init__()
+        self.blocks = nn.Sequential(
+            *[EncoderBlock(dim=dim, hidden_dim=dim, dropout=dropout) for _ in range(num_blocks)]
+        )
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.blocks(x)
+        x = self.proj(x)
+        return x
+
+
+@support_torch_compile
+@default_pooling_type("CLS")
+class EmbedderModel(nn.Module):
+    """
+    vLLM wrapper for HF-trained EmbedderModel
+    (encoder + custom graph head)
+    """
+
+    # HF state_dict keys start with "model."
+    hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ):
+
+        super().__init__()
+        self.hf_config = vllm_config.model_config.hf_config
+        # --------------------------------------------------
+        # Base encoder (identical to training)
+        # --------------------------------------------------
+        self.encoder = NomicBertModel(
+            vllm_config=vllm_config,
+            prefix=f"{prefix}.encoder",
+            add_pooling_layer=False,
+        )
+        # --------------------------------------------------
+        # Custom head (must match HF exactly)
+        # --------------------------------------------------
+        self.head = Head(
+            dim=self.hf_config.hidden_size,
+            num_blocks=self.hf_config.num_blocks,
+            dropout=self.hf_config.dropout,
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.encoder.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # vLLM manages attention & KV internally
+        hidden_states = self.encoder(
+            input_ids=input_ids,
+            positions=positions,
+            inputs_embeds=inputs_embeds,
+            token_type_ids=token_type_ids,
+        )
+        if not self.hf_config.encoder_only:
+            # Head + normalize (same as HF)
+            emb = self.head(hidden_states)
+        emb = F.normalize(emb, dim=-1)
+        return emb