1_SparseUpPooling/sparseup.py

"""SPARSEUP (Linkup) recipe for Sentence Transformers.

Architecture: MLM → log1p(ReLU(x − θ)) → per-position top-k → max-pool
→ case / space-marker folding. The ``[Q] `` / ``[D] `` prefixes
are expected by the backbone but excluded from pooling.

References:
- https://huggingface.co/Linkup-Platform/linkup-sparseup-embed-v1
- https://www.linkup.so/blog/sparseup
"""
from __future__ import annotations

import html as html_lib
import json
import logging
import os
import re
from collections import defaultdict
from pathlib import Path
from typing import Any

import torch
from sentence_transformers.base.modules.module import Module

logger = logging.getLogger(__name__)

QUERY_PREFIX = "[Q] "
DOCUMENT_PREFIX = "[D] "


# ---------------------------------------------------------------------------
# Vocab fold (ingredient #3)
# ---------------------------------------------------------------------------

def _surface_body(token: str) -> str | None:
    """Lowercase body of a BPE token, without the space marker ``Ġ`` / ``▁``."""
    if not token or token.startswith(("[", "<", "Ġ[", "▁[")):
        return None
    body = token[1:] if token[:1] in {"Ġ", "▁"} else token
    body = body.lower()
    return body or None


def _fold_rank(token: str, body: str) -> tuple[int, str]:
    """Prefer the lowercase form with a word-space marker, then the bare form."""
    stripped = token[1:] if token[:1] in {"Ġ", "▁"} else token
    space = token[:1] if token[:1] in {"Ġ", "▁"} else ""
    if space == "▁" and stripped == body:
        return (0, token)
    if space == "Ġ" and stripped == body:
        return (1, token)
    if token == body:
        return (2, token)
    if stripped.lower() == body:
        return (3, token)
    return (4, token)


def build_case_space_fold(tokenizer) -> torch.Tensor:
    """Map each id to a canonical id (lowercase + space marker)."""
    vocab = tokenizer.get_vocab()
    size = max(len(tokenizer), max(vocab.values(), default=-1) + 1)
    fold = torch.arange(size, dtype=torch.long)

    groups: dict[str, list[tuple[str, int]]] = defaultdict(list)
    for token, idx in vocab.items():
        if idx < 0 or idx >= size:
            continue
        body = _surface_body(token)
        if body is None:
            continue
        groups[body].append((token, int(idx)))

    n_merged = 0
    for body, members in groups.items():
        if len(members) < 2:
            continue
        members.sort(key=lambda item: _fold_rank(item[0], body))
        canon = members[0][1]
        for _, idx in members:
            fold[idx] = canon
        n_merged += len(members) - 1

    logger.info(
        "vocab_fold case_space: %s ids → ~%s canonical (%s merged)",
        f"{size:,}",
        f"{size - n_merged:,}",
        f"{n_merged:,}",
    )
    return fold


def prefix_token_ids(tokenizer, prefixes: tuple[str, ...] = (QUERY_PREFIX, DOCUMENT_PREFIX)) -> list[list[int]]:
    ids = []
    for prefix in prefixes:
        enc = tokenizer(prefix, add_special_tokens=False)["input_ids"]
        if enc:
            ids.append([int(x) for x in enc])
    return ids


def pooling_mask(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    prefix_ids: list[list[int]] | None,
    special_ids: list[int] | None,
) -> torch.Tensor:
    """Pooling mask: drop specials + prefix at the start of the sequence."""
    mask = attention_mask.clone()
    if special_ids:
        special = torch.tensor(special_ids, device=input_ids.device, dtype=input_ids.dtype)
        mask = mask.masked_fill(torch.isin(input_ids, special), 0)
    if not prefix_ids:
        return mask

    # If the first token is an already-masked special, the prefix starts at 1.
    starts_with_special = mask[:, 0] == 0
    for pref in prefix_ids:
        n = len(pref)
        if n == 0 or n >= input_ids.size(1):
            continue
        pref_t = torch.tensor(pref, device=input_ids.device, dtype=input_ids.dtype)
        head0 = input_ids[:, :n]
        head1 = input_ids[:, 1 : 1 + n]
        match0 = (head0 == pref_t).all(dim=1) & ~starts_with_special
        match1 = (head1 == pref_t).all(dim=1) & starts_with_special
        if bool(match0.any()):
            mask[match0, :n] = 0
        if bool(match1.any()):
            mask[match1, 1 : 1 + n] = 0
    return mask


# ---------------------------------------------------------------------------
# SPARSEUP pooling
# ---------------------------------------------------------------------------

def _pair_mean_pool(weights: torch.Tensor, mask: torch.Tensor, factor: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Average ``factor`` adjacent tokens (cheap hpool). Ward is not applicable in vocab dim."""
    if factor <= 1:
        return weights, mask
    batch, length, vocab = weights.shape
    pad = (factor - length % factor) % factor
    if pad:
        weights = torch.nn.functional.pad(weights, (0, 0, 0, pad))
        mask = torch.nn.functional.pad(mask, (0, pad))
    groups = weights.size(1) // factor
    w = weights.view(batch, groups, factor, vocab)
    m = mask.view(batch, groups, factor, 1)
    denom = m.sum(dim=2).clamp_min(1.0)
    pooled = (w * m).sum(dim=2) / denom
    new_mask = (m.sum(dim=2).squeeze(-1) > 0).to(dtype=mask.dtype)
    return pooled, new_mask


class SparseUpPooling(Module):
    """log1p(ReLU(logits − θ)) + per-position top-k + max-pool + vocab fold.

    ``pool_factor`` > 1: average adjacent tokens **before** max-pool
    (hpool documents, like KURE / LI). Queries: leave at 1.
    """

    config_keys = [
        "logit_shift",
        "position_top_k",
        "vocab_fold",
        "chunk_size",
        "prefix_token_ids",
        "special_token_ids",
        "embedding_dimension",
        "pool_factor",
    ]

    def __init__(
        self,
        logit_shift: float = 15.0,
        position_top_k: int | None = 12,
        vocab_fold: str | None = "case_space",
        chunk_size: int | None = 32,
        prefix_token_ids: list[list[int]] | None = None,
        special_token_ids: list[int] | None = None,
        embedding_dimension: int | None = None,
        vocab_fold_index: torch.Tensor | None = None,
        pool_factor: int = 1,
    ) -> None:
        super().__init__()
        self.logit_shift = float(logit_shift)
        self.position_top_k = int(position_top_k) if position_top_k else None
        self.vocab_fold = vocab_fold
        self.chunk_size = chunk_size
        self.prefix_token_ids = prefix_token_ids or []
        self.special_token_ids = special_token_ids or []
        self.embedding_dimension = embedding_dimension
        self.pool_factor = max(1, int(pool_factor))
        self._hpool_enabled = False
        if vocab_fold_index is None:
            n = embedding_dimension or 1
            vocab_fold_index = torch.arange(n, dtype=torch.long)
        self.register_buffer("vocab_fold_index", vocab_fold_index.long(), persistent=True)

    def get_embedding_dimension(self) -> int | None:
        return self.embedding_dimension

    def _activate(self, logits: torch.Tensor) -> torch.Tensor:
        weights = torch.log1p(torch.relu(logits - self.logit_shift))
        k = self.position_top_k
        if k is not None and 0 < k < weights.size(-1):
            cutoff = weights.topk(k, dim=-1).values[..., -1:]
            weights = weights * (weights >= cutoff)
        return weights

    def _fold(self, sparse: torch.Tensor) -> torch.Tensor:
        if self.vocab_fold is None or self.vocab_fold_index.numel() != sparse.size(-1):
            return sparse
        index = self.vocab_fold_index.to(device=sparse.device, dtype=torch.long)
        if torch.equal(index, torch.arange(index.numel(), device=index.device)):
            return sparse
        index = index.unsqueeze(0).expand_as(sparse)
        return torch.zeros_like(sparse).scatter_reduce(1, index, sparse, reduce="amax", include_self=False)

    def _pool_mask(self, features: dict[str, torch.Tensor], length: int, device) -> torch.Tensor:
        attn = features.get("attention_mask")
        ids = features.get("input_ids")
        if attn is None:
            return torch.ones(ids.size(0), length, device=device) if ids is not None else None
        if ids is None:
            return attn
        return pooling_mask(ids, attn, self.prefix_token_ids, self.special_token_ids)

    def forward(self, features: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]:
        if "cu_seq_lens_q" in features:
            return self._forward_flattened(features)

        logits = features["token_embeddings"]
        batch, seq_len, vocab = logits.shape
        device = logits.device
        mask = self._pool_mask(features, seq_len, device)
        if mask is None:
            mask = torch.ones(batch, seq_len, device=device, dtype=logits.dtype)
        mask = mask.to(device=device, dtype=logits.dtype)

        pooled = torch.full((batch, vocab), float("-inf"), dtype=logits.dtype, device=device)
        factor = self.pool_factor if self._hpool_enabled else 1
        step = seq_len if not self.chunk_size else max(1, int(self.chunk_size))
        for start in range(0, seq_len, step):
            chunk = logits[:, start : start + step]
            chunk_mask = mask[:, start : start + step]
            weights = self._activate(chunk) * chunk_mask.unsqueeze(-1)
            if factor > 1:
                weights, chunk_mask = _pair_mean_pool(weights, chunk_mask, factor)
                weights = weights * chunk_mask.unsqueeze(-1)
            pooled = torch.maximum(pooled, weights.max(dim=1).values)

        pooled = torch.where(torch.isfinite(pooled), pooled, torch.zeros_like(pooled))
        pooled = self._fold(pooled)
        if self.embedding_dimension is None:
            self.embedding_dimension = pooled.size(1)
        features["sentence_embedding"] = pooled
        return features

    def _forward_flattened(self, features: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        flat = features["token_embeddings"].squeeze(0)
        cu = features["cu_seq_lens_q"].tolist()
        ids = features.get("input_ids")
        attn = features.get("attention_mask")
        rows = []
        for i, (start, end) in enumerate(zip(cu[:-1], cu[1:])):
            weights = self._activate(flat[start:end])
            if ids is not None and attn is not None:
                row_ids = ids[i : i + 1, : end - start]
                row_attn = attn[i : i + 1, : end - start]
                # Unpadded tensors often don't have the padded shape; keep everything.
                _ = (row_ids, row_attn)
            if self._hpool_enabled and self.pool_factor > 1 and weights.dim() == 2:
                w = weights.unsqueeze(0)
                m = torch.ones(1, w.size(1), device=w.device, dtype=w.dtype)
                w, _ = _pair_mean_pool(w, m, self.pool_factor)
                weights = w.squeeze(0)
            rows.append(weights.max(dim=0).values)
        pooled = self._fold(torch.stack(rows))
        if self.embedding_dimension is None:
            self.embedding_dimension = pooled.size(1)
        features["sentence_embedding"] = pooled
        return features

    def save(self, output_path: str, *args, safe_serialization: bool = True, **kwargs) -> None:
        os.makedirs(output_path, exist_ok=True)
        self.save_config(output_path)
        self.save_torch_weights(output_path, safe_serialization=safe_serialization)

    @classmethod
    def load(cls, model_name_or_path: str, subfolder: str = "", **kwargs):
        # Sentence Transformers passes trust_remote_code here; load_config does not accept it.
        config = cls.load_config(
            model_name_or_path,
            subfolder=subfolder,
            token=kwargs.get("token"),
            cache_folder=kwargs.get("cache_folder"),
            revision=kwargs.get("revision"),
            local_files_only=kwargs.get("local_files_only", False),
        )
        module = cls(**config)
        try:
            cls.load_torch_weights(
                model_name_or_path,
                subfolder=subfolder,
                model=module,
                token=kwargs.get("token"),
                cache_folder=kwargs.get("cache_folder"),
                revision=kwargs.get("revision"),
                local_files_only=kwargs.get("local_files_only", False),
            )
        except ValueError:
            logger.warning("No SparseUpPooling weights (identity fold).")
        return module

    def __repr__(self) -> str:
        return f"SparseUpPooling({self.get_config_dict()})"


class hpool_documents:
    """Enable hpool only for ``encode_document`` (queries untouched)."""

    def __init__(self, module: SparseUpPooling | None):
        self.module = module
        self._prev = False

    def __enter__(self):
        if self.module is not None:
            self._prev = self.module._hpool_enabled
            self.module._hpool_enabled = True
        return self

    def __exit__(self, *exc):
        if self.module is not None:
            self.module._hpool_enabled = self._prev
        return False


def build_pooling(
    tokenizer,
    logit_shift: float = 15.0,
    position_top_k: int | None = 12,
    vocab_fold: str | None = "case_space",
    chunk_size: int | None = 32,
) -> SparseUpPooling:
    fold = build_case_space_fold(tokenizer) if vocab_fold == "case_space" else None
    specials = sorted({int(i) for i in tokenizer.all_special_ids})
    dim = int(len(tokenizer))
    if fold is not None and fold.numel() < dim:
        padded = torch.arange(dim, dtype=torch.long)
        padded[: fold.numel()] = fold
        fold = padded
    elif fold is not None:
        dim = int(fold.numel())
    return SparseUpPooling(
        logit_shift=logit_shift,
        position_top_k=position_top_k,
        vocab_fold=vocab_fold,
        chunk_size=chunk_size,
        prefix_token_ids=prefix_token_ids(tokenizer),
        special_token_ids=specials,
        embedding_dimension=dim,
        vocab_fold_index=fold,
    )


# ---------------------------------------------------------------------------
# Grafting the MLM head (backbone ingredient)
# ---------------------------------------------------------------------------

def _encoder_module(model):
    if hasattr(model, "model") and hasattr(model.model, "embeddings"):
        return model.model
    if hasattr(model, "embeddings"):
        return model
    raise TypeError(f"Encoder not found on {type(model)}")


def _copy_matching(dst: torch.nn.Module, src: torch.nn.Module) -> tuple[int, int]:
    dst_sd = dst.state_dict()
    src_sd = src.state_dict()
    merged = dict(dst_sd)
    n_full = n_partial = 0
    for key, src_t in src_sd.items():
        if key not in dst_sd:
            continue
        dst_t = dst_sd[key]
        if dst_t.shape == src_t.shape:
            merged[key] = src_t
            n_full += 1
            continue
        if dst_t.ndim != src_t.ndim:
            continue
        sl = tuple(slice(0, min(a, b)) for a, b in zip(dst_t.shape, src_t.shape))
        cloned = dst_t.clone()
        cloned[sl] = src_t[sl]
        merged[key] = cloned
        n_partial += 1
    dst.load_state_dict(merged, strict=False)
    return n_full, n_partial


def _ensure_qd_tokens(tokenizer, model) -> int:
    """Add ``[Q] `` / ``[D] `` if missing (like LateOn / mLateOn)."""
    extra = [QUERY_PREFIX, DOCUMENT_PREFIX]
    existing = set(tokenizer.all_special_tokens)
    to_add = [t for t in extra if t not in existing]
    if not to_add:
        return 0
    n = tokenizer.add_special_tokens({"additional_special_tokens": to_add})
    if n:
        model.resize_token_embeddings(len(tokenizer))
        logger.info("Added tokens: %s (vocab %s)", to_add, len(tokenizer))
    return n


def prepare_bare_mlm(mlm_name: str, output_dir: str | Path) -> Path:
    """MLM checkpoint + ``[Q] `` / ``[D] `` prefixes, without the LateOn encoder."""
    from transformers import AutoModelForMaskedLM, AutoTokenizer

    output_dir = Path(output_dir)
    marker = output_dir / "graft_meta.json"
    if marker.exists() and (output_dir / "config.json").exists():
        meta = json.loads(marker.read_text())
        if meta.get("late") is None and meta.get("mlm") == mlm_name:
            logger.info("MLM cache already present: %s", output_dir)
            return output_dir

    logger.info("Loading bare MLM %s", mlm_name)
    tokenizer = AutoTokenizer.from_pretrained(mlm_name)
    mlm = AutoModelForMaskedLM.from_pretrained(mlm_name, torch_dtype=torch.float32)
    _ensure_qd_tokens(tokenizer, mlm)
    if hasattr(mlm, "tie_weights"):
        mlm.tie_weights()

    output_dir.mkdir(parents=True, exist_ok=True)
    mlm.config.architectures = [mlm.__class__.__name__]
    mlm.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    marker.write_text(
        json.dumps({"late": None, "mlm": mlm_name, "vocab_size": int(mlm.config.vocab_size)}, indent=2)
    )
    logger.info("Bare MLM → %s (vocab %s)", output_dir, mlm.config.vocab_size)
    del mlm
    return output_dir


def graft_mlm_head(
    late_name: str,
    mlm_name: str,
    output_dir: str | Path,
    model_kwargs: dict[str, Any] | None = None,
) -> Path:
    """Graft the ``mlm_name`` MLM head onto the ``late_name`` LateOn encoder.

    ``decoder.weight`` stays tied to the LateOn embeddings (like SPARSEUP).
    ``head`` + ``decoder.bias`` come from the original MLM checkpoint.

    Both checkpoints must share the same tokenizer (up to the ``[Q]`` / ``[D]``
    tokens). An mmBERT ↔ modernCamemBERT graft is rejected: 256k vs 32k.
    """
    from sentence_transformers import MultiVectorEncoder
    from transformers import AutoModelForMaskedLM

    output_dir = Path(output_dir)
    marker = output_dir / "graft_meta.json"
    if marker.exists() and (output_dir / "config.json").exists():
        meta = json.loads(marker.read_text())
        if meta.get("late") == late_name and meta.get("mlm") == mlm_name:
            logger.info("Graft cache already present: %s", output_dir)
            return output_dir

    logger.info("Loading LateOn %s", late_name)
    late = MultiVectorEncoder(late_name, model_kwargs=model_kwargs or {"torch_dtype": "float32"})
    late_enc = _encoder_module(late[0].auto_model)
    tokenizer = late[0].tokenizer
    late_vocab = int(late_enc.config.vocab_size)

    logger.info("Loading MLM head %s", mlm_name)
    mlm = AutoModelForMaskedLM.from_pretrained(mlm_name, torch_dtype=torch.float32)
    mlm_vocab = int(mlm.config.vocab_size)
    extra = late_vocab - mlm_vocab
    if extra < 0 or extra > 16:
        raise ValueError(
            f"Incompatible vocabularies for a SPARSEUP graft: "
            f"encoder {late_vocab} vs MLM head {mlm_vocab} ({mlm_name}). "
            "The head must come from the same tokenizer as the encoder "
            "(mmBERT-base with mLateOn, modernCamemBERT with a modernCamemBERT PFT). "
            "We do not graft a 32k head onto a 256k encoder: the ids no longer make sense."
        )
    if late_vocab != mlm_vocab:
        logger.info("resize_token_embeddings %s → %s ([Q]/[D])", mlm_vocab, late_vocab)
        mlm.resize_token_embeddings(late_vocab)
        if hasattr(mlm, "decoder") and getattr(mlm.decoder, "bias", None) is not None:
            with torch.no_grad():
                mlm.decoder.bias[mlm_vocab:] = 0.0

    n_full, n_partial = _copy_matching(_encoder_module(mlm), late_enc)
    if hasattr(mlm, "tie_weights"):
        mlm.tie_weights()

    output_dir.mkdir(parents=True, exist_ok=True)
    mlm.config.architectures = [mlm.__class__.__name__]
    mlm.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    marker.write_text(
        json.dumps(
            {
                "late": late_name,
                "mlm": mlm_name,
                "copied_full": n_full,
                "copied_partial": n_partial,
                "vocab_size": late_vocab,
            },
            indent=2,
        )
    )
    logger.info("Graft OK (%s full tensors, %s partial) → %s", n_full, n_partial, output_dir)

    del late, mlm
    return output_dir


def build_sparseup_encoder(
    init_dir: str | Path,
    *,
    logit_shift: float = 15.0,
    position_top_k: int | None = 12,
    vocab_fold: str | None = "case_space",
    chunk_size: int | None = 32,
    query_length: int = 128,
    document_length: int = 512,
    model_kwargs: dict[str, Any] | None = None,
):
    from sentence_transformers.base.modules.transformer import Transformer
    from sentence_transformers.sparse_encoder import SparseEncoder
    from sentence_transformers.sparse_encoder.model_card import SparseEncoderModelCardData

    init_dir = str(init_dir)
    transformer = Transformer(
        init_dir,
        transformer_task="fill-mask",
        query_length=query_length,
        document_length=document_length,
        model_kwargs=model_kwargs or {"torch_dtype": "float32"},
    )
    pooling = build_pooling(
        transformer.tokenizer,
        logit_shift=logit_shift,
        position_top_k=position_top_k,
        vocab_fold=vocab_fold,
        chunk_size=chunk_size,
    )
    model = SparseEncoder(
        modules=[transformer, pooling],
        prompts={"query": QUERY_PREFIX, "document": DOCUMENT_PREFIX},
        similarity_fn_name="dot",
        model_card_data=SparseEncoderModelCardData(
            language=["fr"],
            license="apache-2.0",
            model_name="SPARSEUP-style French sparse encoder",
        ),
    )
    return model


# ---------------------------------------------------------------------------
# Probes and reading bags
# ---------------------------------------------------------------------------

@torch.inference_mode()
def probe_logits(
    model,
    texts: list[str],
    batch_size: int = 1,
    max_length: int = 64,
    max_tokens: int = 48,
    max_values: int = 400_000,
) -> dict[str, float]:
    """Percentiles of the MLM logits (before shift) to calibrate θ.

    We don't materialize the whole vocabulary: a
    ``[B, L, 256k]`` float32 tensor is enough to saturate RAM/GPU and freeze the cell.
    """
    transformer = model[0]
    tokenizer = transformer.tokenizer
    device = next(transformer.parameters()).device
    cap = min(max_length, transformer.document_length or max_length)
    chunks: list[torch.Tensor] = []
    n_seen = 0
    print(f"Probing {len(texts)} texts (max_length={cap}, batch={batch_size})…", flush=True)
    for start in range(0, len(texts), batch_size):
        batch = texts[start : start + batch_size]
        enc = tokenizer(
            [DOCUMENT_PREFIX + t for t in batch],
            padding=True,
            truncation=True,
            max_length=cap,
            return_tensors="pt",
        )
        enc = {k: v.to(device) for k, v in enc.items() if torch.is_tensor(v)}
        logits = transformer.auto_model(**enc).logits
        flat = logits[enc["attention_mask"].bool()].float()
        if flat.size(0) > max_tokens:
            pick = torch.randperm(flat.size(0), device=flat.device)[:max_tokens]
            flat = flat[pick]
        chunks.append(flat.reshape(-1).cpu())
        n_seen += int(flat.size(0))
        print(f"  batch {start // batch_size + 1}: {flat.size(0)} tokens, vocab={logits.size(-1)}", flush=True)
        del logits, flat
    values = torch.cat(chunks)
    if values.numel() > max_values:
        values = values[torch.randperm(values.numel())[:max_values]]
    qs = torch.quantile(values, torch.tensor([0.01, 0.05, 0.25, 0.50, 0.75, 0.95, 0.99]))
    return {
        "n_tokens": n_seen,
        "n_sample": int(values.numel()),
        "min": float(values.min()),
        "p1": float(qs[0]),
        "p5": float(qs[1]),
        "p25": float(qs[2]),
        "median": float(qs[3]),
        "p75": float(qs[4]),
        "p95": float(qs[5]),
        "p99": float(qs[6]),
        "max": float(values.max()),
        "frac>0": float((values > 0).float().mean()),
        "frac>15": float((values > 15).float().mean()),
    }


def l0(embeddings: torch.Tensor) -> torch.Tensor:
    if embeddings.is_sparse:
        embeddings = embeddings.to_dense()
    return (embeddings > 0).sum(dim=-1).float()


def format_bag(pairs: list[tuple[str, float]], top_k: int = 20) -> str:
    if not pairs:
        return "  (empty bag — lower LOGIT_SHIFT or check the pooling)"
    peak = max(pairs[0][1], 1e-6)
    lines = []
    for token, weight in pairs[:top_k]:
        shown = token[1:] if token[:1] in {"Ġ", "▁"} else token
        bar = "█" * max(1, int(round(10 * weight / peak)))
        lines.append(f"  {shown:20s} {weight:6.2f} {bar}")
    return "\n".join(lines)


_WORD = re.compile(r"\S+")


def highlight_html(text: str, bag: list[tuple[str, float]]) -> str:
    """Highlight the document words that (approximately) appear in the bag."""
    weights = {}
    for token, weight in bag:
        body = _surface_body(token) or token.lower()
        weights[body] = max(weights.get(body, 0.0), float(weight))
    peak = max(weights.values()) if weights else 1.0
    parts = []
    last = 0
    for match in _WORD.finditer(text):
        parts.append(html_lib.escape(text[last : match.start()]))
        word = match.group()
        key = word.lower().strip(".,;:!?()[]\"'«»")
        w = weights.get(key, 0.0)
        safe = html_lib.escape(word)
        if w <= 0:
            parts.append(f"<span>{safe}</span>")
        else:
            alpha = min(0.85, 0.18 + 0.67 * (w / peak))
            parts.append(
                f'<span style="background:rgba(53,159,171,{alpha:.2f});padding:1px 3px;'
                f'border-radius:3px" title="{w:.2f}">{safe}</span>'
            )
        last = match.end()
    parts.append(html_lib.escape(text[last:]))
    return "<p style='line-height:1.8;font-family:sans-serif'>" + "".join(parts) + "</p>"