Train Laya for your own task. Then measure it honestly.

{
  "id": "3f9c0a17e2b4d851",
  "state": {
    "queue": { "time": "Monday morning", "backlog": "high" },
    "tickets": {
      "T-4812": { "plan": "business", "channel": "email",
                  "text": "Charged twice and the export still fails. Third time I write.",
                  "customer": { "customer_for": "over two years",
                                "history": ["refund in March", "outage report in July"] } }
    }
  },
  "questions": {
    "team": { "type": "choice",
              "instructions": "Which team should handle `tickets.T-4812`? Judge from its `text` and `plan`.",
              "criteria": { "billing": "Charges, refunds, invoices. Comes first when money was taken wrongly.",
                            "technical": "Something in the product does not work.", "sales": "Upgrades, quotes, new seats.", "account": null } },
    "urgent": { "type": "noul", "instructions": "Is `tickets.T-4812` urgent? Yes if the customer is blocked or losing money now." },
    "angry": { "type": "noul", "instructions": "Is the customer who wrote `tickets.T-4812` angry?" }
  },
  "targets": {
    "team": { "type": "choice",
              "probabilities": { "billing": 0.58, "technical": 0.34, "sales": 0.02, "account": 0.06 } },
    "urgent": { "type": "noul", "p": 0.71 },
    "angry": { "type": "noul", "p": 0.83 }
  },
  "source": { "group": "2026-09-14", "moment": "T-4812", "model": "teacher-2.1.0", "inputs": "exact", "duplicates_merged": 1 }
}

# data.py: logged answers (JSONL) -> training items, split by whole groups
import hashlib
import json

from laya.common import QTYPES, build_sequence, render_options

from load import HEAD_MAX_LEN, MAX_LEN


def split_of(group: str, test_groups: set[str], val_share: float = 0.10) -> str:
    """Whole groups only (a day, a customer). The test groups are a real deployment you chose before training."""
    if group in test_groups:
        return "test"
    h = int(hashlib.sha256(group.encode()).hexdigest()[:8], 16) / 0xFFFFFFFF
    return "val" if h < val_share else "train"


def target_of(q: dict, t: dict) -> list[float]:
    if q["type"] == "choice":
        return [float(t["probabilities"][k]) for k in q["criteria"]]  # in criteria order
    if q["type"] == "noul":
        return [1.0 - float(t["p"]), float(t["p"])]  # Laya shows a yes/no as [false, true]
    raise ValueError(f"unsupported question type {q['type']}")


def to_items(ex: dict, tok) -> tuple[list[dict], int]:
    items, dropped = [], 0
    for qid, q in ex["questions"].items():
        iq = {"t": q["type"], "ins": q["instructions"], "crit": q.get("criteria")}
        ids, markers = build_sequence(tok, ex["state"], iq, MAX_LEN, HEAD_MAX_LEN)
        if len(markers) != len(render_options(iq)):
            dropped += 1  # an option fell off the end: count it, never train or score on it
            continue
        t = target_of(q, ex["targets"][qid])
        s = sum(t)
        items.append({
            "ids": ids, "markers": markers, "qtype": QTYPES[q["type"]],
            "target": [v / s for v in t],
            "group": ex["source"]["group"], "moment": ex["source"].get("moment"),
        })
    return items, dropped


def build(path: str, tok, test_groups: set[str]) -> dict[str, list[dict]]:
    out, seen, dropped = {"train": [], "val": [], "test": []}, {"val": set(), "test": set()}, 0
    rows = [json.loads(line) for line in open(path)]
    key = lambda ex: hashlib.sha256(json.dumps([ex["state"], ex["questions"]], sort_keys=True).encode()).hexdigest()
    for ex in rows:  # first pass: what val and test contain
        s = split_of(ex["source"]["group"], test_groups)
        if s != "train":
            seen[s].add(key(ex))
    for ex in rows:
        s = split_of(ex["source"]["group"], test_groups)
        if s == "train" and (key(ex) in seen["val"] or key(ex) in seen["test"]):
            continue  # the same request in two splits is a leak: drop it from train
        items, d = to_items(ex, tok)
        out[s] += items
        dropped += d
    print({k: len(v) for k, v in out.items()}, "dropped (did not fit):", dropped)
    return out

# load.py: the base checkpoint, the sequence lengths, and loading
import json
import os

from huggingface_hub import snapshot_download
from laya.agent import _fix_tokenizer_config  # a private helper in laya 0.3.4; it patches older tokenizer files
from laya.common import build_model
from safetensors.torch import load_file
from transformers import AutoTokenizer

BASE_REPO = "convaiinnovations/laya"
BASE_REVISION = "1c5edc17a7acd8701df6fc341c0d179f1c62c982"  # the revision we trained from: pin yours
MAX_LEN = 512       # the whole sequence
HEAD_MAX_LEN = 256  # question + options (Laya's default, 192, can cut long option rules short)


def resolve(path: str) -> str:
    return snapshot_download(BASE_REPO, revision=BASE_REVISION) if path == "base" else path


def load(path: str = "base", device: str = "cuda"):
    d = resolve(path)
    _fix_tokenizer_config(d)
    tok = AutoTokenizer.from_pretrained(os.path.join(d, "tokenizer"))
    with open(os.path.join(d, "rl_agent_config.json")) as f:
        cfg = json.load(f)
    model = build_model(cfg, encoder_dir=os.path.join(d, "encoder"))
    model.load_state_dict(load_file(os.path.join(d, "model.safetensors")), strict=True)
    return model.to(device), tok, cfg

# train.py: train Laya on the soft labels (or one-hot labels), one epoch, one GPU
import json
import os
import random
import shutil
import sys

import torch
from laya.common import collate_items
from safetensors.torch import save_file

from data import build
from evaluate import fit_temperature
from load import HEAD_MAX_LEN, MAX_LEN, load, resolve


def batches(items, max_tokens, rng):
    """Length-bucketed batches of about max_tokens padded tokens, in random order."""
    order = sorted(range(len(items)), key=lambda i: (len(items[i]["ids"]), rng.random()))
    out, cur, longest = [], [], 0
    for i in order:
        n = len(items[i]["ids"])
        if cur and max(longest, n) * (len(cur) + 1) > max_tokens:
            out.append(cur)
            cur, longest = [], 0
        cur.append(i)
        longest = max(longest, n)
    if cur:
        out.append(cur)
    rng.shuffle(out)
    return out


def main(logs: str, out_dir: str, test_groups: set[str]):
    torch.manual_seed(7)
    model, tok, cfg = load("base")
    data = build(logs, tok, test_groups)
    items = data["train"]  # repeat a rare question type a few times so it is not drowned out

    model.encoder.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    model.train()
    enc = [p for n, p in model.named_parameters() if n.startswith("encoder.")]
    head = [p for n, p in model.named_parameters() if not n.startswith("encoder.")]
    opt = torch.optim.AdamW([{"params": enc, "lr": 3e-5}, {"params": head, "lr": 1e-4}], weight_decay=0.01, fused=True)
    plan = batches(items, 12288, random.Random(7))  # one epoch; fixed schedule, no early stopping
    warmup = max(1, len(plan) // 50)                # 2% warm-up, then linear decay to 5%
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / warmup) * max(0.05, 1 - s / len(plan)))

    for step, idx in enumerate(plan):
        b = collate_items([[items[i] for i in idx]], tok.pad_token_id)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            logits, act = model(b["input_ids"].cuda(), b["attention_mask"].cuda(),
                                b["marker_pos"].cuda(), b["marker_mask"].cuda(), b["qtype"].cuda())
        mask, target = b["marker_mask"].cuda(), b["target"].cuda()
        logp = torch.log_softmax(logits.float().masked_fill(~mask, -1e4), -1)
        # KL(teacher || student): cross-entropy minus the target's own entropy, so 0 means a perfect copy.
        # 0 * act keeps Laya's unused action head in the graph, as our script does.
        loss = (target * (torch.log(target.clamp_min(1e-9)) - logp) * mask).sum(-1).mean() + 0.0 * act.sum()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        sched.step()
        opt.zero_grad(set_to_none=True)
        if step % 100 == 0:
            print(f"step {step}/{len(plan)}  KL {loss.item():.4f}", flush=True)

    model.eval()
    temps = fit_temperature(model, data["val"], tok.pad_token_id)  # validation only, never test
    os.makedirs(out_dir, exist_ok=True)
    save_file({k: v.detach().to(torch.bfloat16).float().contiguous().cpu() for k, v in model.state_dict().items()},
              os.path.join(out_dir, "model.safetensors"))
    for sub in ("encoder", "tokenizer"):
        shutil.copytree(os.path.join(resolve("base"), sub), os.path.join(out_dir, sub), dirs_exist_ok=True)
    cfg.update({"model_name": "my-laya", "max_len": MAX_LEN, "head_max_len": HEAD_MAX_LEN,
                "temperature": temps, "temperature_by_options": {}})  # drop the base model's own temperatures
    with open(os.path.join(out_dir, "rl_agent_config.json"), "w") as f:
        json.dump(cfg, f, indent=2)
    print("saved", out_dir, "temperatures (choice, score, noul):", temps)


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2], set(sys.argv[3].split(",")))

# evaluate.py: predictions, temperatures (validation only), metrics with intervals over clusters
import numpy as np
import torch
from laya.common import collate_items


@torch.no_grad()
def predict(model, items, pad_id, temperature=None, batch_tokens=16384):
    """Probabilities per item, in input order, batched by length."""
    model.eval()
    order = sorted(range(len(items)), key=lambda i: len(items[i]["ids"]))
    out, i = [None] * len(items), 0
    while i < len(order):
        longest = len(items[order[min(i + 255, len(order) - 1)]]["ids"])
        idx = order[i : i + min(max(1, batch_tokens // longest), 512)]
        b = collate_items([[items[j] for j in idx]], pad_id)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            logits, _ = model(b["input_ids"].cuda(), b["attention_mask"].cuda(),
                              b["marker_pos"].cuda(), b["marker_mask"].cuda(), b["qtype"].cuda())
        logits = logits.float()
        if temperature is not None:
            logits = logits / temperature[b["qtype"].cuda()][:, None]
        p = torch.softmax(logits, -1).cpu().numpy()
        for r, j in enumerate(idx):
            out[j] = p[r, : len(items[j]["markers"])]
        i += len(idx)
    return out


def fit_temperature(model, val_items, pad_id):
    """One temperature per question type (choice, score, noul), minimising cross-entropy against the teacher."""
    by_type = {0: [], 1: [], 2: []}
    for it, p in zip(val_items, predict(model, val_items, pad_id)):
        by_type[it["qtype"]].append((torch.log(torch.tensor(p).clamp_min(1e-9)), torch.tensor(it["target"])))
    temps = [1.0, 1.0, 1.0]
    for qt, sel in by_type.items():
        if len(sel) < 20:
            continue  # too few to fit: leave 1.0
        log_t = torch.zeros(1, requires_grad=True)
        opt = torch.optim.LBFGS([log_t], lr=0.1, max_iter=100)

        def closure():
            opt.zero_grad()
            loss = sum(-(t * torch.log_softmax(z / log_t.exp(), -1)).sum() for z, t in sel) / len(sel)
            loss.backward()
            return loss

        opt.step(closure)
        temps[qt] = float(log_t.detach().exp().clamp(0.2, 5.0))
    return temps


def per_item(items, preds):
    t = [np.asarray(it["target"]) for it in items]
    rows = {
        "tvd": [0.5 * np.abs(p - q).sum() for p, q in zip(preds, t)],       # distance between the two distributions
        "top1": [float(p.argmax() == q.argmax()) for p, q in zip(preds, t)],  # same top answer as the teacher
        "brier": [float(((p - q) ** 2).sum()) for p, q in zip(preds, t)],
    }
    return {k: np.asarray(v) for k, v in rows.items()}


def bootstrap(values, clusters, n=10_000, seed=0):
    """95% interval for the mean, resampling whole clusters (items from one moment are correlated)."""
    keys, inv = np.unique(np.asarray(clusters).astype(str), return_inverse=True)
    sums = np.bincount(inv, weights=values, minlength=len(keys))
    counts = np.bincount(inv, minlength=len(keys)).astype(float)
    rng = np.random.default_rng(seed)
    means = [sums[pick].sum() / counts[pick].sum() for pick in (rng.integers(0, len(keys), len(keys)) for _ in range(n))]
    return [float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5))]


def report(items, preds):
    clusters = [f"{it['group']}:{it['moment']}" for it in items]
    return {k: (float(v.mean()), bootstrap(v, clusters)) for k, v in per_item(items, preds).items()}

# serve.py: the teacher's wire shape, answered by your checkpoint on the GPU
import hashlib
import hmac
import json
import os
import re
from collections import OrderedDict

import torch
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from laya.common import QTYPES, build_sequence, collate_items, render_options

from load import HEAD_MAX_LEN, MAX_LEN, load

MODEL_DIR = os.environ.get("MODEL_DIR", "checkpoints/v1")
SECRET = os.environ.get("LAYA_SECRET")
BATCH_TOKENS = 24576
PATH = re.compile(r"`([A-Za-z_]\w*)\.([^`.\s]+)")  # `collection.key` in a question's instructions, e.g. `tickets.T-4812`

model, tok, cfg = load(MODEL_DIR)
model.eval()
temps = torch.tensor(cfg.get("temperature", [1.0, 1.0, 1.0]), device="cuda")
cache: OrderedDict[str, list[float]] = OrderedDict()  # exact: same (slice, question) -> same answer
app = FastAPI()


def slice_state(state, instructions: str):
    """Cut each collection a question names down to that one entry: the same slicing as the training data."""
    if not isinstance(state, dict):
        return state
    named: dict[str, set[str]] = {}
    for coll, key in PATH.findall(instructions):
        if isinstance(state.get(coll), dict) and key in state[coll]:
            named.setdefault(coll, set()).add(key)
    return {k: ({n: v[n] for n in v if n in named[k]} if k in named else v) for k, v in state.items()}


@torch.no_grad()
def run(batch: list[dict]) -> list[list[float]]:
    b = collate_items([batch], tok.pad_token_id)
    with torch.autocast("cuda", dtype=torch.bfloat16):
        logits, _ = model(b["input_ids"].cuda(), b["attention_mask"].cuda(),
                          b["marker_pos"].cuda(), b["marker_mask"].cuda(), b["qtype"].cuda())
    p = torch.softmax(logits.float() / temps[b["qtype"].cuda()][:, None], -1).cpu()
    return [p[r, : len(it["markers"])].tolist() for r, it in enumerate(batch)]


@app.post("/v1/systemone")
async def systemone(req: Request):
    if SECRET and not hmac.compare_digest(req.headers.get("authorization", "").encode(), f"Bearer {SECRET}".encode()):
        raise HTTPException(401, "unauthorized")
    body = await req.json()
    state, questions = body.get("state"), body.get("questions")
    if not isinstance(questions, dict) or not 0 < len(questions) <= 1000:
        raise HTTPException(400, "questions must be an object with 1 to 1000 entries")
    asked, todo = {}, []
    for qid, q in questions.items():
        crit = q.get("criteria")
        if q.get("type") == "choice" and isinstance(crit, list):
            crit = {c: None for c in crit}
        if q.get("type") not in ("choice", "noul") or (q["type"] == "choice" and not crit):
            raise HTTPException(400, f"{qid}: a choice with criteria, or a noul")
        iq = {"t": q["type"], "ins": str(q["instructions"]), "crit": crit}
        sliced = slice_state(state, iq["ins"])
        key = hashlib.sha256(json.dumps([sliced, iq]).encode()).hexdigest()
        if key not in cache:
            ids, markers = build_sequence(tok, sliced, iq, MAX_LEN, HEAD_MAX_LEN)
            if len(markers) != len(render_options(iq)):
                raise HTTPException(400, f"{qid}: the options do not fit in {HEAD_MAX_LEN} tokens")
            todo.append((key, {"ids": ids, "markers": markers, "qtype": QTYPES[iq["t"]]}))
        asked[qid] = (key, iq)
    todo.sort(key=lambda kv: len(kv[1]["ids"]))  # shortest first, batches of about BATCH_TOKENS padded tokens
    batch: list = []
    for key, it in todo + [(None, None)]:
        if batch and (it is None or len(it["ids"]) * (len(batch) + 1) > BATCH_TOKENS):
            for (k, _), p in zip(batch, run([x for _, x in batch])):
                cache[k] = p
            batch = []
        if it is not None:
            batch.append((key, it))
    answers = {}
    for qid, (key, iq) in asked.items():
        p = cache[key]
        cache.move_to_end(key)
        if iq["t"] == "choice":
            opts = list(iq["crit"])
            answers[qid] = {"type": "choice", "choice": opts[p.index(max(p))], "probabilities": dict(zip(opts, p))}
        else:
            answers[qid] = {"type": "noul", "noul": p[1]}
    while len(cache) > 200_000:
        cache.popitem(last=False)
    return {"model": cfg.get("model_name", "my-laya"), "answers": answers}


if __name__ == "__main__":
    host = os.environ.get("HOST", "127.0.0.1")
    if host not in ("127.0.0.1", "localhost") and not SECRET:
        raise SystemExit("refusing to listen beyond localhost without LAYA_SECRET")
    uvicorn.run(app, host=host, port=int(os.environ.get("PORT", "8765")))