assay-4b

assay-4b

Calibrated typed decisions from one forward pass. Send a state and named typed questions (bool yes/no, choice over 2..255 described options, score over 2..10 ordered levels); get a probability distribution per question, a confidence and an evidence score. No text is generated, so nothing can come back off-schema.

Code, server and training recipe: https://github.com/bgokden/assay

In one line

It abstains. On tasks it has never seen, this model acts on 75% of bool questions at a 10.2% error rate, 74% of choice at 8.0%, and hands the rest over instead of guessing. That is a conformal prediction set against a fitted error rate, not a threshold on a softmax someone eyeballed.

Accuracy on those unseen tasks is 0.803. Every answer carries a calibrated probability, a confidence, and the set of options that could not be ruled out.

The other five models, and the dataset, are in the Assay collection.

How it is built

• Backbone Qwen/Qwen3-4B-Base with a LoRA adapter (r=16, alpha=32, lr=5e-05, 1.0 epoch, batch 8 x 1 accumulation); merged weights are in this repository, the adapter is in adapter/. The merged weights load with transformers like any Qwen checkpoint.
• The answer is read from the model's own next-token logits over option label tokens at a single decision position, so the base model's zero-shot competence is the starting point.
• Questions are isolated branches over a shared state (block attention mask, restarted positions): packed and separate requests agree exactly.
• Trained with cross-entropy against soft targets: human label distributions where the source has them, SORD-smoothed levels for ordinal questions, one-hot otherwise. Choice options are shuffled per example.
• An evidence head (linear on the decision token, assay_head.safetensors) predicts whether the state supports the question, trained on passage-swapped negatives.
• Global temperature 1.263 fitted on the calibration split of the training tasks and applied unchanged everywhere else.

Evaluation

"Unseen tasks" are eleven datasets never used in training (bbc_news, app_reviews, scitail, medical_questions_pairs, tweet_irony, ethos, stance_climate, dream, copa, truthful_qa, hh_rlhf). "kev transfer-v4 dev" is the public suite from jaredpalmer/kev-suites (mmlu, emotion, sciq, tweet_offensive, qnli, paws and synthetic rule holdouts); none of its sources are in the training data. Brier is the multi-class sum of squared errors (0..2), ECE uses 15 bins, confident errors are answers with p >= 0.9 that are wrong.

Latency on one RTX 5090 (bf16, transformers, packed questions over one state versus separate requests):

questions  packed_ms  separate_ms
        1       22.9         22.6
        3       22.9         68.2
        6       33.3        136.0
       12       33.0        271.8
       24       56.5        543.7

Abstention

conformal.json holds per-question-type thresholds fitted on the seen-task calibration split (alpha 0.1, delta 0.05): a prediction-set threshold with coverage at least 1 - alpha and an act threshold on the top probability whose acted-on error rate is at most alpha at confidence 1 - delta, both on inputs distributed like the calibration split. assay.server returns them as act and set on every answer. How they carry over to other tasks:

Score questions get prediction sets but usually no act threshold, because exact-level accuracy is the wrong error notion for ordinal answers. Refit on your own labelled data with python -m assay.conformal for a guarantee about your distribution.

The family

Accuracy / Brier after temperature scaling. Same recipe, same splits, different backbones; per-tier abstention and latency are in docs/models.md.

Serving

python -m assay.server --model Berk/assay-4b --port 8000

POST /v1/decide is the native shape; POST /v1/systemone and /v1/systemone/batch accept the shape other open decision models use (criteria options, noul booleans); POST /v1/decide_graph walks a decision tree in one forward pass; POST /v1/agents registers an agent -- a graph plus the actions its outcomes stand for -- and /v1/agents/<name>/run decides a case. Requests arriving together share a pass, and /health and /metrics are for operations. assay.backends.sglang runs the same model on an SGLang deployment.

Guides: deployment and the full API, agents, runnable examples.

Train one on your own data

python -m assay.pipeline --config <your>.json runs training, temperature calibration, evaluation and the conformal thresholds over your own records, and writes a directory this same server and publisher accept. The repository's examples/ has a configuration per tier and a dataset in the record format; records written for other decision models (criteria options, noul booleans) load unchanged.

Usage

from assay import load_model
from assay.schema import Question

model = load_model("Berk/assay-4b")
answers = model.answer(
    state="My card was charged twice for order A-104.",
    questions={
        "refund": Question(type="bool", instructions="Does the customer ask for money back?"),
        "team": Question(type="choice", instructions="Which team should handle this?",
                         options={"billing": "Charges and refunds", "technical": "Bugs"}),
    },
)
print(answers["team"].probabilities, answers["refund"].p_true, answers["refund"].evidence)

Limitations

Text only, English training data. No arithmetic, counting, date comparison or multi-hop reasoning in one pass; keep those in code. Accuracy drops with unrelated state. The evidence head is trained on coarse swapped-passage negatives. Probabilities are calibrated in aggregate on the evaluated distributions, which is not a guarantee about any single answer or about your data; check calibration on your own labels before acting on thresholds.

Training data

Fifty-five public classification, inference, reading-comprehension and preference datasets rendered as typed questions with described options, plus a synthetic policy-application generator (see assay/data/tasks.py in the repository for the rubrics). Each dataset keeps its own licence; the per-dataset list is in docs/datasets.md. Several sources carry non-commercial or research-only terms; check them before commercial use.

Relationship to other work

Assay is an independent project. Jev and System One are names of TypeSafe AI's products and are mentioned only to describe and compare; kev-suites is Jared Palmer's evaluation data. Assay is not affiliated with or endorsed by either.