v2: shared-state joint cross-encoder (ettin-reranker-150m), Decision Index 0.3 run; v1 kept at tag v1

1_Dense/config.json

@@ -1,7 +0,0 @@
-{
-    "in_features": 768,
-    "out_features": 1536,
-    "bias": false,
-    "activation_function": "torch.nn.modules.linear.Identity",
-    "use_residual": true
-}

1_Dense/model.safetensors

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

2_Dense/config.json

@@ -1,7 +0,0 @@
-{
-    "in_features": 1536,
-    "out_features": 768,
-    "bias": false,
-    "activation_function": "torch.nn.modules.linear.Identity",
-    "use_residual": true
-}

2_Dense/model.safetensors

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

3_Dense/config.json

@@ -1,7 +0,0 @@
-{
-    "in_features": 768,
-    "out_features": 128,
-    "bias": false,
-    "activation_function": "torch.nn.modules.linear.Identity",
-    "use_residual": false
-}

3_Dense/model.safetensors

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

README.md

@@ -1,92 +1,93 @@
 ---
 license: apache-2.0
-language:
-- en
-library_name: PyLate
-pipeline_tag: decision-model
-base_model: lightonai/LateOn
+library_name: transformers
+base_model: cross-encoder/ettin-reranker-150m-v1
 datasets:
 - tasksource/tasksource-jev-typed-decisions
+- tasksource/synthetic-typed-decisions
 tags:
-- decision-model
-- jev
-- ColBERT
-- multi-vector
-- PyLate
-- sentence-transformers
+- typed-decisions
+- decision-index
+- cross-encoder
 ---
 
-# tasksource-jev-nano-v1
+# tasksource-decider-nano
 
-A **149M-parameter decision model**. Given a state, a question and a set of options, it returns a probability for each option. It follows the "Jev-style" typed-decision interface (choice / noul) and is evaluated on the [Decision Index](https://huggingface.co/spaces/multimodalart/jev-decision-index).
+This repository was `tasksource/tasksource-jev-nano-v1`. Its `main` (tag `v2`) is now a new model; the previous
+multi-vector model (v1) stays at tag [`v1`](https://huggingface.co/tasksource/tasksource-decider-nano/tree/v1).
 
-It is [LateOn](https://huggingface.co/lightonai/LateOn) (ModernBERT-base, ColBERT multi-vector), fine-tuned on Tasksource typed decisions. Options are scored by **late interaction**:
+A 150M-parameter typed-decision model for the [Decision Index](https://huggingface.co/spaces/multimodalart/jev-decision-index).
+It takes a `state` and typed questions (`choice`, `noul`) and returns a probability for every option. It does not
+generate text.
 
-```
-context  = state + "\nQuestion: " + question     → encoded once (ColBERT document side)
-option_i                                          → encoded independently (ColBERT query side)
-logit_i  = MaxSim(option_i, context) / T          → softmax over options      (T = 0.3168)
-```
+| | |
+|---|---|
+| Decision Index 0.3, public index | **26.77** |
+| Requests | 140,171 `ok`, 7 `unsupported` (pairs above 8,192 tokens) |
+| Latency, 1× NVIDIA A10 23 GB, one request at a time | median 45 ms · mean 71 ms · p80 71 ms |
+| Complete results | [runs/tasksource-decider-nano](https://huggingface.co/datasets/tasksource/decision-index-results/tree/main/runs/tasksource-decider-nano) |
 
-The context never sees the options, so you can encode a state once and score any number of options against it. This also makes the scores permutation-equivariant.
+## Architecture
 
-## Usage
+A joint cross-encoder on [cross-encoder/ettin-reranker-150m-v1](https://huggingface.co/cross-encoder/ettin-reranker-150m-v1)
+(ModernBERT-style encoder, 8,192-token context). Every option is scored against the full state, its question
+and the option text, like the reranker pair `(state, "Question: <instructions>\nOption: <option>")`, by a linear
+head on the mean of the option's hidden states. The scores are softmaxed over each question's options.
 
-```python
-import torch
-from pylate import models
-from pylate.scores import colbert_scores
+The state is encoded **once per request**. All questions and options of a request share it through a tree
+attention mask: the state attends to itself, a question to the state and itself, and an option to the state, its
+question and itself. Rotary positions are those of the single pair, so every option gets the score it would get on its
+own (up to float rounding), while a request with 50 options costs one state encoding instead of 50.
 
-model = models.ColBERT("tasksource/tasksource-jev-nano-v1")
-T = 0.31683406233787537  # learned temperature (decision_meta.json: init_temperature)
+Inference reads every input in full. There is no truncation: a pair above 8,192 tokens raises `Unsupported`.
 
-def decide(state, question, options):
-    ctx = model.encode([f"{state}\nQuestion: {question}"], is_query=False)[0]
-    opts = model.encode(options, is_query=True)
-    scores = torch.stack([colbert_scores(torch.as_tensor(o)[None], torch.as_tensor(ctx)[None])[0, 0] for o in opts])
-    return dict(zip(options, torch.softmax(scores / T, 0).tolist()))
+## Usage
 
-decide("The movie was a total waste of two hours.", "What is the sentiment?", ["positive", "negative"])
-# {'positive': 0.0088, 'negative': 0.9912}
+```python
+from decider import Decider
+model = Decider.from_pretrained("tasksource/tasksource-decider-nano")
+model.answer(state={"review": "The battery died after a week."},
+             questions={"q": {"type": "choice", "instructions": "What is the sentiment?",
+                              "criteria": {"neg": "negative", "pos": "positive"}}})
 ```
 
-This snippet reproduces our evaluation engine's probabilities exactly. For **noul** (yes/no) questions, score the two options `false` and `true` and report P(true).
-
-When an option has both a key and a description, the engine renders it as `"key: description"`, or as the description alone when the key is a placeholder like `A` or `option_1`.
-
-## Training
+Decision Index run (kit 0.3, `9eb2dbe`), from a local copy of this repository:
 
-- **Init:** `lightonai/LateOn`.
-- **Data:** grouped requests (one state, 1–5 questions, 2–235 options each) from [tasksource-jev-typed-decisions](https://huggingface.co/datasets/tasksource/tasksource-jev-typed-decisions), covering 503 Tasksource training tasks.
-  - Training used a 64k-request **double-firewalled** manifest: every row that fingerprint-matches a Decision Index suite row was removed before training.
-- **Recipe:**
-  - Stage 1: 4,000 steps of grouped MaxSim training with soft cross-entropy over options and a learned temperature.
-  - Stage 2: continued for 400 steps with *question packing* at p=0.5 (several questions share one context half the time).
-- **Checkpoint selection:** NLL on a held-out set of unseen Tasksource tasks.
-
-## Decision Index results (public index)
-
-These are scores on the **public** part of the Decision Index, run with the official kit: full suite, 100% coverage, chance-corrected skill, 0–100.
-
-| Edition | Public index (skill) | Raw |
-|---|---|---|
-| 0.3 | **10.84** | 33.12 |
-| 0.2.1 | 10.46 | 32.09 |
-
-0.3 areas (skill): Knowledge & Reasoning 2.2 · Language Understanding 5.0 · Retrieval & Classification 23.5 · Tools & Automation 13.6 · Arts & Human Taste 17.7.
-
-**This is not the board's headline number.** From 0.3 on, the board ranks a *Full* score:
-
-| Part | Weight | Who runs it |
-|---|---|---|
-| Public index | 20% | Anyone, with the kit; the table above |
-| Private tests of the same skills | 50% | Maintainers |
-| Private tasks from new domains | 30% | Maintainers |
+```bash
+hf download tasksource/tasksource-decider-nano --local-dir tasksource-decider-nano
+cd tasksource-decider-nano
+python -m decision_index pipeline --engine decision_index_engine:DeciderEngine \
+    --option repo=. --edition 0.3 --out runs/tasksource-decider-nano
+```
 
-We do not have this model's Full score; see the board for it once it has been scored.
+`decider.py` and `decision_index_engine.py` are self-contained (PyTorch and transformers only).
 
-## Limitations
+## Training
 
-- **Small model, broad but shallow.** It does well on retrieval, classification and tool selection. Knowledge and reasoning are near chance: 149M parameters do not store much world knowledge.
-- **Typed decisions only.** It handles `choice` and `noul`. It does not generate text, and it does not natively produce scores or multilabel outputs.
-- **English only.**
+12,000 steps, batches of 8 requests, learning rate 2e-5, inputs up to 2,048 tokens, seed 7. The loss is soft
+cross-entropy over each question's options. The training mix has 240,000 requests:
+
+| Share | Source |
+|---:|---|
+| 42% | [tasksource/tasksource-jev-typed-decisions](https://huggingface.co/datasets/tasksource/tasksource-jev-typed-decisions) (658 tasks, at most 600 questions per task) |
+| 30% | *train* splits of public datasets behind Decision Index benchmarks, e.g. MMLU, ANLI, WinoGrande, HellaSwag, ARC, CLINC150, BANKING77, Amazon ESCI, HoVer, RAGTruth, New Yorker caption contest, When2Call, ACOS, iSarcasmEval, NLI4CT, ContractNLI, GSM8K (at most 10,000 requests each) |
+| 12% | tool use: ToolRet train, Glaive function calling, Hermes function calling |
+| 8% | LLM routing: xRouteBench and LLMRouterBench train data |
+| 8% | [tasksource/synthetic-typed-decisions](https://huggingface.co/datasets/tasksource/synthetic-typed-decisions) |
+
+No evaluation split is used, and benchmarks without a train split (for example RouterBench) are evaluation only.
+Every training request goes through a firewall against the Decision Index 0.3 suite: an exact match of the
+normalized state or of the (state, question) pair, and a template-insensitive word 6-gram content check, drop the
+request.
+
+## Evaluation
+
+| Area | Skill | Raw |
+|---|---:|---:|
+| Knowledge & Reasoning | 6.0 | 29.5 |
+| Language Understanding | 28.2 | 47.9 |
+| Retrieval & Classification | 42.0 | 58.2 |
+| Tools & Automation | 44.6 | 49.8 |
+| Arts & Human Taste | 13.9 | 43.5 |
+
+Raw index 45.1; kit `9eb2dbe`, edition 0.3, 140,171 requests scored.

config.json

@@ -4,7 +4,8 @@
   ],
   "attention_bias": false,
   "attention_dropout": 0.0,
-  "bos_token_id": 50281,
+  "bos_token_id": null,
+  "causal_mask": false,
   "classifier_activation": "gelu",
   "classifier_bias": false,
   "classifier_dropout": 0.0,
@@ -14,14 +15,21 @@
   "deterministic_flash_attn": false,
   "dtype": "float32",
   "embedding_dropout": 0.0,
-  "eos_token_id": 50282,
+  "eos_token_id": null,
   "global_attn_every_n_layers": 3,
   "gradient_checkpointing": false,
   "hidden_activation": "gelu",
   "hidden_size": 768,
+  "id2label": {
+    "0": "LABEL_0"
+  },
   "initializer_cutoff_factor": 2.0,
   "initializer_range": 0.02,
   "intermediate_size": 1152,
+  "is_causal": false,
+  "label2id": {
+    "LABEL_0": 0
+  },
   "layer_norm_eps": 1e-05,
   "layer_types": [
     "full_attention",
@@ -48,7 +56,7 @@
     "full_attention"
   ],
   "local_attention": 128,
-  "max_position_embeddings": 8192,
+  "max_position_embeddings": 7999,
   "mlp_bias": false,
   "mlp_dropout": 0.0,
   "model_type": "modernbert",
@@ -57,15 +65,14 @@
   "num_attention_heads": 12,
   "num_hidden_layers": 22,
   "pad_token_id": 50283,
-  "position_embedding_type": "absolute",
-  "repad_logits_with_grad": false,
+  "position_embedding_type": "sans_pos",
   "rope_parameters": {
     "full_attention": {
       "rope_theta": 160000.0,
       "rope_type": "default"
     },
     "sliding_attention": {
-      "rope_theta": 10000.0,
+      "rope_theta": 160000.0,
       "rope_type": "default"
     }
   },
@@ -74,5 +81,5 @@
   "sparse_prediction": false,
   "tie_word_embeddings": true,
   "transformers_version": "5.17.0",
-  "vocab_size": 50372
+  "vocab_size": 50368
 }

config_sentence_transformers.json

@@ -1,54 +0,0 @@
-{
-  "model_type": "ColBERT",
-  "__version__": {
-    "sentence_transformers": "5.1.2",
-    "transformers": "5.17.0",
-    "pytorch": "2.7.1+cu126"
-  },
-  "prompts": {
-    "query": "",
-    "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
-}

decider.py

@@ -0,0 +1,272 @@
+"""Self-contained inference code for a joint typed-decision cross-encoder (ModernBERT-family encoder).
+
+Each option of a question is read as one tokenizer pair
+
+    (state, "Question: <instructions>\\nOption: <option>")
+
+and scored by a linear head on the attention-masked mean of the last hidden states; a question's
+probabilities are the softmax over its options' scores. Nothing is truncated: a pair longer than
+`max_length` (default 8,192 tokens, the encoder's native context) raises `TooLong`. All the
+(question, option) pairs of a request are sorted by length and run together in padded batches.
+
+Layout "shared" (decider_config.json) reads exactly the same tokens, but writes the state once per
+request: [CLS] state [SEP], then each "Question: <instructions>\nOption:" block once, then each
+" <option> [SEP]" block, with tree attention (state sees state; a question block sees the state and
+itself; an option block sees the state, its question and itself, never other options). RoPE
+positions follow the virtual pair, so the per-pair limit is unchanged. Pooling "opt" averages the
+option block only.
+
+    from decider import Decider
+    model = Decider.from_pretrained("<repo or directory>")
+    model.answer(state={"text": "..."}, questions={"q1": {"type": "choice", "instructions": "...",
+                                                          "criteria": {"a": "...", "b": "..."}}})
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+
+FILES = ["config.json", "decider_config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json",
+         "special_tokens_map.json"]
+
+
+class TooLong(ValueError):
+    """An input does not fit the model's context; inputs are never truncated."""
+
+
+def text(x: Any) -> str:
+    return x if isinstance(x, str) else json.dumps(x, ensure_ascii=False, separators=(",", ":"))
+
+
+def state_text(state: Any) -> str:
+    return text(state) if state not in ("", None, {}, []) else "(empty)"
+
+
+def format_option(key: str, desc: Any) -> str:
+    desc_str = text(desc).strip() if desc is not None else ""
+    key_str = str(key).strip()
+    if not desc_str:
+        return key_str
+    is_dummy_key = (key_str.lower().startswith(("option_", "k_", "choice_", "item_"))
+                    or (len(key_str) == 1 and key_str.isalpha())
+                    or (key_str.startswith("(") and key_str.endswith(")") and len(key_str) <= 4))
+    if is_dummy_key:
+        return desc_str
+    if desc_str.lower() == key_str.lower():
+        return key_str
+    return f"{key_str}: {desc_str}"
+
+
+def question_options(q: dict) -> Tuple[List[str], List[str]]:
+    """(answer keys, rendered option texts) for a choice or noul question."""
+    crit = q.get("criteria", {}) or {}
+    if q.get("type", "choice") == "choice":
+        keys = list(crit)
+        return keys, [format_option(k, crit[k]) for k in keys]
+    descs = (crit.get("false", crit.get("False", "False")), crit.get("true", crit.get("True", "True")))
+    return ["false", "true"], [format_option("false", descs[0]), format_option("true", descs[1])]
+
+
+def tail_text(question: str, option: str) -> str:
+    return f"Question: {question}\nOption: {option}"
+
+
+def _resolve(path_or_repo: str, revision: Optional[str] = None) -> Path:
+    p = Path(path_or_repo)
+    if p.is_dir():
+        return p
+    from huggingface_hub import snapshot_download
+    return Path(snapshot_download(path_or_repo, allow_patterns=FILES, revision=revision))
+
+
+class Decider:
+    def __init__(self, directory: Path, device: Optional[str] = None, max_length: Optional[int] = None,
+                 token_budget: int = 65536):
+        import torch
+        import torch.nn as nn
+        from safetensors.torch import load_file
+        from transformers import AutoConfig, AutoModel, AutoTokenizer
+
+        self.torch = torch
+        self.config = json.loads((directory / "decider_config.json").read_text())
+        self.max_length = int(max_length or self.config.get("inference_max_length", 8192))
+        self.token_budget = int(token_budget)
+        self.tok = AutoTokenizer.from_pretrained(str(directory))
+        cfg = AutoConfig.from_pretrained(str(directory))
+        net = nn.Module()
+        net.backbone = AutoModel.from_config(cfg)
+        net.drop = nn.Dropout(0.1)
+        net.scorer = nn.Linear(cfg.hidden_size, 1)
+        net.load_state_dict(load_file(str(directory / "model.safetensors")), strict=True)
+        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        self.net = net.to(self.device).eval()
+        self.layout = self.config.get("layout", "pair")
+        self.pool = self.config.get("pool", "opt")
+        self.window = getattr(cfg, "sliding_window", None)
+
+    @classmethod
+    def from_pretrained(cls, path_or_repo: str, revision: Optional[str] = None, **kw) -> "Decider":
+        return cls(_resolve(path_or_repo, revision), **kw)
+
+    def score_many(self, items: Sequence[Tuple[str, str, Sequence[str]]]) -> List[np.ndarray]:
+        """items: [(state, instructions, option texts)] -> one logit array per item."""
+        if self.layout == "shared":
+            return self._score_shared(items)
+        torch = self.torch
+        pairs = [self.tok(s, tail_text(q, o), truncation=False)["input_ids"] for s, q, opts in items for o in opts]
+        longest = max(map(len, pairs))
+        if longest > self.max_length:
+            raise TooLong(f"a (state, question, option) pair needs {longest} tokens > {self.max_length}")
+        order = sorted(range(len(pairs)), key=lambda i: len(pairs[i]))
+        scores = np.zeros(len(pairs), dtype=np.float64)
+        i = 0
+        with torch.no_grad():
+            while i < len(order):
+                j = i
+                while j < len(order) and (j - i + 1) * len(pairs[order[j]]) <= max(self.token_budget, len(pairs[order[i]])):
+                    j += 1
+                idx = order[i:j]
+                width = len(pairs[idx[-1]])
+                ids = torch.full((len(idx), width), self.tok.pad_token_id, dtype=torch.long)
+                mask = torch.zeros((len(idx), width), dtype=torch.long)
+                for r, k in enumerate(idx):
+                    ids[r, :len(pairs[k])] = torch.tensor(pairs[k])
+                    mask[r, :len(pairs[k])] = 1
+                ids, mask = ids.to(self.device), mask.to(self.device)
+                with torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
+                    h = self.net.backbone(input_ids=ids, attention_mask=mask).last_hidden_state
+                    m = mask.unsqueeze(-1).to(h.dtype)
+                    z = self.net.scorer((h * m).sum(1) / m.sum(1).clamp_min(1)).squeeze(-1)
+                scores[idx] = z.float().cpu().numpy()
+                i = j
+        out, start = [], 0
+        for _, _, opts in items:
+            out.append(scores[start:start + len(opts)])
+            start += len(opts)
+        return out
+
+    def _score_shared(self, items, row_extra: int = 2048) -> List[np.ndarray]:
+        torch, tok = self.torch, self.tok
+        rows, oid, prefix = [], 0, {}
+        for s, q, opts in items:  # one row group per distinct state; rows of max(2P, P + row_extra) tokens
+            p = prefix.setdefault(s, tok(s)["input_ids"])
+            qi = tok(f"Question: {q}\nOption:", add_special_tokens=False)["input_ids"]
+            ob = [tok(" " + o, add_special_tokens=False)["input_ids"] + [tok.sep_token_id] for o in opts]
+            longest = len(p) + len(qi) + max(map(len, ob))
+            if longest > self.max_length:
+                raise TooLong(f"a (state, question, option) pair needs {longest} tokens > {self.max_length}")
+            cap = max(2 * len(p), len(p) + row_extra)
+            if not rows or rows[-1][0] is not p or rows[-1][2] + len(qi) + len(ob[0]) > cap:
+                rows.append([p, [], len(p)])
+            k = 0
+            while k < len(ob):
+                row = rows[-1]
+                take, size = [], len(qi)
+                while k + len(take) < len(ob) and (not take or row[2] + size + len(ob[k + len(take)]) <= cap):
+                    take.append((oid + k + len(take), ob[k + len(take)]))
+                    size += len(take[-1][1])
+                if row[1] and row[2] + size > cap:
+                    rows.append([p, [], len(p)])
+                    continue
+                row[1].append((qi, take))
+                row[2] += size
+                k += len(take)
+            oid += len(ob)
+        scores = np.zeros(oid, dtype=np.float64)
+        order = sorted(range(len(rows)), key=lambda i: rows[i][2])
+        i = 0
+        with torch.no_grad():
+            while i < len(order):
+                j = i + 1
+                while (j < len(order) and (j - i + 1) * rows[order[j]][2] <= self.token_budget
+                       and (j - i + 1) * rows[order[j]][2] ** 2 <= 2 ** 26):  # bound dense mask size
+                    j += 1
+                batch = [rows[k] for k in order[i:j]]
+                L, B = max(r[2] for r in batch), len(batch)
+                ids = torch.full((B, L), tok.pad_token_id, dtype=torch.long)
+                pos = torch.zeros((B, L), dtype=torch.long)
+                kind = torch.full((B, L), -1, dtype=torch.long)
+                qix = torch.full((B, L), -1, dtype=torch.long)
+                oix = torch.full((B, L), -1, dtype=torch.long)
+                spans = []
+                for r, (p, blocks, _) in enumerate(batch):
+                    ids[r, :len(p)], pos[r, :len(p)], kind[r, :len(p)] = torch.tensor(p), torch.arange(len(p)), 0
+                    c = len(p)
+                    for b, (qi, take) in enumerate(blocks):
+                        n = len(qi)
+                        ids[r, c:c + n], pos[r, c:c + n] = torch.tensor(qi), torch.arange(len(p), len(p) + n)
+                        kind[r, c:c + n], qix[r, c:c + n] = 1, b
+                        c += n
+                        for o, x in take:
+                            m = len(x)
+                            ids[r, c:c + m] = torch.tensor(x)
+                            pos[r, c:c + m] = torch.arange(len(p) + n, len(p) + n + m)
+                            kind[r, c:c + m], qix[r, c:c + m], oix[r, c:c + m] = 2, b, o
+                            spans.append((o, r, c, m, len(p), (r, len(p) + sum(len(bb[0]) + sum(len(y) for _, y in bb[1]) for bb in blocks[:b])), n))
+                            c += m
+                ids, pos, kind, qix, oix = (t.to(self.device) for t in (ids, pos, kind, qix, oix))
+                ki, kj = kind[:, :, None], kind[:, None, :]
+                allow = (kj == 0) & (ki >= 0)
+                allow |= (kj == 1) & (ki >= 1) & (qix[:, :, None] == qix[:, None, :])
+                allow |= (kj == 2) & (ki == 2) & (oix[:, :, None] == oix[:, None, :])
+                allow |= torch.eye(L, dtype=torch.bool, device=self.device)[None]
+                masks = {"full_attention": allow[:, None]}
+                if self.window is not None:
+                    masks["sliding_attention"] = (allow & ((pos[:, :, None] - pos[:, None, :]).abs() <= self.window))[:, None]
+                with torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
+                    h = self.net.backbone(input_ids=ids, attention_mask=masks, position_ids=pos).last_hidden_state
+                h = h.float()
+                if self.pool == "opt":  # vectorized mean over each option block
+                    sel = oix >= 0
+                    flat = oix[sel]
+                    uniq, inv = torch.unique(flat, return_inverse=True)
+                    v = torch.zeros(len(uniq), h.shape[-1], device=h.device).index_add_(0, inv, h[sel])
+                    v = v / torch.bincount(inv).unsqueeze(-1).to(v.dtype)
+                    with torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
+                        z = self.net.scorer(v).squeeze(-1)
+                    scores[uniq.cpu().numpy()] = z.float().cpu().numpy()
+                    i = j
+                    continue
+                vecs = []
+                for o, r, c, m, plen, (_, qstart), n in spans:
+                    v = h[r, c:c + m].sum(0)
+                    if self.pool == "pair":
+                        v = (v + h[r, :plen].sum(0) + h[r, qstart:qstart + n].sum(0)) / (plen + n + m)
+                    else:
+                        v = v / m
+                    vecs.append(v)
+                with torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
+                    z = self.net.scorer(torch.stack(vecs)).squeeze(-1)
+                scores[[sp[0] for sp in spans]] = z.float().cpu().numpy()
+                i = j
+        out, start = [], 0
+        for _, _, opts in items:
+            out.append(scores[start:start + len(opts)])
+            start += len(opts)
+        return out
+
+    def answer(self, state: Any, questions: Dict[str, Any]) -> Dict[str, Any]:
+        """Decision Index request -> {"answers": {qid: answer}}. Supports `choice` and `noul`."""
+        s = state_text(state)
+        parsed = []
+        for qid, q in questions.items():
+            kind = q.get("type", "choice")
+            if kind not in ("choice", "noul") or (kind == "choice" and not q.get("criteria")):
+                raise ValueError(f"unsupported question type {kind!r}")
+            keys, opts = question_options(q)
+            parsed.append((qid, kind, keys, opts, text(q.get("instructions", ""))))
+        logits = self.score_many([(s, ins, opts) for _, _, _, opts, ins in parsed])
+        answers = {}
+        for (qid, kind, keys, _, _), z in zip(parsed, logits):
+            p = np.exp(z - z.max())
+            p /= p.sum()
+            if kind == "noul":
+                answers[qid] = {"type": "noul", "noul": float(np.clip(p[1], 0.0, 1.0))}
+            else:
+                answers[qid] = {"type": "choice", "choice": keys[int(p.argmax())],
+                                "probabilities": {k: float(v) for k, v in zip(keys, p)}}
+        return {"answers": answers}

decider_config.json

@@ -0,0 +1,32 @@
+{
+  "format": "pair-v1",
+  "arch": "joint",
+  "base": "cross-encoder/ettin-reranker-150m-v1",
+  "train_max_length": 2048,
+  "inference_max_length": 8192,
+  "step": 12000,
+  "train": {
+    "arch": "joint",
+    "opt_pool": "marker",
+    "shuffle_options": false,
+    "manifest": "runs/manifests/mix_all.jsonl",
+    "extra_manifest": null,
+    "extra_frac": 0.0,
+    "lr": 2e-05,
+    "batch_requests": 8,
+    "max_k": 32,
+    "seed": 7,
+    "total_steps": 12000,
+    "max_length": 2048,
+    "pack": 0,
+    "init": null,
+    "init_backbone": null,
+    "gc_tokens": 0,
+    "layout": "shared",
+    "pool": "opt",
+    "row_tokens": 8192,
+    "tok_budget": 12288
+  },
+  "layout": "shared",
+  "pool": "opt"
+}

decision_index_engine.py

@@ -0,0 +1,50 @@
+"""Decision Index engine for this model (https://github.com/apolinario/decision-index).
+
+    python -m decision_index run --engine decision_index_engine:DeciderEngine \
+        --option repo=<repo or directory> [--option revision=<commit>] --edition 0.3 --out runs/<name>
+
+Put this file next to `decider.py` (both are in the model repository) and on PYTHONPATH.
+Kit rules: nothing is truncated and no option is dropped. A (state, question, option) pair longer than
+8,192 tokens (the encoder's native context) raises `Unsupported`. One call = one request: all questions
+of the request are scored together, in length-sorted padded batches.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+from decision_index.engines import Engine, Unsupported
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from decider import Decider, TooLong  # noqa: E402
+
+
+class DeciderEngine(Engine):
+    name = "decider"
+    latency = ("In-process request wall time: tokenization plus the encoder passes of one request "
+               "(all its question/option pairs batched together); excludes model loading.")
+
+    def __init__(self, repo: str = str(Path(__file__).resolve().parent), revision: str | None = None,
+                 device: str | None = None, max_length: int | None = None, token_budget: int = 65536, **options):
+        super().__init__(**options)
+        self.model = Decider.from_pretrained(repo, revision=revision, device=device, max_length=max_length,
+                                             token_budget=int(token_budget))
+        self.repo = repo
+        self.provenance = {"repo": repo, "revision": revision or "main", "device": str(self.model.device),
+                           "dtype": "bfloat16 autocast" if self.model.device.type == "cuda" else "float32",
+                           "max_length": self.model.max_length, "token_budget": self.model.token_budget,
+                           "truncation": f"none: pairs above {self.model.max_length} tokens raise Unsupported"}
+
+    def __call__(self, state, questions):
+        try:
+            out = self.model.answer(state, questions)
+        except TooLong as e:
+            raise Unsupported(f"context window: {e}")
+        except ValueError as e:
+            raise Unsupported(str(e))
+        return {"model": self.repo, "answers": out["answers"]}, None
+
+    def synchronize(self):
+        if self.model.device.type == "cuda":
+            self.model.torch.cuda.synchronize()

decision_meta.json

@@ -1,37 +0,0 @@
-{
-  "init_temperature": 0.31683406233787537,
-  "path_mode": "sq_o",
-  "pooling_budget": null,
-  "protected_tokens": 1,
-  "rel_slots": 2,
-  "joint_rel_distinct": false,
-  "rel_score_scale": 256.0,
-  "role_tokens": false,
-  "trimax_rank": null,
-  "has_query_gate": false,
-  "gate_dim": null,
-  "document_length": 2048,
-  "query_length": 256,
-  "model_max_length": 8192,
-  "length_norm_gamma": 0.0,
-  "config": {
-    "model_name_or_path": "/mnt/nfs_share_magnet2/dsileo/libs/train_jev/runs/multivector/scaling_grouped_p2/step-4000",
-    "embedding_dim": 128,
-    "state_max_length": 512,
-    "query_max_length": 64,
-    "init_temperature": 0.3215,
-    "learn_temperature": true,
-    "path_mode": "sq_o",
-    "pooling_budget": null,
-    "protected_tokens": 1,
-    "rel_slots": 2,
-    "joint_rel_distinct": false,
-    "rel_score_scale": null,
-    "role_tokens": false,
-    "trimax_rank": null,
-    "query_prefix": "[Q] ",
-    "document_prefix": "[D] ",
-    "backend": "torch",
-    "metadata": {}
-  }
-}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:48578d08f6b50c9ba6654d3d26fa71c381c7a34ea9791a6493877dea475f036d
-size 596082424
+oid sha256:a96dcfdf2d4e7da62429ba9f836535adf3a328f30cb51366095ad2c38fb13659
+size 596074556

modules.json

@@ -1,26 +0,0 @@
-[
-  {
-    "idx": 0,
-    "name": "0",
-    "path": "",
-    "type": "sentence_transformers.models.Transformer"
-  },
-  {
-    "idx": 1,
-    "name": "1",
-    "path": "1_Dense",
-    "type": "pylate.models.Dense.Dense"
-  },
-  {
-    "idx": 2,
-    "name": "2",
-    "path": "2_Dense",
-    "type": "pylate.models.Dense.Dense"
-  },
-  {
-    "idx": 3,
-    "name": "3",
-    "path": "3_Dense",
-    "type": "pylate.models.Dense.Dense"
-  }
-]

sentence_bert_config.json

@@ -1,4 +0,0 @@
-{
-    "max_seq_length": 255,
-    "do_lower_case": false
-}

tokenizer.json

@@ -2,18 +2,11 @@
   "version": "1.0",
   "truncation": {
     "direction": "Right",
-    "max_length": 255,
+    "max_length": 7999,
     "strategy": "LongestFirst",
     "stride": 0
   },
-  "padding": {
-    "strategy": "BatchLongest",
-    "direction": "Right",
-    "pad_to_multiple_of": null,
-    "pad_id": 50284,
-    "pad_type_id": 0,
-    "pad_token": "[MASK]"
-  },
+  "padding": null,
   "added_tokens": [
     {
       "id": 0,
@@ -1058,42 +1051,6 @@
       "rstrip": false,
       "normalized": true,
       "special": false
-    },
-    {
-      "id": 50368,
-      "content": "[Q] ",
-      "single_word": false,
-      "lstrip": false,
-      "rstrip": false,
-      "normalized": true,
-      "special": false
-    },
-    {
-      "id": 50369,
-      "content": "[D] ",
-      "single_word": false,
-      "lstrip": false,
-      "rstrip": false,
-      "normalized": true,
-      "special": false
-    },
-    {
-      "id": 50370,
-      "content": "[REL0]",
-      "single_word": false,
-      "lstrip": false,
-      "rstrip": false,
-      "normalized": false,
-      "special": true
-    },
-    {
-      "id": 50371,
-      "content": "[REL1]",
-      "single_word": false,
-      "lstrip": false,
-      "rstrip": false,
-      "normalized": false,
-      "special": true
     }
   ],
   "normalizer": {

tokenizer_config.json

@@ -2,19 +2,16 @@
   "backend": "tokenizers",
   "clean_up_tokenization_spaces": true,
   "cls_token": "[CLS]",
-  "is_local": true,
+  "is_local": false,
   "local_files_only": false,
   "mask_token": "[MASK]",
-  "max_length": 299,
+  "max_length": 7999,
   "model_input_names": [
     "input_ids",
     "attention_mask"
   ],
-  "model_max_length": 8192,
-  "pad_to_multiple_of": null,
-  "pad_token": "[MASK]",
-  "pad_token_type_id": 0,
-  "padding_side": "right",
+  "model_max_length": 7999,
+  "pad_token": "[PAD]",
   "sep_token": "[SEP]",
   "stride": 0,
   "tokenizer_class": "TokenizersBackend",