agent-task-classifier

(laya loop-taxonomy v1)

A fine-tuned laya checkpoint that classifies a user request to a personal-assistant agent by the agent-loop shape it demands — the skeleton of the high-level process required to fulfil it — rather than by topic.

Same interface as every laya checkpoint: text in, one typed choice decision out, one forward pass, no generation. English, ~0.3 s/query on CPU.

The eleven labels

label the loop it demands
conversation answer/act directly — the default, safe fallback
diagnose find the cause of a specific malfunction or anomaly
research assemble understanding from investigation (multi-source, synthesis)
arrange commit a real-world arrangement — book, reserve, schedule, hire
acquire obtain a purchasable good, and verify it arrived
execute complete a known, externally-defined procedure (renew, file, claim)
make produce an artifact judged against a standard (tests, spec, taste)
tend review/maintain a standing concern (budget, health, home, wardrobe)
learn build the user's own ability over sessions (practice, tutoring)
companion inner-life support with no outcome in the world (reflection, therapy)
react urgent/deteriorating situation — act before full understanding

Boundary rule: two requests share a label iff they demand the same loop skeleton — same phases, same iteration driver, same exit criterion. Topic is irrelevant (conversation spans finance, wardrobe, memories; make spans code, letters, meal plans).

Usage

import laya

agent = laya.load("kumarabd1996/agent-task-classifier", device="cpu")
questions = {
    "loop": {
        "type": "choice",
        "instructions": "What does this request (or the request it continues) demand?",
        "criteria": { ... },  # shipped as questions.json in this repo
    }
}
out = agent.predict("the user's message", questions)
out["answers"]["loop"]["choice"]         # e.g. "arrange"
out["answers"]["loop"]["probabilities"]  # calibrated (temperatures fitted at train time)

Continuations: pass the request being continued in the state, the format the model was trained on:

Earlier request: What are my biggest categories?
New request: Continue the bulk update

Serving

A ready-made HTTP image (CPU, weights baked in) is published from realnighthawk/classifier:

docker run -p 8000:8000 ghcr.io/realnighthawk/classifier:latest
curl localhost:8000/classify -H 'content-type: application/json' \
  -d '{"text":"book me a table for two at 8pm friday"}'
# {"label":"arrange","probabilities":{...}}

POST /classify takes {text, earlier?}; earlier is the request being continued.

Training

  • Base: convaiinnovations/laya (English checkpoint; ModernBERT-large + decision head), Apache-2.0.
  • Corpus: 5,280 items, 8 sources — MultiWOZ 2.1 booking utterances (990, arrange), MASSIVE en-US slices (920, acquire/arrange/conversation), OASST2 prompts + continuations (1,060, teacher-labeled), synthetic items (2,315) covering all eleven shapes (labels by construction, incl. 34 hand-written boundary cases), and 1,056 continuation-form items (Earlier request: … / New request: …). Per-class caps keep the largest class ≤ ~28%. Released as a dataset: kumarabd1996/agent-task-classifier-dataset.
  • Hyperparameters: 4 epochs, effective batch 64 (micro 4 × accum 16), encoder lr 2.5e-5 / head lr 1e-4, RLCD objective (default), option-order augmentation, max-len 256.
  • Hardware: one Mac, CPU only (~2 hours). A frozen-encoder variant failed to learn the taxonomy (37% in-distribution); encoder adaptation is required — matching laya's own findings.
  • Temperatures fitted automatically on a held-out calibration slice (choice temperature ≈ 3.82).

Evaluation

Boundary drills: 24 held-out hand-written cases targeting the taxonomy's hard boundaries (info-vs-procedure, understanding-vs-commitment, urgency-vs-worry, artifact-vs-investigation).

accuracy ECE mean confidence
base convaiinnovations/laya (zero-shot) 0.29 (7/24) 0.54 0.83 (miscalibrated)
this checkpoint 0.83 (20/24) 0.13 0.96

The 4 remaining misses, all at genuine boundaries: "Is Bali worth visiting in December?" (research, predicted arrange), "What does renewing a passport require?" (conversation, predicted execute), "Why have I been so tired lately?" (diagnose, predicted conversation), "Sign me up for the best EV lease deal you find." (arrange, predicted conversation).

Limitations

  • Trained largely on synthetic + public chat corpora; real user traffic is under-represented.
  • English only (the base multilingual checkpoint was not used).
  • The eleven labels are fixed; requests outside them resolve to conversation by design.
  • Continuations require the caller to supply the continued request in the state.
  • Not a safety classifier — do not use for guardrails or moderation.

Related work

Small-model routing exists on two other axes, both public on the Hub:

  • Capability buckets — e.g. ai-systems/agent-task-classifier, a TF-IDF reference model over tool/subsystem categories (browser, retrieval, verification, …). It answers which capability a task needs.
  • Model-choice routing — e.g. subhash4face/agentic-ai-routing-task-classifier, a DeBERTa fine-tune over LLM task types (QA, coding, summarization, long-context, …). It answers which model should handle a prompt.

This model answers a third question: which execution loop the task demands — the only one of the three that decides how the agent's own work should proceed (phases, iteration driver, exit criteria) rather than which resource to engage. Capability and loop are orthogonal — a diagnose task may need retrieval tools, a conversation task may need none — and in an agent harness the loop decision is the routing decision that determines everything downstream.

Training data provenance

MultiWOZ 2.1 (research use, MIT-licensed distribution) · MASSIVE en-US (CC-BY-4.0, Amazon) · OASST2 (Apache-2.0, OpenAssistant) · synthetic/generated items (Apache-2.0 with this release). No user/private data was used in training.

Files

model.safetensors (weights, incl. encoder) · encoder/config.json · tokenizer/ · questions.json (the taxonomy question, ready to pass to predict) · rl_agent_config.json (fitted temperatures, training config) · train_report.json (before/after metrics).

Downloads last month
21
Safetensors
Model size
0.4B params
Tensor type
F16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for kumarabd1996/agent-task-classifier

Finetuned
(165)
this model