Linkup-Platform/linkup-sparseup-embed-v1

1"""SPLADE head over a ModernBERT/LateOn MLM backbone.2 3Sparse vector = fold(max-pool(top_k-gate(log1p(relu(logits - shift))) * pooling_mask))4where the instruction prefix ("[Q] " / "[D] ") is attended by the backbone but5excluded from pooling. Scores are dot products. Load with:6 7    model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)8    q = model.encode(["a query"], kind="query")                # [N, V] float329    d = model.encode(["a document"])                           # [N, V] float3210    model.score(q, d)                                          # [Nq, Nd] dot11    model.encode_to_dict(["a query"], kind="query", top_k=20)  # {token: weight}12    model.attribute(["a query"], kind="query")                 # + winning input token13    print(model.render(["a query"], kind="query"))             # terminal bar chart14    print(model.highlight(["a document"]))                     # text, fired words lit15"""16 17from __future__ import annotations18 19import sys20 21import torch22import torch.nn.functional as F23from transformers import AutoTokenizer, ModernBertConfig, ModernBertForMaskedLM24 25 26class SpladeConfig(ModernBertConfig):27    def __init__(28        self,29        logit_shift: float = 0.0,30        position_top_k: int | None = None,31        vocab_fold: str | None = None,32        query_prefix: str = "",33        document_prefix: str = "",34        query_max_length: int = 128,35        doc_max_length: int = 512,36        **kwargs,37    ):38        super().__init__(**kwargs)39        self.logit_shift = logit_shift40        self.position_top_k = position_top_k41        self.vocab_fold = vocab_fold42        self.query_prefix = query_prefix43        self.document_prefix = document_prefix44        self.query_max_length = query_max_length45        self.doc_max_length = doc_max_length46 47 48class SpladeModel(ModernBertForMaskedLM):49    config_class = SpladeConfig50 51    def __init__(self, config: SpladeConfig):52        super().__init__(config)53        # [V] canonical-id map for vocab folding, computed once at export from54        # the tokenizer and stored in the checkpoint (identity when unused).55        self.register_buffer(56            "vocab_fold_index", torch.arange(config.vocab_size), persistent=True57        )58        self._tokenizer = None59 60    # -- forward path ---------------------------------------------------------61 62    def _token_weights(63        self, input_ids: torch.Tensor, attention_mask: torch.Tensor64    ) -> torch.Tensor:65        """[B, L] tokens -> [B, L, V] per-position activations (pre-pooling)."""66        cfg = self.config67        logits = super().forward(input_ids=input_ids, attention_mask=attention_mask).logits68        weights = torch.log1p(F.relu(logits - cfg.logit_shift))69        if cfg.position_top_k is not None and cfg.position_top_k < weights.shape[-1]:70            # Keep each position's k largest dims (ties keep more than k).71            cutoff = weights.topk(cfg.position_top_k, dim=-1).values[..., -1:]72            weights = weights * (weights >= cutoff)73        return weights74 75    def _fold(76        self, sparse: torch.Tensor, source_indices: torch.Tensor | None = None77    ) -> tuple[torch.Tensor, torch.Tensor | None]:78        """Reroute each fold group's mass onto its canonical vocab dim.79 80        `source_indices` [B, V] (max-pool argmax positions) follows the winning81        group member so attribution keeps pointing at a real input position.82        """83        index = self.vocab_fold_index.unsqueeze(0).expand_as(sparse)84        folded = torch.zeros_like(sparse).scatter_reduce(85            1, index, sparse, reduce="amax", include_self=False86        )87        if source_indices is None:88            return folded, None89        winner = (sparse == folded.gather(1, index)) & (sparse > 0)90        folded_sources = torch.full_like(source_indices, -1).scatter_reduce(91            1, index, torch.where(winner, source_indices, -1), reduce="amax", include_self=False92        )93        return folded, folded_sources.clamp_min_(0)94 95    def forward(96        self,97        input_ids: torch.Tensor,98        attention_mask: torch.Tensor,99        pooling_mask: torch.Tensor | None = None,100        **kwargs,101    ) -> torch.Tensor:102        """[B, L] tokens -> [B, V] sparse activations (dot-product scoring)."""103        if pooling_mask is None:104            pooling_mask = attention_mask105        weights = self._token_weights(input_ids, attention_mask)106        weights = weights * pooling_mask.unsqueeze(-1).to(weights.dtype)107        sparse = weights.max(dim=1).values108        if self.config.vocab_fold is not None:109            sparse, _ = self._fold(sparse)110        return sparse111 112    @staticmethod113    def score(queries: torch.Tensor, documents: torch.Tensor) -> torch.Tensor:114        """Dot-product relevance scores: [Nq, V] x [Nd, V] -> [Nq, Nd]."""115        return queries @ documents.T116 117    # -- tokenization ---------------------------------------------------------118 119    def _get_tokenizer(self):120        if self._tokenizer is None:121            self._tokenizer = AutoTokenizer.from_pretrained(self.config._name_or_path)122        return self._tokenizer123 124    def _encode_args(self, kind: str, max_length: int | None) -> tuple[str, int]:125        cfg = self.config126        if kind == "query":127            return cfg.query_prefix, max_length or cfg.query_max_length128        if kind == "document":129            return cfg.document_prefix, max_length or cfg.doc_max_length130        raise ValueError("`kind` must be 'query' or 'document'.")131 132    def _tokenize(self, texts: list[str], prefix: str, max_length: int):133        """-> (input_ids, attention_mask, pooling_mask, char_offsets)."""134        enc = self._get_tokenizer()(135            [prefix + t for t in texts],136            padding=True,137            truncation=True,138            max_length=max_length,139            return_tensors="pt",140            return_offsets_mapping=True,141            return_special_tokens_mask=True,142        )143        pooling_mask = enc["attention_mask"]144        if prefix:145            # Prefix tokens are attended but dropped from pooling: a token is146            # prefix iff its char span starts inside the prefix string.147            # Specials carry a zero-width (0, 0) span, so guard them.148            in_prefix = (enc["offset_mapping"][..., 0] < len(prefix)) & ~enc[149                "special_tokens_mask"150            ].bool()151            pooling_mask = pooling_mask.masked_fill(in_prefix, 0)152        return enc["input_ids"], enc["attention_mask"], pooling_mask, enc["offset_mapping"]153 154    def _encode_with_sources(self, texts: list[str], prefix: str, max_length: int):155        """One batch through the model, keeping max-pool source positions.156 157        -> (input_ids, pooling_mask, char_offsets, sparse [B, V], sources [B, V])158        """159        device = next(self.parameters()).device160        ids, attn, pool, offsets = self._tokenize(texts, prefix, max_length)161        ids, attn, pool = ids.to(device), attn.to(device), pool.to(device)162        weights = self._token_weights(ids, attn)163        weights = weights * pool.unsqueeze(-1).to(weights.dtype)164        pooled = weights.max(dim=1)165        sparse, sources = pooled.values, pooled.indices166        if self.config.vocab_fold is not None:167            sparse, sources = self._fold(sparse, sources)168        return ids, pool, offsets, sparse, sources169 170    # -- encoding APIs --------------------------------------------------------171 172    @torch.inference_mode()173    def encode(174        self,175        texts: list[str],176        kind: str = "document",177        batch_size: int = 32,178        max_length: int | None = None,179    ) -> torch.Tensor:180        """Encode raw texts -> [N, V] float32 sparse vectors on CPU.181 182        `kind` ("query" | "document") selects the instruction prefix and the183        default max length.184        """185        prefix, max_length = self._encode_args(kind, max_length)186        device = next(self.parameters()).device187        rows = []188        for start in range(0, len(texts), batch_size):189            ids, attn, pool, _ = self._tokenize(190                texts[start : start + batch_size], prefix, max_length191            )192            rows.append(193                self(ids.to(device), attn.to(device), pool.to(device)).float().cpu()194            )195        return torch.cat(rows)196 197    @torch.inference_mode()198    def attribute(199        self,200        texts: list[str],201        kind: str = "document",202        top_k: int | None = 25,203        batch_size: int = 32,204        max_length: int | None = None,205        round_to: int = 4,206    ) -> list[list[dict]]:207        """Encode texts and attribute each output dim to its input subtoken.208 209        Per text, a weight-sorted list of entries210        `{"token", "weight", "source", "position", "expansion"}`:211        `source`/`position` name the input subtoken whose activation won the212        max for that vocab dim (after folding); `expansion` is True when the213        output term is not the source token's own (folded) dim.214        """215        prefix, max_length = self._encode_args(kind, max_length)216        tokenizer = self._get_tokenizer()217        fold_index = self.vocab_fold_index.cpu()218        results = []219        for start in range(0, len(texts), batch_size):220            ids, _, _, sparse, sources = self._encode_with_sources(221                texts[start : start + batch_size], prefix, max_length222            )223            for row, row_sources, row_ids in zip(224                sparse.float().cpu(), sources.cpu(), ids.cpu()225            ):226                dims = torch.nonzero(row, as_tuple=False).flatten()227                order = torch.argsort(row[dims], descending=True)[:top_k]228                entries = []229                for dim in dims[order].tolist():230                    pos = int(row_sources[dim])231                    src_id = int(row_ids[pos])232                    entries.append(233                        {234                            "token": tokenizer.convert_ids_to_tokens(dim),235                            "weight": round(float(row[dim]), round_to),236                            "source": tokenizer.convert_ids_to_tokens(src_id),237                            "position": pos,238                            "expansion": int(fold_index[src_id]) != dim,239                        }240                    )241                results.append(entries)242        return results243 244    def encode_to_dict(245        self, texts: list[str], kind: str = "document", top_k: int | None = None, **kwargs246    ) -> list[dict[str, float]]:247        """Encode texts -> {token: weight} dicts sorted by descending weight."""248        return [249            {e["token"]: e["weight"] for e in entries}250            for entries in self.attribute(texts, kind=kind, top_k=top_k, **kwargs)251        ]252 253    # -- terminal displays ----------------------------------------------------254 255    _FADE = "▓▒░"256    _HEAT = (196, 202, 208, 214, 220, 190, 108, 66, 60, 241)  # ANSI-256, hot -> cold257 258    def _heat(self, ratio: float) -> int:259        return self._HEAT[min(int((1 - ratio) * len(self._HEAT)), len(self._HEAT) - 1)]260 261    @staticmethod262    def _display(token: str) -> str:263        """Strip the Ġ word marker; dot-prefix continuation pieces."""264        if token.startswith("Ġ"):265            return token[1:]266        if token.startswith("["):  # specials: [CLS], [SEP], [Q], [D]267            return token268        return "·" + token269 270    def render(271        self,272        texts: list[str],273        kind: str = "document",274        top_k: int | None = 25,275        width: int = 36,276        color: bool | None = None,277        **attribute_kwargs,278    ) -> str:279        """Terminal bar chart of the sparse expansions, with attributions.280 281        One block per text: an `L0` line with the total number of active dims,282        then bars proportional to weight (peak-normalized), each line ending283        with the input subtoken that produced the dim and `<exp>` for pure284        expansions. `color=None` auto-detects a TTY.285        """286        if color is None:287            color = sys.stdout.isatty()288        blocks = []289        for text, entries in zip(290            texts, self.attribute(texts, kind=kind, top_k=None, **attribute_kwargs)291        ):292            shown = text if len(text) <= 70 else text[:67] + "..."293            if not entries:294                blocks.append(f"{kind} · {shown}\n  (empty vector)")295                continue296            total = len(entries)297            entries = entries[:top_k]298            l0 = f"L0 = {total} active dims"299            if len(entries) < total:300                l0 += f" (showing {len(entries)})"301            peak = entries[0]["weight"]302            name_width = max(len(self._display(e["token"])) for e in entries)303            lines = [f"{kind} · {shown}", f"\033[2m{l0}\033[0m" if color else l0]304            for e in entries:305                ratio = e["weight"] / peak306                cells = max(1, round(ratio * width))307                bar = ("█" * cells)[:-3] + self._FADE if cells > 3 else self._FADE[3 - cells :]308                pad = " " * (width - cells)309                token = self._display(e["token"]).rjust(name_width)310                attrib = f"<- {self._display(e['source'])}@{e['position']}"311                if e["expansion"]:312                    attrib += " <exp>"313                if color:314                    heat = self._heat(ratio)315                    token = f"\033[38;5;{heat}m{token}\033[0m"316                    bar = f"\033[38;5;{heat}m{bar}\033[0m"317                    attrib = f"\033[2m{attrib}\033[0m"318                lines.append(f"{token} {bar}{pad} {e['weight']:>6.2f}  {attrib}")319            blocks.append("\n".join(lines))320        return "\n\n".join(blocks)321 322    @torch.inference_mode()323    def highlight(324        self,325        texts: list[str],326        kind: str = "document",327        reduce: str = "sum",328        color: bool | None = None,329        batch_size: int = 32,330        max_length: int | None = None,331    ) -> str:332        """Render each text with its firing words lit up.333 334        A word fires when one of its subtokens wins the max for at least one335        output dim; intensity is the total mass it contributes (`reduce="sum"`,336        default) or the largest single weight it wins (`reduce="max"`). TTY:337        reverse-video heat colors. Plain: tiered markers `⟦strong⟧ «mid» ‹weak›`.338        """339        if reduce not in ("sum", "max"):340            raise ValueError("`reduce` must be 'sum' or 'max'.")341        if color is None:342            color = sys.stdout.isatty()343        prefix, max_length = self._encode_args(kind, max_length)344        blocks = []345        for start in range(0, len(texts), batch_size):346            batch = texts[start : start + batch_size]347            ids, pool, offsets, sparse, sources = self._encode_with_sources(348                batch, prefix, max_length349            )350            # [B, L] per-position intensity, reduced over the dims each position351            # won. Zero-weight dims carry a clamped position 0 but contribute 0,352            # so they can't corrupt the reduction.353            intensity = torch.zeros(354                ids.shape, dtype=sparse.dtype, device=sparse.device355            ).scatter_reduce(356                1, sources, sparse, reduce="sum" if reduce == "sum" else "amax", include_self=False357            )358            for text, row_int, row_off, row_pool in zip(359                batch, intensity.float().cpu(), offsets, pool.cpu()360            ):361                blocks.append(362                    self._paint(text, row_int, row_off, row_pool, len(prefix), color, reduce)363                )364        return "\n\n".join(blocks)365 366    def _paint(self, text, intensity, offsets, pooling_mask, prefix_len, color, reduce) -> str:367        """Wrap fired char spans of `text` in intensity markers."""368        # Byte-BPE offsets include the word's leading space: trim it, so merging369        # only fuses glued subtokens of the same word (one span per word).370        spans: list[list] = []371        for pos in range(len(offsets)):372            w = intensity[pos].item()373            if w <= 0 or pooling_mask[pos] == 0:374                continue375            s, e = int(offsets[pos][0]) - prefix_len, int(offsets[pos][1]) - prefix_len376            s = max(s, 0)377            while s < e and text[s].isspace():378                s += 1379            if e <= s:  # zero-width specials / whitespace-only380                continue381            if spans and s == spans[-1][1]:382                spans[-1][1] = e383                spans[-1][2] = spans[-1][2] + w if reduce == "sum" else max(spans[-1][2], w)384            else:385                spans.append([s, e, w])386        if not spans:387            return text388        peak = max(w for _, _, w in spans)  # after merging, so summed words stay <= 1389        out, cursor = [], 0390        for s, e, w in spans:391            ratio = w / peak392            out.append(text[cursor:s])393            if color:394                # Reverse video with heat foreground = heat-colored highlighter.395                out.append(f"\033[7;38;5;{self._heat(ratio)}m{text[s:e]}\033[0m")396            else:397                marks = "⟦⟧" if ratio > 0.66 else "«»" if ratio > 0.33 else "‹›"398                out.append(f"{marks[0]}{text[s:e]}{marks[1]}")399            cursor = e400        out.append(text[cursor:])401        return "".join(out)402