raya / README.md
cderinbogaz's picture
Add training kit: train your own System-1 model
48c8658 verified
|
Raw
History Blame Contribute Delete
8.95 kB
metadata
license: apache-2.0
library_name: laya
pipeline_tag: text-classification
base_model: convaiinnovations/laya
datasets:
  - allenai/WildChat-1M
language:
  - en
  - de
  - fr
  - es
  - it
  - pt
  - nl
  - pl
  - tr
  - ru
  - ar
  - zh
  - ja
  - ko
tags:
  - routing
  - llm-routing
  - model-selection
  - prompt-difficulty
  - laya
  - system-one
  - calibrated-decisions
  - multilingual

Raya β€” a fast multilingual LLM router (Laya fine-tune)

Raya picks which model tier should answer a prompt β€” a small fast model, a mid-tier model, or a frontier model β€” in one ~17 ms forward pass, about 20Γ— faster than a hosted routing API.

  • Robust to how you phrase the routing question. Raya scores 80–81% on all three question styles we tested. On the difficulty-score style it beats TypeSafe Jev by ~10 points (80.3% vs 70.5%, paired McNemar p < 0.001).
  • Matches or beats Jev on the minimal routing question in Japanese (90 vs 87), Polish (86 vs 83), Spanish (87 vs 85), German (74 vs 72), English (86 vs 86), and Arabic (73 vs 73).
  • Open and self-hostable. A 300M-parameter model with no per-call cost, served through Laya's Jev-compatible HTTP API.

Raya is Laya's multilingual decision model (mmBERT-base encoder), fine-tuned end-to-end for 3-tier routing. It is not affiliated with Convai Innovations or TypeSafe.

Quick start

import laya  # pip install laya  (tested with laya 0.3.7 and 0.3.20; identical results)

raya = laya.Agent("TextCortex/raya", device="cuda")   # or "mps" / "cpu"

ROUTE = {
    "type": "choice",
    "instructions": "Route this prompt to a model.",
    "criteria": {
        "small_model": "simple requests",
        "medium_model": "moderately complex requests",
        "frontier_model": "very hard requests",
    },
}

out = raya.system_one({"prompt": "Schreibe eine professionelle E-Mail an einen Kunden …"}, {"route": ROUTE})
print(out["answers"]["route"])
# {'choice': 'medium_model', 'probabilities': {'small_model': 0.24, 'medium_model': 0.74, 'frontier_model': 0.03}, ...}

Raya was trained on three routing questions: the minimal choice above, a detailed rubric choice, and a 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/, 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.
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.

from huggingface_hub import snapshot_download
from laya.onnx_agent import ONNXAgent   # pip install laya onnxruntime

FILE = "onnx/raya.onnx"                 # see the table below
path = snapshot_download("TextCortex/raya", allow_patterns=["rl_agent_config.json", "tokenizer/*", "encoder/*", FILE])
raya = ONNXAgent(path, onnx_path=f"{path}/{FILE}")
raya.cfg["max_len"] = 512
print(raya.system_one({"prompt": "Schreibe eine professionelle E-Mail …"}, {"route": ROUTE})["answers"]["route"])  # ROUTE as above

Which file? If you're not sure, use raya.onnx. At the same token budget it matches the PyTorch model exactly (0 of 563 choices changed, probabilities within 0.0001), on any CPU. On CPUs with VNNI int8 instructions (Intel Cascade Lake or Alder Lake and newer, AMD Zen 4 and newer), raya-int8-blockwise.onnx keeps accuracy and was about 10–15% faster on our Intel i9-13900. The other int8 files are faster still, but lose about 2 accuracy points.

File Size Accuracy Choices changed vs PyTorch Intel i9-13900, VNNI (p50 / p95) AMD EPYC 7502P, no VNNI (p50 / p95)
raya.onnx (fp32) 1.23 GB 81.2% 3 / 563 38 / 297 ms 60 / 487 ms
raya-int8-blockwise.onnx 0.89 GB 81.5% 5 / 563 34 / 282 ms 127 / 975 ms
raya-int8.onnx 0.92 GB 79.0% 53 / 563 24 / 200 ms not recommended
raya-int8-emb.onnx 0.35 GB 79.4% 47 / 563 24 / 201 ms not recommended
raya-int8-mixed.onnx 1.05 GB 79.2% 41 / 563 31 / 242 ms not recommended
PyTorch reference (1,024-token budget) β€” 81.0% β€” 42 / 421 ms 89 / 526 ms

All rows use the same 563-prompt benchmark below, with the minimal routing question and single requests on 8 threads. The fp32 and block-wise rows include the switch from a 1,024- to a 512-token budget, which accounts for 3 of their changed choices. Per-tensor int8 (the last three files) overflows on CPUs without VNNI; on the EPYC it dropped to 54–80% depending on settings. 16 threads was slower than 8 on both CPUs. onnx/SHA256SUMS has checksums and onnx/benchmark_results.jsonl has the raw results.

Benchmark: 3-tier routing on real multilingual prompts

Test set. 563 first-turn prompts from WildChat-1M (shards never used for training), 14 languages, ~40 per language, 75% natural traffic + 25% hard-leaning. Gold tier = the tier two independent blind annotators (Claude Opus and Claude Sonnet, same written rubric) both assigned; the annotators agree on 78% of prompts, which is the practical ceiling. Tier mix: 225 small / 317 medium / 21 frontier.

Question style Raya TypeSafe Jev 1.13 Laya 0.3.7 (stock) Von 1.1 (stock)
Difficulty score (3 levels) 80.3% 70.5% 61.6% 42.3%
Rubric choice 81.0% 84.2% 47.1% 58.8%
Minimal choice 80.8% 84.5% 55.2% 54.5%
Latency (p50) 17 ms (1Γ— GPU, in-process) ~330 ms (hosted API) β€” β€”

Always answering "medium" scores 56.3%.

Per language, minimal choice (accuracy %, n = 31–53 per language):

ja pl es de en ar it ko nl ru fr zh pt tr
Raya 90 86 87 74 86 73 87 88 78 82 78 79 76 67
Jev 87 83 85 72 86 73 90 91 81 89 86 88 90 82

Confusion, minimal choice (rows = gold small/medium/frontier, cols = Raya): [[169, 56, 0], [30, 279, 8], [0, 14, 7]].

Limitations of these numbers. On the choice questions Jev is ahead overall by 3–4 points. That gap is not statistically significant at these sample sizes (p = 0.07 and 0.13), but it is consistent. Raya is weakest in Turkish, Portuguese, and Chinese. Differences of less than ~10 points within a single language are within noise. Like Jev, Raya rarely routes to the frontier tier (7 of 21). Gold labels come from LLM annotators, not from measured downstream answer quality.

Training

  • Architecture: Laya decision model β€” mmBERT-base encoder + 2-layer transformer decision head with per-option [MASK] scoring, initialised from convaiinnovations/laya (multilingual) and fine-tuned end-to-end.
  • Data (soft labels from two blind annotators; 50/50 target where they disagreed): WildChat-1M first-turn prompts (ODC-BY); synthetic hard / hard-looking prompts across the 14 languages (kept only if a blind annotator confirmed the intended tier); additional in-house routing data.
  • Procedure: soft-target cross-entropy, options shuffled, AdamW (encoder 2e-5, head 1e-4), batch 32, sequences capped at 512 tokens, bf16, best epoch chosen by validation accuracy (86.7% on 659 held-out 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/ reproduces this procedure on your own data. The training data itself is not published.

Intended use and limitations

For choosing among 2–3 model tiers for chat/assistant traffic. Not for safety moderation or for more than 3 tiers. Not validated on multi-turn context. Calibrate a confidence threshold on your own traffic before auto-escalating.

Acknowledgements

Laya architecture, runtime, and base checkpoint by Convai Innovations (Apache-2.0). mmBERT by JHU CLSP (MIT). WildChat-1M by AI2 (ODC-BY). Jev is a product of TypeSafe AI; its scores come from our own runs via its API (September 2026).