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,266 @@
+---
+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) |
+
+---
+
+## 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.
+
+### Production Preload & Memory
+
+A cold checkpoint build costs seconds; language detection costs microseconds. At the default `max_loaded=1`, traffic that alternates languages rebuilds a model on *every* request (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 1 hot, LRU eviction)
+router = Router(max_loaded=2)       # keep two hot
+router.unload()                     # free memory
+```
+
+| Deployment Mode | Per-Request Latency | Model Reloads |
+|---|---|---|
+| `Router()` (lazy, `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`.
+
+---
+
+## 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).
+- **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]"
+}