test api

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:005e215841154f89015dc063bc0eec38388d97838a198df4172e73d1b7d654e9
+oid sha256:0f6e509537aa47a23abdd3cd174dc0a8f4398c603439041181e7ef804deb07f4
 size 565859400

vllm_configuration_embedder.py

@@ -0,0 +1,92 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig  # ty: ignore[unresolved-import]
+
+
+from vllm.model_executor.models.config import VerifyAndUpdateConfig  # ty: ignore[unresolved-import]
+
+
+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  # ty: ignore[unresolved-import]
+
+        config = vllm_config.model_config.hf_config
+        assert config.__class__.__name__ == "EmbedderConfig"  # nosec B101
+        assert config.activation_function in ["swiglu", "gelu"]  # nosec B101
+        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  # nosec B101
+        config.bias = config.qkv_proj_bias
+
+        assert config.rotary_emb_scale_base is None  # nosec B101
+        assert not config.rotary_emb_interleaved  # nosec B101
+
+        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 cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.compilation.decorators import support_torch_compile  # ty: ignore[unresolved-import]
+from vllm.config import VllmConfig  # ty: ignore[unresolved-import]
+from vllm.model_executor.models.bert_with_rope import NomicBertModel  # ty: ignore[unresolved-import]
+from vllm.model_executor.models.interfaces_base import default_pooling_type  # ty: ignore[unresolved-import]
+from vllm.model_executor.models.utils import WeightsMapper  # ty: ignore[unresolved-import]
+
+
+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: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        token_type_ids: torch.Tensor | None = 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,
+        )
+        emb = hidden_states
+        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