Add UIBK-Embedding-v1 (weights, tokenizer, model code, model card)

1_HybridTokens/config.json

@@ -0,0 +1,17 @@
+{
+    "dim": 768,
+    "token_dim": 128,
+    "slots": 3,
+    "local_weight": 0.9,
+    "slot_weight": 0.225,
+    "num_latents": 512,
+    "num_heads": 8,
+    "ff_mult": 4,
+    "query_marker_id": 50368,
+    "query_local_cap": 64,
+    "query_local_ref": 0,
+    "query_local_power": 1.0,
+    "query_slot_weight": -1.0,
+    "doc_slot_weight": -1.0,
+    "slot_copies": 1
+}

1_HybridTokens/model.safetensors

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

README.md

@@ -1,3 +1,152 @@
 ---
 license: apache-2.0
+language:
+- en
+library_name: PyLate
+pipeline_tag: sentence-similarity
+base_model: answerdotai/ModernBERT-base
+tags:
+- ColBERT
+- PyLate
+- sentence-transformers
+- late-interaction
+- retrieval
+- feature-extraction
+- ModernBERT
 ---
+
+# UIBK-Embedding-v1
+
+**UIBK-Embedding-v1** is an English retrieval model that holds a **late-interaction (ColBERT) retriever and a dense
+retriever in one 179M-parameter encoder**, and stores both in one standard ColBERT index.
+
+- One encoder pass per text. Output: a list of 128-dimensional vectors, one per token plus 3 vectors that carry the
+  dense sentence embedding.
+- Plain MaxSim over these vectors equals **token MaxSim + 0.75 x dense cosine**. No special index or scorer is needed:
+  the model loads as a [PyLate](https://github.com/lightonai/pylate) ColBERT model and works with PLAID.
+- 179M parameters (ModernBERT-base trunk of 18 shared layers, 4 own layers per head), 2048-token documents,
+  256-token queries.
+
+## Results
+
+nDCG@10, [MTEB](https://github.com/embeddings-benchmark/mteb) 2.21.6 with PyLate and a PLAID index (default settings).
+
+| | UIBK-Embedding-v1 | LateOn (149M) |
+|---|---|---|
+| **BEIR (15 datasets), average** | **59.74** | 57.22 |
+| **NanoBEIR (13 datasets), average** | **70.55** | - |
+
+| BEIR dataset | UIBK-Embedding-v1 | LateOn |
+|---|---|---|
+| ArguAna | 56.78 | 50.52 |
+| CQADupStack | 48.61 | 47.36 |
+| ClimateFEVER | 44.29 | 39.67 |
+| DBPedia | 49.41 | 45.99 |
+| FEVER | 93.23 | 92.02 |
+| FiQA | 53.61 | 53.12 |
+| HotpotQA | 81.60 | 79.98 |
+| MSMARCO | 46.71 | 45.67 |
+| NFCorpus | 38.91 | 37.79 |
+| NQ | 67.66 | 63.91 |
+| Quora | 89.90 | 89.67 |
+| SCIDOCS | 23.50 | 21.90 |
+| SciFact | 77.98 | 76.61 |
+| TREC-COVID | 83.73 | 83.60 |
+| Touche-2020 | 40.21 | 30.52 |
+
+LateOn's numbers are those of its model card.
+
+| NanoBEIR dataset | UIBK-Embedding-v1 |
+|---|---|
+| NanoArguAna | 59.84 |
+| NanoClimateFever | 51.65 |
+| NanoDBPedia | 69.20 |
+| NanoFEVER | 96.57 |
+| NanoFiQA2018 | 64.34 |
+| NanoHotpotQA | 92.72 |
+| NanoMSMARCO | 68.99 |
+| NanoNFCorpus | 39.22 |
+| NanoNQ | 82.14 |
+| NanoQuora | 96.99 |
+| NanoSCIDOCS | 45.40 |
+| NanoSciFact | 83.93 |
+| NanoTouche2020 | 66.14 |
+
+## Usage
+
+The model needs `trust_remote_code=True` (two small modules in this repository define the two-top encoder and the
+output layer) and **the prompts `query` / `document` must be passed** when encoding.
+
+```bash
+pip install -U pylate "transformers>=5.3" "sentence-transformers>=5.3"
+```
+
+```python
+from pylate import indexes, models, retrieve
+
+model = models.ColBERT("DataScience-UIBK/uibk-embedding-v1", trust_remote_code=True)
+
+documents = [
+    "ColBERT scores a document by matching every query token with its most similar document token.",
+    "Comprehensive smoke-free laws were followed by falls in hospital admissions for heart attacks.",
+]
+documents_embeddings = model.encode(documents, is_query=False, prompt_name="document")
+
+index = indexes.PLAID(index_folder="pylate-index", index_name="index", override=True)
+index.add_documents(documents_ids=["d1", "d2"], documents_embeddings=documents_embeddings)
+
+queries_embeddings = model.encode(["how does late interaction score documents"], is_query=True, prompt_name="query")
+results = retrieve.ColBERT(index=index).retrieve(queries_embeddings=queries_embeddings, k=10)
+```
+
+Re-ranking a list of candidates without an index:
+
+```python
+from pylate import rank
+
+reranked = rank.rerank(
+    documents_ids=[["d1", "d2"]],
+    queries_embeddings=queries_embeddings,
+    documents_embeddings=[documents_embeddings],
+)
+```
+
+## How it works
+
+```
+text ("query: ..." / "document: ...", with a [Q] / [D] marker)
+  -> shared trunk: ModernBERT embeddings + layers 0-17
+       -> token top (layers 18-21)  -> linear 768 -> 128 per token          = token vectors
+       -> dense top (own layers 18-21) -> latent-attention pooling -> 768-dim sentence vector
+                                       -> rotated, cut into 3 slices of 127 dims = 3 "slot" vectors
+```
+
+Every output vector has 128 dimensions and unit norm. Dimension 0 is a tag with opposite signs for token vectors and
+slot vectors, so in MaxSim a token only ever matches a token and a slot only its own slot. One MaxSim pass therefore
+computes `token MaxSim + 0.75 x dense cosine` (up to a constant per query). The token head rewards exact term and
+phrase matches; the dense head rewards overall meaning; their sum is better than either.
+
+## Training
+
+- **Backbone**: ModernBERT-base, contrastive pre-training on LightOn's public pre-training collection.
+- **Fine-tuning** (1 epoch): contrastive losses on both heads with hard negatives, distillation from a cross-encoder
+  (`mixedbread-ai/mxbai-rerank-large-v2`) on the dense head and on the fused score, and feature anchors to our earlier
+  single-head models.
+- **Data**: LightOn's public English fine-tuning set (MS MARCO, NQ, HotpotQA, FEVER, FiQA, SQuAD v2, TriviaQA, MIRACL)
+  plus about 560k mined pairs from the pre-training collection (Quora, StackExchange, S2ORC, Reddit). Training
+  queries that equal a BEIR test query were removed.
+- **Teachers**, for transparency: the cross-encoder above; one of the earlier single-head models was itself trained
+  with a ranking loss to a fine-tune of `dunzhang/stella_en_400M_v5`.
+
+The model is **not zero-shot on the BEIR datasets whose training splits are in the data** (MS MARCO, NQ, HotpotQA,
+FEVER, FiQA), like other models trained on this collection.
+
+## Limitations
+
+- English only.
+- Multi-vector output: a ColBERT index is larger than a single-vector index.
+- Needs `trust_remote_code=True`, `transformers >= 5.3`, and the prompts above.
+
+## Citation
+
+To be added.

branched_transformer.py

@@ -0,0 +1,99 @@
+"""UIBK-Embedding-v1 encoder: ModernBERT with a shared trunk and two tops.
+
+(This file ships with the model and is loaded from the model repository: `trust_remote_code=True`.)
+
+The embeddings and the first `shared_layers` transformer layers are shared; each retrieval head has its own copy of the
+remaining layers and final norm, so that the dense head and the token (ColBERT) head make different errors:
+
+    token_embeddings = final_norm(layers[K:](trunk))            the wrapped model's own top: the ColBERT branch
+    dense_states     = dense_norm(dense_layers(trunk))          the second top: read by the dense pooling
+
+`auto_model` is a complete ModernBertModel (trunk + ColBERT top), so PyLate and the tokenizer handling see an ordinary
+Transformer module; the dense top lives in dense_layers / dense_norm and is saved to dense_branch.safetensors.
+18 shared layers and 4 own layers per branch: 149.0M + 20.1M parameters (+ 9.9M in the output layer).
+Requires transformers >= 5.3: the forward pass follows ModernBertModel.forward of that version.
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+
+import torch
+from safetensors.torch import load_file, save_file
+from sentence_transformers.models import Transformer
+from torch import nn
+from transformers.masking_utils import create_bidirectional_mask, create_bidirectional_sliding_window_mask
+
+BRANCH_FILE = "dense_branch.safetensors"
+
+
+class BranchedTransformer(Transformer):
+    config_keys: list[str] = ["max_seq_length", "do_lower_case", "shared_layers"]
+
+    def __init__(self, model_name_or_path: str, *args, shared_layers: int = 18, **kwargs) -> None:
+        super().__init__(model_name_or_path, *args, **kwargs)
+        self.shared_layers = shared_layers
+        model = self.auto_model
+        self.dense_layers = nn.ModuleList(copy.deepcopy(layer) for layer in model.layers[shared_layers:])
+        self.dense_norm = copy.deepcopy(model.final_norm)
+        for module in self.dense_layers.modules():  # one config object: the attention implementation is switched on it
+            if hasattr(module, "config"):
+                module.config = model.config
+        path = os.path.join(model_name_or_path, BRANCH_FILE)
+        if os.path.exists(path):
+            self.load_dense_branch(load_file(path))
+
+    def load_dense_branch(self, state: dict[str, torch.Tensor]) -> None:
+        """`state`: dense_layers.* / dense_norm.* tensors, or a full encoder state dict (its layers from K on are used)."""
+        if not any(k.startswith("dense_layers.") for k in state):
+            k0 = self.shared_layers
+            renamed = {}
+            for name, tensor in state.items():
+                if name.startswith("final_norm."):
+                    renamed["dense_norm." + name[len("final_norm."):]] = tensor
+                elif name.startswith("layers.") and int(name.split(".")[1]) >= k0:
+                    parts = name.split(".")
+                    renamed[".".join(["dense_layers", str(int(parts[1]) - k0), *parts[2:]])] = tensor
+            state = renamed
+        own = {k: v for k, v in self.state_dict().items() if k.startswith(("dense_layers.", "dense_norm."))}
+        if set(own) != set(state):
+            raise ValueError(f"dense branch weights do not match: {sorted(set(own) ^ set(state))[:4]}")
+        self.load_state_dict(state, strict=False)
+
+    def branched(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None,
+                 position_ids: torch.Tensor | None = None, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
+        """ModernBertModel.forward (transformers 5.3) with the layer loop split after the trunk."""
+        model = self.auto_model
+        kwargs.pop("return_dict", None)
+        if position_ids is None:
+            position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0)
+        hidden = model.embeddings(input_ids=input_ids, inputs_embeds=None)
+        masks = attention_mask
+        if not isinstance(masks, dict):
+            mask_kwargs = {"config": model.config, "inputs_embeds": hidden, "attention_mask": attention_mask}
+            masks = {"full_attention": create_bidirectional_mask(**mask_kwargs),
+                     "sliding_attention": create_bidirectional_sliding_window_mask(**mask_kwargs)}
+        positions = {kind: model.rotary_emb(hidden, position_ids, kind) for kind in model.config.layer_types}
+
+        def run(layers, states):
+            for layer in layers:
+                states = layer(states, attention_mask=masks[layer.attention_type],
+                               position_embeddings=positions[layer.attention_type], **kwargs)
+            return states
+
+        trunk = run(model.layers[: self.shared_layers], hidden)
+        tokens = model.final_norm(run(model.layers[self.shared_layers:], trunk))
+        dense = self.dense_norm(run(self.dense_layers, trunk))
+        return tokens, dense
+
+    def forward(self, features: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]:
+        inputs = {k: features[k] for k in ("input_ids", "attention_mask", "position_ids") if k in features}
+        features["token_embeddings"], features["dense_states"] = self.branched(**inputs, **kwargs)
+        return features
+
+    def save(self, output_path: str, safe_serialization: bool = True, **kwargs) -> None:
+        super().save(output_path, safe_serialization=safe_serialization, **kwargs)
+        own = {k: v.detach().cpu().contiguous() for k, v in self.state_dict().items()
+               if k.startswith(("dense_layers.", "dense_norm."))}
+        save_file(own, os.path.join(output_path, BRANCH_FILE), metadata={"format": "pt"})

config.json

@@ -0,0 +1,77 @@
+{
+  "architectures": [
+    "ModernBertModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "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,
+  "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,
+  "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",
+  "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": true,
+  "transformers_version": "5.3.0",
+  "vocab_size": 50370
+}

config_sentence_transformers.json

@@ -0,0 +1,75 @@
+{
+  "model_type": "ColBERT",
+  "__version__": {
+    "sentence_transformers": "5.3.0",
+    "transformers": "5.3.0",
+    "pytorch": "2.9.0+cu128"
+  },
+  "prompts": {
+    "query": "query: ",
+    "document": "document: "
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": "MaxSim",
+  "query_prefix": "[Q] ",
+  "document_prefix": "[D] ",
+  "query_length": 256,
+  "document_length": 2048,
+  "attend_to_expansion_tokens": false,
+  "skiplist_words": [
+    "!",
+    "\"",
+    "#",
+    "$",
+    "%",
+    "&",
+    "'",
+    "(",
+    ")",
+    "*",
+    "+",
+    ",",
+    "-",
+    ".",
+    "/",
+    ":",
+    ";",
+    "<",
+    "=",
+    ">",
+    "?",
+    "@",
+    "[",
+    "\\",
+    "]",
+    "^",
+    "_",
+    "`",
+    "{",
+    "|",
+    "}",
+    "~"
+  ],
+  "do_query_expansion": false,
+  "hybrid_head": {
+    "dim": 768,
+    "token_dim": 128,
+    "slots": 3,
+    "local_weight": 0.9,
+    "slot_weight": 0.225,
+    "num_latents": 512,
+    "num_heads": 8,
+    "ff_mult": 4,
+    "query_marker_id": 50368,
+    "query_local_cap": 64,
+    "query_local_ref": 0,
+    "query_local_power": 1.0,
+    "query_slot_weight": -1.0,
+    "doc_slot_weight": -1.0,
+    "slot_copies": 1,
+    "sha256": "c42363bb43bcfadef8154bc5d13cdc14231242e8b57676c5242590bde2dfbc6f",
+    "code_sha256": "0f0c540e201da7150b2dbc2447914d2f68a4e6e1d66280ab8843e829a62db80d",
+    "branch_sha256": "1e924278f9c2ec9c7dc326892a19ade45feb72cf2a780bef77006e0bc07ccac6",
+    "branch_code_sha256": "38023f3eaaeac0e7da5a92d19059a3ebbb5e644841c13ff3aa05adf3e6d8d9e1"
+  }
+}

dense_branch.safetensors

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

hybrid_tokens.py

@@ -0,0 +1,224 @@
+"""UIBK-Embedding-v1 output layer: a dense retriever and a late-interaction (ColBERT) retriever in one list of vectors.
+
+(This file ships with the model and is loaded from the model repository: `trust_remote_code=True`.)
+
+One encoder pass gives the token states H. From them:
+  local tokens   u_t = normalize(W H_t)              the ColBERT projection (768 -> 128), rotated down to 127 dims
+  dense vector   g   = mean_t LatentAttention(H)_t   latent-attention pooling of the dense branch's states
+The 128-dim output vectors (unit norm) are
+  a local token      [ +sqrt(1 - b2) ; sqrt(b2) * u_t ]      b2 = local_weight
+  dense slot s       [ -sqrt(1 - d2) ; sqrt(d2) * g_s ]      d2 = slot_weight; g rotated to `slots` slices of 127
+                                                             dims, each slice normalised
+and the slots overwrite the first `slots` token positions of every query and document ([CLS], the [Q] / [D] marker,
+the prompt word), which are never punctuation, so PyLate keeps them. Dimension 0 is a tag: it adds a constant to pairs
+of the same kind and subtracts one from mixed pairs, so a vector never picks a partner of the other kind:
+  local  x local   = (1 - b2) + b2 * cos(u, u')
+  slot s x slot s  = (1 - d2) + d2 * cos(g_s, g'_s)
+  slot   x local   = -sqrt((1 - b2)(1 - d2)) + sqrt(b2 d2) * (g_s . u)     (unrelated spaces: the last term is ~0)
+so plain MaxSim(q, d) = const(q) + b2 * MaxSim_local(q, d) + d2 * sum_s cos(g_s, g'_s): the fixed-weight sum of the
+token score and the dense score, MaxSim_local + lam * (mean slice cosine) with lam = slots * d2 / b2 (0.75 in the
+released configuration). Any ColBERT index and scorer works unchanged.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from pylate.models import Dense
+from torch import nn
+
+
+class GEGLU(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, gates = x.chunk(2, dim=-1)
+        return x * F.gelu(gates)
+
+
+def orthonormal(rows: int, cols: int, seed: int) -> torch.Tensor:
+    """(rows, cols) matrix with orthonormal columns, the same for a given seed on every machine."""
+    q, _ = torch.linalg.qr(torch.randn(rows, rows, generator=torch.Generator().manual_seed(seed), dtype=torch.float64))
+    return q[:, :cols].float()
+
+
+class HybridTokens(Dense):
+    config_keys: list[str] = ["dim", "token_dim", "slots", "local_weight", "slot_weight", "num_latents", "num_heads",
+                              "ff_mult", "query_marker_id", "query_local_cap", "query_local_ref", "query_local_power",
+                              "query_slot_weight", "doc_slot_weight", "slot_copies"]
+    refine_rows = 32768  # token rows per call of the pooling block (memory only)
+
+    def __init__(self, dim: int = 768, token_dim: int = 128, slots: int = 3, local_weight: float = 0.8,
+                 slot_weight: float = 0.8, num_latents: int = 512, num_heads: int = 8, ff_mult: int = 4,
+                 query_marker_id: int = -1, query_local_cap: int = 0, query_local_ref: int = 0,
+                 query_local_power: float = 1.0, query_slot_weight: float = -1.0,
+                 doc_slot_weight: float = -1.0, slot_copies: int = 1) -> None:
+        super().__init__(in_features=dim, out_features=token_dim, bias=False)  # self.linear: the ColBERT projection
+        self.dim, self.token_dim, self.slots = dim, token_dim, slots
+        self.local_weight, self.slot_weight = local_weight, slot_weight
+        # Long queries: the dense vector pools the whole query, but only its first `query_local_cap` tokens take part
+        # in the token-level MaxSim (a ColBERT reads 32-64 query tokens): the later query tokens are masked out of the
+        # output, so PyLate's encode drops them. A query is recognised by its [Q] marker at position 1.
+        self.query_marker_id, self.query_local_cap = query_marker_id, query_local_cap
+        # Optional, off in the released model (query_local_ref = 0): with r > 0, a query with n > r local tokens gets
+        # the local weight b2 * f^2, f = (r / n) ** query_local_power, which scales its token score by f.
+        self.query_local_ref, self.query_local_power = query_local_ref, query_local_power
+        # Optional, off in the released model (both -1 = slot_weight on both sides): separate slot weights for queries
+        # (q2) and documents (p2); slot x slot = const + sqrt(q2 p2) cos, lam = slots * sqrt(q2 p2) / b2.
+        self.query_slot_weight, self.doc_slot_weight = query_slot_weight, doc_slot_weight
+        # Optional, off in the released model (slot_copies = 1): with k > 1 (padded batches) every text carries its
+        # dense vector k times, copy c rotated by its own fixed orthogonal matrix before slicing; the extra slot
+        # vectors are appended as new columns after the token columns (mask 1, the id of position 0).
+        self.slot_copies = slot_copies
+        self._copy_rotations: dict[int, torch.Tensor] = {}
+        self.num_latents, self.num_heads, self.ff_mult = num_latents, num_heads, ff_mult
+        content = token_dim - 1
+        self.register_buffer("local_rotation", orthonormal(token_dim, content, seed=1))
+        self.register_buffer("dense_rotation", orthonormal(dim, slots * content, seed=0))
+        # the dense pooling block (names as in LatentAttentionPooling)
+        self.latents = nn.Parameter(torch.randn(num_latents, dim) * dim ** -0.5)
+        self.norm_x = nn.LayerNorm(dim)
+        self.norm_latents = nn.LayerNorm(dim)
+        self.to_q = nn.Linear(dim, dim, bias=False)
+        self.to_kv = nn.Linear(dim, 2 * dim, bias=False)
+        self.to_out = nn.Linear(dim, dim)
+        self.norm_ff = nn.LayerNorm(dim)
+        self.ff = nn.Sequential(nn.Linear(dim, 2 * ff_mult * dim), GEGLU(), nn.Linear(ff_mult * dim, dim))
+
+    def refine(self, x: torch.Tensor) -> torch.Tensor:
+        """The latent-attention block of our dense models on token states (..., T, dim), token by token."""
+        b, t, d = x.shape
+        h, dh = self.num_heads, d // self.num_heads
+        q = self.to_q(self.norm_x(x)).view(b, t, h, dh).transpose(1, 2)
+        k, v = self.to_kv(self.norm_latents(self.latents)).chunk(2, dim=-1)
+        k = k.view(self.num_latents, h, dh).transpose(0, 1).unsqueeze(0).expand(b, -1, -1, -1).to(q.dtype)
+        v = v.view(self.num_latents, h, dh).transpose(0, 1).unsqueeze(0).expand(b, -1, -1, -1).to(q.dtype)
+        attn = F.scaled_dot_product_attention(q, k, v)
+        x = x + self.to_out(attn.transpose(1, 2).reshape(b, t, d))
+        return x + self.ff(self.norm_ff(x))
+
+    def layout(self, features: dict[str, torch.Tensor]):
+        """Token rows (1, rows, dim) of a padded or packed batch, with each row's sequence, position and realness."""
+        states = features["token_embeddings"]
+        device = states.device
+        if "cu_seq_lens_q" in features:  # packed batch (training): (1, rows, dim), sequences back to back
+            bounds = features["cu_seq_lens_q"].long()
+            lengths = bounds[1:] - bounds[:-1]
+            n_seq = lengths.numel()
+            seq = torch.repeat_interleave(torch.arange(n_seq, device=device), lengths)
+            pos = features["position_ids"][0]
+            if "kv_index" in features:  # query expansion: the [MASK] rows are not part of the pooled vector
+                real = torch.zeros(seq.numel(), device=device).index_fill_(0, features["kv_index"], 1.0)
+            else:
+                real = torch.ones(seq.numel(), device=device)
+        else:  # padded batch (B, T, dim)
+            n_seq, t = states.shape[:2]
+            seq = torch.arange(n_seq, device=device).repeat_interleave(t)
+            pos = torch.arange(t, device=device).repeat(n_seq)
+            real = features["attention_mask"].reshape(-1).float()
+        return states.reshape(1, -1, states.size(-1)), seq, pos, real, n_seq
+
+    def dense_vectors(self, features: dict[str, torch.Tensor]) -> torch.Tensor:
+        """The dense models' sentence embedding of every sequence, normalised: (sequences, dim), float32."""
+        rows, seq, _, real, n_seq = self.layout(features)
+        if "dense_states" in features:  # two-top encoder (branched_transformer.py): the dense branch's states
+            rows = features["dense_states"].reshape(1, -1, rows.size(-1))
+        # The block works token by token, so only the rows that enter the pooled vector are refined, in chunks: the
+        # same result with bounded memory for long documents in large batches.
+        keep = real > 0
+        if not bool(keep.all()):
+            idx = keep.nonzero(as_tuple=True)[0]
+            rows, seq, real = rows[:, idx], seq[idx], real[idx]
+        refined = torch.cat([self.refine(chunk)[0].float() for chunk in rows.split(self.refine_rows, dim=1)])
+        refined = refined * real.unsqueeze(-1)
+        pooled = torch.zeros(n_seq, refined.size(-1), device=rows.device).index_add_(0, seq, refined)
+        pooled = pooled / torch.zeros(n_seq, device=rows.device).index_add_(0, seq, real).clamp(min=1.0).unsqueeze(-1)
+        return F.normalize(pooled, dim=-1)
+
+    def forward(self, features: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        states = features["token_embeddings"]
+        rows, seq, pos, real, n_seq = self.layout(features)
+        content = self.token_dim - 1
+        is_query = None
+        asymmetric = self.query_slot_weight >= 0 and self.doc_slot_weight >= 0
+        if self.query_local_cap > 0 or self.query_local_ref > 0 or asymmetric:
+            ids = features["input_ids"].reshape(-1)
+            marker = pos == 1
+            is_query = torch.zeros(n_seq, dtype=torch.bool, device=rows.device)
+            is_query[seq[marker]] = ids[marker] == self.query_marker_id
+
+        local = F.normalize(self.linear(rows[0]), dim=-1) @ self.local_rotation.to(rows.dtype)
+        local = F.normalize(local, dim=-1)
+        if self.query_local_ref > 0:
+            counted = real * (pos >= self.slots)
+            if self.query_local_cap > 0:
+                counted = counted * (pos < self.query_local_cap)
+            n = torch.zeros(n_seq, device=rows.device).index_add_(0, seq, counted.float()).clamp(min=1.0)
+            f = (self.query_local_ref / n).clamp(max=1.0) ** self.query_local_power
+            f = torch.where(is_query, f, torch.ones_like(f))
+            weight = (self.local_weight * f[seq] ** 2).unsqueeze(-1)  # per row
+            out = torch.cat([(1.0 - weight) ** 0.5, local.float() * weight ** 0.5], dim=-1).to(local.dtype)
+        else:
+            local = local * self.local_weight ** 0.5
+            out = torch.cat([local.new_full((local.size(0), 1), (1.0 - self.local_weight) ** 0.5), local], dim=-1)
+
+        dense = self.dense_vectors(features) @ self.dense_rotation.float()
+        dense = F.normalize(dense.view(n_seq, self.slots, content), dim=-1)
+        unit_slices = dense  # (sequences, slots, content), float32, each slice of unit norm
+        if not asymmetric:
+            dense = dense.to(out.dtype) * self.slot_weight ** 0.5
+        if self.query_local_cap > 0:
+            capped = is_query[seq] & (pos >= self.query_local_cap)
+            if "attention_mask" in features:  # padded batch: pylate keeps the query tokens whose mask is 1
+                keep = ~capped.view(features["attention_mask"].shape)
+                features["attention_mask"] = features["attention_mask"] * keep.to(features["attention_mask"].dtype)
+        slot = pos < self.slots
+        if asymmetric:
+            weight = torch.where(is_query[seq[slot]], self.query_slot_weight, self.doc_slot_weight)
+            weight = weight.to(torch.float32).unsqueeze(-1)
+            out[slot] = torch.cat([-(1.0 - weight) ** 0.5, dense[seq[slot], pos[slot]] * weight ** 0.5],
+                                  dim=-1).to(out.dtype)
+        else:
+            tag = dense.new_full((int(slot.sum()), 1), -((1.0 - self.slot_weight) ** 0.5))
+            out[slot] = torch.cat([tag, dense[seq[slot], pos[slot]]], dim=-1)
+        features["token_embeddings"] = out.view(*states.shape[:-1], self.token_dim)
+        if self.slot_copies > 1 and "cu_seq_lens_q" not in features and "attention_mask" in features:
+            extra = []
+            for c in range(1, self.slot_copies):
+                if c not in self._copy_rotations:
+                    self._copy_rotations[c] = orthonormal(self.slots * content, self.slots * content, seed=1000 + c)
+                rotation = self._copy_rotations[c].to(unit_slices.device)
+                copy = F.normalize((unit_slices.reshape(n_seq, -1) @ rotation).view(n_seq, self.slots, content), dim=-1)
+                if asymmetric:
+                    weight = torch.where(is_query, self.query_slot_weight, self.doc_slot_weight)
+                    weight = weight.to(torch.float32).view(n_seq, 1, 1).expand(n_seq, self.slots, 1)
+                else:
+                    weight = copy.new_full((n_seq, self.slots, 1), self.slot_weight)
+                extra.append(torch.cat([-(1.0 - weight) ** 0.5, copy * weight ** 0.5], dim=-1))
+            extra = torch.cat(extra, dim=1).to(out.dtype)  # (sequences, slots x (k - 1), token_dim)
+            n_extra = extra.size(1)
+            features["token_embeddings"] = torch.cat([features["token_embeddings"], extra], dim=1)
+            mask, ids = features["attention_mask"], features["input_ids"]
+            features["attention_mask"] = torch.cat([mask, mask.new_ones(n_seq, n_extra)], dim=1)
+            features["input_ids"] = torch.cat([ids, ids[:, :1].expand(n_seq, n_extra)], dim=1)
+            if "token_type_ids" in features:
+                types = features["token_type_ids"]
+                features["token_type_ids"] = torch.cat([types, types.new_zeros(n_seq, n_extra)], dim=1)
+        return features
+
+    def get_config_dict(self) -> dict:
+        return {k: getattr(self, k) for k in self.config_keys}
+
+    def save(self, output_path: str, *args, safe_serialization: bool = True, **kwargs) -> None:
+        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 = "", token=None, cache_folder=None, revision=None,
+             local_files_only: bool = False, **kwargs):
+        hub = {"subfolder": subfolder, "token": token, "cache_folder": cache_folder, "revision": revision,
+               "local_files_only": local_files_only}
+        model = cls(**cls.load_config(model_name_or_path=model_name_or_path, **hub))
+        return cls.load_torch_weights(model_name_or_path=model_name_or_path, model=model, **hub)
+
+    def __repr__(self) -> str:
+        n = sum(p.numel() for p in self.parameters())
+        return f"HybridTokens({self.get_config_dict()}, params={n / 1e6:.1f}M)"

model.safetensors

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

modules.json

@@ -0,0 +1,14 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "branched_transformer.BranchedTransformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_HybridTokens",
+    "type": "hybrid_tokens.HybridTokens"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,5 @@
+{
+    "max_seq_length": 299,
+    "do_lower_case": false,
+    "shared_layers": 18
+}

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": 512,
+  "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]"
+}