Duplicate from convaiinnovations/laya

Co-authored-by: Convai Innovations <convaiinnovations@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,41 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+eval/benchmark_comparison.png filter=lfs diff=lfs merge=lfs -text
+multilingual/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+assets/laya_vs_jev.png filter=lfs diff=lfs merge=lfs -text
+assets/laya_benchmark.png filter=lfs diff=lfs merge=lfs -text
+assets/laya_benchmark_common.png filter=lfs diff=lfs merge=lfs -text
+assets/laya_vs_jev_full.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,323 @@
+---
+license: apache-2.0
+library_name: transformers
+pipeline_tag: text-classification
+tags: [laya, system-one, calibrated-decisions, rlcd, classification, routing, scoring, guardrails, moderation, reinforcement-learning, commercial-use]
+---
+
+# Laya
+
+**Multilingual, non-autoregressive System 1 decision model.** Give it a **state** (text, email, ticket, or JSON) and **typed questions**; it returns typed answers with mathematically calibrated probabilities in a single forward pass (~33 ms) across 100+ languages. Trained with reinforcement learning against strictly proper scoring rules (**RLCD**), so reporting honest probabilities is the only way to maximise reward. It never generates text, so there is nothing to parse and nothing to hallucinate.
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/NandhaKishorM/laya/main/assets/laya_vs_jev_full.png" alt="Laya versus TypeSafe Jev: accuracy, every application workflow, all 51 languages, speed, calibration and routing cost" width="100%" />
+</p>
+
+**This repo holds all three checkpoints** and is the hub for the family. The English checkpoint is at the repo root; the other two are bundled subfolders, and only the one you request is downloaded:
+
+| Checkpoint | Backbone Encoder | Params | Context | Best at |
+|---|---|---|---|---|
+| **`convaiinnovations/laya`** (this repo root) | ModernBERT-large | 421M | 512 | English text, guardrails, email triage |
+| [`convaiinnovations/laya-multilingual`](https://huggingface.co/convaiinnovations/laya-multilingual) | mmBERT-base | 322M | 1024 (up to 8k) | 100+ languages, ~2.2x faster |
+| [`convaiinnovations/laya-typed-decisions`](https://huggingface.co/convaiinnovations/laya-typed-decisions) | ModernBERT-large | 421M | 1024 | the four typed-decisions workflows (0.766 acc) |
+
+## What's new in laya 0.3.11
+
+`pip install -U laya` for all of this; everything below is new since 0.3.6. The checkpoints themselves are unchanged.
+
+* **About 10x faster loading.** Checkpoints are built without the throwaway random weight initialisation, so `laya.load()` drops from about 22 s to about 2 s on CPU, with bit-identical answers. This also skips the pass that crashed on Windows with Python 3.14.
+* **`import laya` no longer loads torch.** Routing, language detection and e-mail cleaning work in lightweight processes.
+* **Batch scoring.** `agent.predict_batch(states, questions)` scores many states in shared forward passes, with answers identical to calling `predict` one state at a time.
+* **Routed batches.** `Router.predict_batch(requests)` routes each request, groups them by checkpoint and question set, and scores each group in shared forward passes, with answers identical to one `predict` call per request.
+* **Prediction hooks.** Opt-in hooks run around every decision, to audit, trace, redact, cache or gate results. With no hooks set, answers are unchanged.
+* **transformers 4.x and Apple GPUs.** Checkpoints load with the right RoPE settings on transformers 4.x, and `predict()` no longer crashes on MPS builds without an autocast backend.
+* **Faster paths, all opt-in.** `laya.load(..., fast=True)` uses a TileLang GPU fast path that matches the stock bf16 forward within rounding. `Agent(compile=True)` enables `torch.compile`, and `laya.onnx_agent.ONNXAgent` runs an exported model on ONNX Runtime.
+* **Run it your way, locally.** A self-hosted Jev-compatible HTTP server (`pip install "laya[serve]"`, then `laya-serve`, see [Self-hosting](#self-hosting-jev-compatible-http-server)), a `laya` command for quick local tests, an optional MCP server (`pip install "laya[mcp]"`), LangChain and LangGraph integrations (`pip install "laya[langchain]"`), and [`laya-ts`](https://github.com/NandhaKishorM/laya/tree/main/laya-ts), a TypeScript package for Node and the browser that gives the same answers as the Python package.
+* **Better routing.** Plain-ASCII Spanish, Italian, Portuguese, French and German, Brazilian Portuguese support text, CJK text with Latin brand names, romanized Bangla and Azerbaijani now reach the multilingual checkpoint. Checked on 20,000 English texts: at most 5 English sentences move, all quoting long native-script names.
+* **Router defaults and hooks.** `Router()` keeps two checkpoints resident, you can pass your own language guess with `lang_guess=`, and `Router`/`Agent` work as context managers.
+* **Correctness and clearer errors.** Long conversation lists keep the newest turn when truncated, non-ASCII instructions reach the model as text, and a malformed question is rejected with a message naming the question and what to fix. A `noul` `criteria` dict keyed anything other than `true`/`false` is rejected rather than silently replaced; use `labels` to change the wording.
+
+---
+
+## Quickstart: Route Mode (Recommended)
+
+Laya's built-in **`Router`** is the recommended way to use Laya in production. It evaluates any state in any language, automatically detects scripts and languages in sub-milliseconds, and dispatches to the optimal checkpoint in a single forward pass.
+
+```bash
+pip install laya
+```
+
+```python
+import laya
+from laya import Router
+
+# Preload checkpoints into memory for instant sub-35ms routing
+router = Router(preload=True)
+
+state = {
+    "from": "user@acme.com",
+    "subject": "Duplicate charge on invoice #4411",
+    "body": "Hi, we were billed twice for March. Please refund the duplicate today or we will cancel our plan."
+}
+
+questions = {
+    "department": {
+        "type": "choice",
+        "instructions": "Which department should handle this request?",
+        "criteria": {
+            "billing": "invoices, payments, refunds",
+            "technical": "bugs, outages, system errors",
+            "sales": "pricing, new contracts",
+            "other": "everything else"
+        }
+    },
+    "urgency": {
+        "type": "score",
+        "instructions": "How urgent is this request?",
+        "criteria": ["not urgent", "soon", "critical deadline or blocking issue"]
+    },
+    "churn_risk": {
+        "type": "noul",
+        "instructions": "Does the user threaten to cancel or leave?"
+    },
+    "refund_requested": {
+        "type": "noul",
+        "instructions": "Does the user explicitly request a refund?"
+    }
+}
+
+# 1. English state -> automatically routed to ModernBERT-large (39.5 ms)
+res_en = router.predict(state, questions)
+print("Department :", res_en["answers"]["department"]["choice"])  # -> billing (confidence: 0.94)
+print("Routing    :", res_en["routing"]["model"])                 # -> english
+
+# 2. Hindi state -> automatically routed to mmBERT-base (100+ languages, 32.8 ms)
+res_hi = router.predict({"body": "मुझसे दो बार शुल्क लिया गया, कृपया पैसे वापस करें।"}, questions)
+print("Department :", res_hi["answers"]["department"]["choice"])  # -> billing (confidence: 0.86)
+print("Routing    :", res_hi["routing"]["model"])                 # -> multilingual
+
+# 3. Explicit override when you already know the checkpoint
+res_td = router.predict(state, questions, model="typed-decisions")
+```
+
+Every result carries full routing metadata explaining why the choice was made:
+
+```python
+res_hi["routing"]
+# {
+#   'model': 'multilingual',
+#   'repo': 'convaiinnovations/laya/multilingual',
+#   'reason': 'non-Latin script (devanagari, 100% of letters); the English checkpoint cannot read it'
+# }
+```
+
+### Why Route: The Evidence
+
+On a shared benchmark (17,416 questions, one T4 GPU, identical questions per model):
+
+| Benchmark / Task | English (`laya`) | Multilingual (`laya-multilingual`) | `Router` (Routed) |
+|---|---|---|---|
+| MASSIVE intent, English | **0.783** | 0.657 | **0.783** |
+| MASSIVE intent, 13 other languages | 0.306 | **0.451** | **0.451** |
+| XNLI, English | **0.860** | 0.843 | **0.860** |
+| XNLI, 14 other languages | 0.521 | **0.731** | **0.731** |
+| Languages usable (>3x random) | 23 / 51 | 45 / 51 | **45 / 51** |
+| Latency, 1 question (T4 GPU) | 39.5 ms | **32.8 ms** | **32.8 ms** |
+| Latency, 10 questions batched | 158.6 ms | **72.3 ms** | **72.3 ms** |
+
+The English checkpoint collapses on non-Latin scripts (Khmer scores **0.000 accuracy at 0.952 confidence**). Because the model stays confident while being wrong, confidence gating cannot save you. `Router` detects the script in <0.5 ms pure Python before the forward pass.
+
+### Supplying your own language detection
+
+If you already run a language-identification model, pass its answer instead of relying on the built-in heuristic. `lang_guess` takes a language code or a callable, is checked after an explicit `lang=` and before detection, and a callable that returns `None` falls through to detection:
+
+```python
+router.predict(state, questions, lang_guess="ro")              # a code you already know
+router = Router(preload=True, lang_guess=my_lid)               # or install one for every request
+```
+
+### Production Preload & Memory
+
+A cold checkpoint build costs seconds; language detection costs microseconds. Since laya 0.3.11 the lazy default keeps **two** checkpoints resident (`english` and `multilingual`, the only two automatic routing chooses between), so after each language's first load a switch costs detection only. A single-language deployment never builds the second. `max_loaded=1` rebuilds on every switch (measured at a 7.4 s median reload on CPU and 10.3 s on T4).
+
+For a server or a demo, preload:
+
+```python
+# Every checkpoint resident in memory; language flips cost detection only (<1 ms)
+router = Router(preload=True)
+router = Router(preload=True, device="cuda")
+
+# Or preload only the specific checkpoints you serve:
+router.preload(["english", "multilingual"])
+
+# If your app already built an agent, attach it to avoid duplicate VRAM:
+router.attach("english", existing_agent)
+
+# Manage resident memory (default keeps two hot: english + multilingual, LRU eviction)
+router = Router(max_loaded=3)       # all three hot, e.g. with auto_task_detection
+router = Router(max_loaded=1)       # memory-constrained host, reloads on every switch
+router.unload()                     # free memory
+
+with Router() as r:                 # releases the models when the block ends
+    r.predict(state, questions)
+```
+
+| Deployment Mode | Per-Request Latency | Model Reloads |
+|---|---|---|
+| `Router()` (lazy, `max_loaded=2`) | detection only (<1 ms) after each language's first load | 1 the first time a language appears |
+| `Router(max_loaded=1)` | 7 to 10 s on every language switch | 1 per switch |
+| `Router(preload=True)` | **32.8 ms (GPU) / 193–464 ms (CPU)** | **none** |
+
+---
+
+## Single-Model Mode (Direct SDK)
+
+If you only need a single checkpoint for a dedicated pipeline:
+
+```python
+import laya
+
+# 1. Load from the repo root or subfolders (downloads only the requested weights)
+agent = laya.load("convaiinnovations/laya")                           # English root (~808 MB)
+agent_ml = laya.load("convaiinnovations/laya", subfolder="multilingual") # 100+ languages (~647 MB)
+agent_td = laya.load("convaiinnovations/laya", subfolder="typed-decisions")
+
+# 2. Run all questions in ONE single forward pass (~35 ms on GPU)
+result = agent.predict(state, questions)
+answers = result["answers"]
+
+print("Department :", answers["department"]["choice"])   # -> billing (confidence: 0.94)
+print("Urgency    :", answers["urgency"]["score"])        # -> 1.84 / 2.0
+print("Churn Risk :", answers["churn_risk"]["noul"])       # -> 0.892 (89.2% probability)
+```
+
+> **If `laya.load()` hangs:** `transformers` probes for TensorFlow at import, and when TF is
+> installed its abseil runtime can deadlock model construction. Run with `USE_TF=0`.
+
+---
+
+## Self-hosting: Jev-compatible HTTP server
+
+`laya-serve` exposes the `Router` on the same `POST /v1/systemone` request and response shape as TypeSafe Jev, so existing TypeSafe clients work by changing their base URL:
+
+```bash
+pip install "laya[serve]"
+LAYA_DEVICE=cuda LAYA_PRELOAD=1 laya-serve        # 0.0.0.0:8000, preloads the checkpoints
+```
+
+```bash
+curl -s localhost:8000/v1/systemone -H 'Content-Type: application/json' -d '{
+  "state": {"document": "I was charged twice. Please fix this ASAP."},
+  "questions": {"billing": {"type": "noul", "instructions": "Is this ticket about billing?"}}
+}'
+```
+
+It accepts every question shape the Jev API does (for example `criteria` as a list), ignores unknown fields, and returns a 422 naming the problem for a malformed question. It binds `0.0.0.0` with no authentication unless `LAYA_API_KEY` is set, in which case it requires `Authorization: Bearer <key>`.
+
+---
+
+## Architecture
+
+- **Backbone:** ModernBERT-large (395M, bidirectional, fully fine-tuned) + a decision head trained from scratch: 2 transformer layers, an option-marker scorer, and an act/escalate head. 421M total. (Multilingual uses mmBERT-base, 22 layers, 256k vocab, 322M total).
+- **Option markers:** Every option is scored at its own `[MASK]` token, then softmaxed over that question's options. The answer space is defined at request time, so new schemas need no retraining.
+- **Budget:** 512 tokens per question for English (`head_max_len = 192`); 1024 tokens for multilingual (`head_max_len = 256`).
+- **Batching:** Every question in a call is answered in one single forward pass.
+
+---
+
+## Training
+
+**RLCD (Reinforcement Learning for Calibrated Decisions).** The policy reports a distribution; exploration adds zero-mean Gaussian noise to the logits; the reward is a strictly proper scoring rule (log + spherical, plus ranked probability score for ordinal questions). Expected reward is maximised only by reporting honest probabilities. Updates are REINFORCE with a group-mean baseline (GRPO-style). Multi-turn conversations use TD(λ=1.0) over prefix slices.
+
+---
+
+## Benchmarks
+
+Measured on a Tesla T4; every checkpoint answered byte-identical questions in the same run.
+
+### Speed
+
+| questions per call | `laya` | `laya-multilingual` |
+|---|---|---|
+| 1 | 39.5 ms | **32.8 ms** |
+| 5 | 84.5 ms | **40.1 ms** |
+| 10 | 158.6 ms (15.9 ms/q) | **72.3 ms (7.2 ms/q)** |
+| 50 | 771 ms | **337 ms (6.8 ms/q)** |
+
+103–332 questions/sec batched on a single T4. For reference, TypeSafe Jev has been independently measured at 236–276 ms p50 ([AbdelStark](https://github.com/AbdelStark/jev-benchmarks), [nibzard](https://github.com/nibzard/decision-model-benchmark)), so Laya answers a single question roughly **6–8× faster**.
+
+### Laya (with routing) vs TypeSafe Jev
+
+Every Laya figure is what `Router().predict(...)` returns — the checkpoint the router selects for that input. Jev figures are **third-party published, never measured here** (no TypeSafe API access); sample sizes and prompts differ.
+
+| Benchmark / Metric | TypeSafe Jev 1.13.0 | Laya (routed) | Comparison |
+|---|---|---|---|
+| typed-decisions, 2,000 decisions | 0.727 | **0.766** | +0.039 (beats 0.735 teacher ceiling) |
+| AG News, 4 labels | 0.910 | **0.950** | +0.040 |
+| DAIR Emotion, 6 labels | 0.480 | **0.595** | +0.115 |
+| Banking77 (72 vs 77 labels) | **0.870** | 0.425 | Jev leads on >20 options |
+| ECE *(lower better)* | 0.246 | **0.081** | 3× better (post-temperature) |
+| p50 latency, 1 question | 236–276 ms | **32.8 ms** | 7.8× faster |
+| Languages usable (>3x random) | *no published benchmark* | **45 of 51** | Global language coverage |
+| Weights | closed API | **Apache 2.0** | Open weights, on-premise capable |
+| Cost | $0.042 / 1M tokens | **$0 self-hosted** | 100% free |
+
+On DAIR Emotion, Jev assigned zero probability to the true label on 16% of examples.
+
+#### Where Jev leads
+
+* **High-cardinality label spaces (>20 options at default settings):** On Banking77, Jev scores 0.870 (on 72 labels) while Laya scores 0.425 (on 77 labels at default 256-token head budget). Options share a fixed `head_max_len` budget (192 tokens on English, 256 on multilingual), so 77 options receive only ~3 to 4 tokens per label, causing text to become indistinguishable. Jev supports up to 255 options out-of-the-box. While `laya-multilingual` supports 1,024 context (and up to 8,192 in the encoder) and you can raise `agent.cfg["head_max_len"] = 512` at runtime, Jev is currently better suited for 50+ options in a single prompt without tuning.
+* **Soft distribution matching:** On typed-decisions, while Laya achieves higher argmax accuracy (0.766 vs 0.727), Jev achieves higher soft accuracy (0.580 vs 0.471) against the teacher's full probability distributions.
+* **Out-of-the-box raw calibration:** Before temperature scaling, the base checkpoint has higher raw ECE (0.213 vs 0.144). Laya achieves its 0.081 ECE after domain temperature fitting.
+
+Full report: [BENCHMARKS.md](https://github.com/NandhaKishorM/laya/blob/main/BENCHMARKS.md).
+
+### typed-decisions, measured on all three checkpoints
+
+400 cases, 2,000 decisions, four workflows — measured here.
+
+| model | accuracy | soft acc | Brier | ECE | score MAE |
+|---|---|---|---|---|---|
+| **`laya-typed-decisions`** | **0.766** | 0.471 | **0.062** | 0.213 | **0.242** |
+| `laya` | 0.362 | 0.332 | 0.316 | 0.175 | 0.694 |
+| `laya-multilingual` | 0.342 | 0.326 | 0.439 | 0.285 | 0.687 |
+| *Jev 1.13.0 (published)* | *0.727* | *0.580* | *0.148* | *0.144* | *0.391* |
+| *teacher self-agreement ceiling* | *0.735* | | | | |
+| *per-question majority class* | *0.461* | | | | |
+
+The fine-tuned checkpoint clears the teacher ceiling and wins all four workflows: invoice processing 0.804, security incidents 0.766, customer service 0.764, agent-trace observability 0.730. By primitive: `noul` 0.857, `choice` 0.733, `score` 0.723.
+
+The base checkpoints sit below the majority-class baseline here — the capability on this benchmark comes from fine-tuning, which is what the [fine-tuning notebook](https://github.com/NandhaKishorM/laya/blob/main/notebooks/laya_finetune_typed_decisions_2xT4_kaggle.ipynb) is for.
+
+---
+
+## Honest Limits
+
+- **Base checkpoints are near chance on typed-decisions zero-shot** — 0.362 here and 0.352 for multilingual, against a 0.318 random and 0.461 majority-class baseline. The 0.766 belongs to the checkpoint fine-tuned on that benchmark's own training split. Laya is a fast base to specialise, not a zero-shot decision engine.
+- **High-cardinality choice questions and token budgets:** Sequences split into an option prompt budget (`head_max_len`) and the remaining document/state budget (`max_len - head_max_len`):
+  * `laya` (English) defaults to 512 context (`head_max_len = 192`, ~320 tokens for state).
+  * `laya-multilingual` and `laya-typed-decisions` default to 1,024 context (`head_max_len = 256`, ~768 tokens for state; mmBERT-base encoder supports up to 8,192 with RoPE).
+  At default settings, a 77-option question like Banking77 allocates only `(256 - 16) // 77` ≈ 3–4 tokens per label, causing accuracy to fall off sharply (0.425 vs Jev's 0.870). If evaluating 50+ options in a single question:
+  1. Raise `agent.cfg["head_max_len"] = 512` and `agent.cfg["max_len"] = 1024` (or up to 2048 / 4096 / 8192) so every option has enough tokens to remain distinct.
+  2. Or split large option sets into a two-step coarse-to-fine hierarchical choice.
+- **Ordinal `score` questions are the weakest primitive** (SST-5 0.372).
+- **`noul` can follow its option labels instead of the state, most strongly on this English checkpoint.** `noul` renders its two options as `false:` / `true:`, and here that label pair can dominate the answer, returning a confident "no" for clearly positive input ([#156](https://github.com/NandhaKishorM/laya/issues/156)). Check `noul` answers on your own data. If they look stuck, ask the same question as a two-option `choice` with neutral keys and your yes/no wording as the descriptions:
+
+  ```python
+  {"type": "choice", "instructions": "Is this review positive?",
+   "criteria": {"A": "yes, the review is positive", "B": "no, the review is negative"}}
+  ```
+- **`action.act_probability` carries no usable signal yet** ([#185](https://github.com/NandhaKishorM/laya/issues/185)). It reads 1.0 for almost every input, and its raw logits run against correctness (AUROC 0.30 on 396 labelled decisions). Gate on `confidence` instead, which reaches an AUROC of 0.77 on the same items.
+- **Ships over-confident:** Refitting one temperature per (question type, option count) moves mean ECE **0.466 → 0.081** (`laya`) and **0.314 → 0.106** (`laya-multilingual`). Do this on your own data before trusting the probabilities.
+- **English only on root:** Use `laya-multilingual` for anything outside English.
+
+---
+
+## Links
+
+- **GitHub:** https://github.com/NandhaKishorM/laya
+- **PyPI:** https://pypi.org/project/laya/
+- **Live Demo:** https://huggingface.co/spaces/convaiinnovations/laya-demo
+- **Write-up:** [Read on Dev.to](https://dev.to/nandakishor_m_6cc0adfde9f/i-built-non-autoregressive-decision-models-a-year-ago-then-a-frontier-lab-called-it-a-18me)
+
+Apache 2.0 · Convai Innovations

email_utils.py

@@ -0,0 +1,71 @@
+"""Email helpers for RL Agent: clean raw emails into a compact state and a ready-made set of email questions.
+
+Jev-style models lose accuracy on long, noisy state, and RL Agent reads at most max_len (512) tokens,
+so strip quoted replies, signatures and disclaimers in code before asking questions.
+"""
+import re
+
+_QUOTE_HEADERS = [
+    re.compile(r"^\s*On .{0,300}wrote:\s*$", re.I),
+    re.compile(r"^\s*-{2,}\s*(Original|Forwarded) Message\s*-{2,}", re.I),
+    re.compile(r"^\s*_{8,}\s*$"),
+    re.compile(r"^\s*From:\s.+$", re.I),
+]
+_SIGNATURE_MARKERS = [
+    re.compile(r"^\s*--\s*$"),
+    re.compile(r"^\s*(best|kind|warm|many thanks|thanks|thank you|regards|cheers|sincerely)[\w ,!.]*$", re.I),
+    re.compile(r"^\s*sent from my (iphone|android|mobile|ipad)", re.I),
+]
+_DISCLAIMER = re.compile(r"(confidential|intended (solely )?for the (use of the )?(named )?(addressee|recipient)|"
+                         r"if you (have )?received this (e-?mail|message) in error)", re.I)
+
+
+def clean_email_body(body, max_chars=3000):
+    """Remove quoted history, signature and legal disclaimer; collapse whitespace; truncate."""
+    text = (body or "").replace("\r\n", "\n").replace("\r", "\n").replace("\\n", "\n")
+    lines = []
+    for line in text.split("\n"):
+        if any(p.match(line) for p in _QUOTE_HEADERS) and lines:
+            break  # everything below is the previous thread
+        if line.lstrip().startswith(">"):
+            continue
+        lines.append(line.rstrip())
+    # a sign-off only counts near the end (last 40%, or last 8 lines of a short email) and must be a short line
+    cut = len(lines)
+    for i in range(max(1, min(int(len(lines) * 0.6), len(lines) - 8)), len(lines)):
+        if len(lines[i].strip()) <= 40 and any(p.match(lines[i]) for p in _SIGNATURE_MARKERS):
+            cut = i
+            break
+    lines = lines[:cut]
+    paragraphs = [p for p in re.split(r"\n\s*\n", "\n".join(lines)) if not _DISCLAIMER.search(p)]
+    text = re.sub(r"[ \t]+", " ", "\n\n".join(p.strip() for p in paragraphs if p.strip()))
+    return text[:max_chars]
+
+
+def email_state(subject, body, sender=None, clean=True, **extra):
+    """Build the state dict the email questions refer to (`subject`, `body`, optional `from`)."""
+    state = {"subject": (subject or "").strip(), "body": clean_email_body(body) if clean else (body or "")}
+    if sender:
+        state["from"] = sender
+    state.update({k: v for k, v in extra.items() if v is not None})
+    return state
+
+
+def email_questions(categories=None):
+    """A default fan-out of email questions. `categories` = {key: description} for your own routing labels."""
+    categories = categories or {
+        "billing": "invoices, payments, refunds", "technical": "bugs, outages, integrations",
+        "sales": "pricing, demos, new purchases", "account": "login, access, profile changes",
+        "hr": "hiring, leave, payroll", "other": "none of the above",
+    }
+    return {
+        "category": {"type": "choice", "instructions": "Which team should handle the email in `body`?", "criteria": categories},
+        "is_spam": {"type": "noul", "instructions": "Is this email unsolicited spam or bulk marketing?"},
+        "is_phishing": {"type": "noul", "instructions": "Is this email a phishing or scam attempt to steal money, credentials, or personal data?",
+                        "criteria": {"true": "phishing, scam, or fraud", "false": "a legitimate email"}},
+        "urgency": {"type": "score", "instructions": "How urgent is the issue described in `body`?",
+                    "criteria": ["no time pressure", "needs attention soon", "blocking issue or hard deadline"]},
+        "needs_reply": {"type": "noul", "instructions": "Does the sender expect a reply?"},
+        "sentiment": {"type": "score", "instructions": "What is the sender's tone in `body`?",
+                      "criteria": ["angry or very negative", "negative", "neutral", "positive"]},
+    }

encoder/config.json

@@ -0,0 +1,84 @@
+{
+  "architectures": [
+    "ModernBertForMaskedLM"
+  ],
+  "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": 1024,
+  "initializer_cutoff_factor": 2.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 2624,
+  "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",
+    "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": 16,
+  "num_hidden_layers": 28,
+  "pad_token_id": 50283,
+  "position_embedding_type": "absolute",
+  "repad_logits_with_grad": false,
+  "rope_parameters": {
+    "full_attention": {
+      "rope_theta": 160000.0,
+      "rope_type": "default"
+    },
+    "sliding_attention": {
+      "rope_theta": 10000.0,
+      "rope_type": "default"
+    }
+  },
+  "sep_token_id": 50282,
+  "sparse_pred_ignore_index": -100,
+  "sparse_prediction": false,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.0.0",
+  "vocab_size": 50368
+}

eval/results.json

@@ -0,0 +1,142 @@
+{
+  "model": "rl-agent",
+  "questions_evaluated": {
+    "eval_in": 23024,
+    "eval_zs": 2400
+  },
+  "by_task_family": {
+    "eval_in": {
+      "conversation outcomes": {
+        "n": 3600,
+        "accuracy": 0.4817,
+        "ece": 0.0193,
+        "nll": 0.6933
+      },
+      "email triage": {
+        "n": 2691,
+        "accuracy": 0.7321,
+        "ece": 0.0172,
+        "nll": 0.5954
+      },
+      "emotion and tone": {
+        "n": 1825,
+        "accuracy": 0.9058,
+        "ece": 0.0183,
+        "nll": 0.2382
+      },
+      "inference and fact checking": {
+        "n": 3022,
+        "accuracy": 0.8832,
+        "ece": 0.054,
+        "nll": 0.3404
+      },
+      "instruction-following tasks": {
+        "n": 600,
+        "accuracy": 0.8783,
+        "ece": 0.0465,
+        "nll": 0.3021
+      },
+      "intent and routing": {
+        "n": 1475,
+        "accuracy": 0.9912,
+        "ece": 0.0085,
+        "nll": 0.1811
+      },
+      "moderation and safety": {
+        "n": 2708,
+        "accuracy": 0.9671,
+        "ece": 0.0613,
+        "nll": 0.1527
+      },
+      "reading comprehension": {
+        "n": 770,
+        "accuracy": 0.8468,
+        "ece": 0.0833,
+        "nll": 0.4086
+      },
+      "response quality scoring": {
+        "n": 3146,
+        "accuracy": 0.5814,
+        "ece": 0.023,
+        "nll": 1.0087
+      },
+      "robustness checks": {
+        "n": 744,
+        "accuracy": 0.8508,
+        "ece": 0.1085,
+        "nll": 1.0576
+      },
+      "search relevance": {
+        "n": 733,
+        "accuracy": 0.6276,
+        "ece": 0.066,
+        "nll": 0.7281
+      },
+      "sentiment and rating": {
+        "n": 961,
+        "accuracy": 0.4422,
+        "ece": 0.4384,
+        "nll": 3.545
+      },
+      "topic classification": {
+        "n": 749,
+        "accuracy": 0.9386,
+        "ece": 0.0285,
+        "nll": 0.1957
+      }
+    },
+    "eval_zs": {
+      "emotion and tone": {
+        "n": 600,
+        "accuracy": 0.5833,
+        "ece": 0.3178,
+        "nll": 1.9761
+      },
+      "instruction-following tasks": {
+        "n": 600,
+        "accuracy": 0.8633,
+        "ece": 0.0455,
+        "nll": 0.3187
+      },
+      "moderation and safety": {
+        "n": 600,
+        "accuracy": 0.7967,
+        "ece": 0.1713,
+        "nll": 1.4151
+      },
+      "sentiment and rating": {
+        "n": 600,
+        "accuracy": 0.3617,
+        "ece": 0.2915,
+        "nll": 1.7985
+      }
+    }
+  },
+  "calibration_temperature": [
+    1.6369030475616455,
+    1.2514300346374512,
+    1.983399510383606
+  ],
+  "latency_ms": {
+    "1_questions": {
+      "p50_ms": 38.4,
+      "p95_ms": 42.1
+    },
+    "10_questions": {
+      "p50_ms": 156.0,
+      "p95_ms": 158.4
+    },
+    "50_questions": {
+      "p50_ms": 721.4,
+      "p95_ms": 733.0
+    }
+  },
+  "act_policy": {
+    "eval_in": {
+      "automation_rate": 1.0,
+      "accuracy_when_acting": 0.8032331136738056,
+      "accuracy_when_escalating": null,
+      "accuracy_all": 0.8032331136738056
+    }
+  }
+}

eval/results.md

@@ -0,0 +1,53 @@
+# RL Agent evaluation
+
+Metrics after calibration. Zero-shot = task families held out of training entirely.
+
+## In-task
+
+| task family | questions | accuracy | ECE | NLL |
+|---|---|---|---|---|
+| conversation outcomes | 3600 | 0.482 | 0.019 | 0.693 |
+| email triage | 2691 | 0.732 | 0.017 | 0.595 |
+| emotion and tone | 1825 | 0.906 | 0.018 | 0.238 |
+| inference and fact checking | 3022 | 0.883 | 0.054 | 0.340 |
+| instruction-following tasks | 600 | 0.878 | 0.046 | 0.302 |
+| intent and routing | 1475 | 0.991 | 0.009 | 0.181 |
+| moderation and safety | 2708 | 0.967 | 0.061 | 0.153 |
+| reading comprehension | 770 | 0.847 | 0.083 | 0.409 |
+| response quality scoring | 3146 | 0.581 | 0.023 | 1.009 |
+| robustness checks | 744 | 0.851 | 0.108 | 1.058 |
+| search relevance | 733 | 0.628 | 0.066 | 0.728 |
+| sentiment and rating | 961 | 0.442 | 0.438 | 3.545 |
+| topic classification | 749 | 0.939 | 0.029 | 0.196 |
+
+Overall: accuracy 0.753, ECE 0.030, Brier 0.308, accuracy at 50% coverage 0.947
+
+## Zero-shot
+
+| task family | questions | accuracy | ECE | NLL |
+|---|---|---|---|---|
+| emotion and tone | 600 | 0.583 | 0.318 | 1.976 |
+| instruction-following tasks | 600 | 0.863 | 0.045 | 0.319 |
+| moderation and safety | 600 | 0.797 | 0.171 | 1.415 |
+| sentiment and rating | 600 | 0.362 | 0.291 | 1.798 |
+
+Overall: accuracy 0.651, ECE 0.204, Brier 0.532, accuracy at 50% coverage 0.818
+
+## Latency
+
+```
+{
+  "1_questions": {
+    "p50_ms": 38.4,
+    "p95_ms": 42.1
+  },
+  "10_questions": {
+    "p50_ms": 156.0,
+    "p95_ms": 158.4
+  },
+  "50_questions": {
+    "p50_ms": 721.4,
+    "p95_ms": 733.0
+  }
+}
+```

model.safetensors

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

multilingual/encoder/config.json

@@ -0,0 +1,79 @@
+{
+  "architectures": [
+    "ModernBertForMaskedLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 2,
+  "classifier_activation": "gelu",
+  "classifier_bias": false,
+  "classifier_dropout": 0.0,
+  "classifier_pooling": "mean",
+  "cls_token_id": 1,
+  "decoder_bias": true,
+  "deterministic_flash_attn": false,
+  "dtype": "float32",
+  "embedding_dropout": 0.0,
+  "eos_token_id": 1,
+  "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,
+  "mask_token_id": 4,
+  "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": 0,
+  "position_embedding_type": "sans_pos",
+  "repad_logits_with_grad": false,
+  "rope_parameters": {
+    "full_attention": {
+      "rope_theta": 160000,
+      "rope_type": "default"
+    },
+    "sliding_attention": {
+      "rope_theta": 160000,
+      "rope_type": "default"
+    }
+  },
+  "sep_token_id": 1,
+  "sparse_pred_ignore_index": -100,
+  "sparse_prediction": false,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.0.0",
+  "vocab_size": 256000
+}

multilingual/model.safetensors

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

multilingual/rl_agent_config.json

@@ -0,0 +1,26 @@
+{
+  "encoder": "jhu-clsp/mmBERT-base",
+  "head_layers": 2,
+  "max_len": 1024,
+  "head_max_len": 256,
+  "max_prefixes": 6,
+  "act_costs": {
+    "escalate": 0.5
+  },
+  "cost_wrong_act": 3.0,
+  "amp_dtype": "bf16",
+  "model_name": "rl-agent",
+  "temperature": [
+    1.0,
+    1.0,
+    1.0
+  ],
+  "temperature_by_options": {},
+  "training": {
+    "updates": 15987,
+    "epochs_completed": 4,
+    "hours": 4.97,
+    "world_size": 1,
+    "fine_tuned_from_checkpoint": false
+  }
+}

multilingual/tokenizer/tokenizer.json

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

multilingual/tokenizer/tokenizer_config.json

@@ -0,0 +1,22 @@
+{
+  "bos_token": "<bos>",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "<bos>",
+  "eos_token": "<eos>",
+  "extra_special_tokens": {
+    "extra_0": "<start_of_turn>",
+    "extra_1": "<end_of_turn>"
+  },
+  "mask_token": "<mask>",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 8192,
+  "pad_token": "<pad>",
+  "padding_side": "right",
+  "sep_token": "<eos>",
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "<unk>"
+}

rl_agent_api.py

@@ -0,0 +1,78 @@
+"""Jev-compatible inference for a saved RL Agent model: system_one(state, questions) -> typed answers."""
+import json
+import math
+import os
+
+import numpy as np
+import torch
+
+from rl_common import (QTYPES, amp_dtype, build_model, build_sequence, collate_items, confidence_from_probs,
+                       render_options, temp_bucket)
+
+
+class RLAgent:
+    def __init__(self, model_dir, device=None):
+        from safetensors.torch import load_file
+        from transformers import AutoTokenizer
+        with open(os.path.join(model_dir, "rl_agent_config.json")) as f:
+            self.cfg = json.load(f)
+        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        self.tok = AutoTokenizer.from_pretrained(os.path.join(model_dir, "tokenizer"))
+        self.model = build_model(self.cfg, encoder_dir=os.path.join(model_dir, "encoder"))
+        self.model.load_state_dict(load_file(os.path.join(model_dir, "model.safetensors")), strict=True)
+        self.model.to(self.device).eval()
+        self.model.encoder.config.reference_compile = False  # torch.compile is a loss on small batches / few SMs (T4)
+        self.temperature = self.cfg.get("temperature", [1.0, 1.0, 1.0])
+        self.temperature_by_options = self.cfg.get("temperature_by_options", {})
+        self.dtype = amp_dtype(self.cfg.get("amp_dtype", "fp16"))
+        if self.device.type == "cuda" and torch.cuda.get_device_capability(self.device)[0] < 8:
+            self.dtype = torch.float16  # e.g. a bf16-trained model evaluated on a T4
+
+    @staticmethod
+    def _to_internal(qdef):
+        t = qdef["type"]
+        crit = qdef.get("criteria")
+        if t == "choice" and isinstance(crit, list):
+            crit = {c: None for c in crit}
+        return {"t": t, "ins": qdef["instructions"] if isinstance(qdef["instructions"], str) else json.dumps(qdef["instructions"]),
+                "crit": crit}
+
+    @torch.no_grad()
+    def system_one(self, state, questions):
+        """questions: {id: {"type": "choice"|"score"|"noul", "instructions": ..., "criteria": ...}} (Jev request shape)."""
+        ids, items = list(questions.keys()), []
+        for qid in ids:
+            q = self._to_internal(questions[qid])
+            seq, markers = build_sequence(self.tok, state, q, self.cfg["max_len"], self.cfg["head_max_len"])
+            if len(markers) != len(render_options(q)):
+                raise ValueError("question %r: options do not fit in head_max_len=%d tokens" % (qid, self.cfg["head_max_len"]))
+            items.append({"ids": seq, "markers": markers, "qtype": QTYPES[q["t"]], "target": [0.0] * len(markers), "label": -1,
+                          "episode": 0, "ep_step": 0, "ep_len": 1, "src": "api"})
+        b = collate_items([items], self.tok.pad_token_id)
+        use_amp = self.device.type == "cuda"
+        with torch.autocast(device_type=self.device.type, dtype=self.dtype, enabled=use_amp):
+            logits, act = self.model(b["input_ids"].to(self.device), b["attention_mask"].to(self.device),
+                                     b["marker_pos"].to(self.device), b["marker_mask"].to(self.device), b["qtype"].to(self.device))
+        logits, act = logits.float().cpu().numpy(), torch.softmax(act.float(), -1).cpu().numpy()
+        answers, n_tokens = {}, int(b["attention_mask"].sum())
+        for r, qid in enumerate(ids):
+            q = self._to_internal(questions[qid])
+            k = len(items[r]["markers"])
+            qt = QTYPES[q["t"]]
+            z = logits[r, :k] / self.temperature_by_options.get(temp_bucket(qt, k), self.temperature[qt])
+            p = np.exp(z - z.max())
+            p = p / p.sum()
+            ext = {"act_probability": float(act[r, 0])}
+            if q["t"] == "choice":
+                keys = list(q["crit"].keys())
+                answers[qid] = {"type": "choice", "choice": keys[int(p.argmax())],
+                                "probabilities": {kk: round(float(v), 4) for kk, v in zip(keys, p)},
+                                "confidence": round(confidence_from_probs(p, k), 4), "rl_agent": ext}
+            elif q["t"] == "score":
+                answers[qid] = {"type": "score", "score": round(float((np.arange(k) * p).sum()), 4),
+                                "legend": {str(i): c for i, c in enumerate(q["crit"])},
+                                "probabilities": {str(i): round(float(v), 4) for i, v in enumerate(p)},
+                                "confidence": round(confidence_from_probs(p, k), 4), "rl_agent": ext}
+            else:
+                answers[qid] = {"type": "noul", "noul": round(float(p[1]), 4), "rl_agent": ext}
+        return {"model": "rl-agent", "answers": answers, "usage": {"input_tokens": n_tokens, "output_tokens": 0}}

rl_agent_config.json

@@ -0,0 +1,33 @@
+{
+  "encoder": "answerdotai/ModernBERT-large",
+  "head_layers": 2,
+  "max_len": 512,
+  "head_max_len": 192,
+  "max_prefixes": 6,
+  "act_costs": {
+    "escalate": 0.5
+  },
+  "cost_wrong_act": 3.0,
+  "amp_dtype": "bf16",
+  "model_name": "rl-agent",
+  "temperature": [
+    1.6369030475616455,
+    1.2514300346374512,
+    1.983399510383606
+  ],
+  "temperature_by_options": {
+    "choice:3-5": 1.7601518630981445,
+    "choice:6-10": 1.0000158548355103,
+    "score:3-5": 1.2514300346374512,
+    "noul:2": 1.983399510383606,
+    "choice:11+": 0.10058280825614929,
+    "choice:2": 1.9063563346862793
+  },
+  "training": {
+    "updates": 7313,
+    "epochs_completed": 1,
+    "hours": 1.96,
+    "world_size": 1,
+    "fine_tuned_from_checkpoint": true
+  }
+}

rl_common.py

@@ -0,0 +1,408 @@
+"""RL Agent shared code: config, Jev-style question rendering, model, proper-scoring rewards, metrics.
+
+Kept Python 3.9 compatible so the same file runs on Kaggle and on a laptop smoke test.
+"""
+import json
+import math
+import os
+import random
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+QTYPES = {"choice": 0, "score": 1, "noul": 2}
+QTYPE_NAMES = {v: k for k, v in QTYPES.items()}
+
+
+# ----------------------------------------------------------------------------- config
+def load_cfg(path: Optional[str] = None) -> Dict:
+    path = path or os.environ.get("RL_AGENT_CFG", "rl_agent_config.json")
+    with open(path) as f:
+        return json.load(f)
+
+
+# ----------------------------------------------------------------------------- rendering
+def serialize_state(state) -> str:
+    if isinstance(state, str):
+        return state
+    return json.dumps(state, ensure_ascii=False)
+
+
+def render_options(q: Dict) -> List[str]:
+    """Option texts in label-index order. Noul is always [false, true] so p[1] == noul."""
+    t, crit = q["t"], q.get("crit")
+    if t == "choice":
+        return [k if not v else "%s: %s" % (k, v) for k, v in crit.items()]
+    if t == "score":
+        return ["level %d: %s" % (i, c) for i, c in enumerate(crit)]
+    crit = crit or {}
+    return ["false: " + (crit.get("false") or "no, the statement does not hold"),
+            "true: " + (crit.get("true") or "yes, the statement holds")]
+
+
+def build_sequence(tok, state, q: Dict, max_len: int, head_max_len: int,
+                   option_order: Optional[List[int]] = None, truncate_left: bool = False):
+    """[CLS] <type> instructions [SEP] [MASK] opt0 [MASK] opt1 ... [SEP] state [SEP].
+
+    Returns input_ids and the positions of the per-option [MASK] markers (in the given option order).
+    """
+    mask_tok = tok.mask_token
+    opts = render_options(q)
+    order = option_order if option_order is not None else list(range(len(opts)))
+    ins = str(q["ins"]).replace(mask_tok, " ")
+    head_ids = tok("%s question: %s" % (q["t"], ins), add_special_tokens=False)["input_ids"]
+    opt_ids = []
+    for i in order:
+        opt_ids.append([tok.mask_token_id] + tok(" " + opts[i].replace(mask_tok, " "), add_special_tokens=False)["input_ids"][:48])
+    opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+    if opt_budget < 16:  # too many / too long options: shrink every option text evenly
+        per = max(4, (head_max_len - 16) // max(1, len(opt_ids)))
+        opt_ids = [o[:per] for o in opt_ids]
+        opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+    head_ids = head_ids[:max(8, opt_budget)]
+    ids = [tok.cls_token_id] + head_ids + [tok.sep_token_id]
+    markers = []
+    for o in opt_ids:
+        markers.append(len(ids))
+        ids.extend(o)
+    ids.append(tok.sep_token_id)
+    room = max(0, max_len - len(ids) - 1)
+    st = tok(serialize_state(state).replace(mask_tok, " "), add_special_tokens=False)["input_ids"]
+    st = st[-room:] if truncate_left else st[:room]
+    ids = ids + st + [tok.sep_token_id]
+    return ids[:max_len], [m for m in markers if m < max_len]
+
+
+# ----------------------------------------------------------------------------- model
+class DecisionModel(nn.Module):
+    """Pretrained bidirectional encoder (no LLM, no LoRA) + from-scratch decision head.
+
+    Each option gets a [MASK] marker; the head scores markers -> softmax over the question's options.
+    """
+
+    def __init__(self, encoder: nn.Module, head_layers: int = 2, n_act: int = 2, dropout: float = 0.1):
+        super().__init__()
+        self.encoder = encoder
+        d = encoder.config.hidden_size
+        nhead = max(1, d // 64)
+        layer = nn.TransformerEncoderLayer(d, nhead, 4 * d, dropout, batch_first=True, norm_first=True)
+        self.head = nn.TransformerEncoder(layer, head_layers, enable_nested_tensor=False) if head_layers > 0 else None
+        self.type_emb = nn.Embedding(3, d)
+        self.scorer = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1))
+        self.act_head = nn.Sequential(nn.Linear(d + 4, 256), nn.GELU(), nn.Linear(256, n_act))
+        self.register_buffer("temperature", torch.ones(3))  # per qtype, fitted post-hoc in evaluate.py
+        self.head_checkpointing = False
+
+    def forward(self, input_ids, attention_mask, marker_pos, marker_mask, qtype, detach_encoder: bool = False):
+        h = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+        if detach_encoder:
+            h = h.detach()
+        h = h + self.type_emb(qtype)[:, None, :]
+        if self.head is not None:
+            pad = ~attention_mask.bool()
+            for layer in self.head.layers:
+                if self.head_checkpointing and self.training and torch.is_grad_enabled():
+                    h = torch.utils.checkpoint.checkpoint(layer, h, None, pad, use_reentrant=False)
+                else:
+                    h = layer(h, src_key_padding_mask=pad)
+        idx = marker_pos.clamp(min=0)[:, :, None].expand(-1, -1, h.size(-1))
+        m = torch.gather(h, 1, idx)
+        logits = self.scorer(m).squeeze(-1).float()
+        logits = logits.masked_fill(~marker_mask, -1e4)
+        # act head sees the pooled sequence + detached summary of its own answer distribution
+        p = torch.softmax(logits.detach(), -1)
+        k = marker_mask.sum(-1).clamp(min=2).float()
+        ent = -(p * torch.log(p.clamp_min(1e-9))).sum(-1) / torch.log(k)
+        top2 = p.topk(2, -1).values
+        feats = torch.stack([top2[:, 0], top2[:, 0] - top2[:, 1], ent, k / 255.0], -1)
+        pooled = h[:, 0].float()
+        act_logits = self.act_head(torch.cat([pooled, feats], -1))
+        return logits, act_logits
+
+
+def build_model(cfg: Dict, encoder_dir: Optional[str] = None) -> DecisionModel:
+    from transformers import AutoConfig, AutoModel
+    if encoder_dir:  # offline: architecture only, weights come from the saved state dict
+        ecfg = AutoConfig.from_pretrained(encoder_dir)
+        enc = AutoModel.from_config(ecfg, attn_implementation="sdpa")
+    else:
+        enc = AutoModel.from_pretrained(cfg["encoder"], attn_implementation="sdpa")
+    return DecisionModel(enc, cfg["head_layers"], len(cfg["act_costs"]) + 1)
+
+
+# ----------------------------------------------------------------------------- rewards (strictly proper)
+def proper_reward(q: torch.Tensor, target: torch.Tensor, qtype: torch.Tensor, mask: torch.Tensor,
+                  w_sph: float = 0.5, w_rps: float = 1.0, log_floor: float = -9.21) -> torch.Tensor:
+    """q: [..., N, K] reported distributions, target: [N, K] (one-hot or soft) -> reward [..., N].
+
+    log score + spherical score for all types, + ranked probability score for ordinal (score) questions.
+    All three are strictly proper, so the only way to maximize reward is to report honest probabilities.
+    """
+    q = q * mask
+    logq = torch.log(q.clamp_min(1e-12)).clamp_min(log_floor)
+    log_score = (target * logq).sum(-1)
+    sph = (target * q).sum(-1) / q.norm(dim=-1).clamp_min(1e-9)
+    r = log_score + w_sph * sph
+    is_score = (qtype == QTYPES["score"]).float()
+    if is_score.any():
+        k = mask.sum(-1).clamp(min=2).float()
+        cdf_q = torch.cumsum(q, -1)
+        cdf_t = torch.cumsum(target, -1)
+        rps = (((cdf_q - cdf_t) ** 2) * mask).sum(-1) / (k - 1)
+        r = r - w_rps * rps * is_score
+    return r
+
+
+# ----------------------------------------------------------------------------- metrics (numpy, no sklearn)
+def ece_score(conf: np.ndarray, correct: np.ndarray, bins: int = 15) -> float:
+    if len(conf) == 0:
+        return float("nan")
+    edges = np.linspace(0, 1, bins + 1)
+    e = 0.0
+    for lo, hi in zip(edges[:-1], edges[1:]):
+        sel = (conf > lo) & (conf <= hi)
+        if sel.any():
+            e += sel.mean() * abs(conf[sel].mean() - correct[sel].mean())
+    return float(e)
+
+
+def auroc(scores: np.ndarray, labels: np.ndarray) -> float:
+    pos, neg = labels == 1, labels == 0
+    if pos.sum() == 0 or neg.sum() == 0:
+        return float("nan")
+    order = np.argsort(scores)
+    ranks = np.empty(len(scores))
+    ranks[order] = np.arange(1, len(scores) + 1)
+    # average ties
+    s_sorted = scores[order]
+    i = 0
+    while i < len(s_sorted):
+        j = i
+        while j + 1 < len(s_sorted) and s_sorted[j + 1] == s_sorted[i]:
+            j += 1
+        if j > i:
+            ranks[order[i:j + 1]] = (i + j + 2) / 2.0
+        i = j + 1
+    return float((ranks[pos].sum() - pos.sum() * (pos.sum() + 1) / 2) / (pos.sum() * neg.sum()))
+
+
+def spearman(a: np.ndarray, b: np.ndarray) -> float:
+    if len(a) < 3:
+        return float("nan")
+    ra = np.argsort(np.argsort(a)).astype(float)
+    rb = np.argsort(np.argsort(b)).astype(float)
+    if ra.std() == 0 or rb.std() == 0:
+        return float("nan")
+    return float(np.corrcoef(ra, rb)[0, 1])
+
+
+def aurc(conf: np.ndarray, correct: np.ndarray) -> float:
+    """Area under the risk-coverage curve (lower is better)."""
+    if len(conf) == 0:
+        return float("nan")
+    order = np.argsort(-conf)
+    err = 1 - correct[order]
+    return float((np.cumsum(err) / np.arange(1, len(err) + 1)).mean())
+
+
+def confidence_from_probs(p: np.ndarray, k: int) -> float:
+    """Jev-style confidence: 1 - normalized entropy of the answer distribution."""
+    if k < 2:
+        return 1.0
+    p = p[:k]
+    ent = -(p * np.log(np.clip(p, 1e-12, 1))).sum()
+    return float(1 - ent / math.log(k))
+
+
+def seed_all(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+# ----------------------------------------------------------------------------- record -> model inputs
+def episode_prefix_lengths(n_turns: int, max_prefixes: int) -> List[int]:
+    if n_turns <= max_prefixes:
+        return list(range(1, n_turns + 1))
+    return sorted(set(int(round(x)) for x in np.linspace(1, n_turns, max_prefixes)))
+
+
+def encode_record(rec: Dict, tok, cfg: Dict, rng: Optional[random.Random], train: bool) -> List[Dict]:
+    """One stored record -> list of model sequences (one per question, or one per conversation prefix)."""
+    items = []
+    if rec.get("kind") == "episode":
+        ep, q = rec["ep"], rec["qs"][0]
+        lens = episode_prefix_lengths(len(ep["turns"]), cfg["max_prefixes"])
+        for step, t in enumerate(lens):
+            state = dict(ep["ctx"], conversation=ep["turns"][:t])
+            ids, markers = build_sequence(tok, state, q, cfg["max_len"], cfg["head_max_len"], truncate_left=True)
+            if len(markers) != 2:
+                continue
+            items.append({"ids": ids, "markers": markers, "qtype": QTYPES["noul"], "target": [1.0 - ep["y"], float(ep["y"])],
+                          "label": int(ep["y"]), "episode": 1, "ep_step": step, "ep_len": len(lens), "src": rec.get("src", ""),
+                          "prefix_frac": t / float(len(ep["turns"]))})
+        return items
+    for qi, q in enumerate(rec["qs"]):
+        k = len(render_options(q))
+        target = list(q["soft"]) if q.get("soft") else [1.0 if i == q["y"] else 0.0 for i in range(k)]
+        order = list(range(k))
+        if train and rng is not None and q["t"] != "score":
+            rng.shuffle(order)
+        ids, markers = build_sequence(tok, rec["state"], q, cfg["max_len"], cfg["head_max_len"], option_order=order)
+        if len(markers) != k:
+            continue  # options did not fit; skip rather than train on a truncated answer space
+        target = [target[i] for i in order]
+        label = order.index(q["y"]) if q.get("y") is not None else -1
+        items.append({"ids": ids, "markers": markers, "qtype": QTYPES[q["t"]], "target": target, "label": label,
+                      "episode": 0, "ep_step": 0, "ep_len": 1, "src": rec.get("src", ""), "q_index": qi, "order": order})
+    return items
+
+
+def collate_items(batch, pad_id: int):
+    items = [it for group in batch for it in group]
+    if not items:
+        return None
+    n, L = len(items), max(len(it["ids"]) for it in items)
+    kmax = max(len(it["markers"]) for it in items)
+    ids = torch.full((n, L), pad_id, dtype=torch.long)
+    att = torch.zeros((n, L), dtype=torch.long)
+    mpos = torch.zeros((n, kmax), dtype=torch.long)
+    mmask = torch.zeros((n, kmax), dtype=torch.bool)
+    target = torch.zeros((n, kmax), dtype=torch.float32)
+    ep_group = torch.full((n,), -1, dtype=torch.long)
+    group_of = {}
+    for i, it in enumerate(items):
+        ids[i, :len(it["ids"])] = torch.tensor(it["ids"])
+        att[i, :len(it["ids"])] = 1
+        k = len(it["markers"])
+        mpos[i, :k] = torch.tensor(it["markers"])
+        mmask[i, :k] = True
+        target[i, :k] = torch.tensor(it["target"], dtype=torch.float32)
+    # episodes: all prefixes of the same record share a group id (used for TD(lambda) targets)
+    for i, it in enumerate(items):
+        if it["episode"]:
+            ep_group[i] = group_of.setdefault(it.get("rec_uid", -1 - i), len(group_of))
+    return {"input_ids": ids, "attention_mask": att, "marker_pos": mpos, "marker_mask": mmask, "target": target,
+            "qtype": torch.tensor([it["qtype"] for it in items]), "label": torch.tensor([it["label"] for it in items]),
+            "episode": torch.tensor([it["episode"] for it in items], dtype=torch.bool), "ep_group": ep_group,
+            "ep_step": torch.tensor([it["ep_step"] for it in items]), "meta": [{k: it[k] for k in it if k not in ("ids", "markers", "target")} for it in items],
+            "n_tokens": int(att.sum())}
+
+
+def pack_groups(groups: List[List[Dict]], max_tokens: int, max_seqs: int) -> List[List[List[Dict]]]:
+    """Split one sampled batch into sub-batches using the *real* tokenized lengths, so padded tokens never exceed
+    max_tokens (the index only stores estimates). A record's items stay together (TD targets need all prefixes)."""
+    groups = sorted([g for g in groups if g], key=lambda g: max(len(it["ids"]) for it in g))
+    subs, cur, cur_max, cur_n = [], [], 0, 0
+    for g in groups:
+        g_max, g_n = max(len(it["ids"]) for it in g), len(g)
+        if g_max * g_n > max_tokens:  # one record bigger than the budget (only if max_tokens < max_len * n_items)
+            step = max(1, max_tokens // g_max)
+            for s in range(0, g_n, step):
+                subs.append([g[s:s + step]])
+            continue
+        new_max, new_n = max(cur_max, g_max), cur_n + g_n
+        if cur and (new_max * new_n > max_tokens or new_n > max_seqs):
+            subs.append(cur)
+            cur, new_max, new_n = [], g_max, g_n
+        cur.append(g)
+        cur_max, cur_n = new_max, new_n
+    if cur:
+        subs.append(cur)
+    return subs
+
+
+def td_lambda_targets(p_true: torch.Tensor, batch: Dict, lam: float) -> torch.Tensor:
+    """TD(lambda) soft targets for conversation prefixes: G_last = outcome, G_t = (1-lam) V_{t+1} + lam G_{t+1}."""
+    target = batch["target"].clone()
+    groups = batch["ep_group"]
+    for g in torch.unique(groups[groups >= 0]).tolist():
+        idx = (groups == g).nonzero(as_tuple=True)[0]
+        idx = idx[torch.argsort(batch["ep_step"][idx])]
+        y = batch["target"][idx[-1], 1]
+        G = y
+        for j in range(len(idx) - 1, -1, -1):
+            if j < len(idx) - 1:
+                G = (1 - lam) * p_true[idx[j + 1]] + lam * G
+            target[idx[j], 0], target[idx[j], 1] = 1 - G, G
+    return target
+
+
+def make_token_batches(lengths: np.ndarray, nseq: np.ndarray, max_tokens: int, max_seqs: int, rng: np.random.RandomState,
+                       chunk: int = 4096) -> List[List[int]]:
+    """Length-bucketed batches of record indices under a padded-token budget."""
+    order = rng.permutation(len(lengths))
+    batches = []
+    for s in range(0, len(order), chunk):
+        part = order[s:s + chunk]
+        part = part[np.argsort(lengths[part])]
+        cur, cur_max, cur_n = [], 0, 0
+        for i in part:
+            ln, ns = int(lengths[i]), int(nseq[i])
+            new_max, new_n = max(cur_max, ln), cur_n + ns
+            if cur and (new_max * new_n > max_tokens or new_n > max_seqs):
+                batches.append(cur)
+                cur, new_max, new_n = [], ln, ns
+            cur.append(int(i))
+            cur_max, cur_n = new_max, new_n
+        if cur:
+            batches.append(cur)
+    rng.shuffle(batches)
+    return batches
+
+
+def temp_bucket(qtype: int, k: int) -> str:
+    """Key for per-cardinality temperature fitting: a 2-option noul and a 20-option choice need different scaling."""
+    size = "2" if k <= 2 else "3-5" if k <= 5 else "6-10" if k <= 10 else "11+"
+    return "%s:%s" % (QTYPE_NAMES[int(qtype)], size)
+
+
+def amp_dtype(name: Optional[str]) -> torch.dtype:
+    """'bf16' on GPUs that support it (Ampere+, e.g. RTX 6000 Pro); 'fp16' on T4."""
+    return torch.bfloat16 if name == "bf16" else torch.float16
+
+
+@torch.no_grad()
+def predict_items(model, items: List[Dict], pad_id: int = 0, device=None, max_tokens: int = 16384, use_amp: bool = True,
+                  dtype: torch.dtype = torch.float16, max_seqs: int = 256, progress: str = ""):
+    """Run the model over pre-encoded items; returns list of dicts with probs/logits (uncalibrated) and act probs."""
+    import sys
+    import time as _time
+    model.eval()
+    out = []
+    t0, done_tok = _time.time(), 0
+    order = sorted(range(len(items)), key=lambda i: len(items[i]["ids"]))
+    i = 0
+    while i < len(order):
+        j, L = i, 0
+        while j < len(order) and j - i < max_seqs and max(L, len(items[order[j]]["ids"])) * (j - i + 1) <= max_tokens:
+            L = max(L, len(items[order[j]]["ids"]))
+            j += 1
+        j = max(j, i + 1)
+        sel = [items[order[t]] for t in range(i, j)]
+        b = collate_items([sel], pad_id)
+        with torch.autocast(device_type=device.type, dtype=dtype, enabled=use_amp and device.type == "cuda"):
+            logits, act = model(b["input_ids"].to(device), b["attention_mask"].to(device), b["marker_pos"].to(device),
+                                b["marker_mask"].to(device), b["qtype"].to(device))
+        logits, act = logits.float().cpu(), torch.softmax(act.float(), -1).cpu()
+        done_tok += int(b["attention_mask"].sum())
+        if progress and (j % max(1, len(order) // 2000) == 0 or j >= len(order)):
+            el = _time.time() - t0
+            eta = el * (len(order) - j) / max(1, j)
+            sys.stdout.write("\r  [%s] %d/%d sequences | %.1fk tok/s | ETA %dm%02ds   " %
+                             (progress, j, len(order), done_tok / max(el, 1e-9) / 1000, int(eta // 60), int(eta % 60)))
+            sys.stdout.flush()
+        for r, it in enumerate(sel):
+            k = len(it["markers"])
+            out.append((order[i + r], {"logits": logits[r, :k].detach().numpy(), "act": act[r].detach().numpy()}))
+        i = j
+    if progress:
+        print("\r  [%s] %d sequences in %.0fs (%.1fk tok/s)%s" % (progress, len(order), _time.time() - t0,
+                                                                  done_tok / max(_time.time() - t0, 1e-9) / 1000, " " * 20))
+    out.sort(key=lambda x: x[0])
+    model.train()
+    return [o for _, o in out]

tokenizer/tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 8192,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}

typed-decisions/encoder/config.json

@@ -0,0 +1,84 @@
+{
+  "architectures": [
+    "ModernBertForMaskedLM"
+  ],
+  "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": 1024,
+  "initializer_cutoff_factor": 2.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 2624,
+  "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",
+    "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": 16,
+  "num_hidden_layers": 28,
+  "pad_token_id": 50283,
+  "position_embedding_type": "absolute",
+  "repad_logits_with_grad": false,
+  "rope_parameters": {
+    "full_attention": {
+      "rope_theta": 160000.0,
+      "rope_type": "default"
+    },
+    "sliding_attention": {
+      "rope_theta": 10000.0,
+      "rope_type": "default"
+    }
+  },
+  "sep_token_id": 50282,
+  "sparse_pred_ignore_index": -100,
+  "sparse_prediction": false,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "vocab_size": 50368
+}

typed-decisions/model.safetensors

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

typed-decisions/rl_agent_config.json

@@ -0,0 +1,36 @@
+{
+  "encoder": "answerdotai/ModernBERT-large",
+  "head_layers": 2,
+  "max_len": 1024,
+  "head_max_len": 256,
+  "max_prefixes": 6,
+  "act_costs": {
+    "escalate": 0.5
+  },
+  "cost_wrong_act": 3.0,
+  "amp_dtype": "bf16",
+  "model_name": "laya-typed-decisions",
+  "temperature": [
+    1.0148024559020996,
+    1.0374259948730469,
+    1.0575125217437744
+  ],
+  "temperature_by_options": {
+    "choice:3-5": 1.7601518630981445,
+    "choice:6-10": 1.0000158548355103,
+    "score:3-5": 1.2514300346374512,
+    "noul:2": 1.983399510383606,
+    "choice:11+": 0.10058280825614929,
+    "choice:2": 1.9063563346862793
+  },
+  "training": {
+    "updates": 7313,
+    "epochs_completed": 1,
+    "hours": 1.96,
+    "world_size": 1,
+    "fine_tuned_from_checkpoint": true
+  },
+  "gradient_checkpointing": true,
+  "max_tokens_per_batch": 4096,
+  "fine_tuned": true
+}

typed-decisions/tokenizer/tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 8192,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}