Add training kit: train your own System-1 model

Label (two independent LLM annotators), train, evaluate and export (ONNX fp32 + block-wise int8) scripts generalised from Raya's training. Model files unchanged.

README.md

@@ -52,6 +52,22 @@ Raya was trained on three routing questions: the minimal choice above, a detaile
 3-level difficulty score. Use one of those. Option order does not matter because options were shuffled in
 training. Raya serves through Laya's Jev-compatible HTTP server (`POST /v1/systemone`).
 
+## Train your own
+
+The code that trained Raya is in [`training/`](training/README.md), generalised so you can train a fast
+decision model on your own data:
+
+- **Any choice or score question:** LLM routing like Raya, ticket triage, intent detection, escalation.
+- **Labelling included:** labels come from two independent LLM annotators, which is how Raya's labels
+  were made.
+- **Adapt Raya:** start from Raya to fit it to your own traffic (`--base TextCortex/raya`).
+- **Serving:** exports to ONNX for CPU serving, checked against PyTorch.
+
+```bash
+pip install -r training/requirements.txt
+python training/train.py --task training/task.example.json --data my_labelled.jsonl --out my-router
+```
+
 ## CPU inference (ONNX)
 
 No GPU? `onnx/` has ONNX Runtime builds of Raya that run through Laya's own `ONNXAgent`, which gives the same answer format as `laya.Agent`. Use a **512-token** input budget, which is what Raya was trained on.
@@ -124,6 +140,8 @@ annotators, not from measured downstream answer quality.
   validation prompts, minimal choice), per-question temperature fitted on validation. The seed was also
   chosen on validation only.
 - **Compute:** one NVIDIA RTX A6000, ~6 minutes.
+- **Code:** [`training/`](training/README.md) reproduces this procedure on your own data. The training
+  data itself is not published.
 
 ## Intended use and limitations
 

training/README.md

@@ -0,0 +1,215 @@
+# Train your own System-1 model
+
+This folder is the recipe we used to train [Raya](../README.md), packaged so you can train your own
+fast decision model on your own data in an afternoon.
+
+A **System-1 model** answers one well-defined question about an input, instantly and with calibrated
+probabilities: *which model should answer this prompt?*, *does this ticket need a human?*, *which team
+owns this request?* It is a small encoder (a [Laya](https://huggingface.co/convaiinnovations/laya)
+decision model), not an LLM. It runs in tens of milliseconds on a CPU, costs nothing per call, and you
+can host it anywhere.
+
+## The whole pipeline
+
+```bash
+pip install -r requirements.txt
+
+# 1. Describe your task: the labels and a few ways of asking the question
+cp task.example.json my_task.json
+
+# 2. Label your inputs with two independent LLMs (skip if you already have labels)
+python label.py --task my_task.json --data prompts.jsonl --out labelled.jsonl \
+    --annotator <model-a> --annotator <model-b>@https://api.anthropic.com/v1/#ANTHROPIC_API_KEY
+
+# 3. Train
+python train.py --task my_task.json --data labelled.jsonl --out my-model
+
+# 4. Check it on data it has never seen
+python evaluate.py --model my-model --task my_task.json --data test.jsonl
+
+# 5. (optional) Export for fast CPU serving, then publish
+pip install -r requirements-onnx.txt
+python export_onnx.py --model my-model --task my_task.json --data test.jsonl
+python train.py ... --push-to-hub your-name/my-model      # or upload my-model/ yourself
+```
+
+**Try it in five minutes** with the bundled toy data (44 hand-written routing prompts, just enough to
+see every step run; it is far too small to train a useful model):
+
+```bash
+python train.py --task task.example.json --data data/example.jsonl --out my-router --epochs 2
+python evaluate.py --model my-router --task task.example.json --data data/example_test.jsonl
+```
+
+Then use it like any Laya model:
+
+```python
+import json, laya
+
+model = laya.Agent("my-router")                       # a local dir or a Hub repo id
+question = json.load(open("my-router/task.json"))["questions"][0]
+print(model.system_one({"prompt": "Prove that √2 is irrational."}, {"route": question})["answers"]["route"])
+# {'choice': 'frontier_model', 'probabilities': {...}, ...}
+```
+
+## 1. Describe the task (`task.json`)
+
+```json
+{
+  "labels": ["small_model", "medium_model", "frontier_model"],
+  "questions": [
+    {"type": "choice", "instructions": "Route this prompt to a model.",
+     "criteria": {"small_model": "simple requests", "medium_model": "moderately complex requests",
+                  "frontier_model": "very hard requests"}},
+    {"type": "score", "instructions": "How difficult is this prompt for an AI assistant to answer well?",
+     "criteria": ["simple: a small model answers it perfectly", "moderate: needs a capable general model",
+                  "hard: needs the strongest frontier model"]}
+  ],
+  "rubric": "Detailed labelling instructions for label.py (optional)."
+}
+```
+
+- **`labels`**: the possible answers, 2 or more.
+- **`questions`**: one or more *phrasings* of the same decision, written as normal Laya questions. The
+  model trains on all of them, so it learns the decision rather than one exact wording. Raya was
+  trained on three phrasings and scores 80–81% on each.
+  - A `choice` question's `criteria` keys must be exactly your labels. Their order doesn't matter,
+    because options are shuffled during training.
+  - A `score` question is ordinal: one criterion per label, **in label order** (lowest first).
+- **`rubric`**: what `label.py` shows the annotators. Be specific, with examples per label: the model
+  can only be as consistent as its labels.
+
+`task.example.json` is Raya's exact routing task.
+
+## 2. Get labels (`label.py`)
+
+Your data is JSON Lines, one input per line:
+
+```json
+{"prompt": "hi there!", "label": "small_model"}
+{"prompt": "Review this contract …", "labels": ["medium_model", "frontier_model"]}
+{"state": {"ticket": "You charged me twice!", "plan": "enterprise"}, "label": "human_agent"}
+{"prompt": "…", "label": "…", "split": "val"}
+```
+
+- `label` is one gold answer. `labels` holds several annotators' votes: disagreements become soft
+  targets (50/50 above), which is better than forcing a hard label on a genuinely ambiguous input.
+- Use `state` instead of `prompt` for structured inputs (any JSON object).
+- Mark rows `"split": "val"` to fix your validation set; otherwise 10% is held out at random.
+- Rows labelled `"exclude"` are skipped.
+
+**No labels yet?** `label.py` has two (or more) different LLMs label every input independently from
+your rubric, which is how Raya's labels were made (Claude Opus and Claude Sonnet, blind to each other).
+It speaks the OpenAI chat-completions API, so it works with OpenAI, Anthropic's OpenAI-compatible
+endpoint, OpenRouter, vLLM, Ollama and others. An annotator is `MODEL[@BASE_URL][#API_KEY_ENV]`.
+It resumes where it stopped, and prints how often the annotators agree: that agreement rate is
+roughly the ceiling your model can reach against these labels (78% for Raya's test set).
+
+**How much data?**
+
+- 1,000–5,000 is a good target for a first model.
+- Raya used about 10,000.
+- Use **real inputs** from your product where you can, in the languages you serve.
+- Don't balance the classes artificially: `train.py` already up-weights rare labels.
+- Keep a separate test set you never train or validate on.
+
+## 3. Train (`train.py`)
+
+```bash
+python train.py --task my_task.json --data labelled.jsonl --out my-model
+```
+
+What it does, in the same way as Raya's training:
+
+- **Soft targets:** each label's share of the annotator votes, learned with cross-entropy.
+- **Every question phrasing:** trained on each one, with choice options shuffled every epoch.
+- **Class weights:** about 1/√(label frequency), so rare labels still count.
+- **Memory:** the token-embedding table stays frozen.
+- **Best epoch:** picked on validation accuracy (or `--select nll`).
+- **Calibration:** a temperature is then fitted per question on validation, so that 0.9 means about 90%.
+
+The output folder is a normal Laya checkpoint, plus `task.json` and `training_log.json`.
+
+**Pick a starting point:**
+
+| Flag | Starts from | When |
+|---|---|---|
+| *(default)* | Laya multilingual (mmBERT-base, 300M) | Most tasks, multilingual input (Raya was tested on 14 languages) |
+| `--base TextCortex/raya` | Raya | LLM routing on your own traffic: adapt Raya instead of starting over |
+| `--subfolder .` | Laya English (ModernBERT-large) | English-only input; a larger encoder, so slower |
+| `--encoder <hf-id>` | Any Hugging Face encoder, fresh decision head | e.g. a larger multilingual encoder |
+| `--base <dir or repo>` | Any Laya checkpoint | Continue from a model you trained before |
+
+A bigger encoder was the biggest single lever in our experiments: with the same data, a large encoder
+reached about 84% on Raya's benchmark where mmBERT-base topped out around 81–82%. Adding more data
+barely moved the smaller model.
+
+**Hardware:**
+
+- **GPU:** Raya trained in about 6 minutes on one 48 GB RTX A6000 (batch 32, bf16). With the default
+  batch size of 16 it should fit on a 24 GB GPU.
+- **Apple Silicon or CPU:** fine for a few thousand examples. `train.py` picks CUDA, then MPS, then CPU
+  automatically, and turns on gradient checkpointing off-GPU to save memory.
+
+Useful flags: `--epochs` (3), `--batch-size` (16), `--lr-encoder` (2e-5), `--lr-head` (1e-4, or 3e-4
+for a fresh head), `--max-tokens` (512), `--val-data`, `--seed`. Raya used these learning rates and
+token budget with `--batch-size 32`, and its seed was chosen on validation only.
+
+## 4. Evaluate (`evaluate.py`)
+
+```bash
+python evaluate.py --model my-model --task my_task.json --data test.jsonl [--out predictions.jsonl]
+```
+
+For each question phrasing this prints accuracy, macro-F1 and a confusion matrix on rows with a single
+gold label. It also prints the "always answer the most common label" baseline, which is the number to
+beat, and per-decision latency. Pass `--onnx <file>` to score an exported model.
+
+## 5. Export and serve (`export_onnx.py`)
+
+```bash
+python export_onnx.py --model my-model --task my_task.json --data test.jsonl
+```
+
+This writes `my-model/onnx/model.onnx` (fp32) and `model-int8-blockwise.onnx`, then checks both against
+PyTorch on your data. The export fails if any choice changes or probabilities drift by more than 0.001
+(fp32) or 0.05 (int8). Serve either with Laya's `ONNXAgent`, which uses the same call and answer format:
+
+```python
+from laya.onnx_agent import ONNXAgent
+
+model = ONNXAgent("my-model", onnx_path="my-model/onnx/model.onnx")
+model.cfg["max_len"] = 512   # match --max-tokens
+```
+
+**Which ONNX file?** `model.onnx` matches PyTorch on any CPU. The block-wise int8 file kept Raya's accuracy
+and was about 10–15% faster on x86 CPUs with VNNI instructions (Intel Cascade Lake or Alder Lake and
+newer, AMD Zen 4 and newer), but slower on CPUs without VNNI and on ARM. Measure on your own hardware
+before choosing it. In our tests 8 CPU threads were faster than 16.
+
+## Tips
+
+- **Evaluate on data you never trained on.** Always, and keep the split fixed when comparing runs.
+- **Don't tune on the test set.** Choose epochs, seeds and ensembles on validation only.
+- **Match the serving token budget to training** (`--max-tokens`, default 512).
+- **Your labels are the ceiling.** If two good annotators agree only 75% of the time, a 90% score
+  means the model learned your annotator's quirks. Tighten the rubric first.
+- **Calibrate your action threshold on real traffic.** Before acting automatically on a prediction
+  (for example, only escalate when p(frontier) > 0.6), set the threshold from what your real traffic
+  looks like.
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `task.example.json` | Raya's routing task, a template for yours |
+| `common.py` | Task and data format, validation (read its docstring for the full format) |
+| `label.py` | Label inputs with independent LLM annotators |
+| `train.py` | Fine-tune and calibrate |
+| `evaluate.py` | Accuracy, macro-F1, confusion, latency |
+| `export_onnx.py` | ONNX export (fp32 + block-wise int8) with equivalence checks |
+| `data/example.jsonl`, `data/example_test.jsonl` | Tiny hand-written demo data for training and evaluation (not Raya's training data) |
+| `data/unlabelled.jsonl` | A few unlabelled prompts to try `label.py` on |
+
+Tested with laya 0.3.20, torch 2.14, transformers 5.17 and onnxruntime 1.30. The training data behind
+Raya is not published. This code is Apache-2.0, like Raya and Laya.

training/common.py

@@ -0,0 +1,125 @@
+"""Task and data loading shared by label.py, train.py, evaluate.py and export_onnx.py.
+
+A *task* (``task.json``) fixes the label set and the questions the model learns to answer:
+
+    {
+      "labels": ["small_model", "medium_model", "frontier_model"],
+      "questions": [ <Laya question>, ... ],   # 1+ phrasings of the same decision
+      "rubric": "..."                          # optional, used by label.py
+    }
+
+Each question is a regular Laya question. A ``choice`` question's ``criteria`` keys must be
+exactly the labels (any order: options are shuffled in training). A ``score`` question's
+``criteria`` list is ordinal and must have one entry per label, in label order.
+
+*Data* is JSON Lines, one example per line:
+
+    {"prompt": "hi!", "label": "small_model"}                       # one gold label
+    {"prompt": "...", "labels": ["medium_model", "frontier_model"]}  # several annotators -> soft target
+    {"state": {"ticket": "...", "plan": "pro"}, "label": "..."}      # any Laya state instead of a prompt
+    {"prompt": "...", "label": "...", "split": "val"}                # optional fixed split
+
+Rows labelled ``"exclude"`` (by any annotator) are skipped.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+EXCLUDE = "exclude"
+
+
+class TaskError(ValueError):
+    """Raised when task.json or a data file is malformed."""
+
+
+def load_task(path: str | Path) -> dict[str, Any]:
+    """Read and validate a task file; returns it with ``questions`` checked against ``labels``."""
+    task = json.loads(Path(path).read_text(encoding="utf-8"))
+    labels = task.get("labels")
+    if not isinstance(labels, list) or len(labels) < 2 or len(set(labels)) != len(labels):
+        raise TaskError("'labels' must be a list of 2+ distinct strings")
+    if EXCLUDE in labels:
+        raise TaskError(f"'{EXCLUDE}' is reserved for skipping rows; rename that label")
+    questions = task.get("questions")
+    if not isinstance(questions, list) or not questions:
+        raise TaskError("'questions' must be a non-empty list of Laya questions")
+    for i, q in enumerate(questions):
+        where = f"questions[{i}]"
+        if q.get("type") == "choice":
+            crit = q.get("criteria")
+            keys = list(crit) if isinstance(crit, (dict, list)) else None
+            if keys is None or set(keys) != set(labels) or len(keys) != len(labels):
+                raise TaskError(f"{where}: a choice question's criteria keys must be exactly the labels")
+        elif q.get("type") == "score":
+            crit = q.get("criteria")
+            if not isinstance(crit, list) or len(crit) != len(labels):
+                raise TaskError(f"{where}: a score question needs one criterion per label, in label order")
+        else:
+            raise TaskError(f"{where}: type must be 'choice' or 'score'")
+        if not isinstance(q.get("instructions"), str) or not q["instructions"].strip():
+            raise TaskError(f"{where}: 'instructions' must be a non-empty string")
+    return task
+
+
+def row_state(row: dict[str, Any]) -> Any:
+    """The Laya state for a data row: its ``state`` if given, else ``{"prompt": prompt}``."""
+    if "state" in row:
+        return row["state"]
+    if "prompt" in row:
+        return {"prompt": row["prompt"]}
+    raise TaskError("each row needs a 'prompt' or a 'state'")
+
+
+def row_votes(row: dict[str, Any]) -> list[str]:
+    """All label votes on a row (``label`` and/or ``labels``)."""
+    votes = list(row.get("labels") or [])
+    if row.get("label") is not None:
+        votes.append(row["label"])
+    return votes
+
+
+def load_rows(path: str | Path, labels: list[str], *, require_labels: bool = True) -> list[dict[str, Any]]:
+    """Read a JSONL data file into rows with a soft ``target`` over ``labels``.
+
+    The target is each label's share of the votes, so two annotators who disagree give a 50/50
+    target. ``gold`` is the label when every vote agrees, else ``None`` (such rows still train,
+    but accuracy is only reported on rows with a gold label).
+    """
+    rows = []
+    for n, line in enumerate(Path(path).read_text(encoding="utf-8").splitlines(), 1):
+        if not line.strip():
+            continue
+        try:
+            row = json.loads(line)
+        except json.JSONDecodeError as exc:
+            raise TaskError(f"{path}:{n}: not valid JSON ({exc.msg})") from exc
+        row_state(row)  # validates prompt/state presence
+        row.setdefault("id", n)
+        votes = row_votes(row)
+        if EXCLUDE in votes:
+            continue
+        if not votes:
+            if require_labels:
+                raise TaskError(f"{path}:{n}: no 'label' or 'labels'")
+            rows.append(row)
+            continue
+        unknown = sorted(set(votes) - set(labels))
+        if unknown:
+            raise TaskError(f"{path}:{n}: unknown label(s) {unknown}; expected one of {labels}")
+        row["target"] = [votes.count(label) / len(votes) for label in labels]
+        row["gold"] = votes[0] if len(set(votes)) == 1 else None
+        rows.append(row)
+    if not rows:
+        raise TaskError(f"{path}: no usable rows")
+    return rows
+
+
+def answer_probs(answer: dict[str, Any], question: dict[str, Any], labels: list[str]) -> list[float]:
+    """Map a Laya answer's probabilities back to label order (score options are ordinal)."""
+    probs = answer.get("probabilities") or {}
+    if question["type"] == "choice":
+        return [float(probs.get(label, 0.0)) for label in labels]
+    return [float(probs.get(str(i), probs.get(i, 0.0))) for i in range(len(labels))]

training/data/example.jsonl

@@ -0,0 +1,44 @@
+{"prompt": "Derive the Black-Scholes PDE from first principles using Itô's lemma and a delta-hedged portfolio.", "label": "frontier_model"}
+{"prompt": "Write a two-line birthday wish for my sister.", "label": "small_model"}
+{"prompt": "Write a cover letter for a senior data scientist role at a bank.", "labels": ["medium_model", "frontier_model"]}
+{"prompt": "hi there!", "label": "small_model"}
+{"prompt": "What are the pros and cons of remote work for a 50-person agency?", "label": "medium_model"}
+{"prompt": "Rédige une lettre de motivation pour un poste de chef de projet marketing.", "label": "medium_model"}
+{"prompt": "Ciao! Come stai?", "label": "small_model"}
+{"prompt": "Give me a synonym for 'happy'.", "label": "small_model"}
+{"prompt": "Write a Python function that parses a CSV of orders and returns total revenue per customer.", "label": "medium_model"}
+{"prompt": "Create a one-week meal plan for a vegetarian who trains for a marathon.", "label": "medium_model"}
+{"prompt": "Convert 5 miles to kilometers.", "label": "small_model"}
+{"prompt": "How do I reverse a list in Python?", "label": "small_model"}
+{"prompt": "thanks, that's all", "label": "small_model"}
+{"prompt": "What's the capital of Canada?", "label": "small_model"}
+{"prompt": "Draft a LinkedIn post announcing our new office in Munich.", "label": "medium_model"}
+{"prompt": "Write a polite email to a client explaining that their order will arrive two weeks late.", "label": "medium_model"}
+{"prompt": "¿Cuántos días tiene febrero en un año bisiesto?", "label": "small_model"}
+{"prompt": "Explain quantum entanglement.", "labels": ["small_model", "medium_model"]}
+{"prompt": "Refactor this React component to use hooks instead of a class: [80-line class component with state and lifecycle methods]", "label": "medium_model"}
+{"prompt": "Build a discounted cash flow model for a SaaS company with 40% growth, negative free cash flow and a 2-year path to breakeven; justify every assumption.", "label": "frontier_model"}
+{"prompt": "Scrivi una recensione di un ristorante giapponese immaginario, circa 300 parole.", "label": "medium_model"}
+{"prompt": "Optimize this SQL query that takes 40 seconds on a 200M-row table: [query with three joins and a window function]", "labels": ["medium_model", "frontier_model"]}
+{"prompt": "Find the race condition in this lock-free queue implementation and prove your fix is linearizable: [120 lines of Rust using atomics]", "label": "frontier_model"}
+{"prompt": "Explain how HTTPS works to a non-technical colleague.", "label": "medium_model"}
+{"prompt": "Écris un compilateur minimal pour un langage à la Lisp en Rust, avec analyse lexicale, analyse syntaxique et évaluation.", "label": "frontier_model"}
+{"prompt": "Compare three treatment strategies for a patient with resistant hypertension and chronic kidney disease, citing the relevant trial evidence.", "label": "frontier_model"}
+{"prompt": "Analysiere die kartellrechtlichen Risiken einer Fusion zweier Marktführer im deutschen Zementmarkt und schlage Abhilfemaßnahmen vor.", "label": "frontier_model"}
+{"prompt": "Fix the typo: 'I recieved your mesage'", "label": "small_model"}
+{"prompt": "Translate this paragraph into French: [one short paragraph]", "labels": ["small_model", "medium_model"]}
+{"prompt": "What is 17 * 23?", "label": "small_model"}
+{"prompt": "Bonjour, comment ça va ?", "label": "small_model"}
+{"prompt": "Wie spät ist es in Tokio, wenn es in Berlin 9 Uhr ist?", "label": "small_model"}
+{"prompt": "Develop a go-to-market strategy for entering the Japanese B2B software market, including channel partners, pricing and a 3-year plan.", "label": "frontier_model"}
+{"prompt": "What does HTML stand for?", "label": "small_model"}
+{"prompt": "Translate this product description into German and keep the marketing tone: [four-paragraph description of a smart thermostat]", "label": "medium_model"}
+{"prompt": "Prove that there are infinitely many primes of the form 4k+3.", "label": "frontier_model"}
+{"prompt": "Schreibe eine professionelle E-Mail an einen Kunden über eine Lieferverzögerung von zwei Wochen.", "label": "medium_model"}
+{"prompt": "Escribe un artículo de blog de 600 palabras sobre consejos para ahorrar energía en casa.", "label": "medium_model"}
+{"prompt": "Design a multi-region, active-active payment system that guarantees exactly-once settlement, and explain how you handle network partitions.", "label": "frontier_model"}
+{"prompt": "Review this 40-page share purchase agreement and list every clause that shifts liability to the buyer, with the risk each one creates.", "label": "frontier_model"}
+{"prompt": "Summarize the key points of this meeting transcript in bullet points: [transcript of a 30-minute product sync about launch timing, pricing and support staffing]", "label": "medium_model"}
+{"prompt": "Translate 'good morning' into Spanish.", "label": "small_model"}
+{"prompt": "Explain the difference between a stock and a bond with examples.", "label": "medium_model"}
+{"prompt": "Write a SQL query that finds the top 5 products by revenue for each month of 2025.", "label": "medium_model"}

training/data/example_test.jsonl

@@ -0,0 +1,13 @@
+{"prompt": "hello!", "label": "small_model"}
+{"prompt": "What is the boiling point of water in Fahrenheit?", "label": "small_model"}
+{"prompt": "Translate 'thank you' into Japanese.", "label": "small_model"}
+{"prompt": "Was ist 12 plus 30?", "label": "small_model"}
+{"prompt": "Rewrite this sentence more formally: 'gonna be late, sorry'", "label": "small_model"}
+{"prompt": "Write a short product update email for our customers about a new export feature.", "label": "medium_model"}
+{"prompt": "Explain what a REST API is with a simple example.", "label": "medium_model"}
+{"prompt": "Summarize the plot of Hamlet in one paragraph.", "label": "medium_model"}
+{"prompt": "Écris un e-mail pour reporter une réunion à la semaine prochaine.", "label": "medium_model"}
+{"prompt": "Write a JavaScript function that debounces another function.", "label": "medium_model"}
+{"prompt": "Prove that the square root of 2 is irrational and generalize the argument to any non-square integer.", "label": "frontier_model"}
+{"prompt": "Design the data model and consistency strategy for a global inventory system with offline-capable warehouses.", "label": "frontier_model"}
+{"prompt": "Evaluate the tax implications of relocating a German GmbH's headquarters to the Netherlands, including exit taxation.", "label": "frontier_model"}

training/data/unlabelled.jsonl

@@ -0,0 +1,4 @@
+{"prompt": "hello!"}
+{"prompt": "What is the boiling point of water in Fahrenheit?"}
+{"prompt": "Translate 'thank you' into Japanese."}
+{"prompt": "Was ist 12 plus 30?"}

training/evaluate.py

@@ -0,0 +1,97 @@
+"""Score a trained model on a labelled JSONL file: accuracy, macro-F1, confusion and latency.
+
+    python evaluate.py --model my-router --task task.example.json --data data/example.jsonl
+    python evaluate.py --model my-router --onnx my-router/onnx/model-int8-blockwise.onnx ...
+
+Accuracy counts rows whose annotators all agreed (a single gold label). Always test on data the
+model never trained or validated on. The majority-label baseline is printed for comparison.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+
+os.environ.setdefault("USE_TF", "0")
+
+import numpy as np
+
+from common import answer_probs, load_rows, load_task, row_state
+
+
+def load_model(path: str, onnx_path: str | None, device: str | None, max_tokens: int):
+    if onnx_path:
+        from laya.onnx_agent import ONNXAgent
+
+        model = ONNXAgent(path, onnx_path=onnx_path)
+    else:
+        import laya
+
+        model = laya.Agent(path, device=device)
+    model.cfg["max_len"] = max_tokens
+    return model
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--model", required=True, help="checkpoint dir or Hub repo id")
+    ap.add_argument("--task", required=True)
+    ap.add_argument("--data", required=True, help="labelled JSONL the model has not seen")
+    ap.add_argument("--onnx", help="score this ONNX file (with the checkpoint's config/tokenizer) instead")
+    ap.add_argument("--device", default=None)
+    ap.add_argument("--max-tokens", type=int, default=512)
+    ap.add_argument("--out", help="write per-row predictions to this JSONL")
+    args = ap.parse_args()
+
+    task = load_task(args.task)
+    labels, questions = task["labels"], task["questions"]
+    rows = [r for r in load_rows(args.data, labels) if r["gold"] is not None]
+    model = load_model(args.model, args.onnx, args.device, args.max_tokens)
+
+    preds = {qi: [] for qi in range(len(questions))}
+    latencies = []
+    out = open(args.out, "w", encoding="utf-8") if args.out else None
+    for row in rows:
+        record = {"id": row["id"], "gold": row["gold"], "predictions": {}}
+        for qi, q in enumerate(questions):
+            t = time.perf_counter()
+            answer = model.system_one(row_state(row), {"q": q})["answers"]["q"]
+            latencies.append((time.perf_counter() - t) * 1000)
+            probs = answer_probs(answer, q, labels)
+            preds[qi].append(int(np.argmax(probs)))
+            record["predictions"][qi] = dict(zip(labels, [round(p, 4) for p in probs]))
+        if out:
+            out.write(json.dumps(record, ensure_ascii=False) + "\n")
+    if out:
+        out.close()
+
+    golds = [labels.index(r["gold"]) for r in rows]
+    majority = max(labels, key=lambda l: sum(r["gold"] == l for r in rows))
+    print(f"{len(rows)} rows with a gold label; always answering {majority!r} scores "
+          f"{sum(r['gold'] == majority for r in rows) / len(rows):.1%}")
+    for qi, q in enumerate(questions):
+        p = preds[qi]
+        acc = np.mean([a == b for a, b in zip(p, golds)])
+        f1s = []
+        for k in range(len(labels)):
+            tp = sum(a == k == b for a, b in zip(p, golds))
+            fp = sum(a == k != b for a, b in zip(p, golds))
+            fn = sum(b == k != a for a, b in zip(p, golds))
+            f1s.append(2 * tp / (2 * tp + fp + fn) if tp else 0.0)
+        confusion = [[sum(g == i and a == j for a, g in zip(p, golds)) for j in range(len(labels))]
+                     for i in range(len(labels))]
+        print(f"\nquestion {qi} ({q['type']}): {q['instructions'][:70]!r}")
+        print(f"  accuracy {acc:.1%}   macro-F1 {np.mean(f1s):.3f}")
+        print("  confusion (rows = gold, columns = predicted):")
+        width = max(len(l) for l in labels)
+        print("  " + " " * width + "  " + "  ".join(f"{l:>{width}}" for l in labels))
+        for label, line in zip(labels, confusion):
+            print(f"  {label:>{width}}  " + "  ".join(f"{n:>{width}}" for n in line))
+    lat = np.array(latencies)
+    print(f"\nlatency per decision: p50 {np.percentile(lat, 50):.0f} ms, p95 {np.percentile(lat, 95):.0f} ms")
+
+
+if __name__ == "__main__":
+    main()

training/export_onnx.py

@@ -0,0 +1,151 @@
+"""Export a trained checkpoint to ONNX for fast CPU serving, and check it against PyTorch.
+
+    python export_onnx.py --model my-router --task task.example.json --data data/example.jsonl
+
+Writes into ``<model>/onnx/``:
+
+* ``model.onnx``: fp32, matches PyTorch on any CPU (the safe default);
+* ``model-int8-blockwise.onnx``: 8-bit block-wise encoder weights (ONNX Runtime MatMulNBits,
+  block 32). It kept Raya's accuracy and ran ~10-15% faster on CPUs with VNNI int8 instructions
+  (Intel Cascade Lake/Alder Lake+, AMD Zen 4+), but slower without VNNI. Skip with --fp32-only.
+
+Serve either with Laya's ONNXAgent (same answers and format as laya.Agent):
+
+    from laya.onnx_agent import ONNXAgent
+    agent = ONNXAgent("my-router", onnx_path="my-router/onnx/model.onnx"); agent.cfg["max_len"] = 512
+
+Every exported file is checked on up to 20 rows of --data: the choice must not change and
+probabilities must stay within 0.001 (fp32) or 0.05 (int8) of PyTorch.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+os.environ.setdefault("USE_TF", "0")
+
+import numpy as np
+import onnx
+import torch
+from laya import Agent
+from laya.onnx_agent import ONNXAgent
+
+from common import answer_probs, load_rows, load_task, row_state
+
+GRAPH_INPUTS = ("input_ids", "attention_mask", "marker_pos", "marker_mask", "qtype")
+
+
+def capture_inputs(agent: Agent, state, question: dict) -> dict[str, np.ndarray]:
+    """The network inputs Laya builds for one decision (used as the export example)."""
+    captured: dict[str, np.ndarray] = {}
+    model = agent.model
+
+    class Capture(torch.nn.Module):
+        def forward(self, *tensors):
+            captured.update({n: t.numpy() for n, t in zip(GRAPH_INPUTS, tensors)})
+            return model(*tensors)
+
+    agent.model = Capture()
+    try:
+        agent.system_one(state, {"q": question})
+    finally:
+        agent.model = model
+    return captured
+
+
+def export_fp32(agent: Agent, state, question: dict, path: Path) -> None:
+    example = capture_inputs(agent, state, question)
+    seq = torch.export.Dim("seq", min=8, max=4096)
+    opts = torch.export.Dim("opts", min=1, max=16)
+    torch.onnx.export(agent.model.eval(), tuple(torch.from_numpy(example[k]) for k in GRAPH_INPUTS), str(path),
+                      input_names=list(GRAPH_INPUTS), output_names=["logits", "act_logits"],
+                      dynamic_shapes=({1: seq}, {1: seq}, {1: opts}, {1: opts}, None),
+                      dynamo=True, external_data=False)
+    # Shapes recorded by the exporter contradict ONNX shape inference during quantization;
+    # drop them (ONNX Runtime re-infers them at load).
+    model = onnx.load(str(path))
+    del model.graph.value_info[:]
+    onnx.save(model, str(path))
+
+
+def encoder_weight_matmuls(path: Path, n_layers: int) -> list[str] | None:
+    """Names of the encoder's weight matmuls (4 per ModernBERT/mmBERT layer), or None if not found."""
+    model = onnx.load(str(path), load_external_data=False)
+    constants = {i.name for i in model.graph.initializer}
+    producers = {o: n for n in model.graph.node for o in n.output}
+
+    def is_constant(name: str) -> bool:
+        p = producers.get(name)
+        return name in constants or (p is not None and p.op_type in ("Transpose", "Cast", "Identity")
+                                      and all(is_constant(i) for i in p.input))
+
+    names = [node.name for node in model.graph.node
+             if node.op_type == "MatMul" and is_constant(node.input[1])
+             and "encoder.layers." in {p.key: p.value for p in node.metadata_props}.get("pkg.torch.onnx.name_scopes", "")]
+    return names if len(names) == 4 * n_layers else None
+
+
+def quantize_blockwise(src: Path, dst: Path, nodes: list[str]) -> None:
+    from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer
+
+    quantizer = MatMulNBitsQuantizer(onnx.load(str(src)), bits=8, block_size=32, is_symmetric=True,
+                                     accuracy_level=4, nodes_to_include=nodes)
+    quantizer.process()
+    quantizer.model.save_model_to_file(str(dst), use_external_data_format=False)
+
+
+def check(agent: Agent, model_dir: str, onnx_path: Path, rows, question, labels, max_tokens, tolerance) -> float:
+    served = ONNXAgent(model_dir, onnx_path=str(onnx_path))
+    served.cfg["max_len"] = max_tokens
+    worst = 0.0
+    for row in rows:
+        state = row_state(row)
+        want = answer_probs(agent.system_one(state, {"q": question})["answers"]["q"], question, labels)
+        got = answer_probs(served.system_one(state, {"q": question})["answers"]["q"], question, labels)
+        delta = max(abs(a - b) for a, b in zip(want, got))
+        worst = max(worst, delta)
+        if int(np.argmax(want)) != int(np.argmax(got)) or delta > tolerance:
+            raise SystemExit(f"{onnx_path.name} disagrees with PyTorch on row {row['id']} (delta {delta:.4f})")
+    return worst
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--model", required=True, help="trained checkpoint directory")
+    ap.add_argument("--task", required=True)
+    ap.add_argument("--data", required=True, help="JSONL rows used to check the export (labels not needed)")
+    ap.add_argument("--max-tokens", type=int, default=512, help="serving token budget (match training)")
+    ap.add_argument("--fp32-only", action="store_true")
+    args = ap.parse_args()
+
+    task = load_task(args.task)
+    labels, question = task["labels"], task["questions"][0]
+    rows = load_rows(args.data, labels, require_labels=False)[:20]
+    agent = Agent(args.model, device="cpu")
+    agent.cfg["max_len"] = args.max_tokens
+
+    out_dir = Path(args.model) / "onnx"
+    out_dir.mkdir(exist_ok=True)
+    fp32 = out_dir / "model.onnx"
+    export_fp32(agent, row_state(rows[0]), question, fp32)
+    print(f"{fp32}: max probability difference vs PyTorch "
+          f"{check(agent, args.model, fp32, rows, question, labels, args.max_tokens, 1e-3):.5f}")
+    if args.fp32_only:
+        return
+
+    n_layers = json.loads((Path(args.model) / "encoder" / "config.json").read_text())["num_hidden_layers"]
+    nodes = encoder_weight_matmuls(fp32, n_layers)
+    if nodes is None:
+        print("int8: encoder is not ModernBERT/mmBERT-shaped; skipping block-wise quantization")
+        return
+    int8 = out_dir / "model-int8-blockwise.onnx"
+    quantize_blockwise(fp32, int8, nodes)
+    print(f"{int8}: max probability difference vs PyTorch "
+          f"{check(agent, args.model, int8, rows, question, labels, args.max_tokens, 0.05):.5f}")
+
+
+if __name__ == "__main__":
+    main()

training/label.py

@@ -0,0 +1,135 @@
+"""Label your prompts with two (or more) independent LLM annotators.
+
+Raya's labels were made this way: two different strong models (Claude Opus and Claude Sonnet)
+read the same written rubric and labelled every prompt without seeing each other's answers.
+Where they agree you get a clean label; where they disagree the example keeps both votes and
+train.py learns a 50/50 target instead of a wrong hard label.
+
+    export OPENAI_API_KEY=...        # or any OpenAI-compatible provider
+    python label.py --task task.example.json --data prompts.jsonl --out labelled.jsonl \\
+        --annotator <model-a> \\
+        --annotator <model-b>@https://api.anthropic.com/v1/#ANTHROPIC_API_KEY
+
+An annotator is ``MODEL[@BASE_URL][#API_KEY_ENV]``: BASE_URL defaults to
+https://api.openai.com/v1/ and API_KEY_ENV to OPENAI_API_KEY. Use models from different
+families so their mistakes are independent. Input rows need a ``prompt`` (or ``state``); any
+existing labels are ignored. Re-running resumes: rows already labelled in --out are skipped.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import threading
+import time
+import urllib.error
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+from common import EXCLUDE, load_rows, load_task, row_state
+
+DEFAULT_BASE = "https://api.openai.com/v1/"
+
+
+def parse_annotator(spec: str) -> dict:
+    spec, _, key_env = spec.partition("#")
+    model, _, base = spec.partition("@")
+    return {"model": model, "base": (base or DEFAULT_BASE).rstrip("/") + "/",
+            "key_env": key_env or "OPENAI_API_KEY", "name": model}
+
+
+def system_prompt(task: dict) -> str:
+    labels = task["labels"]
+    rubric = task.get("rubric") or "\n".join(
+        f"- {label}: {task['questions'][0]['criteria'][label]}" for label in labels
+        if isinstance(task["questions"][0].get("criteria"), dict))
+    return (
+        "You are an independent annotator building a training set.\n"
+        f"{task['questions'][0]['instructions']}\n\n{rubric}\n\n"
+        f'Use "{EXCLUDE}" only if the input is unusable (empty, gibberish, no discernible request).\n'
+        "Judge the input yourself; do not guess from keywords. "
+        f'Answer with JSON only: {{"label": one of {labels + [EXCLUDE]}}}'
+    )
+
+
+def call(annotator: dict, system: str, user: str, labels: list[str], retries: int = 6) -> str:
+    key = os.environ.get(annotator["key_env"])
+    if not key:
+        raise SystemExit(f"set {annotator['key_env']} for annotator {annotator['model']}")
+    body = json.dumps({"model": annotator["model"], "temperature": 0, "max_tokens": 50,
+                       "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}]})
+    request = urllib.request.Request(annotator["base"] + "chat/completions", data=body.encode(), headers={
+        "Content-Type": "application/json", "Authorization": f"Bearer {key}", "x-api-key": key,
+        "anthropic-version": "2023-06-01"})
+    for attempt in range(retries):
+        try:
+            with urllib.request.urlopen(request, timeout=120) as response:
+                text = json.load(response)["choices"][0]["message"]["content"] or ""
+            match = re.search(r"\{.*?\}", text, re.S)
+            label = json.loads(match.group(0)).get("label") if match else None
+            if label in labels or label == EXCLUDE:
+                return label
+            raise ValueError(f"unexpected answer {text[:80]!r}")
+        except urllib.error.HTTPError as exc:
+            if exc.code not in (408, 409, 429) and exc.code < 500:
+                raise SystemExit(f"{annotator['model']}: HTTP {exc.code} {exc.read()[:200]!r}") from exc
+            error = exc
+        except (urllib.error.URLError, TimeoutError, ValueError, KeyError) as exc:
+            error = exc
+        time.sleep(min(60, 2 ** attempt))
+    raise RuntimeError(f"{annotator['model']} failed after {retries} attempts: {error}")
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--task", required=True)
+    ap.add_argument("--data", required=True, help="JSONL of rows with a 'prompt' or 'state'")
+    ap.add_argument("--out", required=True, help="labelled JSONL (appended; re-run to resume)")
+    ap.add_argument("--annotator", action="append", required=True, help="MODEL[@BASE_URL][#API_KEY_ENV]; repeat")
+    ap.add_argument("--workers", type=int, default=8)
+    ap.add_argument("--max-chars", type=int, default=6000, help="truncate long inputs sent to annotators")
+    args = ap.parse_args()
+
+    task = load_task(args.task)
+    labels = task["labels"]
+    annotators = [parse_annotator(a) for a in args.annotator]
+    if len(annotators) < 2:
+        print("warning: one annotator gives hard labels only; two different models are recommended")
+    system = system_prompt(task)
+    rows = load_rows(args.data, labels, require_labels=False)
+    out = Path(args.out)
+    done = {json.loads(l)["id"] for l in out.read_text(encoding="utf-8").splitlines() if l.strip()} if out.exists() else set()
+    todo = [r for r in rows if r["id"] not in done]
+    print(f"{len(rows)} rows, {len(done)} already labelled, {len(todo)} to go, {len(annotators)} annotators")
+
+    lock = threading.Lock()
+
+    def label_row(row: dict) -> dict:
+        state = row_state(row)
+        user = state["prompt"] if isinstance(state, dict) and set(state) == {"prompt"} else json.dumps(state, ensure_ascii=False)
+        if len(user) > args.max_chars:
+            user = user[:args.max_chars] + " …[truncated]"
+        votes = {a["name"]: call(a, system, user, labels) for a in annotators}
+        clean = {k: v for k, v in row.items() if k not in ("label", "labels", "target", "gold")}
+        return dict(clean, labels=list(votes.values()), annotators=votes)
+
+    agree = labelled = 0
+    with ThreadPoolExecutor(args.workers) as pool, out.open("a", encoding="utf-8") as sink:
+        for future in as_completed(pool.submit(label_row, r) for r in todo):
+            result = future.result()
+            with lock:
+                sink.write(json.dumps(result, ensure_ascii=False) + "\n")
+                sink.flush()
+                labelled += 1
+                agree += len(set(result["labels"])) == 1
+                if labelled % 50 == 0:
+                    print(f"{labelled}/{len(todo)} labelled, annotators agree on {agree / labelled:.0%}", flush=True)
+    if labelled:
+        print(f"done: {labelled} labelled, annotators agree on {agree / labelled:.0%}")
+
+
+if __name__ == "__main__":
+    main()

training/requirements-onnx.txt

@@ -0,0 +1,5 @@
+# Only needed for export_onnx.py and CPU serving. Tested: onnx 1.23.0, onnxscript 0.7.2, onnxruntime 1.30.0.
+-r requirements.txt
+onnx>=1.17
+onnxscript>=0.2
+onnxruntime>=1.20

training/requirements.txt

@@ -0,0 +1,8 @@
+# Tested with these exact versions (September 2026). laya is pinned because its tokenisation
+# and temperature handling decide the model's inputs; the rest are minimums.
+laya==0.3.20
+torch>=2.4
+transformers>=4.48
+safetensors>=0.4
+numpy>=1.26
+huggingface_hub>=0.26

training/task.example.json

@@ -0,0 +1,34 @@
+{
+  "description": "3-tier LLM routing, the task Raya was trained on. Copy this file and edit it for your own task.",
+  "labels": ["small_model", "medium_model", "frontier_model"],
+  "questions": [
+    {
+      "type": "choice",
+      "instructions": "Route this prompt to a model.",
+      "criteria": {
+        "small_model": "simple requests",
+        "medium_model": "moderately complex requests",
+        "frontier_model": "very hard requests"
+      }
+    },
+    {
+      "type": "choice",
+      "instructions": "Which model tier should answer this user prompt? Pick the cheapest tier that will still answer it well.",
+      "criteria": {
+        "small_model": "greetings, trivial facts, single-sentence translation or rewrite, typo fixes, simple arithmetic, one-liner code questions, very short simple creative requests",
+        "medium_model": "normal multi-paragraph writing, emails, essays, standard coding tasks, summarizing/translating/rewriting a provided text, explaining well-known concepts, routine analysis",
+        "frontier_model": "hard multi-step reasoning, non-trivial math or proofs, complex system design or large tricky code, expert legal/financial/medical/scientific analysis, research-grade or strategy work"
+      }
+    },
+    {
+      "type": "score",
+      "instructions": "How difficult is this prompt for an AI assistant to answer well?",
+      "criteria": [
+        "simple: a small fast model answers it perfectly",
+        "moderate: needs a capable general model",
+        "hard: needs the strongest frontier model"
+      ]
+    }
+  ],
+  "rubric": "Decide which model tier should answer the prompt for a production assistant that wants the CHEAPEST model that will still answer well.\n- small_model: greetings, trivial facts, single-sentence translation/rewrite/typo fix, simple arithmetic, one-liner code questions, very short simple creative requests.\n- medium_model: normal multi-paragraph writing, emails, essays, standard coding tasks, summarizing/translating/rewriting a provided text, explanations of well-known concepts, routine analysis.\n- frontier_model: hard multi-step reasoning, non-trivial math/proofs, complex system design or large/tricky code, expert legal/financial/medical/scientific analysis, long research-grade or strategy work synthesizing many sources.\nJudge the prompt in its own language. Do not aim for class balance."
+}

training/train.py

@@ -0,0 +1,354 @@
+"""Fine-tune a Laya decision model (a "System-1" model) on your own labelled data.
+
+This is the recipe behind TextCortex/raya, generalised to any choice/score task:
+
+* soft targets: each label's share of the annotator votes (a 50/50 target where two disagree);
+* every question phrasing in the task is trained on, and choice options are shuffled, so the
+  model learns the content rather than the wording or the option order;
+* soft-target cross-entropy with 1/sqrt(frequency) class weights, so rare labels still count;
+* the best epoch is picked on validation, then a temperature per question is fitted there so
+  the output probabilities are calibrated;
+* the result is a normal Laya checkpoint: ``laya.Agent("<out>")`` loads it.
+
+Quick start (CPU/Apple Silicon works for small data; a GPU is much faster):
+
+    python train.py --task task.example.json --data data/example.jsonl --out my-router
+
+Start from Raya instead of stock Laya to adapt the router to your own traffic:
+
+    python train.py --task task.example.json --data my_data.jsonl --base TextCortex/raya --out my-router
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import shutil
+import time
+from pathlib import Path
+
+os.environ.setdefault("USE_TF", "0")
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from safetensors.torch import save_file
+
+import laya
+from laya.common import QTYPES, build_sequence, collate_items, temp_bucket
+
+from common import load_rows, load_task, row_state
+
+BASE_FILES = ["rl_agent_config.json", "model.safetensors", "tokenizer/*", "encoder/*"]
+
+
+def resolve_base(base: str, subfolder: str | None) -> Path:
+    """Local directory holding the base checkpoint (downloads it from the Hub if needed)."""
+    if Path(base).is_dir():
+        return Path(base) / subfolder if subfolder else Path(base)
+    from huggingface_hub import snapshot_download
+
+    patterns = [f"{subfolder}/{p}" for p in BASE_FILES] if subfolder else BASE_FILES
+    local = Path(snapshot_download(base, allow_patterns=patterns))
+    return local / subfolder if subfolder else local
+
+
+def internal(question: dict) -> dict:
+    return laya.Agent._to_internal(question)
+
+
+def make_item(tok, cfg, state, qi: int, question: dict, labels: list[str], target: list[float], shuffle: bool):
+    """Tokenise one (example, question) pair and align its target with the option order."""
+    q = internal(question)
+    if q["t"] == "choice":
+        keys = list(q["crit"].keys())
+        by_label = dict(zip(labels, target))
+        order = list(range(len(keys)))
+        if shuffle:
+            random.shuffle(order)
+        seq, markers = build_sequence(tok, state, q, cfg["max_len"], cfg["head_max_len"], option_order=order)
+        tgt = [by_label[keys[j]] for j in order]
+    else:  # ordinal score: the option order is the meaning, never shuffle
+        seq, markers = build_sequence(tok, state, q, cfg["max_len"], cfg["head_max_len"])
+        tgt = list(target)
+    return {"ids": seq, "markers": markers, "qtype": QTYPES[q["t"]], "target": tgt,
+            "question": qi, "label_target": list(target)}
+
+
+def length_batches(items, bs: int, shuffle: bool):
+    """Batches of similar length (less padding), in random order when training."""
+    idx = sorted(range(len(items)), key=lambda i: len(items[i]["ids"]))
+    chunks = [idx[i:i + bs] for i in range(0, len(idx), bs)]
+    if shuffle:
+        random.shuffle(chunks)
+    return chunks
+
+
+def forward(model, batch, dev):
+    with torch.autocast(device_type=dev.type, dtype=torch.bfloat16, enabled=dev.type == "cuda"):
+        logits, _ = model(batch["input_ids"].to(dev), batch["attention_mask"].to(dev), batch["marker_pos"].to(dev),
+                          batch["marker_mask"].to(dev), batch["qtype"].to(dev))
+    return logits.float()
+
+
+@torch.no_grad()
+def predict_logits(model, items, pad_id, dev, bs=16):
+    model.eval()
+    out = [None] * len(items)
+    for chunk in length_batches(items, bs, shuffle=False):
+        batch = collate_items([[items[i] for i in chunk]], pad_id)
+        logits = forward(model, batch, dev).cpu()
+        for j, i in enumerate(chunk):
+            out[i] = logits[j, :len(items[i]["markers"])].numpy()
+    return out
+
+
+def softmax(z: np.ndarray, t: float = 1.0) -> np.ndarray:
+    z = z / t
+    p = np.exp(z - z.max())
+    return p / p.sum()
+
+
+def evaluate(logits, items, rows_of_items, n_questions: int, n_labels: int, temps=None):
+    """Accuracy and macro-F1 per question, on examples with a single gold label."""
+    res = {}
+    for qi in range(n_questions):
+        preds, golds = [], []
+        for lg, it, row in zip(logits, items, rows_of_items):
+            if it["question"] != qi or row["gold"] is None:
+                continue
+            preds.append(int(softmax(lg, (temps or {}).get(qi, 1.0)).argmax()))
+            golds.append(int(np.argmax(it["target"])))
+        if not golds:
+            continue
+        f1s = []
+        for k in range(n_labels):
+            tp = sum(p == k == g for p, g in zip(preds, golds))
+            fp = sum(p == k != g for p, g in zip(preds, golds))
+            fn = sum(g == k != p for p, g in zip(preds, golds))
+            f1s.append(2 * tp / (2 * tp + fp + fn) if tp else 0.0)
+        res[qi] = {"acc": round(float(np.mean([p == g for p, g in zip(preds, golds)])), 4),
+                   "macro_f1": round(float(np.mean(f1s)), 4), "n": len(golds)}
+    return res
+
+
+def soft_nll(logits, items) -> float:
+    total = 0.0
+    for lg, it in zip(logits, items):
+        total -= float((np.array(it["target"]) * np.log(softmax(lg) + 1e-12)).sum())
+    return total / len(items)
+
+
+def fit_temperature(logits, items, qi: int) -> float:
+    """Temperature minimising the soft-target NLL of one question on validation."""
+    sel = [(lg, np.array(it["target"])) for lg, it in zip(logits, items) if it["question"] == qi]
+    best_t, best_nll = 1.0, float("inf")
+    for t in np.arange(0.5, 5.01, 0.05):
+        nll = -sum(float((tg * np.log(softmax(lg, t) + 1e-12)).sum()) for lg, tg in sel)
+        if nll < best_nll:
+            best_t, best_nll = round(float(t), 2), nll
+    return best_t
+
+
+def split_rows(rows, val_frac: float, seed: int):
+    """Rows marked ``"split": "val"`` are validation; otherwise a random ``val_frac`` share is."""
+    marked = [r for r in rows if r.get("split") == "val"]
+    if marked:
+        return [r for r in rows if r.get("split") != "val"], marked
+    rows = list(rows)
+    random.Random(seed).shuffle(rows)
+    n_val = max(1, int(len(rows) * val_frac))
+    return rows[n_val:], rows[:n_val]
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--task", required=True, help="task.json: labels + question phrasings")
+    ap.add_argument("--data", required=True, help="training JSONL (see common.py for the format)")
+    ap.add_argument("--val-data", help="optional separate validation JSONL")
+    ap.add_argument("--out", required=True, help="output checkpoint directory")
+    ap.add_argument("--base", default="convaiinnovations/laya",
+                    help="base checkpoint: a Hub repo id or local dir (e.g. TextCortex/raya to adapt Raya)")
+    ap.add_argument("--subfolder", default=None,
+                    help="checkpoint subfolder; defaults to 'multilingual' for convaiinnovations/laya "
+                         "(pass --subfolder . for Laya's English ModernBERT-large checkpoint)")
+    ap.add_argument("--encoder", default=None,
+                    help="build a fresh decision head on this HF encoder (e.g. jhu-clsp/mmBERT-small) instead of --base")
+    ap.add_argument("--epochs", type=int, default=3)
+    ap.add_argument("--batch-size", type=int, default=16)
+    ap.add_argument("--lr-encoder", type=float, default=2e-5)
+    ap.add_argument("--lr-head", type=float, default=None, help="default 1e-4, or 3e-4 for a fresh head")
+    ap.add_argument("--max-tokens", type=int, default=512, help="input token budget while training")
+    ap.add_argument("--val-frac", type=float, default=0.1)
+    ap.add_argument("--select", choices=["acc", "nll"], default="acc", help="best-epoch criterion on validation")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--device", default=None, help="cuda | mps | cpu (default: best available)")
+    ap.add_argument("--train-embeddings", action="store_true",
+                    help="also train the token-embedding table (frozen by default: saves memory, rarely helps)")
+    ap.add_argument("--push-to-hub", metavar="REPO_ID", help="upload the result to this Hugging Face model repo")
+    ap.add_argument("--private", action="store_true", help="with --push-to-hub: create the repo as private")
+    args = ap.parse_args()
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    task = load_task(args.task)
+    labels, questions = task["labels"], task["questions"]
+    rows = load_rows(args.data, labels)
+    if args.val_data:
+        train_rows, val_rows = rows, load_rows(args.val_data, labels)
+    else:
+        train_rows, val_rows = split_rows(rows, args.val_frac, args.seed)
+    mass = np.array([sum(r["target"][i] for r in train_rows) for i in range(len(labels))])
+    print(f"train {len(train_rows)} rows, validation {len(val_rows)} rows, {len(questions)} question phrasing(s)")
+    print("train label mass:", dict(zip(labels, mass.round(1).tolist())))
+    if (mass == 0).any():
+        raise SystemExit(f"no training examples for: {[l for l, m in zip(labels, mass) if m == 0]}")
+
+    dev = torch.device(args.device or ("cuda" if torch.cuda.is_available()
+                                        else "mps" if torch.backends.mps.is_available() else "cpu"))
+    if args.encoder:
+        from transformers import AutoTokenizer
+        from laya.common import build_model
+
+        # Borrow Laya's head/config layout, but train the decision head from scratch on a new encoder.
+        ref_dir = resolve_base("convaiinnovations/laya", "multilingual")
+        cfg = dict(json.loads((ref_dir / "rl_agent_config.json").read_text()), encoder=args.encoder,
+                   max_len=args.max_tokens, head_max_len=192, temperature=[1.0, 1.0, 1.0], temperature_by_options={})
+        tok = AutoTokenizer.from_pretrained(args.encoder)
+        model = build_model(cfg, pretrained=True).to(dev)
+        base_dir = None
+        lr_head = args.lr_head or 3e-4
+        print(f"device {dev}: fresh decision head on {args.encoder}")
+    else:
+        subfolder = args.subfolder if args.subfolder is not None else (
+            "multilingual" if args.base == "convaiinnovations/laya" else None)
+        subfolder = None if subfolder in ("", ".") else subfolder
+        base_dir = resolve_base(args.base, subfolder)
+        agent = laya.Agent(str(base_dir), device=str(dev))
+        model, tok, cfg = agent.model, agent.tok, agent.cfg
+        lr_head = args.lr_head or 1e-4
+        print(f"device {dev}: fine-tuning {args.base}{'/' + subfolder if subfolder else ''}")
+
+    # Weight each example by its target's class weight (~1/sqrt(frequency), mean 1).
+    cw = (mass.sum() / mass) ** 0.5
+    cw = torch.tensor(cw / cw.mean(), dtype=torch.float32)
+    print("class weights:", {l: round(float(w), 2) for l, w in zip(labels, cw)})
+
+    if not args.train_embeddings:
+        for p in model.encoder.get_input_embeddings().parameters():
+            p.requires_grad_(False)
+    if dev.type != "cuda":
+        model.encoder.gradient_checkpointing_enable()  # trade speed for memory off-GPU
+    train_cfg = dict(cfg, max_len=min(cfg["max_len"], args.max_tokens))
+
+    def items_for(rs, shuffle, c):
+        its, owners = [], []
+        for r in rs:
+            state = row_state(r)
+            for qi, q in enumerate(questions):
+                its.append(make_item(tok, c, state, qi, q, labels, r["target"], shuffle))
+                owners.append(r)
+        return its, owners
+
+    val_items, val_owners = items_for(val_rows, False, train_cfg)
+    before = predict_logits(model, val_items, tok.pad_token_id, dev)
+    print("validation before training:", json.dumps(evaluate(before, val_items, val_owners, len(questions), len(labels))))
+
+    enc_params = [p for p in model.encoder.parameters() if p.requires_grad]
+    head_params = [p for n, p in model.named_parameters() if not n.startswith("encoder.") and p.requires_grad]
+    opt = torch.optim.AdamW([{"params": enc_params, "lr": args.lr_encoder}, {"params": head_params, "lr": lr_head}],
+                            weight_decay=0.01)
+    steps_per_epoch = math.ceil(len(train_rows) * len(questions) / args.batch_size)
+    total = steps_per_epoch * args.epochs
+    warm = max(1, int(0.1 * total))
+    sched = torch.optim.lr_scheduler.LambdaLR(
+        opt, lambda s: min(1.0, (s + 1) / warm) * max(0.0, (total - s) / max(1, total - warm)))
+
+    best_crit, best_state, best_logits = float("inf"), None, None
+    history, step, t0 = [], 0, time.time()
+    for epoch in range(args.epochs):
+        model.train()
+        items, _ = items_for(train_rows, True, train_cfg)  # fresh option shuffle every epoch
+        running = []
+        for chunk in length_batches(items, args.batch_size, shuffle=True):
+            batch = collate_items([[items[i] for i in chunk]], tok.pad_token_id)
+            logits = forward(model, batch, dev)
+            tgt = batch["target"][:, :logits.size(1)].to(dev)
+            logp = F.log_softmax(logits.masked_fill(~batch["marker_mask"].to(dev), -1e4), -1)
+            per_example = -(tgt * logp).sum(-1)
+            w = torch.tensor([float((torch.tensor(items[i]["label_target"]) * cw).sum()) for i in chunk], device=dev)
+            loss = (per_example * w).sum() / w.sum()
+            opt.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step()
+            sched.step()
+            step += 1
+            running.append(loss.item())
+            if dev.type == "mps" and step % 25 == 0:
+                torch.mps.empty_cache()
+            if step % 50 == 0:
+                print(f"epoch {epoch} step {step}/{total} loss {np.mean(running[-50:]):.4f} ({time.time() - t0:.0f}s)",
+                      flush=True)
+        val_logits = predict_logits(model, val_items, tok.pad_token_id, dev)
+        ev = evaluate(val_logits, val_items, val_owners, len(questions), len(labels))
+        mean_acc = float(np.mean([v["acc"] for v in ev.values()])) if ev else 0.0
+        nll = soft_nll(val_logits, val_items)
+        crit = -mean_acc if args.select == "acc" else nll
+        history.append({"epoch": epoch, "val_mean_acc": round(mean_acc, 4), "val_soft_nll": round(nll, 4),
+                        "val": ev, "seconds": round(time.time() - t0)})
+        print(f"epoch {epoch}: validation {json.dumps(ev)} mean acc {mean_acc:.4f} soft NLL {nll:.4f}", flush=True)
+        if crit < best_crit:
+            best_crit = crit
+            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
+            best_logits = val_logits
+
+    temps = {qi: fit_temperature(best_logits, val_items, qi) for qi in range(len(questions))}
+    final = evaluate(best_logits, val_items, val_owners, len(questions), len(labels), temps)
+    print("fitted temperatures:", temps, "validation after calibration:", json.dumps(final))
+
+    out = Path(args.out)
+    if out.exists():
+        shutil.rmtree(out)
+    if base_dir is None:
+        out.mkdir(parents=True)
+        tok.save_pretrained(out / "tokenizer")
+        model.encoder.config.save_pretrained(out / "encoder")
+    else:
+        shutil.copytree(base_dir, out, ignore=shutil.ignore_patterns(
+            "model.safetensors", "*.onnx", "onnx", "multilingual", "typed-decisions", "assets", "*.md", ".git*", ".cache"))
+    save_file({k: v.contiguous() for k, v in best_state.items()}, str(out / "model.safetensors"))
+
+    # Laya applies one temperature per (question type, option count) bucket; average within a bucket.
+    buckets: dict[str, list[float]] = {}
+    for qi, q in enumerate(questions):
+        iq = internal(q)
+        buckets.setdefault(temp_bucket(QTYPES[iq["t"]], len(iq["crit"])), []).append(temps[qi])
+    new_cfg = dict(cfg)
+    new_cfg["temperature_by_options"] = {**cfg.get("temperature_by_options", {}),
+                                         **{b: round(float(np.mean(ts)), 2) for b, ts in buckets.items()}}
+    new_cfg["training"] = dict(cfg.get("training", {}), fine_tuned={
+        "base": args.encoder or args.base, "train_rows": len(train_rows), "val_rows": len(val_rows),
+        "epochs": args.epochs, "select": args.select, "seed": args.seed, "labels": labels})
+    (out / "rl_agent_config.json").write_text(json.dumps(new_cfg, indent=2))
+    (out / "task.json").write_text(json.dumps(task, indent=2, ensure_ascii=False))
+    (out / "training_log.json").write_text(json.dumps(
+        {"args": vars(args), "history": history, "temperatures": temps, "validation": final}, indent=2))
+    print(f"saved {out}  (load it with laya.Agent({str(out)!r}))")
+
+    if args.push_to_hub:
+        from huggingface_hub import HfApi
+
+        api = HfApi()
+        api.create_repo(args.push_to_hub, private=args.private, exist_ok=True)
+        api.upload_folder(folder_path=str(out), repo_id=args.push_to_hub)
+        print(f"uploaded to https://huggingface.co/{args.push_to_hub}")
+
+
+if __name__ == "__main__":
+    main()