DEV Community

tomek7667
tomek7667

Posted on AI-assisted

Ask ten questions, read the text once: making a decision model 6.7x faster

cbjev vs Laya vs TypeSafe Jev

A lot of "AI" in production isn't generation at all. It's a pile of small decisions about a piece of text: which team should handle this ticket? how urgent is it? is the customer threatening to leave? is this e-mail phishing? Sending each of those to an LLM and parsing the answer is slow, expensive and occasionally creative in ways you didn't ask for.

Laya is an open model for exactly this job: you send a state (a ticket, an e-mail, a JSON document) and typed questions — choice, score, or a yes/no noul — and a ModernBERT encoder returns calibrated probabilities in one forward pass. I built cbjev on top of it. This post is about the one idea that made it fast, how it could reuse Laya's weights, and what didn't work.

The problem: the document is read once per question

Laya builds one sequence per question:

[CLS] question 1 [SEP] options [SEP] document [SEP]
[CLS] question 2 [SEP] options [SEP] document [SEP]
...
Enter fullscreen mode Exit fullscreen mode

Ask ten questions about a 500-token document and the encoder processes the document ten times — about 5,500 tokens for a call whose unique content is about 1,000.

The fix: one row, and an attention mask that does the separating

cbjev packs a whole call into one row:

[CLS] q1 | q2 | ... | q10 | document [SEP]
Enter fullscreen mode Exit fullscreen mode

and shapes the attention mask so that:

  • questions never see each other (each segment attends only to itself and the document),
  • the document reads every question (so its encoding is still question-aware, like in Laya),
  • every question segment restarts its positions right after [CLS], and the document starts after the longest segment.

That last detail is the important one. With a single question, the row is token for token and position for position exactly what Laya sees. So a Laya checkpoint dropped into this layout already works (I measured 0.749 on the 5-question typed-decisions benchmark before any training, against 0.768 in its own layout), and fine-tuning starts from Laya's full ability instead of relearning the task.

My first attempt didn't have this: the document came first and could not read the questions. It was just as fast, but fine-tuning had to rebuild skills Laya already had, and it kept losing on half the benchmarks. Switching to the "shared" layout above is what made the accuracy numbers work.

Making the rest of the call cheap

Once the token count is down, a small call is dominated by overhead, not math:

  • a from-scratch ModernBERT forward (no transformers model at runtime) with bf16 matmuls over an fp32 residual stream,
  • each layer fused with torch.compile (one dynamic-shape compile, a few seconds, cached),
  • the whole forward replayed as a CUDA graph per 32-token shape bucket — one launch instead of ~300,
  • the document tokenized once per call, question text tokenized once and cached.

Numbers

Measured side by side with Laya on one RTX 4090, same cases, through both libraries' public predict APIs:

cbjev Laya (better checkpoint)
10 questions, 500-token document 11.4 ms 75.8 ms
1 question, short ticket 3.0 ms 5.4 ms (TileLang fast path)
mean accuracy, 15 English suites 0.741 0.710
typed-decisions, 2,000 decisions 0.783 0.768
answers that change when options are reordered 0.2 % 7.8 %
MASSIVE intent, 51 languages 0.436 0.401

The option-order number comes from a cheap trick the packed layout makes almost free: every choice and score question is also asked with its options reversed, and the two answers are averaged. That's one extra short segment, not another pass.

What didn't work

  • Adding a prompt-injection dataset to training lowered the prompt-injection benchmark by 10 points. The benchmark is half German and short; the added data taught the model that instruction-looking text is usually benign.
  • Weight soups (averaging several fine-tuned runs, or blending with Laya's original weights) nudged the mean up slightly but never closed the specific gaps.
  • Seven training rounds later, cbjev still trails Laya on four suites: support triage (−4.0), prompt injection (−3.5), DAIR emotion (−2.5) and AG News (−0.8, three cases out of 400). They're listed in the README rather than tuned away.

Try it

pip install "cbjev[serve] @ git+https://github.com/tomek7667/cbjev"
Enter fullscreen mode Exit fullscreen mode
import cbjev

agent = cbjev.load()   # weights download from Hugging Face
res = agent.predict(
    {"body": "Billed twice for March. Refund it today or we cancel."},
    {"team": {"type": "choice", "instructions": "Which team should handle this?",
              "criteria": {"billing": "invoices, refunds", "technical": "bugs", "other": "anything else"}},
     "churn": {"type": "noul", "instructions": "Does the customer threaten to cancel their subscription?"}},
)
Enter fullscreen mode Exit fullscreen mode

It also ships a server that speaks TypeSafe Jev's /v1/systemone wire format, so an existing Jev client can point at it.

Thanks to Convai Innovations for releasing Laya openly; cbjev is fine-tuned from their Apache-2.0 checkpoints. I'd love to hear where it breaks on your data.

Top comments (0)