File size: 2,524 Bytes
b7875cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """Sentence Transformers module for the V-SPLADE inference-free query encoder.
Referenced from ``router_config.json``: the "query" route uses
:class:`VSPLADEStaticEmbedding`, whose weights are the precomputed Li-LSR
lookup table ``softplus(projection(embedding))`` extracted from the
``query_encoder.*`` tensors in ``model.safetensors`` (with the special tokens
[UNK]/[CLS]/[SEP]/[PAD]/[MASK] zeroed out).
"""
from __future__ import annotations
import torch
try:
# sentence-transformers >= 5.6
from sentence_transformers.sparse_encoder.modules import SparseStaticEmbedding
except ImportError:
from sentence_transformers.sparse_encoder.models import SparseStaticEmbedding
class VSPLADEStaticEmbedding(SparseStaticEmbedding):
"""Inference-free Li-LSR query encoder for V-SPLADE.
Behaves like :class:`SparseStaticEmbedding` with two differences, matching
``InferenceFreeQueryEncoder.encode_with_lookup`` from
https://github.com/naver/v-splade:
* repeated query tokens accumulate their weight (scatter-add) instead of
being counted once;
* token ids outside the lookup table (the 40 added vision tokens, e.g.
``<image>``) contribute nothing instead of raising an index error;
* the lookup table covers the base (MLM) vocabulary (50368 entries), which
is smaller than the full tokenizer vocabulary, so its size is stored in
the module config (``num_dimensions``) for loading.
"""
config_keys: list[str] = ["frozen", "num_dimensions"]
def __init__(self, tokenizer, weight: torch.Tensor | None = None, frozen: bool = False, num_dimensions: int | None = None):
if weight is None and num_dimensions is not None:
weight = torch.zeros(num_dimensions)
super().__init__(tokenizer=tokenizer, weight=weight, frozen=frozen)
def forward(self, features: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
input_ids = features["input_ids"]
attention_mask = features["attention_mask"]
valid = (input_ids < self.num_dimensions) & (attention_mask > 0)
safe_ids = input_ids.clamp(max=self.num_dimensions - 1)
scores = self.weight[safe_ids] * valid.to(self.weight.dtype)
embeddings = torch.zeros(
input_ids.size(0), self.num_dimensions, device=input_ids.device, dtype=self.weight.dtype
)
embeddings.scatter_add_(1, safe_ids, scores)
features["sentence_embedding"] = embeddings
return features
__all__ = ["VSPLADEStaticEmbedding"]
|