Upload folder using huggingface_hub

config.json

@@ -0,0 +1,89 @@
+{
+  "architectures": [
+    "SpladeModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "modeling_splade.SpladeConfig",
+    "AutoModel": "modeling_splade.SpladeModel"
+  },
+  "bos_token_id": 50281,
+  "classifier_activation": "gelu",
+  "classifier_bias": false,
+  "classifier_dropout": 0.0,
+  "classifier_pooling": "mean",
+  "cls_token_id": 50281,
+  "decoder_bias": true,
+  "deterministic_flash_attn": false,
+  "doc_max_length": 512,
+  "document_prefix": "[D] ",
+  "dtype": "float32",
+  "embedding_dropout": 0.0,
+  "eos_token_id": 50282,
+  "global_attn_every_n_layers": 3,
+  "gradient_checkpointing": false,
+  "hidden_activation": "gelu",
+  "hidden_size": 768,
+  "initializer_cutoff_factor": 2.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 1152,
+  "layer_norm_eps": 1e-05,
+  "layer_types": [
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "local_attention": 128,
+  "logit_shift": 15,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "mlp_dropout": 0.0,
+  "model_type": "modernbert",
+  "norm_bias": false,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 22,
+  "pad_token_id": 50283,
+  "position_embedding_type": "absolute",
+  "position_top_k": 12,
+  "query_max_length": 128,
+  "query_prefix": "[Q] ",
+  "repad_logits_with_grad": false,
+  "rope_parameters": {
+    "full_attention": {
+      "rope_theta": 160000.0,
+      "rope_type": "default"
+    },
+    "sliding_attention": {
+      "rope_theta": 10000.0,
+      "rope_type": "default"
+    }
+  },
+  "sep_token_id": 50282,
+  "sparse_pred_ignore_index": -100,
+  "sparse_prediction": false,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "vocab_fold": "case_space",
+  "vocab_size": 50370
+}

config_sentence_transformers.json

@@ -0,0 +1,6 @@
+{
+  "model_type": "SparseEncoder",
+  "prompts": {"query": "[Q] ", "document": "[D] "},
+  "default_prompt_name": null,
+  "similarity_fn_name": "dot"
+}

custom_st.py

@@ -0,0 +1,26 @@
+from sentence_transformers.base.modules import InputModule
+from transformers import AutoModel, AutoTokenizer
+class SpladeSTModule(InputModule):
+    save_in_root = True
+    def __init__(self, model_name_or_path: str, **kwargs):
+        super().__init__()
+        self.model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)  # for SparseEncoder.decode
+    def preprocess(self, inputs, prompt=None, **kwargs):
+        prefix = prompt or ""
+        cfg = self.model.config
+        max_length = cfg.query_max_length if prefix == cfg.query_prefix else cfg.doc_max_length
+        ids, attn, pool, _ = self.model._tokenize(list(inputs), prefix, max_length)
+        return {"input_ids": ids, "attention_mask": attn, "pooling_mask": pool}
+    def forward(self, features, **kwargs):
+        features["sentence_embedding"] = self.model(
+            features["input_ids"], features["attention_mask"], features["pooling_mask"]
+        )
+        return features
+    def get_embedding_dimension(self):
+        return self.model.config.vocab_size
+    @classmethod
+    def load(cls, model_name_or_path, **kwargs):
+        return cls(model_name_or_path)
+    def save(self, output_path, **kwargs):
+        pass  # repo is assembled by export.py, never by ST save

model.safetensors

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

modeling_splade.py

@@ -0,0 +1,387 @@
+"""SPLADE head over a ModernBERT/LateOn MLM backbone.
+
+Sparse vector = fold(max-pool(top_k-gate(log1p(relu(logits - shift))) * pooling_mask))
+where the instruction prefix ("[Q] " / "[D] ") is attended by the backbone but
+excluded from pooling. Scores are dot products. Load with:
+
+    model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
+    q = model.encode(["a query"], kind="query")                # [N, V] float32
+    d = model.encode(["a document"])                           # [N, V] float32
+    model.score(q, d)                                          # [Nq, Nd] dot
+    model.encode_to_dict(["a query"], kind="query", top_k=20)  # {token: weight}
+    model.attribute(["a query"], kind="query")                 # + winning input token
+    print(model.render(["a query"], kind="query"))             # terminal bar chart
+    print(model.highlight(["a document"]))                     # text, fired words lit
+"""
+
+from __future__ import annotations
+
+import sys
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, ModernBertConfig, ModernBertForMaskedLM
+
+
+class SpladeConfig(ModernBertConfig):
+    def __init__(
+        self,
+        logit_shift: float = 0.0,
+        position_top_k: int | None = None,
+        vocab_fold: str | None = None,
+        query_prefix: str = "",
+        document_prefix: str = "",
+        query_max_length: int = 128,
+        doc_max_length: int = 512,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.logit_shift = logit_shift
+        self.position_top_k = position_top_k
+        self.vocab_fold = vocab_fold
+        self.query_prefix = query_prefix
+        self.document_prefix = document_prefix
+        self.query_max_length = query_max_length
+        self.doc_max_length = doc_max_length
+
+
+class SpladeModel(ModernBertForMaskedLM):
+    config_class = SpladeConfig
+
+    def __init__(self, config: SpladeConfig):
+        super().__init__(config)
+        # [V] canonical-id map for vocab folding, computed once at export from
+        # the tokenizer and stored in the checkpoint (identity when unused).
+        self.register_buffer(
+            "vocab_fold_index", torch.arange(config.vocab_size), persistent=True
+        )
+        self._tokenizer = None
+
+    # -- forward path ---------------------------------------------------------
+
+    def _token_weights(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
+    ) -> torch.Tensor:
+        """[B, L] tokens -> [B, L, V] per-position activations (pre-pooling)."""
+        cfg = self.config
+        logits = super().forward(input_ids=input_ids, attention_mask=attention_mask).logits
+        weights = torch.log1p(F.relu(logits - cfg.logit_shift))
+        if cfg.position_top_k is not None and cfg.position_top_k < weights.shape[-1]:
+            # Keep each position's k largest dims (ties keep more than k).
+            cutoff = weights.topk(cfg.position_top_k, dim=-1).values[..., -1:]
+            weights = weights * (weights >= cutoff)
+        return weights
+
+    def _fold(
+        self, sparse: torch.Tensor, source_indices: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Reroute each fold group's mass onto its canonical vocab dim.
+
+        `source_indices` [B, V] (max-pool argmax positions) follows the winning
+        group member so attribution keeps pointing at a real input position.
+        """
+        index = self.vocab_fold_index.unsqueeze(0).expand_as(sparse)
+        folded = torch.zeros_like(sparse).scatter_reduce(
+            1, index, sparse, reduce="amax", include_self=False
+        )
+        if source_indices is None:
+            return folded, None
+        winner = (sparse == folded.gather(1, index)) & (sparse > 0)
+        folded_sources = torch.full_like(source_indices, -1).scatter_reduce(
+            1, index, torch.where(winner, source_indices, -1), reduce="amax", include_self=False
+        )
+        return folded, folded_sources.clamp_min_(0)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        pooling_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """[B, L] tokens -> [B, V] sparse activations (dot-product scoring)."""
+        if pooling_mask is None:
+            pooling_mask = attention_mask
+        weights = self._token_weights(input_ids, attention_mask)
+        weights = weights * pooling_mask.unsqueeze(-1).to(weights.dtype)
+        sparse = weights.max(dim=1).values
+        if self.config.vocab_fold is not None:
+            sparse, _ = self._fold(sparse)
+        return sparse
+
+    @staticmethod
+    def score(queries: torch.Tensor, documents: torch.Tensor) -> torch.Tensor:
+        """Dot-product relevance scores: [Nq, V] x [Nd, V] -> [Nq, Nd]."""
+        return queries @ documents.T
+
+    # -- tokenization ---------------------------------------------------------
+
+    def _get_tokenizer(self):
+        if self._tokenizer is None:
+            self._tokenizer = AutoTokenizer.from_pretrained(self.config._name_or_path)
+        return self._tokenizer
+
+    def _encode_args(self, kind: str, max_length: int | None) -> tuple[str, int]:
+        cfg = self.config
+        if kind == "query":
+            return cfg.query_prefix, max_length or cfg.query_max_length
+        if kind == "document":
+            return cfg.document_prefix, max_length or cfg.doc_max_length
+        raise ValueError("`kind` must be 'query' or 'document'.")
+
+    def _tokenize(self, texts: list[str], prefix: str, max_length: int):
+        """-> (input_ids, attention_mask, pooling_mask, char_offsets)."""
+        enc = self._get_tokenizer()(
+            [prefix + t for t in texts],
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+            return_tensors="pt",
+            return_offsets_mapping=True,
+            return_special_tokens_mask=True,
+        )
+        pooling_mask = enc["attention_mask"]
+        if prefix:
+            # Prefix tokens are attended but dropped from pooling: a token is
+            # prefix iff its char span starts inside the prefix string.
+            # Specials carry a zero-width (0, 0) span, so guard them.
+            in_prefix = (enc["offset_mapping"][..., 0] < len(prefix)) & ~enc[
+                "special_tokens_mask"
+            ].bool()
+            pooling_mask = pooling_mask.masked_fill(in_prefix, 0)
+        return enc["input_ids"], enc["attention_mask"], pooling_mask, enc["offset_mapping"]
+
+    def _encode_with_sources(self, texts: list[str], prefix: str, max_length: int):
+        """One batch through the model, keeping max-pool source positions.
+
+        -> (input_ids, pooling_mask, char_offsets, sparse [B, V], sources [B, V])
+        """
+        device = next(self.parameters()).device
+        ids, attn, pool, offsets = self._tokenize(texts, prefix, max_length)
+        ids, attn, pool = ids.to(device), attn.to(device), pool.to(device)
+        weights = self._token_weights(ids, attn)
+        weights = weights * pool.unsqueeze(-1).to(weights.dtype)
+        pooled = weights.max(dim=1)
+        sparse, sources = pooled.values, pooled.indices
+        if self.config.vocab_fold is not None:
+            sparse, sources = self._fold(sparse, sources)
+        return ids, pool, offsets, sparse, sources
+
+    # -- encoding APIs --------------------------------------------------------
+
+    @torch.inference_mode()
+    def encode(
+        self,
+        texts: list[str],
+        kind: str = "document",
+        batch_size: int = 32,
+        max_length: int | None = None,
+    ) -> torch.Tensor:
+        """Encode raw texts -> [N, V] float32 sparse vectors on CPU.
+
+        `kind` ("query" | "document") selects the instruction prefix and the
+        default max length.
+        """
+        prefix, max_length = self._encode_args(kind, max_length)
+        device = next(self.parameters()).device
+        rows = []
+        for start in range(0, len(texts), batch_size):
+            ids, attn, pool, _ = self._tokenize(
+                texts[start : start + batch_size], prefix, max_length
+            )
+            rows.append(
+                self(ids.to(device), attn.to(device), pool.to(device)).float().cpu()
+            )
+        return torch.cat(rows)
+
+    @torch.inference_mode()
+    def attribute(
+        self,
+        texts: list[str],
+        kind: str = "document",
+        top_k: int | None = 25,
+        batch_size: int = 32,
+        max_length: int | None = None,
+        round_to: int = 4,
+    ) -> list[list[dict]]:
+        """Encode texts and attribute each output dim to its input subtoken.
+
+        Per text, a weight-sorted list of entries
+        `{"token", "weight", "source", "position", "expansion"}`:
+        `source`/`position` name the input subtoken whose activation won the
+        max for that vocab dim (after folding); `expansion` is True when the
+        output term is not the source token's own (folded) dim.
+        """
+        prefix, max_length = self._encode_args(kind, max_length)
+        tokenizer = self._get_tokenizer()
+        fold_index = self.vocab_fold_index.cpu()
+        results = []
+        for start in range(0, len(texts), batch_size):
+            ids, _, _, sparse, sources = self._encode_with_sources(
+                texts[start : start + batch_size], prefix, max_length
+            )
+            for row, row_sources, row_ids in zip(
+                sparse.float().cpu(), sources.cpu(), ids.cpu()
+            ):
+                dims = torch.nonzero(row, as_tuple=False).flatten()
+                order = torch.argsort(row[dims], descending=True)[:top_k]
+                entries = []
+                for dim in dims[order].tolist():
+                    pos = int(row_sources[dim])
+                    src_id = int(row_ids[pos])
+                    entries.append(
+                        {
+                            "token": tokenizer.convert_ids_to_tokens(dim),
+                            "weight": round(float(row[dim]), round_to),
+                            "source": tokenizer.convert_ids_to_tokens(src_id),
+                            "position": pos,
+                            "expansion": int(fold_index[src_id]) != dim,
+                        }
+                    )
+                results.append(entries)
+        return results
+
+    def encode_to_dict(
+        self, texts: list[str], kind: str = "document", top_k: int | None = None, **kwargs
+    ) -> list[dict[str, float]]:
+        """Encode texts -> {token: weight} dicts sorted by descending weight."""
+        return [
+            {e["token"]: e["weight"] for e in entries}
+            for entries in self.attribute(texts, kind=kind, top_k=top_k, **kwargs)
+        ]
+
+    # -- terminal displays ----------------------------------------------------
+
+    _FADE = "▓▒░"
+    _HEAT = (196, 202, 208, 214, 220, 190, 108, 66, 60, 241)  # ANSI-256, hot -> cold
+
+    def _heat(self, ratio: float) -> int:
+        return self._HEAT[min(int((1 - ratio) * len(self._HEAT)), len(self._HEAT) - 1)]
+
+    @staticmethod
+    def _display(token: str) -> str:
+        """Strip the Ġ word marker; dot-prefix continuation pieces."""
+        if token.startswith("Ġ"):
+            return token[1:]
+        if token.startswith("["):  # specials: [CLS], [SEP], [Q], [D]
+            return token
+        return "·" + token
+
+    def render(
+        self,
+        texts: list[str],
+        kind: str = "document",
+        top_k: int | None = 25,
+        width: int = 36,
+        color: bool | None = None,
+        **attribute_kwargs,
+    ) -> str:
+        """Terminal bar chart of the sparse expansions, with attributions.
+
+        One block per text: bars proportional to weight (peak-normalized), each
+        line ending with the input subtoken that produced the dim and `<exp>`
+        for pure expansions. `color=None` auto-detects a TTY.
+        """
+        if color is None:
+            color = sys.stdout.isatty()
+        blocks = []
+        for text, entries in zip(
+            texts, self.attribute(texts, kind=kind, top_k=top_k, **attribute_kwargs)
+        ):
+            shown = text if len(text) <= 70 else text[:67] + "..."
+            if not entries:
+                blocks.append(f"{kind} · {shown}\n  (empty vector)")
+                continue
+            peak = entries[0]["weight"]
+            name_width = max(len(self._display(e["token"])) for e in entries)
+            lines = [f"{kind} · {shown}"]
+            for e in entries:
+                ratio = e["weight"] / peak
+                cells = max(1, round(ratio * width))
+                bar = ("█" * cells)[:-3] + self._FADE if cells > 3 else self._FADE[3 - cells :]
+                pad = " " * (width - cells)
+                token = self._display(e["token"]).rjust(name_width)
+                attrib = f"<- {self._display(e['source'])}@{e['position']}"
+                if e["expansion"]:
+                    attrib += " <exp>"
+                if color:
+                    heat = self._heat(ratio)
+                    token = f"\033[38;5;{heat}m{token}\033[0m"
+                    bar = f"\033[38;5;{heat}m{bar}\033[0m"
+                    attrib = f"\033[2m{attrib}\033[0m"
+                lines.append(f"{token} {bar}{pad} {e['weight']:>6.2f}  {attrib}")
+            blocks.append("\n".join(lines))
+        return "\n\n".join(blocks)
+
+    @torch.inference_mode()
+    def highlight(
+        self,
+        texts: list[str],
+        kind: str = "document",
+        color: bool | None = None,
+        batch_size: int = 32,
+        max_length: int | None = None,
+    ) -> str:
+        """Render each text with its firing words lit up.
+
+        A word fires when one of its subtokens wins the max for at least one
+        output dim; intensity is the largest weight it wins. TTY: reverse-video
+        heat colors. Plain: tiered markers `⟦strong⟧ «mid» ‹weak›`.
+        """
+        if color is None:
+            color = sys.stdout.isatty()
+        prefix, max_length = self._encode_args(kind, max_length)
+        blocks = []
+        for start in range(0, len(texts), batch_size):
+            batch = texts[start : start + batch_size]
+            ids, pool, offsets, sparse, sources = self._encode_with_sources(
+                batch, prefix, max_length
+            )
+            # [B, L] per-position intensity: max weight over the dims each
+            # position won. Zero-weight dims carry a clamped position 0 but
+            # contribute 0, so they can't corrupt the amax.
+            intensity = torch.zeros(
+                ids.shape, dtype=sparse.dtype, device=sparse.device
+            ).scatter_reduce(1, sources, sparse, reduce="amax", include_self=False)
+            for text, row_int, row_off, row_pool in zip(
+                batch, intensity.float().cpu(), offsets, pool.cpu()
+            ):
+                blocks.append(self._paint(text, row_int, row_off, row_pool, len(prefix), color))
+        return "\n".join(blocks)
+
+    def _paint(self, text, intensity, offsets, pooling_mask, prefix_len, color) -> str:
+        """Wrap fired char spans of `text` in intensity markers."""
+        peak = intensity.max().item()
+        if peak <= 0:
+            return text
+        # Byte-BPE offsets include the word's leading space: trim it, so merging
+        # only fuses glued subtokens of the same word (one span per word).
+        spans: list[list] = []
+        for pos in range(len(offsets)):
+            w = intensity[pos].item()
+            if w <= 0 or pooling_mask[pos] == 0:
+                continue
+            s, e = int(offsets[pos][0]) - prefix_len, int(offsets[pos][1]) - prefix_len
+            s = max(s, 0)
+            while s < e and text[s].isspace():
+                s += 1
+            if e <= s:  # zero-width specials / whitespace-only
+                continue
+            if spans and s == spans[-1][1]:
+                spans[-1][1] = e
+                spans[-1][2] = max(spans[-1][2], w)
+            else:
+                spans.append([s, e, w])
+        out, cursor = [], 0
+        for s, e, w in spans:
+            ratio = w / peak
+            out.append(text[cursor:s])
+            if color:
+                # Reverse video with heat foreground = heat-colored highlighter.
+                out.append(f"\033[7;38;5;{self._heat(ratio)}m{text[s:e]}\033[0m")
+            else:
+                marks = "⟦⟧" if ratio > 0.66 else "«»" if ratio > 0.33 else "‹›"
+                out.append(f"{marks[0]}{text[s:e]}{marks[1]}")
+            cursor = e
+        out.append(text[cursor:])
+        return "".join(out)

modules.json

@@ -0,0 +1 @@
+[{"idx": 0, "name": "0", "path": "", "type": "custom_st.SpladeSTModule"}]

tokenizer_config.json

@@ -0,0 +1,23 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "is_local": true,
+  "mask_token": "[MASK]",
+  "max_length": 299,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 299,
+  "pad_to_multiple_of": null,
+  "pad_token": "[MASK]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "tokenizer_class": "TokenizersBackend",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}