training/README.md

Train your own System-1 model

This folder is the recipe we used to train Raya, 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 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

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):

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:

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)

{
  "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:

{"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)

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)

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)

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:

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.