DEV Community

Take Sunbath
Take Sunbath

Posted on

Not every AI call needs to generate text: ej, an 11 MB model for typed decisions

A lot of "AI" in production pipelines isn't writing anything. It's deciding.

Which queue does this ticket go to? Does this message need a human? How urgent is this request: low, medium or high? The
answer is one item from a list you already know. Yet the usual way to get it is to send the text to a language model,
ask it to reply in JSON, parse the reply, retry when the JSON is broken, and hope the word "high" means the same thing
every time.

That works, but it's an odd fit. You pay for generation you don't need, you get text where you wanted a number, and you
have no clean way to say "only automate this when the model is at least 90% sure."

I built ej to do that one job differently. It's an open model that takes a piece of text plus a set of typed
questions, and returns one probability distribution per question, in one forward pass, without generating a single
token
. The whole model is one 11 MB file that runs on a CPU.

This is version 0.0.1. Below is what it does, how I measured it, and, just as important, where it is weak.

What a "typed decision" looks like

You give ej a state (free text, or a JSON object written as text) and any number of questions. Each question has one
of three types:

Type You give You get
choice instructions + a list of option texts, chosen when you call it a probability for each option
noul a yes/no statement P(false), P(true)
score instructions + ordered levels a distribution over the levels

The options are plain text you write at call time. There is no fixed label set baked in at training.

import ej

model = ej.load("5ak3t/ej", revision="v0.0.1")   # downloads model.ejpack and checks its SHA-256

record = {
    "state": '{"customer_tier": "gold", "message": "The blender arrived with a cracked jug. Replace it before Friday."}',
    "questions": {
        "route": {"type": "choice", "instructions": "Which team should handle this request?",
                  "options": [{"key": "returns", "text": "Returns and replacements for damaged or wrong items"},
                              {"key": "billing", "text": "Billing, invoices and payment problems"}]},
        "needs_human": {"type": "noul",
                        "instructions": "The customer is upset enough that a human agent should reply.",
                        "options": ej.NOUL_OPTIONS},
        "urgency": {"type": "score", "instructions": "How urgent is this request?",
                    "options": [{"key": "0", "text": "Low: can wait a week"},
                                {"key": "1", "text": "Medium: answer within two days"},
                                {"key": "2", "text": "High: answer today"}]},
    },
}

(probs,) = model.predict([record])
print(probs["route"])        # [p_returns, p_billing], sums to 1
print(probs["needs_human"])  # [p_false, p_true]
print(probs["urgency"])      # [p_low, p_medium, p_high]
Enter fullscreen mode Exit fullscreen mode

All three questions are answered in the same pass. Because the output is a distribution, the routing logic in your code
is a threshold, not a string parser:

if max(probs["route"]) >= threshold:
    auto_route(record, probs["route"])
else:
    send_to_human(record)
Enter fullscreen mode Exit fullscreen mode

(Pick that threshold on labelled examples of your workflow. More on why below.)

Why not just use an LLM, or an existing classifier?

There are three common ways to make this kind of decision today. ej is a different trade-off from each, not a strict
upgrade.

Prompting an LLM. It's flexible and strong on new tasks. But it generates text you have to parse, it needs either an
API call (cost, latency, your data leaving your servers) or a local model that is hundreds of MB to several GB, and its
reply isn't a probability you can threshold. ej skips generation entirely and is small enough to ship inside a service.

NLI zero-shot classifiers (the "zero-shot-classification" pipeline). These are the closest relatives. They usually
run one forward pass per (text, option) pair, so cost grows with the number of options, and each question type needs its
own setup. ej reads the text once and answers choice, yes/no and score questions together. To be clear, though, a good
NLI model beats ej on workflows ej has never seen (numbers below).

Fine-tuning your own classifier (BERT, SetFit and friends). It's accurate on its own task, but the label set is fixed
at training, and every change means new labelled data and a retrain. With ej the options are input. It also has an
adapt method that takes a handful of labelled records:

adapted = model.adapt(examples=labelled[:8])   # per-option offsets; the weights don't change
probs = adapted.predict(new_records)
Enter fullscreen mode Exit fullscreen mode

How it fits in 11 MB

Briefly (a technical report with the full method is coming):

  • The text encoder is a distilled version of e5-small-v2, quantised to 2-bit weights and 3-bit attention, with a trimmed vocabulary.
  • On top sit a set of small expert heads, combined per question type, plus a selective head that knows when it isn't sure.
  • Bigger "teacher" models were used only during training and are distilled into those heads. They don't ship.
  • Everything, including the tokenizer, is packed into one file, model.ejpack: 11,384,312 bytes, bit-packed codes with a SHA-256 check on every section. It contains no pickles, so loading it never executes code from the file.

On the test box (a shared 4-core Xeon, one thread), a warm call takes roughly 200 to 360 ms per record, depending on
the run and the inputs, and the process peaks at about 558 MB of RAM. Most of that is torch itself: the model's own tensors peak at 128 MB, or 35 MB with
ej.load(..., low_memory=True), which is about 1.5x slower and gives identical predictions.

Results, including the bad ones

I tried hard to make these numbers trustworthy rather than flattering:

  • Every final test suite was scored once, for this release. Development suites were used to choose between candidates, so their numbers are optimistic, and they aren't reported as results.
  • Every number has a 95% confidence interval, and I only claim "A beats B" when the paired interval excludes zero.
  • Six other models were benchmarked on the same blind inputs, each with a flag for whether it might have seen the test data in training.
Test suite ej 0.0.1 accuracy [95% CI]
Typed Decisions, workflows seen in training .721 [.695, .746]
MASSIVE intents (leak-free split) .832 [.808, .857]
Support tickets (in-house) .712 [.679, .746]
Support tickets, new writing styles .683 [.639, .726]
105 workflows never seen in training .419 [.379, .458]
One held-out Typed Decisions workflow (also used during model selection) .422 [.388, .452]

How that compares with the six rivals:

  • Where ej is good: on the two ticket suites, ej's log loss (how good its probabilities are, not just its top pick) is lower than every rival's. Bear in mind that ej was trained on tickets from the same synthetic source, and the rivals weren't. On seen Typed Decisions workflows it is level with or above five of the six rivals.
  • Where ej is behind: on MASSIVE it trails four of the six. On the held-out Typed Decisions workflow it is below all six.
  • The big caveat: on workflows it has never seen, ej gets about 42% right. That's the number to expect on a brand-new task before you adapt it. During development, a 71M-parameter NLI model scored about six points higher on this kind of task.

A few more things you should know before using it:

  • Calibration is measured per suite, not promised. The probabilities were well calibrated on the two ticket suites and not on the other four. Check on your own data before trusting "0.9" to mean 90%.
  • Adaptation is modest. Eight labelled records lift accuracy on unseen workflows by about three points, and the interval includes zero. Mostly what it learns is how common each label is in your workflow.
  • The ticket data is synthetic. The in-house tickets were written by Claude Haiku, not by real customers, and that corpus isn't released.
  • It's English only.
  • Small differences aren't meaningful. Six training runs of the identical recipe differ by several points on some suites, so I don't read anything into differences that size.

Who should try it

ej makes sense if all of these are true:

  1. You make the same kinds of decisions over and over on text records: routing, triage, yes/no checks, priority levels.
  2. You'd rather get a probability you can set a threshold on than a string you have to parse.
  3. You need it to run on CPU, on your own servers, or offline, where hosting or calling an LLM is awkward.
  4. You can label a few dozen of your own records to check accuracy and pick a threshold before automating anything.

If your task is brand-new and you have no labelled data at all, a prompted LLM or an NLI zero-shot model will probably
serve you better today. I'd rather say that here than have you find out in production.

Try it, and break it

pip install ejai
Enter fullscreen mode Exit fullscreen mode
import ej
model = ej.load("5ak3t/ej", revision="v0.0.1")
Enter fullscreen mode Exit fullscreen mode

The repository also contains the training code (python -m ej.train), the evaluator that produced every number above
(python -m ej.eval), and the benchmark runner with the rival adapters, so you can check my numbers or train your own
pack.

If you run it on your own workflow, I'd love to hear what accuracy you get, good or bad. Results on real workflows are
exactly what version 0.0.2 needs.

Top comments (0)