--- license: apache-2.0 language: - en base_model: answerdotai/ModernBERT-base datasets: - danielrosehill/Agent-Task-Break-Detector pipeline_tag: text-classification library_name: transformers tags: - agents - context-management - topic-shift - cross-encoder - onnx --- # Agent Task Break Detector (ModernBERT-base) A small cross-encoder that answers one question on every user turn of an AI agent session: > **Does this new message need the conversation so far, or is it a new task that a fresh > session could handle just as well?** It outputs `p(new_task)`. A harness hook can use it to start a fresh session (`/clear`, `/new`) automatically when the user drops an unrelated task into a long session. It is meant to run **locally on every turn**, so it is a 150M-parameter encoder rather than an API call. **Status: v0.1, 2026-10-24. Experimental.** Trained only on synthetic data. It has not yet been evaluated on real, hand-labelled sessions. ## Files | Path | What | |---|---| | `model.safetensors`, `config.json`, `tokenizer*` | Fine-tuned weights (fp32), for `transformers` | | `new_task_split.json` | Recommended int8 threshold and input settings (read by the New-Task-Split plugin) | | `onnx/model_int8.onnx` + tokenizer | Dynamic int8 quantisation (onnxruntime `quantize_dynamic`), 151 MB, for CPU inference | ## Input format The model was trained on pairs `(message, context)`: - **First segment:** the new user message. - **Second segment:** the last 4 exchanges, **newest first**, each written as `AGENT: \nUSER: `, joined by newlines. - `truncation="only_second"`, `max_length=512`. Label 1 = `new_task`, 0 = `continue`. ```python from transformers import AutoTokenizer, AutoModelForSequenceClassification import torch repo = "danielrosehill/Agent-Task-Break-Detector-ModernBERT" tok = AutoTokenizer.from_pretrained(repo) model = AutoModelForSequenceClassification.from_pretrained(repo).eval() turns = [{"user": "the payments migration fails on the currency column", "agent": "Added a default of 'USD' in migration 0042 and re-ran it; it passes."}] ctx = "\n".join(f"AGENT: {t['agent']}\nUSER: {t['user']}" for t in turns[-4:][::-1]) enc = tok("now backfill the currency column for old rows", ctx, truncation="only_second", max_length=512, return_tensors="pt") with torch.no_grad(): p_new_task = model(**enc).logits.softmax(-1)[0, 1].item() ``` ONNX (inputs `input_ids`, `attention_mask`; output `logits`): ```python import onnxruntime as ort, numpy as np from transformers import AutoTokenizer from huggingface_hub import hf_hub_download repo = "danielrosehill/Agent-Task-Break-Detector-ModernBERT" tok = AutoTokenizer.from_pretrained(repo, subfolder="onnx") sess = ort.InferenceSession(hf_hub_download(repo, "onnx/model_int8.onnx")) enc = tok(message, ctx, truncation="only_second", max_length=512, return_tensors="np") logits = sess.run(None, {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]})[0] p_new_task = float(np.exp(logits[0, 1]) / np.exp(logits[0]).sum()) ``` ## Training - Base: `answerdotai/ModernBERT-base`. - Data: [danielrosehill/Agent-Task-Break-Detector](https://huggingface.co/datasets/danielrosehill/Agent-Task-Break-Detector) v0.2, rows with `agreed == true` (5,104 train rows, including 1,025 matched pairs). - 3 epochs, lr 5e-5, batch 16, bf16, warmup 10%, best epoch by validation AUC-PR. One NVIDIA L4 on Modal, 6.9 minutes. Run `atbd-r3-modernbert-full3ep`. - Code: [danielrosehill/Agent-Context-Harness](https://github.com/danielrosehill/Agent-Context-Harness), `model/train.py`. ## Results (synthetic test split, fp32) n = 608 agreed rows, 115 matched pairs. A matched pair is the same message once as a new task and once, in a different context, as a continuation, so only the context can separate them. | | Accuracy | AUC-PR | `new_task` recall @ precision 0.95 | Pair accuracy | |---|---|---|---|---| | **This model** | 0.951 | 0.986 | 0.902 | 0.791 | | TF-IDF, message only | 0.793 | 0.758 | 0.037 | — | | TF-IDF, message + context | 0.760 | 0.781 | 0.140 | — | The weakest cases are a message that looks self-contained but continues the work (`counterfactual_continue`, 0.896) and a different task in the same workspace (`same_workspace`, 0.867). **int8 ONNX** (validation n = 697, test n = 608; tokenised with `tokenizers`, identical ids to training): | | Accuracy | AUC-PR | `new_task` recall @ precision 0.95 | |---|---|---|---| | validation | 0.937 | 0.988 | 0.925 | | test | 0.929 | 0.982 | 0.888 | int8 probabilities sit lower than fp32 (mean shift 0.14), so **use threshold 0.76 for int8**, not the fp32 value. At the validation threshold of 0.751, test precision is 0.941 and recall 0.893. This is recorded in `new_task_split.json`. Session load takes about 1.2 s on CPU; per-call latency has not been measured yet. ## Intended use and limitations - The operating point should be **high precision on `new_task`**: a wrong split throws away the agent's working context, while a missed split only costs extra context. Pick a threshold for precision ≥ 0.95 on your own data. - Synthetic training data from one generator (Claude Haiku 5.5), verified by Claude Sonnet 5.5. Real sessions are longer and messier. English only. - The training data is about 50/50; in real use most turns are continuations. ## Licence Apache-2.0 (as the base model). Training data CC BY 4.0.