# 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.