From d2ae81c63190f2a618a8773b5e4beaf7fbd098c4 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 07:44:23 +0000
Subject: [PATCH 01/16] Modernize packaging and add deployment guide

Packaging:
- Move all metadata into a PEP 621 [project] table in pyproject.toml and
  drop setup.py and requirements.txt; version is read from
  adaptive_classifier.__version__ so it lives in one place
- Require Python >=3.10 (3.8/3.9 are EOL) and add 3.11-3.13 classifiers
- Drop unused core dependencies (tqdm, setuptools, wheel); tqdm and the
  other script-only packages move to a `scripts` extra
- CI: test on Python 3.10-3.13, use actions v4/v5 with pip caching,
  install via the `test` extra, build releases with `python -m build`

Docs:
- Add docs/deployment.md: multilingual base models, the saved file
  layout, the exact inference algorithm, and C#/JS notes
- Add examples/portable_inference.py, a NumPy + onnxruntime reference
  that reproduces predict() without PyTorch or FAISS; verified against
  AdaptiveClassifier.predict for mean and cls pooling

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 .github/workflows/publish.yml  |   7 +-
 .github/workflows/test.yml     |  26 +++---
 CONTRIBUTING.md                |   3 +-
 README.md                      |  11 +--
 docs/deployment.md             | 149 +++++++++++++++++++++++++++++++++
 examples/portable_inference.py | 147 ++++++++++++++++++++++++++++++++
 pyproject.toml                 |  69 +++++++++++++--
 requirements.txt               |  11 ---
 setup.py                       |  47 -----------
 9 files changed, 383 insertions(+), 87 deletions(-)
 create mode 100644 docs/deployment.md
 create mode 100644 examples/portable_inference.py
 delete mode 100644 requirements.txt
 delete mode 100644 setup.py

diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index 6a27b8d..19e2c1f 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -21,16 +21,15 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install setuptools wheel
+          pip install build
       - name: Run tests
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          pip install pytest pytest-cov pytest-randomly psutil
-          pip install -e .
+          pip install -e ".[test]"
           pytest tests/
       - name: Build package
         run: |
-          python setup.py sdist bdist_wheel
+          python -m build
       - name: Publish package distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 8193114..1deaf4a 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -9,20 +9,24 @@ on:
 jobs:
   unit-tests:
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13']
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
 
-    - name: Set up Python 3.12
-      uses: actions/setup-python@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
       with:
-        python-version: '3.12'
+        python-version: ${{ matrix.python-version }}
+        cache: pip
 
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -e .
-        pip install pytest pytest-cov psutil
+        pip install -e ".[test]"
 
     - name: Run unit tests
       env:
@@ -35,7 +39,7 @@ jobs:
       with:
         file: ./coverage.xml
         flags: unittests
-        name: codecov-umbrella
+        name: codecov-${{ matrix.python-version }}
         fail_ci_if_error: false
 
   integration-tests:
@@ -44,18 +48,18 @@ jobs:
     timeout-minutes: 30  # Generous timeout for model downloads
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
 
     - name: Set up Python 3.12
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: '3.12'
+        cache: pip
 
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -e .
-        pip install pytest psutil
+        pip install -e ".[test]"
 
     - name: Run integration tests
       env:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index d104beb..3f95e1f 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -61,8 +61,7 @@ source venv/bin/activate  # On Windows: venv\Scripts\activate
 
 3. Install dependencies:
 ```bash
-pip install -r requirements.txt
-pip install -e .
+pip install -e ".[test]"
 ```
 
 4. Run tests:
diff --git a/README.md b/README.md
index d884410..527c24b 100644
--- a/README.md
+++ b/README.md
@@ -108,11 +108,8 @@ pip install adaptive-classifier
 git clone https://github.com/codelion/adaptive-classifier.git
 cd adaptive-classifier
 
-# Install in development mode
-pip install -e .
-
-# Install test dependencies (optional)
-pip install pytest pytest-cov pytest-randomly
+# Install in development mode with test dependencies
+pip install -e ".[test]"
 ```
 
 ---
@@ -315,6 +312,10 @@ ONNX:     2.1ms/query  (4.0x faster) ✓
 
 > **Note:** ONNX optimization is included by default. For GPU inference, PyTorch is automatically used for best performance.
 
+#### Multilingual models and other runtimes (.NET, JavaScript, ...)
+
+Any multilingual encoder works as the base model, and a saved classifier can be run from C#, Node or the browser through an ONNX runtime. See the [deployment guide](docs/deployment.md) and the dependency-light reference implementation in [`examples/portable_inference.py`](examples/portable_inference.py).
+
 ## Advanced Usage
 
 ### Adding New Classes Dynamically
diff --git a/docs/deployment.md b/docs/deployment.md
new file mode 100644
index 0000000..85adb86
--- /dev/null
+++ b/docs/deployment.md
@@ -0,0 +1,149 @@
+# Deployment guide: multilingual models and non-Python runtimes
+
+This guide covers two questions that come up often:
+
+1. [Can I use it for languages other than English?](#multilingual-classification)
+2. [Can I train in Python and run inference in C#, JavaScript or anything else with an ONNX runtime?](#running-outside-python)
+
+Short answers: yes, and yes. You train and update in Python. Inference needs only the encoder in ONNX plus a few lines of arithmetic, which this guide spells out and `examples/portable_inference.py` implements.
+
+---
+
+## Multilingual classification
+
+Nothing in the library is English-specific. The language support comes entirely from the encoder you pick, so swap the base model:
+
+```python
+from adaptive_classifier import AdaptiveClassifier
+
+classifier = AdaptiveClassifier("intfloat/multilingual-e5-small")
+
+classifier.add_examples(
+    ["Das Produkt ist großartig", "Producto terrible, no lo recomiendo", "商品很好用"],
+    ["positive", "negative", "positive"],
+)
+classifier.predict("Muy buen servicio")
+```
+
+Encoders worth trying (all published on the Hugging Face Hub):
+
+| Model | Notes |
+|---|---|
+| `intfloat/multilingual-e5-small` / `-base` | Small and fast, about 100 languages. Mean pooling. |
+| `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2` | Small, trained for sentence similarity. Mean pooling. |
+| `BAAI/bge-m3` | Larger and stronger, long inputs, 100+ languages. CLS pooling. |
+| `xlm-roberta-base` | Multilingual, but not trained for sentence similarity; expect weaker prototypes than the models above. |
+
+Things that matter:
+
+- **Pooling.** The default `pooling: "auto"` reads the model's own sentence-transformers config, so the table above is handled for you in Python. See the "Pooling" section of the README.
+- **Prefixes.** Some models expect a prefix on every input (E5 wants `"query: "`). Add it yourself to training and prediction texts, consistently.
+- **Mixed-language data works.** Good multilingual encoders place translations of the same sentence close together, so a class trained mostly on English examples will often recognise other languages. Add a few examples in each language you care about to firm it up.
+- **Always evaluate on your own data.** Multilingual quality varies a lot by language and domain; the library does not ship per-language benchmarks.
+
+---
+
+## Running outside Python
+
+### What gets saved
+
+`classifier.save("./my_classifier")` (or `push_to_hub`) writes:
+
+| File | Contents |
+|---|---|
+| `config.json` | Label maps (`id_to_label`), per-class `training_history`, and the `config` block: `pooling`, `max_length`, blend weights |
+| `model.safetensors` | `prototype_<label>` tensors (one vector per class) and `adaptive_head_model.{0,3,6}.{weight,bias}` for the neural head |
+| `examples.json` | A few representative training examples per class. Not needed for inference |
+| `onnx/model.onnx`, `onnx/model_quantized.onnx` | The transformer encoder only. It returns token embeddings, not a sentence vector or class scores |
+
+The encoder is the only large file. Everything else is tiny and framework-neutral (JSON plus safetensors, which has readers in most languages).
+
+Two files are **not** written and you must provide them:
+
+- **`tokenizer.json`**. Create it once from the base model:
+  ```python
+  from transformers import AutoTokenizer
+  AutoTokenizer.from_pretrained("<base model name>").save_pretrained("./my_classifier")
+  ```
+- **A fixed pooling mode**, if `config.json` says `"pooling": "auto"`. Auto-detection happens in Python at run time by reading the base model's config, so the saved value is the literal string `auto`. Look up what your base model uses (the table above, or its `1_Pooling/config.json`) and hard-code it. Classifiers saved before 0.2.0 use `cls`.
+
+### The inference algorithm
+
+Given the files above, `predict(text)` is:
+
+1. **Tokenize** with the base model's tokenizer, truncating to `config.max_length` (default 512).
+2. **Run the ONNX encoder** to get `last_hidden_state`, shape `(1, tokens, dim)`. Feed `input_ids` and `attention_mask`, plus `token_type_ids` if the graph declares that input (BERT-family models do).
+3. **Pool** to one vector:
+   - `mean`: sum of token vectors weighted by `attention_mask`, divided by the number of real tokens.
+   - `cls`: the first token's vector.
+4. **L2-normalise** the vector.
+5. **Prototype scores.** For each class prototype `p`, compute `exp(-||p - v||²)` (squared Euclidean distance), then take a softmax across classes.
+6. **Neural head scores** (skip if `model.safetensors` has no head). Compute `h1 = ReLU(W0·v + b0)`, `h2 = ReLU(W3·h1 + b3)`, `logits = W6·h2 + b6`, then softmax. Dropout is inactive at inference. Weights are stored as `(out, in)`.
+7. **Blend per class.** `score = proto_weight · proto_score + neural_weight · head_score`. A class with fewer than `new_class_example_threshold` (default 10) entries in `training_history` uses the `new_class_*` weights (default 0.3 / 0.7) instead of `prototype_weight` / `neural_weight` (default 0.7 / 0.3).
+8. **Normalise** the blended scores to sum to 1 and sort descending.
+
+This applies to the default (non-strategic) mode. Strategic mode adds extra steps that are not covered here.
+
+`examples/portable_inference.py` implements exactly these steps with NumPy, `onnxruntime` and `tokenizers`, with no PyTorch or FAISS:
+
+```bash
+pip install onnxruntime numpy safetensors tokenizers
+python examples/portable_inference.py ./my_classifier "text to classify"
+```
+
+It was checked against `AdaptiveClassifier.predict` on the same saved directory with both `mean` and `cls` pooling, and agrees to about 1e-7 (float32 rounding). Use it as the ground truth when you port: run the same inputs through both and compare scores.
+
+### Quantized or full-precision encoder
+
+`onnx/model_quantized.onnx` is INT8 and about 4x smaller. Its scores differ slightly from the full-precision graph, so check that your class ranking survives on your own data. `onnx/model.onnx` is the full-precision graph.
+
+### .NET (C#)
+
+Load the encoder with the `Microsoft.ML.OnnxRuntime` NuGet package, and tokenize with `Microsoft.ML.Tokenizers` (or any library that can read your model's `vocab.txt` or `tokenizer.json`).
+
+```csharp
+using var session = new InferenceSession("my_classifier/onnx/model_quantized.onnx");
+var inputs = new List<NamedOnnxValue> {
+    NamedOnnxValue.CreateFromTensor("input_ids",      new DenseTensor<long>(ids,  new[] { 1, ids.Length })),
+    NamedOnnxValue.CreateFromTensor("attention_mask", new DenseTensor<long>(mask, new[] { 1, ids.Length })),
+    // add "token_type_ids" (all zeros for single sentences) if session.InputMetadata contains it
+};
+using var results = session.Run(inputs);
+var hidden = results.First().AsTensor<float>();   // (1, tokens, dim)
+// then pool, normalise, and score as in steps 3-8 above
+```
+
+The scoring is a few dozen lines of array maths. Read `model.safetensors` with a safetensors reader (the format is an 8-byte header length, a JSON header, then raw little-endian float32 data), or export the prototypes and head weights to JSON once in Python and ship that.
+
+### JavaScript / Node / browser
+
+Use `onnxruntime-node` (server) or `onnxruntime-web` (browser) for the encoder and `@huggingface/transformers` for tokenization.
+
+```js
+import * as ort from "onnxruntime-node";
+const session = await ort.InferenceSession.create("my_classifier/onnx/model_quantized.onnx");
+const out = await session.run({
+  input_ids:      new ort.Tensor("int64", BigInt64Array.from(ids),  [1, ids.length]),
+  attention_mask: new ort.Tensor("int64", BigInt64Array.from(mask), [1, ids.length]),
+  // token_type_ids too if session.inputNames includes it
+});
+const hidden = out.last_hidden_state;   // Float32Array data, dims [1, tokens, dim]
+// then pool, normalise, and score as in steps 3-8 above
+```
+
+For the browser, check the size of `model_quantized.onnx` first: multilingual encoders carry a large vocabulary embedding and can be well over 100 MB even when quantized.
+
+The C# and JavaScript snippets above are sketches of the API calls and have **not** been run; the Python reference script is the tested artifact.
+
+### Updating a deployed model
+
+Inference outside Python is read-only. The workflow that works well:
+
+1. Keep the Python classifier as the source of truth. Add examples and classes there (`add_examples`) as new data arrives.
+2. `save()` again and ship the updated `config.json` and `model.safetensors`. If the encoder is unchanged, `onnx/` does not need to be redistributed.
+3. Reload them in the non-Python service.
+
+### Known limitations
+
+- Strategic-mode predictions and multi-label classifiers (`MultiLabelAdaptiveClassifier`) have different scoring and are not covered by the reference implementation.
+- The encoder must be one that exports to ONNX with a `last_hidden_state` output, which is the case for standard BERT-family encoders. Custom architectures that need `trust_remote_code` may not export cleanly.
diff --git a/examples/portable_inference.py b/examples/portable_inference.py
new file mode 100644
index 0000000..94488cc
--- /dev/null
+++ b/examples/portable_inference.py
@@ -0,0 +1,147 @@
+"""Run a saved AdaptiveClassifier without PyTorch, FAISS or this library.
+
+A saved classifier directory is plain files:
+
+    config.json        label maps, per-class training counts, blend weights
+    model.safetensors  class prototypes + the neural head weights
+    onnx/model*.onnx   the transformer encoder (token embeddings out)
+
+so inference outside Python only needs three pieces: an ONNX encoder, a
+tokenizer, and the small amount of math below. This script is the reference
+for that math. Port it to C#, Node, Go or anything else with an ONNX runtime;
+docs/deployment.md explains each step.
+
+Usage:
+    pip install onnxruntime numpy safetensors tokenizers
+    python examples/portable_inference.py ./my_classifier "text to classify"
+
+It needs `tokenizer.json` next to the model. `AdaptiveClassifier.save` does not
+write one, so create it once with:
+
+    from transformers import AutoTokenizer
+    AutoTokenizer.from_pretrained("<base model>").save_pretrained("./my_classifier")
+"""
+
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+from safetensors.numpy import load_file
+from tokenizers import Tokenizer
+
+
+def softmax(x):
+    e = np.exp(x - x.max())
+    return e / e.sum()
+
+
+class PortableClassifier:
+    def __init__(self, directory, pooling=None, prefer_quantized=True):
+        root = Path(directory)
+        self.cfg = json.loads((root / "config.json").read_text(encoding="utf-8"))
+        settings = self.cfg.get("config") or {}
+        self.max_length = settings.get("max_length", 512)
+        self.prototype_weight = settings.get("prototype_weight", 0.7)
+        self.neural_weight = settings.get("neural_weight", 0.3)
+        self.new_threshold = settings.get("new_class_example_threshold", 10)
+        self.new_prototype_weight = settings.get("new_class_prototype_weight", 0.3)
+        self.new_neural_weight = settings.get("new_class_neural_weight", 0.7)
+
+        # 'auto' is resolved at run time in Python (it reads the base model's
+        # sentence-transformers config), so it is never written to disk as
+        # 'mean' or 'cls'. Pass `pooling` yourself when config.json says 'auto'
+        # or has no pooling key: use what the base model was trained with
+        # ('mean' for most sentence-transformers models, 'cls' for BGE). Classifiers
+        # saved before 0.2.0 have no pooling key and behave as 'cls'.
+        self.pooling = pooling or settings.get("pooling")
+        if self.pooling not in ("mean", "cls"):
+            raise ValueError(
+                "config.json does not fix a pooling mode; pass pooling='mean' or 'cls' "
+                "to match the base model"
+            )
+
+        tensors = load_file(str(root / "model.safetensors"))
+        self.labels = [self.cfg["id_to_label"][str(i)] for i in range(len(self.cfg["id_to_label"]))]
+        self.prototypes = {
+            label: tensors[f"prototype_{label}"].astype(np.float32)
+            for label in self.labels
+            if f"prototype_{label}" in tensors
+        }
+        # AdaptiveHead is Linear-ReLU-Dropout x2 then Linear. Dropout is a no-op
+        # at inference, leaving three affine layers at indices 0, 3 and 6.
+        self.head = None
+        if "adaptive_head_model.0.weight" in tensors:
+            self.head = [
+                (tensors[f"adaptive_head_model.{i}.weight"], tensors[f"adaptive_head_model.{i}.bias"])
+                for i in (0, 3, 6)
+            ]
+
+        self.tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
+        self.tokenizer.enable_truncation(max_length=self.max_length)
+        self.tokenizer.enable_padding()
+
+        onnx_dir = root / "onnx"
+        quantized, full = onnx_dir / "model_quantized.onnx", onnx_dir / "model.onnx"
+        model_file = quantized if (prefer_quantized and quantized.exists()) else full
+        self.session = ort.InferenceSession(str(model_file), providers=["CPUExecutionProvider"])
+        self.input_names = {i.name for i in self.session.get_inputs()}
+
+    def embed(self, text):
+        enc = self.tokenizer.encode(text)
+        feeds = {
+            "input_ids": np.array([enc.ids], dtype=np.int64),
+            "attention_mask": np.array([enc.attention_mask], dtype=np.int64),
+        }
+        if "token_type_ids" in self.input_names:
+            feeds["token_type_ids"] = np.array([enc.type_ids], dtype=np.int64)
+        hidden = self.session.run(None, feeds)[0]  # (1, tokens, dim)
+
+        if self.pooling == "cls":
+            vec = hidden[:, 0, :]
+        else:
+            mask = feeds["attention_mask"][..., None].astype(hidden.dtype)
+            vec = (hidden * mask).sum(axis=1) / np.clip(mask.sum(axis=1), 1e-9, None)
+        vec = vec[0]
+        return vec / max(np.linalg.norm(vec), 1e-12)  # L2-normalise
+
+    def _head_probs(self, emb):
+        x = emb
+        for n, (w, b) in enumerate(self.head):
+            x = w @ x + b
+            if n < len(self.head) - 1:
+                x = np.maximum(x, 0.0)
+        return softmax(x)
+
+    def predict(self, text, k=5):
+        emb = self.embed(text)
+
+        # Prototype score: exp(-squared L2 distance), softmaxed over classes.
+        labels = list(self.prototypes)
+        sims = np.array([np.exp(-np.sum((p - emb) ** 2)) for p in self.prototypes.values()])
+        proto_scores = dict(zip(labels, softmax(sims)))
+        head_scores = (
+            dict(zip(self.labels, self._head_probs(emb))) if self.head is not None else {}
+        )
+
+        history = self.cfg.get("training_history", {})
+        combined = {}
+        for label in self.labels:
+            if history.get(label, 0) < self.new_threshold:
+                wp, wn = self.new_prototype_weight, self.new_neural_weight
+            else:
+                wp, wn = self.prototype_weight, self.neural_weight
+            combined[label] = proto_scores.get(label, 0.0) * wp + head_scores.get(label, 0.0) * wn
+
+        total = sum(combined.values())
+        ranked = sorted(combined.items(), key=lambda kv: kv[1], reverse=True)
+        return [(label, score / total) for label, score in ranked[:k]]
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 3:
+        sys.exit(__doc__)
+    clf = PortableClassifier(sys.argv[1])
+    for label, score in clf.predict(" ".join(sys.argv[2:])):
+        print(f"{label}\t{score:.4f}")
diff --git a/pyproject.toml b/pyproject.toml
index f80b696..96cd5f9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,12 +1,67 @@
 [build-system]
-requires = [
-    "setuptools>=64.0.0",
-    "wheel",
-    "packaging>=20.0",
-    "pip>=19.3.1",
-]
+requires = ["setuptools>=64.0.0", "wheel"]
 build-backend = "setuptools.build_meta"
 
+[project]
+name = "adaptive-classifier"
+dynamic = ["version"]
+description = "A flexible, adaptive classification system for dynamic text classification"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "Apache-2.0" }
+authors = [{ name = "codelion", email = "codelion@okyasoft.com" }]
+keywords = ["text-classification", "continuous-learning", "online-learning", "few-shot", "onnx", "transformers"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+# ONNX Runtime (via optimum) stays in the core install: the default CPU path
+# uses it, and the README advertises faster CPU inference out of the box.
+dependencies = [
+    "torch>=2.0.0",
+    "transformers>=4.30.0",
+    "safetensors>=0.3.1",
+    "faiss-cpu>=1.7.4",  # use faiss-gpu for GPU support
+    "numpy>=1.24.0",
+    "scikit-learn>=1.3.0",
+    "huggingface_hub>=0.17.0",
+    "optimum[onnxruntime]>=1.14.0",
+]
+
+[project.optional-dependencies]
+test = [
+    "pytest>=7.0.0",
+    "pytest-cov>=4.0.0",
+    "pytest-randomly>=3.12.0",
+    "psutil",
+]
+# Only needed to run the scripts/ benchmarks and data generators.
+scripts = [
+    "tqdm>=4.65.0",
+    "textdistance",
+    "datasets",
+    "sentence-transformers",
+    "openai",
+]
+
+[project.urls]
+Homepage = "https://github.com/codelion/adaptive-classifier"
+"Bug Tracker" = "https://github.com/codelion/adaptive-classifier/issues"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.dynamic]
+version = { attr = "adaptive_classifier.__version__" }
+
 [tool.pytest.ini_options]
 minversion = "7.0"
 addopts = "-ra -q"
@@ -19,4 +74,4 @@ markers = [
 filterwarnings = [
     "ignore::DeprecationWarning",
     "ignore::UserWarning",
-]
\ No newline at end of file
+]
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 49b29c6..0000000
--- a/requirements.txt
+++ /dev/null
@@ -1,11 +0,0 @@
-torch>=2.0.0
-transformers>=4.30.0
-safetensors>=0.3.1
-faiss-cpu>=1.7.4  # Use faiss-gpu for GPU support
-numpy>=1.24.0
-tqdm>=4.65.0
-setuptools>=65.0.0
-wheel>=0.40.0
-scikit-learn
-huggingface_hub>=0.17.0
-optimum[onnxruntime]>=1.14.0
\ No newline at end of file
diff --git a/setup.py b/setup.py
deleted file mode 100644
index d4c3d5e..0000000
--- a/setup.py
+++ /dev/null
@@ -1,47 +0,0 @@
-from setuptools import setup, find_packages
-
-with open("README.md", "r", encoding="utf-8") as fh:
-    long_description = fh.read()
-
-with open("requirements.txt", "r", encoding="utf-8") as fh:
-    requirements = [line.strip() for line in fh if line.strip() and not line.startswith("#")]
-
-test_requirements = [
-    "pytest>=7.0.0",
-    "pytest-cov>=4.0.0",
-    "pytest-randomly>=3.12.0",
-    "psutil",
-]
-
-setup(
-    name="adaptive-classifier",
-    version="0.2.0",
-    author="codelion",
-    author_email="codelion@okyasoft.com",
-    description="A flexible, adaptive classification system for dynamic text classification",
-    long_description=long_description,
-    long_description_content_type="text/markdown",
-    url="https://github.com/codelion/adaptive-classifier",
-    project_urls={
-        "Bug Tracker": "https://github.com/codelion/adaptive-classifier/issues",
-    },
-    classifiers=[
-        "Development Status :: 4 - Beta",
-        "Intended Audience :: Science/Research",
-        "License :: OSI Approved :: Apache Software License",
-        "Operating System :: OS Independent",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.10",
-        "Topic :: Scientific/Engineering :: Artificial Intelligence",
-    ],
-    package_dir={"": "src"},
-    packages=find_packages(where="src"),
-    python_requires=">=3.8",
-    install_requires=requirements,
-    extras_require={
-        "test": test_requirements,
-    },
-    include_package_data=True,
-)

From 2276951bd8f74aa2d41f3552ca757066e8fd4f7e Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 07:51:59 +0000
Subject: [PATCH 02/16] Save resolved pooling and tokenizer files with the
 model

- save()/push_to_hub() now persist the resolved pooling mode ('mean' or
  'cls') instead of the literal 'auto', so reloads and non-Python runtimes
  no longer depend on a Hub lookup of the base model's config
- save() also writes the tokenizer files (tokenizer.json etc.) so a saved
  directory is self-contained; a failure to write them only warns
- Add tests/test_portable_export.py (offline, tiny local BERT) covering
  persisted pooling, tokenizer output, legacy-config handling, reload
  round-trips, and parity of examples/portable_inference.py with predict()
- Update docs/deployment.md and the reference script: the two manual steps
  are no longer needed for newly saved classifiers

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 docs/deployment.md                    |  16 +-
 examples/portable_inference.py        |  14 +-
 src/adaptive_classifier/classifier.py |  15 +-
 tests/test_portable_export.py         | 213 ++++++++++++++++++++++++++
 4 files changed, 238 insertions(+), 20 deletions(-)
 create mode 100644 tests/test_portable_export.py

diff --git a/docs/deployment.md b/docs/deployment.md
index 85adb86..b9b7afd 100644
--- a/docs/deployment.md
+++ b/docs/deployment.md
@@ -36,7 +36,7 @@ Encoders worth trying (all published on the Hugging Face Hub):
 
 Things that matter:
 
-- **Pooling.** The default `pooling: "auto"` reads the model's own sentence-transformers config, so the table above is handled for you in Python. See the "Pooling" section of the README.
+- **Pooling.** The default `pooling: "auto"` reads the model's own sentence-transformers config, so the table above is handled for you. The result is saved with the classifier. See the "Pooling" section of the README.
 - **Prefixes.** Some models expect a prefix on every input (E5 wants `"query: "`). Add it yourself to training and prediction texts, consistently.
 - **Mixed-language data works.** Good multilingual encoders place translations of the same sentence close together, so a class trained mostly on English examples will often recognise other languages. Add a few examples in each language you care about to firm it up.
 - **Always evaluate on your own data.** Multilingual quality varies a lot by language and domain; the library does not ship per-language benchmarks.
@@ -51,21 +51,17 @@ Things that matter:
 
 | File | Contents |
 |---|---|
-| `config.json` | Label maps (`id_to_label`), per-class `training_history`, and the `config` block: `pooling`, `max_length`, blend weights |
+| `config.json` | Label maps (`id_to_label`), per-class `training_history`, and the `config` block: resolved `pooling`, `max_length`, blend weights |
+| `tokenizer.json` and friends | The base model's tokenizer |
 | `model.safetensors` | `prototype_<label>` tensors (one vector per class) and `adaptive_head_model.{0,3,6}.{weight,bias}` for the neural head |
 | `examples.json` | A few representative training examples per class. Not needed for inference |
 | `onnx/model.onnx`, `onnx/model_quantized.onnx` | The transformer encoder only. It returns token embeddings, not a sentence vector or class scores |
 
 The encoder is the only large file. Everything else is tiny and framework-neutral (JSON plus safetensors, which has readers in most languages).
 
-Two files are **not** written and you must provide them:
+`save()` writes the tokenizer files (`tokenizer.json`, `tokenizer_config.json`, ...) next to the model, and records the **resolved** pooling mode (`mean` or `cls`) in `config.json`, so the directory is self-contained. Nothing needs to be looked up from the base model at inference time.
 
-- **`tokenizer.json`**. Create it once from the base model:
-  ```python
-  from transformers import AutoTokenizer
-  AutoTokenizer.from_pretrained("<base model name>").save_pretrained("./my_classifier")
-  ```
-- **A fixed pooling mode**, if `config.json` says `"pooling": "auto"`. Auto-detection happens in Python at run time by reading the base model's config, so the saved value is the literal string `auto`. Look up what your base model uses (the table above, or its `1_Pooling/config.json`) and hard-code it. Classifiers saved before 0.2.0 use `cls`.
+Directories saved by 0.2.0 or earlier are different. They have no tokenizer files, and their `config.json` has either no `pooling` key (meaning `cls`) or the literal `"auto"`. Re-save them with the current version, or generate the tokenizer once with `AutoTokenizer.from_pretrained("<base model>").save_pretrained("./my_classifier")` and pass the pooling mode to the reference script yourself (`mean` for most sentence-transformers models, `cls` for BGE).
 
 ### The inference algorithm
 
@@ -91,7 +87,7 @@ pip install onnxruntime numpy safetensors tokenizers
 python examples/portable_inference.py ./my_classifier "text to classify"
 ```
 
-It was checked against `AdaptiveClassifier.predict` on the same saved directory with both `mean` and `cls` pooling, and agrees to about 1e-7 (float32 rounding). Use it as the ground truth when you port: run the same inputs through both and compare scores.
+It is checked against `AdaptiveClassifier.predict` on the same saved directory with both `mean` and `cls` pooling (`tests/test_portable_export.py`), and agrees to float32 rounding. Use it as the ground truth when you port: run the same inputs through both and compare scores.
 
 ### Quantized or full-precision encoder
 
diff --git a/examples/portable_inference.py b/examples/portable_inference.py
index 94488cc..7c0be20 100644
--- a/examples/portable_inference.py
+++ b/examples/portable_inference.py
@@ -15,11 +15,8 @@
     pip install onnxruntime numpy safetensors tokenizers
     python examples/portable_inference.py ./my_classifier "text to classify"
 
-It needs `tokenizer.json` next to the model. `AdaptiveClassifier.save` does not
-write one, so create it once with:
-
-    from transformers import AutoTokenizer
-    AutoTokenizer.from_pretrained("<base model>").save_pretrained("./my_classifier")
+`AdaptiveClassifier.save` writes `tokenizer.json` and the resolved pooling mode.
+For directories saved by 0.2.0 or earlier, see docs/deployment.md.
 """
 
 import json
@@ -49,10 +46,9 @@ def __init__(self, directory, pooling=None, prefer_quantized=True):
         self.new_prototype_weight = settings.get("new_class_prototype_weight", 0.3)
         self.new_neural_weight = settings.get("new_class_neural_weight", 0.7)
 
-        # 'auto' is resolved at run time in Python (it reads the base model's
-        # sentence-transformers config), so it is never written to disk as
-        # 'mean' or 'cls'. Pass `pooling` yourself when config.json says 'auto'
-        # or has no pooling key: use what the base model was trained with
+        # Current versions save the resolved mode. Older directories say 'auto'
+        # or have no pooling key; pass `pooling` yourself for those: use what
+        # the base model was trained with
         # ('mean' for most sentence-transformers models, 'cls' for BGE). Classifiers
         # saved before 0.2.0 have no pooling key and behave as 'cls'.
         self.pooling = pooling or settings.get("pooling")
diff --git a/src/adaptive_classifier/classifier.py b/src/adaptive_classifier/classifier.py
index 8ae3659..2f95c40 100644
--- a/src/adaptive_classifier/classifier.py
+++ b/src/adaptive_classifier/classifier.py
@@ -534,6 +534,12 @@ def _save_pretrained(
         save_directory = Path(save_directory)
         os.makedirs(save_directory, exist_ok=True)
 
+        # 'auto' is resolved against the base model at run time, so persist the
+        # outcome. Other runtimes (ONNX in .NET/JS, ...) cannot repeat that
+        # lookup, and a reload should not depend on the Hub being reachable.
+        saved_settings = self.config.to_dict()
+        saved_settings['pooling'] = self._resolve_pooling()
+
         # Save configuration and metadata
         config_dict = {
             'model_name': self.model.config._name_or_path,
@@ -542,7 +548,7 @@ def _save_pretrained(
             'id_to_label': {str(k): v for k, v in self.id_to_label.items()},
             'train_steps': self.train_steps,
             'training_history': self.training_history,  # Save cumulative training counts
-            'config': self.config.to_dict(),
+            'config': saved_settings,
             'library_name': 'adaptive-classifier'  # Tell HuggingFace Hub this requires the adaptive-classifier library
         }
 
@@ -580,6 +586,13 @@ def _save_pretrained(
 
         save_file(tensor_dict, tensors_file)
 
+        # Ship the tokenizer so the directory is self-contained for runtimes
+        # that cannot call AutoTokenizer (tokenizer.json is the portable form).
+        try:
+            self.tokenizer.save_pretrained(save_directory)
+        except Exception as e:
+            logger.warning(f"Could not save tokenizer files: {e}")
+
         # Generate model card if it doesn't exist
         model_card_path = save_directory / "README.md"
         if not model_card_path.exists():
diff --git a/tests/test_portable_export.py b/tests/test_portable_export.py
new file mode 100644
index 0000000..edfaecc
--- /dev/null
+++ b/tests/test_portable_export.py
@@ -0,0 +1,213 @@
+"""A saved classifier must be self-describing and runnable outside PyTorch.
+
+These tests use a tiny randomly initialised BERT built on the fly, so they need
+no network access.
+
+Covered here:
+  - `save()` persists the *resolved* pooling mode instead of the literal 'auto'
+  - `save()` writes tokenizer files, including tokenizer.json
+  - the dependency-light reference in examples/portable_inference.py reproduces
+    `AdaptiveClassifier.predict` from the saved files alone
+"""
+
+import importlib.util
+import json
+from pathlib import Path
+
+import pytest
+import torch
+
+from adaptive_classifier import AdaptiveClassifier
+
+WORDS = ("great terrible okay product service love hate awful fine good bad "
+         "slow fast cheap broke works the a is it very").split()
+TEXTS = ["great product love it", "works very good fast", "terrible awful service",
+         "hate it broke bad", "okay fine the product", "it is okay fine"] * 3
+LABELS = ["pos", "pos", "neg", "neg", "neu", "neu"] * 3
+QUERIES = ["great fast product", "awful slow broke", "okay it is fine", "love the service"]
+
+
+@pytest.fixture(scope="module")
+def base_model(tmp_path_factory):
+    from transformers import BertConfig, BertModel, BertTokenizerFast
+
+    path = tmp_path_factory.mktemp("tiny_bert")
+    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + WORDS
+    (path / "vocab.txt").write_text("\n".join(vocab), encoding="utf-8")
+    torch.manual_seed(0)
+    BertModel(BertConfig(
+        vocab_size=len(vocab), hidden_size=32, num_hidden_layers=2,
+        num_attention_heads=2, intermediate_size=64,
+    )).save_pretrained(path)
+    BertTokenizerFast(str(path / "vocab.txt"), do_lower_case=True).save_pretrained(path)
+    return str(path)
+
+
+def _trained(base_model, **config):
+    clf = AdaptiveClassifier(base_model, config=config, use_onnx=False, device="cpu")
+    clf.add_examples(TEXTS, LABELS)
+    return clf
+
+
+def _saved_config(directory):
+    with open(Path(directory) / "config.json", encoding="utf-8") as handle:
+        return json.load(handle)["config"]
+
+
+def _load_reference():
+    path = Path(__file__).resolve().parent.parent / "examples" / "portable_inference.py"
+    spec = importlib.util.spec_from_file_location("portable_inference", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.PortableClassifier
+
+
+def _onnx_available():
+    try:
+        import onnxruntime  # noqa: F401
+        import optimum.onnxruntime  # noqa: F401
+        return True
+    except ImportError:
+        return False
+
+
+# --- resolved pooling is persisted -------------------------------------------
+
+def test_auto_pooling_is_saved_as_the_resolved_mode(base_model, tmp_path):
+    clf = _trained(base_model)  # default is 'auto'; a local model has no ST config -> mean
+    clf.save(str(tmp_path), include_onnx=False)
+    assert _saved_config(tmp_path)["pooling"] == "mean"
+
+
+def test_auto_pooling_saves_whatever_the_model_resolves_to(base_model, tmp_path, monkeypatch):
+    clf = _trained(base_model)
+    monkeypatch.setattr(clf, "_resolve_pooling", lambda: "cls")
+    clf.save(str(tmp_path), include_onnx=False)
+    assert _saved_config(tmp_path)["pooling"] == "cls"
+
+
+@pytest.mark.parametrize("pooling", ["mean", "cls"])
+def test_explicit_pooling_is_saved_unchanged(base_model, tmp_path, pooling):
+    _trained(base_model, pooling=pooling).save(str(tmp_path), include_onnx=False)
+    assert _saved_config(tmp_path)["pooling"] == pooling
+
+
+def test_saving_does_not_change_the_live_config(base_model, tmp_path):
+    clf = _trained(base_model)
+    clf.save(str(tmp_path), include_onnx=False)
+    assert clf.config.pooling == "auto"
+
+
+def test_reload_uses_saved_pooling_without_a_hub_lookup(base_model, tmp_path, monkeypatch):
+    clf = _trained(base_model, pooling="cls")
+    clf.save(str(tmp_path), include_onnx=False)
+
+    def no_network(*args, **kwargs):
+        raise AssertionError("pooling must come from config.json, not the Hub")
+
+    monkeypatch.setattr("huggingface_hub.hf_hub_download", no_network)
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    assert loaded._resolve_pooling() == "cls"
+
+
+def test_reload_round_trips_predictions(base_model, tmp_path):
+    clf = _trained(base_model)
+    clf.save(str(tmp_path), include_onnx=False)
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    for query in QUERIES:
+        before, after = dict(clf.predict(query)), dict(loaded.predict(query))
+        assert before.keys() == after.keys()
+        for label in before:
+            assert after[label] == pytest.approx(before[label], abs=1e-5)
+
+
+def test_pre_0_2_0_directories_still_load_as_cls(base_model, tmp_path):
+    """A config without a 'pooling' key predates 0.2.0 and means CLS pooling."""
+    _trained(base_model, pooling="cls").save(str(tmp_path), include_onnx=False)
+    config_path = tmp_path / "config.json"
+    data = json.loads(config_path.read_text(encoding="utf-8"))
+    del data["config"]["pooling"]
+    config_path.write_text(json.dumps(data), encoding="utf-8")
+
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    assert loaded._resolve_pooling() == "cls"
+
+
+# --- tokenizer ships with the model ------------------------------------------
+
+def test_save_writes_tokenizer_files(base_model, tmp_path):
+    _trained(base_model).save(str(tmp_path), include_onnx=False)
+    assert (tmp_path / "tokenizer.json").exists()
+    assert (tmp_path / "tokenizer_config.json").exists()
+
+
+def test_saved_tokenizer_matches_the_original(base_model, tmp_path):
+    from transformers import AutoTokenizer
+
+    clf = _trained(base_model)
+    clf.save(str(tmp_path), include_onnx=False)
+    reloaded = AutoTokenizer.from_pretrained(str(tmp_path))
+    for query in QUERIES:
+        assert reloaded(query)["input_ids"] == clf.tokenizer(query)["input_ids"]
+
+
+def test_save_survives_a_tokenizer_that_cannot_be_written(base_model, tmp_path, monkeypatch):
+    clf = _trained(base_model)
+
+    def broken(*args, **kwargs):
+        raise OSError("disk full")
+
+    monkeypatch.setattr(clf.tokenizer, "save_pretrained", broken)
+    clf.save(str(tmp_path), include_onnx=False)  # must not raise
+    assert (tmp_path / "config.json").exists()
+    assert (tmp_path / "model.safetensors").exists()
+
+
+# --- reference implementation -------------------------------------------------
+
+@pytest.mark.skipif(not _onnx_available(), reason="optimum[onnxruntime] not installed")
+@pytest.mark.parametrize("pooling", ["mean", "cls"])
+def test_portable_reference_matches_predict(base_model, tmp_path, pooling):
+    pytest.importorskip("tokenizers")
+    clf = _trained(base_model, pooling=pooling)
+    clf.save(str(tmp_path), include_onnx=True, quantize_onnx=False)
+
+    portable = _load_reference()(tmp_path, prefer_quantized=False)
+    assert portable.pooling == pooling
+    for query in QUERIES:
+        expected = dict(clf.predict(query, k=3))
+        actual = dict(portable.predict(query, k=3))
+        assert actual.keys() == expected.keys()
+        for label in expected:
+            assert float(actual[label]) == pytest.approx(expected[label], abs=1e-4)
+
+
+@pytest.mark.skipif(not _onnx_available(), reason="optimum[onnxruntime] not installed")
+def test_portable_reference_ranks_like_predict(base_model, tmp_path):
+    pytest.importorskip("tokenizers")
+    clf = _trained(base_model)
+    clf.save(str(tmp_path), include_onnx=True, quantize_onnx=False)
+    portable = _load_reference()(tmp_path, prefer_quantized=False)
+    for query in QUERIES:
+        assert [label for label, _ in portable.predict(query)] == [label for label, _ in clf.predict(query)]
+
+
+@pytest.mark.skipif(not _onnx_available(), reason="optimum[onnxruntime] not installed")
+def test_portable_reference_needs_a_pooling_mode_for_legacy_configs(base_model, tmp_path):
+    pytest.importorskip("tokenizers")
+    clf = _trained(base_model, pooling="cls")
+    clf.save(str(tmp_path), include_onnx=True, quantize_onnx=False)
+    config_path = tmp_path / "config.json"
+    data = json.loads(config_path.read_text(encoding="utf-8"))
+    del data["config"]["pooling"]
+    config_path.write_text(json.dumps(data), encoding="utf-8")
+
+    Portable = _load_reference()
+    with pytest.raises(ValueError, match="pooling"):
+        Portable(tmp_path, prefer_quantized=False)
+
+    portable = Portable(tmp_path, pooling="cls", prefer_quantized=False)
+    expected = dict(clf.predict(QUERIES[0], k=3))
+    actual = dict(portable.predict(QUERIES[0], k=3))
+    for label in expected:
+        assert float(actual[label]) == pytest.approx(expected[label], abs=1e-4)

From 6af5f5ab1d8c48ea8c23eb618ce5a878c9eea774 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 07:59:44 +0000
Subject: [PATCH 03/16] Add forget, remove_examples and out-of-distribution
 detection

- forget(labels) removes whole classes from memory, label maps, training
  history and the neural head (keeping other classes' learned weights);
  clear_memory(labels) now delegates to it, fixing a gap where cleared
  classes still showed up in predictions through the head
- remove_examples(texts, label=None, retrain=True) removes examples by
  text, recomputes prototypes, drops emptied classes and optionally
  fine-tunes the head
- ood_score()/is_ood() score a text by its distance to the nearest
  prototype relative to that class's radius; radii are saved with the
  model so reloaded classifiers keep their spread
- predict() gains abstain_below and abstain_ood, both off by default
- New config keys ood_threshold and ood_min_radius
- Share the offline tiny-BERT fixture in tests/conftest.py and add
  tests/test_forget_and_ood.py (30 tests)
- Document the new methods in docs/API.md and the README

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 README.md                             |  18 ++
 docs/API.md                           |  41 ++-
 src/adaptive_classifier/classifier.py | 182 +++++++++++--
 src/adaptive_classifier/memory.py     |  59 ++++
 src/adaptive_classifier/models.py     |  27 ++
 tests/conftest.py                     |  28 ++
 tests/test_forget_and_ood.py          | 370 ++++++++++++++++++++++++++
 tests/test_portable_export.py         |  18 --
 8 files changed, 705 insertions(+), 38 deletions(-)
 create mode 100644 tests/conftest.py
 create mode 100644 tests/test_forget_and_ood.py

diff --git a/README.md b/README.md
index 527c24b..e03c5c0 100644
--- a/README.md
+++ b/README.md
@@ -318,6 +318,24 @@ Any multilingual encoder works as the base model, and a saved classifier can be
 
 ## Advanced Usage
 
+### Fixing Mistakes and Knowing When to Abstain
+
+```python
+# Remove a wrongly labelled example, or a whole class
+classifier.remove_examples(["Refund my order"], label="technical")
+classifier.forget("obsolete_class")
+
+# Detect inputs unlike anything the classifier has seen
+classifier.is_ood("completely unrelated text")        # True / False
+classifier.ood_score("completely unrelated text")     # ~1 or below = familiar, higher = further out
+
+# Abstain instead of guessing (an empty list means "don't know")
+classifier.predict(text, abstain_below=0.6)           # low confidence
+classifier.predict(text, abstain_ood=True)            # out of distribution
+```
+
+Out-of-distribution scores compare a text's distance from the nearest class prototype to that class's own spread, so no per-model tuning of absolute distances is needed. The default `ood_threshold` of 1.25 is a starting point; check it on held-out data.
+
 ### Adding New Classes Dynamically
 
 ```python
diff --git a/docs/API.md b/docs/API.md
index 1ba82ac..e82b9c7 100644
--- a/docs/API.md
+++ b/docs/API.md
@@ -128,7 +128,46 @@ def clear_memory(labels: Optional[List[str]] = None)
 Clear stored examples and prototypes.
 
 Parameters:
-- `labels`: Optional list of labels to clear. If None, clears all memory
+- `labels`: Optional list of labels to clear. If given, this is the same as `forget(labels)` (unknown labels are ignored). If None, clears all stored examples and prototypes
+
+#### forget
+
+```python
+def forget(labels: Union[str, List[str]])
+```
+
+Remove whole classes: their examples, prototype, label ids, training history and neural-head output. The remaining classes keep their learned head weights and no retraining happens. Raises `ValueError` for an unknown label.
+
+#### remove_examples
+
+```python
+def remove_examples(texts: List[str], label: Optional[str] = None, retrain: bool = True) -> int
+```
+
+Remove training examples by exact text, for example to fix a wrong label. Prototypes are recomputed; a class left with no examples is removed. Returns the number of examples removed.
+
+Parameters:
+- `texts`: Texts to remove
+- `label`: Restrict removal to one class; searches all classes if None
+- `retrain`: Fine-tune the neural head on the remaining examples. This is fine-tuning, not retraining from scratch, so the head may keep some influence of the removed text
+
+A classifier loaded from disk keeps only a few representative examples per class, so removing from it recomputes the prototype from those and shifts it more than removing from a live classifier (a warning is logged).
+
+#### ood_score
+
+```python
+def ood_score(text: str) -> float
+```
+
+How far `text` is from the known classes: the distance to the nearest class prototype divided by that class's radius (the distance from its prototype to its farthest training example). Around 1 or below looks like known data; larger is further out. Returns `inf` for a classifier with no classes. Class radii are saved with the model.
+
+#### is_ood
+
+```python
+def is_ood(text: str, threshold: Optional[float] = None) -> bool
+```
+
+True when `ood_score(text)` exceeds `threshold`, which defaults to the `ood_threshold` config value (1.25). `ood_min_radius` (default 0.05) sets a floor on class radii so single-example classes do not divide by zero. The right threshold depends on your encoder and data, so check it on held-out examples.
 
 #### merge_classifiers
 
diff --git a/src/adaptive_classifier/classifier.py b/src/adaptive_classifier/classifier.py
index 2f95c40..16c2e3a 100644
--- a/src/adaptive_classifier/classifier.py
+++ b/src/adaptive_classifier/classifier.py
@@ -391,29 +391,169 @@ def _perform_strategic_training(self):
             
             logger.debug("Performed strategic training step")
     
-    def predict(self, text: str, k: int = 5) -> List[Tuple[str, float]]:
+    def predict(
+        self,
+        text: str,
+        k: int = 5,
+        abstain_below: Optional[float] = None,
+        abstain_ood: bool = False,
+    ) -> List[Tuple[str, float]]:
         """Predict with dual prediction system - blends strategic and regular predictions.
-        
+
         If no cost function is provided, uses existing prediction logic (zero changes).
         If cost function is provided, blends strategic and regular predictions.
-        
+
         Args:
             text: Input text to classify
             k: Number of top predictions to return
-            
+            abstain_below: If set, return an empty list when the top
+                prediction's confidence is below this value
+            abstain_ood: If True, return an empty list when the text is out of
+                distribution (see `is_ood`). Costs one extra embedding pass.
+
         Returns:
-            List of (label, confidence) tuples
+            List of (label, confidence) tuples. An empty list means the
+            classifier abstained.
         """
         if not text:
             raise ValueError("Empty input text")
-        
+
+        if abstain_ood and self.is_ood(text):
+            return []
+
         # If strategic mode is not enabled, use regular prediction
         if not self.strategic_mode:
-            return self._predict_regular(text, k)
-        
-        # Dual prediction system: blend strategic and regular predictions
-        return self._predict_dual(text, k)
-    
+            predictions = self._predict_regular(text, k)
+        else:
+            # Dual prediction system: blend strategic and regular predictions
+            predictions = self._predict_dual(text, k)
+
+        if abstain_below is not None and (not predictions or predictions[0][1] < abstain_below):
+            return []
+        return predictions
+
+    def ood_score(self, text: str) -> float:
+        """How far `text` sits outside the known classes.
+
+        The score is the distance to the nearest class prototype divided by
+        that class's radius (the distance from its prototype to its farthest
+        training example). Around 1.0 or below means the text looks like
+        something the class has seen; larger means further out. Returns
+        infinity when the classifier has no classes.
+        """
+        if not text:
+            raise ValueError("Empty input text")
+        with torch.no_grad():
+            embedding = self._get_embeddings([text])[0]
+        score, _ = self.memory.ood_score(embedding)
+        return score
+
+    def is_ood(self, text: str, threshold: Optional[float] = None) -> bool:
+        """True when `text` scores above the out-of-distribution threshold.
+
+        Args:
+            text: Input text
+            threshold: Overrides the `ood_threshold` config value (default 1.25)
+        """
+        limit = self.config.ood_threshold if threshold is None else threshold
+        return self.ood_score(text) > limit
+
+    def forget(self, labels: Union[str, List[str]]):
+        """Remove one or more classes entirely.
+
+        The classes disappear from prototypes, examples, label maps, training
+        history and the neural head. The remaining classes keep their learned
+        head weights, and no retraining happens.
+
+        Raises:
+            ValueError: If a label is not a known class
+        """
+        if isinstance(labels, str):
+            labels = [labels]
+        unknown = sorted(set(labels) - set(self.label_to_id))
+        if unknown:
+            raise ValueError(f"Unknown labels: {unknown}")
+        self._drop_classes(set(labels))
+
+    def remove_examples(
+        self,
+        texts: List[str],
+        label: Optional[str] = None,
+        retrain: bool = True,
+    ) -> int:
+        """Remove training examples by exact text, e.g. to fix a bad label.
+
+        Args:
+            texts: Texts of the examples to remove
+            label: Only remove from this class; searches all classes if None
+            retrain: Fine-tune the neural head on the remaining examples.
+                Prototypes are always recomputed. The head is only
+                fine-tuned, not retrained from scratch, so it can retain some
+                influence of the removed text; rebuild the classifier if you
+                need guaranteed erasure.
+
+        Returns:
+            Number of examples removed. A class left with no examples is
+            removed, as with `forget`.
+
+        Note:
+            A classifier loaded from disk holds only a few representative
+            examples per class, not the full set its prototype was built from.
+            Removing from it recomputes the prototype from the examples that
+            are present, which shifts it more than removing from a live
+            classifier would.
+        """
+        if not texts:
+            raise ValueError("Empty input list")
+        if label is not None and label not in self.label_to_id:
+            raise ValueError(f"Unknown label: {label}")
+
+        targets = [label] if label is not None else sorted(self.memory.examples)
+        total = 0
+        emptied = set()
+        for lbl in targets:
+            held = len(self.memory.examples.get(lbl, []))
+            removed = self.memory.remove_examples(lbl, texts)
+            if not removed:
+                continue
+            if held < self.training_history.get(lbl, 0):
+                logger.warning(
+                    f"Class '{lbl}' holds {held} examples but was trained on "
+                    f"{self.training_history[lbl]}; its prototype is now "
+                    "recomputed from the examples that remain."
+                )
+            total += removed
+            self.training_history[lbl] = max(0, self.training_history.get(lbl, 0) - removed)
+            if lbl not in self.memory.prototypes:
+                emptied.add(lbl)
+
+        if emptied:
+            self._drop_classes(emptied)
+        if total:
+            self.memory._rebuild_index()
+            if retrain and self.adaptive_head is not None and self.memory.examples:
+                self._train_adaptive_head()
+        return total
+
+    def _drop_classes(self, drop: Set[str]):
+        """Remove classes from every structure that tracks them."""
+        ordered = sorted(self.label_to_id.items(), key=lambda item: item[1])
+        keep = [(label, idx) for label, idx in ordered if label not in drop]
+
+        for label in drop:
+            self.memory.remove_class(label)
+            self.training_history.pop(label, None)
+
+        self.label_to_id = {label: new for new, (label, _) in enumerate(keep)}
+        self.id_to_label = {new: label for label, new in self.label_to_id.items()}
+
+        if self.adaptive_head is not None:
+            if keep:
+                self.adaptive_head.remove_classes([old for _, old in keep])
+                self.adaptive_head = self.adaptive_head.to(self.device)
+            else:
+                self.adaptive_head = None
+
     def _predict_regular(self, text: str, k: int = 5) -> List[Tuple[str, float]]:
         """Regular prediction logic (original implementation)."""
         # Ensure deterministic behavior
@@ -548,6 +688,7 @@ def _save_pretrained(
             'id_to_label': {str(k): v for k, v in self.id_to_label.items()},
             'train_steps': self.train_steps,
             'training_history': self.training_history,  # Save cumulative training counts
+            'ood_radii': {label: self.memory.class_radius(label) for label in self.memory.prototypes},
             'config': saved_settings,
             'library_name': 'adaptive-classifier'  # Tell HuggingFace Hub this requires the adaptive-classifier library
         }
@@ -905,6 +1046,7 @@ def _from_pretrained(
 
         # Rebuild memory system
         classifier.memory._restore_from_save()
+        classifier.memory.saved_radii = dict(config_dict.get('ood_radii', {}))
 
         # Restore adaptive head if it exists
         adaptive_head_params = {
@@ -1463,16 +1605,18 @@ def predict_batch(
         return all_predictions
 
     def clear_memory(self, labels: Optional[List[str]] = None):
-        """Clear memory for specified labels or all if none specified."""
+        """Clear memory for specified labels or all if none specified.
+
+        With labels, this is `forget`: the classes are removed everywhere, so
+        they no longer show up in predictions. Labels that are not known
+        classes are ignored.
+        """
         if labels is None:
             self.memory.clear()
-        else:
-            for label in labels:
-                if label in self.memory.examples:
-                    del self.memory.examples[label]
-                if label in self.memory.prototypes:
-                    del self.memory.prototypes[label]
-            self.memory._rebuild_index()
+            return
+        known = [label for label in labels if label in self.label_to_id]
+        if known:
+            self.forget(known)
 
     def merge_classifiers(self, other: 'AdaptiveClassifier') -> 'AdaptiveClassifier':
         """Merge another classifier into this one."""
diff --git a/src/adaptive_classifier/memory.py b/src/adaptive_classifier/memory.py
index 0d3ff63..ddcb74d 100644
--- a/src/adaptive_classifier/memory.py
+++ b/src/adaptive_classifier/memory.py
@@ -29,6 +29,9 @@ def __init__(
         self.examples = defaultdict(list)  # label -> List[Example]
         self.prototypes = {}  # label -> tensor
         self.strategic_prototypes = {}  # label -> strategic prototype tensor
+        # Class spread recorded when a classifier was saved. A reloaded memory
+        # keeps only a few representative examples, which understate the spread.
+        self.saved_radii = {}  # label -> float
         
         # Initialize FAISS index for fast similarity search
         self.index = faiss.IndexFlatL2(embedding_dim)
@@ -235,10 +238,66 @@ def get_stats(self) -> Dict[str, Any]:
             'updates_since_rebuild': self.updates_since_rebuild
         }
     
+    def remove_class(self, label: str):
+        """Drop a class and everything stored for it."""
+        self.examples.pop(label, None)
+        self.prototypes.pop(label, None)
+        self.strategic_prototypes.pop(label, None)
+        self.saved_radii.pop(label, None)
+        self._rebuild_index()
+
+    def remove_examples(self, label: str, texts) -> int:
+        """Remove examples of `label` whose text is in `texts`.
+
+        The prototype is recomputed from what remains; a class left with no
+        examples is removed. Returns the number of examples removed.
+        """
+        wanted = set(texts)
+        before = self.examples.get(label, [])
+        kept = [ex for ex in before if ex.text not in wanted]
+        removed = len(before) - len(kept)
+        if removed == 0:
+            return 0
+        if not kept:
+            self.remove_class(label)
+            return removed
+        self.examples[label] = kept
+        self.prototypes[label] = torch.mean(torch.stack([ex.embedding for ex in kept]), dim=0)
+        self.strategic_prototypes.pop(label, None)
+        # The saved spread described the old membership.
+        self.saved_radii.pop(label, None)
+        self._rebuild_index()
+        return removed
+
+    def class_radius(self, label: str) -> float:
+        """Distance from the prototype to the farthest known example."""
+        radius = self.saved_radii.get(label, 0.0)
+        examples = self.examples.get(label)
+        if examples and label in self.prototypes:
+            embeddings = torch.stack([ex.embedding for ex in examples])
+            distances = torch.norm(embeddings - self.prototypes[label], dim=1)
+            radius = max(radius, float(distances.max()))
+        return radius
+
+    def ood_score(self, query_embedding: torch.Tensor):
+        """Return (score, label): the smallest distance-to-radius ratio over classes.
+
+        1.0 means the query is exactly as far from the prototype as the most
+        remote training example. Returns (inf, None) when the memory is empty.
+        """
+        best_score, best_label = float('inf'), None
+        for label, prototype in self.prototypes.items():
+            radius = max(self.class_radius(label), self.config.ood_min_radius)
+            score = float(torch.norm(query_embedding - prototype)) / radius
+            if score < best_score:
+                best_score, best_label = score, label
+        return best_score, best_label
+
     def clear(self):
         """Clear all memory."""
         self.examples.clear()
         self.prototypes.clear()
+        self.saved_radii.clear()
         self.index = faiss.IndexFlatL2(self.embedding_dim)
         self.label_to_index.clear()
         self.index_to_label.clear()
diff --git a/src/adaptive_classifier/models.py b/src/adaptive_classifier/models.py
index 4710557..14b44d1 100644
--- a/src/adaptive_classifier/models.py
+++ b/src/adaptive_classifier/models.py
@@ -97,6 +97,22 @@ def update_num_classes(self, num_classes: int):
                 
             self.model[-1] = new_layer
 
+    def remove_classes(self, keep_indices: list):
+        """Keep only the given output rows, in the given order.
+
+        The retained classes keep their learned weights, so removing a class
+        does not disturb what the head knows about the others.
+        """
+        if not keep_indices:
+            raise ValueError("Cannot remove every class from the head")
+        layer = self.model[-1]
+        new_layer = nn.Linear(layer.in_features, len(keep_indices))
+        with torch.no_grad():
+            index = torch.tensor(keep_indices, dtype=torch.long, device=layer.weight.device)
+            new_layer.weight.copy_(layer.weight[index])
+            new_layer.bias.copy_(layer.bias[index])
+        self.model[-1] = new_layer.to(layer.weight.device)
+
 class ModelConfig:
     """Configuration for the adaptive classifier."""
     
@@ -151,6 +167,15 @@ def __init__(self, config: Optional[Dict[str, Any]] = None):
         self.new_class_prototype_weight = self.config.get('new_class_prototype_weight', 0.3)
         self.new_class_neural_weight = self.config.get('new_class_neural_weight', 0.7)
         self.min_confidence = self.config.get('min_confidence', 0.1)
+
+        # Out-of-distribution detection. A query is scored by how far it sits
+        # from its nearest class prototype relative to that class's own spread
+        # (the farthest training example from the prototype). A score above
+        # `ood_threshold` means "further out than anything this class has seen
+        # by that factor". `ood_min_radius` stops one-example classes from
+        # having a zero radius.
+        self.ood_threshold = self.config.get('ood_threshold', 1.25)
+        self.ood_min_radius = self.config.get('ood_min_radius', 0.05)
         
         # Device settings
         self.device_map = self.config.get('device_map', 'auto')
@@ -202,6 +227,8 @@ def to_dict(self) -> Dict[str, Any]:
             'new_class_prototype_weight': self.new_class_prototype_weight,
             'new_class_neural_weight': self.new_class_neural_weight,
             'min_confidence': self.min_confidence,
+            'ood_threshold': self.ood_threshold,
+            'ood_min_radius': self.ood_min_radius,
             'device_map': self.device_map,
             'quantization': self.quantization,
             'gradient_checkpointing': self.gradient_checkpointing,
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..7bbde3a
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,28 @@
+"""Shared fixtures.
+
+`base_model` is a tiny randomly initialised BERT built on the fly, so tests
+that use it need no network access. Its embeddings carry no meaning; tests
+that need meaningful geometry should supply their own vectors.
+"""
+
+import pytest
+import torch
+
+WORDS = ("great terrible okay product service love hate awful fine good bad "
+         "slow fast cheap broke works the a is it very").split()
+
+
+@pytest.fixture(scope="session")
+def base_model(tmp_path_factory):
+    from transformers import BertConfig, BertModel, BertTokenizerFast
+
+    path = tmp_path_factory.mktemp("tiny_bert")
+    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + WORDS
+    (path / "vocab.txt").write_text("\n".join(vocab), encoding="utf-8")
+    torch.manual_seed(0)
+    BertModel(BertConfig(
+        vocab_size=len(vocab), hidden_size=32, num_hidden_layers=2,
+        num_attention_heads=2, intermediate_size=64,
+    )).save_pretrained(path)
+    BertTokenizerFast(str(path / "vocab.txt"), do_lower_case=True).save_pretrained(path)
+    return str(path)
diff --git a/tests/test_forget_and_ood.py b/tests/test_forget_and_ood.py
new file mode 100644
index 0000000..430f7cb
--- /dev/null
+++ b/tests/test_forget_and_ood.py
@@ -0,0 +1,370 @@
+"""Removing classes/examples, and telling when an input is out of distribution.
+
+Uses the tiny offline model from conftest.py. Tests that depend on embedding
+geometry replace `_get_embeddings` with hand-made unit vectors, because a
+randomly initialised encoder places every text in nearly the same spot.
+"""
+
+import copy
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from adaptive_classifier import AdaptiveClassifier
+
+DIM = 32
+
+
+def unit(axis, noise=0.0, seed=0):
+    """A unit vector near axis `axis`, jittered by `noise`."""
+    v = torch.zeros(DIM)
+    v[axis] = 1.0
+    if noise:
+        g = torch.Generator().manual_seed(seed)
+        v = v + noise * torch.randn(DIM, generator=g)
+    return F.normalize(v, dim=0)
+
+
+class Geometry:
+    """Maps text to a fixed embedding: 'a3' -> axis 0 jitter 3, 'b1' -> axis 1, 'z' -> axis 9."""
+
+    AXES = {"a": 0, "b": 1, "c": 2, "z": 9}
+
+    def __init__(self, noise=0.15):
+        self.noise = noise
+
+    def __call__(self, texts):
+        out = []
+        for text in texts:
+            axis = self.AXES[text[0]]
+            seed = int(text[1:] or 0)
+            out.append(unit(axis, self.noise if seed else 0.0, seed))
+        return out
+
+
+def make(base_model, geometry=None, **config):
+    clf = AdaptiveClassifier(base_model, config=config, use_onnx=False, device="cpu")
+    if geometry is not False:
+        clf._get_embeddings = geometry or Geometry()
+    return clf
+
+
+def trained(base_model, n=8, geometry=None, **config):
+    clf = make(base_model, geometry, **config)
+    texts, labels = [], []
+    for label, key in (("A", "a"), ("B", "b"), ("C", "c")):
+        for i in range(1, n + 1):
+            texts.append(f"{key}{i}")
+            labels.append(label)
+    clf.add_examples(texts, labels)
+    return clf
+
+
+def head_logits(clf, embedding):
+    clf.adaptive_head.eval()
+    with torch.no_grad():
+        return clf.adaptive_head(embedding.unsqueeze(0))[0]
+
+
+# --- forget -------------------------------------------------------------------
+
+def test_forget_removes_class_everywhere(base_model):
+    clf = trained(base_model)
+    clf.forget("B")
+
+    assert "B" not in clf.label_to_id
+    assert "B" not in clf.memory.examples
+    assert "B" not in clf.memory.prototypes
+    assert "B" not in clf.training_history
+    assert clf.memory.index.ntotal == 2
+    assert sorted(clf.id_to_label.values()) == ["A", "C"]
+
+
+def test_forget_leaves_contiguous_ids_and_a_matching_head(base_model):
+    clf = trained(base_model)
+    clf.forget("A")
+
+    assert sorted(clf.label_to_id.values()) == [0, 1]
+    assert {clf.id_to_label[i] for i in (0, 1)} == {"B", "C"}
+    assert clf.adaptive_head.model[-1].out_features == 2
+
+
+def test_forgotten_class_never_appears_in_predictions(base_model):
+    clf = trained(base_model)
+    clf.forget("B")
+    for text in ("a1", "b1", "c1", "z"):
+        assert "B" not in [label for label, _ in clf.predict(text, k=10)]
+        batch = clf.predict_batch([text], k=10)[0]
+        assert "B" not in [label for label, _ in batch]
+
+
+def test_forget_preserves_what_the_head_learned_about_other_classes(base_model):
+    clf = trained(base_model)
+    probe = unit(0, 0.15, 1)
+    old = head_logits(clf, probe)
+    old_ids = dict(clf.label_to_id)
+
+    clf.forget("B")
+
+    new = head_logits(clf, probe)
+    for label, new_id in clf.label_to_id.items():
+        assert new[new_id].item() == pytest.approx(old[old_ids[label]].item(), abs=1e-6)
+
+
+def test_forget_accepts_a_list_and_checks_labels(base_model):
+    clf = trained(base_model)
+    with pytest.raises(ValueError, match="Unknown"):
+        clf.forget("nope")
+    assert set(clf.label_to_id) == {"A", "B", "C"}   # a bad call changes nothing
+
+    clf.forget(["A", "C"])
+    assert list(clf.label_to_id) == ["B"]
+    assert clf.id_to_label == {0: "B"}
+
+
+def test_forgetting_every_class_gives_a_usable_empty_classifier(base_model):
+    clf = trained(base_model)
+    clf.forget(["A", "B", "C"])
+
+    assert clf.label_to_id == {} and clf.adaptive_head is None
+    assert clf.predict("a1") == []
+    clf.add_examples(["a1", "a2", "b1", "b2"], ["A", "A", "B", "B"])
+    assert {label for label, _ in clf.predict("a1")} == {"A", "B"}
+
+
+def test_classes_can_be_added_after_forgetting(base_model):
+    clf = trained(base_model)
+    clf.forget("A")
+    clf.add_examples(["z", "z"], ["Z", "Z"])
+
+    assert sorted(clf.label_to_id.values()) == [0, 1, 2]
+    assert clf.adaptive_head.model[-1].out_features == 3
+    assert "Z" in [label for label, _ in clf.predict("z", k=10)]
+
+
+def test_forget_survives_save_and_reload(base_model, tmp_path):
+    clf = trained(base_model)
+    clf.forget("B")
+    clf.save(str(tmp_path), include_onnx=False)
+
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    loaded._get_embeddings = Geometry()
+    assert set(loaded.label_to_id) == {"A", "C"}
+    assert "B" not in [label for label, _ in loaded.predict("b1", k=10)]
+
+
+def test_clear_memory_with_labels_now_forgets_them(base_model):
+    clf = trained(base_model)
+    clf.clear_memory(["B", "not-a-class"])
+    assert set(clf.label_to_id) == {"A", "C"}
+    assert "B" not in [label for label, _ in clf.predict("b1", k=10)]
+
+
+# --- remove_examples -----------------------------------------------------------
+
+def test_remove_examples_recomputes_the_prototype(base_model):
+    clf = trained(base_model)
+    assert clf.remove_examples(["a1", "a2"]) == 2
+
+    remaining = torch.stack([ex.embedding for ex in clf.memory.examples["A"]])
+    assert len(remaining) == 6
+    assert torch.allclose(clf.memory.prototypes["A"], remaining.mean(dim=0), atol=1e-6)
+    assert clf.training_history["A"] == 6
+
+
+def test_remove_examples_can_be_limited_to_one_label(base_model):
+    clf = trained(base_model)
+    clf.add_examples(["a1"], ["B"])        # same text under another label
+    assert clf.remove_examples(["a1"], label="B") == 1
+    assert len(clf.memory.examples["A"]) == 8   # untouched
+
+
+def test_remove_examples_reports_zero_for_unknown_text(base_model):
+    clf = trained(base_model)
+    assert clf.remove_examples(["never seen"]) == 0
+    assert clf.training_history["A"] == 8
+
+
+def test_remove_examples_validates_arguments(base_model):
+    clf = trained(base_model)
+    with pytest.raises(ValueError):
+        clf.remove_examples([])
+    with pytest.raises(ValueError, match="Unknown"):
+        clf.remove_examples(["a1"], label="nope")
+
+
+def test_removing_the_last_example_drops_the_class(base_model):
+    clf = make(base_model)
+    clf.add_examples(["a1", "a2", "b1", "b2", "c1"], ["A", "A", "B", "B", "C"])
+    assert clf.remove_examples(["c1"]) == 1
+
+    assert "C" not in clf.label_to_id
+    assert clf.adaptive_head.model[-1].out_features == 2
+    assert "C" not in [label for label, _ in clf.predict("c1", k=10)]
+
+
+def test_removing_a_mislabelled_example_repairs_the_prototype(base_model):
+    clf = trained(base_model)
+    clean = clf.memory.prototypes["A"].clone()
+
+    clf.add_examples(["z"], ["A"])                      # a bad label
+    assert not torch.allclose(clf.memory.prototypes["A"], clean, atol=1e-3)
+
+    clf.remove_examples(["z"], label="A")
+    assert torch.allclose(clf.memory.prototypes["A"], clean, atol=1e-5)
+
+
+def test_retrain_false_leaves_the_head_alone(base_model):
+    clf = trained(base_model)
+    before = copy.deepcopy(clf.adaptive_head.state_dict())
+    clf.remove_examples(["a1"], retrain=False)
+    after = clf.adaptive_head.state_dict()
+    assert all(torch.equal(before[k], after[k]) for k in before)
+
+
+def test_retrain_true_finetunes_the_head(base_model):
+    clf = trained(base_model)
+    before = copy.deepcopy(clf.adaptive_head.state_dict())
+    clf.remove_examples(["a1"], retrain=True)
+    after = clf.adaptive_head.state_dict()
+    assert any(not torch.equal(before[k], after[k]) for k in before)
+
+
+def test_remove_examples_after_reload_warns_about_partial_memory(base_model, tmp_path, caplog):
+    clf = trained(base_model, n=12)
+    clf.save(str(tmp_path), include_onnx=False)
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+
+    held = loaded.memory.examples["A"][0].text
+    with caplog.at_level("WARNING"):
+        loaded.remove_examples([held], label="A", retrain=False)
+    assert "recomputed from the examples that remain" in caplog.text
+
+
+# --- out-of-distribution detection ----------------------------------------------
+
+def test_in_distribution_text_scores_low_and_far_text_scores_high(base_model):
+    clf = trained(base_model)
+    assert clf.ood_score("a3") < 1.0                  # a training example
+    assert clf.ood_score("z") > clf.config.ood_threshold   # an unrelated direction
+    assert clf.ood_score("z") > 1.5 * clf.ood_score("a3")
+    assert not clf.is_ood("a3")
+    assert clf.is_ood("z")
+
+
+def test_score_is_relative_to_each_class_spread(base_model):
+    """The same absolute distance is normal for a diffuse class and alien for a tight one."""
+    geometry = Geometry(noise=0.0)
+
+    def wide(axis):                                    # spread-out class
+        return [unit(axis, 0.6, s) for s in range(1, 9)]
+
+    def tight(axis):
+        return [unit(axis, 0.02, s) for s in range(1, 9)]
+
+    clf = make(base_model, geometry=False)
+    vectors = {}
+    for label, axis, maker in (("wide", 0, wide), ("tight", 1, tight)):
+        for i, v in enumerate(maker(axis)):
+            vectors[f"{label}{i}"] = v
+    clf._get_embeddings = lambda texts: [vectors[t] for t in texts]
+    clf.add_examples(list(vectors), [t.rstrip("0123456789") for t in vectors])
+
+    wide_proto, tight_proto = clf.memory.prototypes["wide"], clf.memory.prototypes["tight"]
+    offset = unit(5)
+    probe_wide = F.normalize(wide_proto + 0.35 * offset, dim=0)
+    probe_tight = F.normalize(tight_proto + 0.35 * offset, dim=0)
+    clf._get_embeddings = lambda texts: [probe_wide if t == "w" else probe_tight for t in texts]
+
+    assert clf.ood_score("w") < clf.ood_score("t")
+
+
+def test_threshold_override_and_config(base_model):
+    clf = trained(base_model)
+    score = clf.ood_score("a3")
+    assert clf.is_ood("a3", threshold=score / 2)
+    assert not clf.is_ood("a3", threshold=score * 2)
+
+    strict = trained(base_model, ood_threshold=0.0)
+    assert strict.is_ood("a3")
+
+
+def test_ood_score_of_an_empty_classifier_is_infinite(base_model):
+    clf = make(base_model)
+    assert clf.ood_score("a1") == float("inf")
+    assert clf.is_ood("a1")
+
+
+def test_single_example_class_uses_the_minimum_radius(base_model):
+    clf = make(base_model)
+    clf.add_examples(["a1", "b1", "b2"], ["A", "B", "B"])
+    assert clf.memory.class_radius("A") == 0.0
+    assert clf.ood_score("a1") == pytest.approx(0.0)
+    assert 0 < clf.ood_score("a2") < float("inf")       # no ZeroDivisionError
+
+
+def test_abstain_ood_returns_nothing_for_foreign_text(base_model):
+    clf = trained(base_model)
+    assert clf.predict("z", abstain_ood=True) == []
+    assert clf.predict("a3", abstain_ood=True) != []
+
+
+def test_abstain_below_uses_the_top_confidence(base_model):
+    clf = trained(base_model)
+    top = clf.predict("a3")[0][1]
+    assert clf.predict("a3", abstain_below=top + 0.01) == []
+    assert clf.predict("a3", abstain_below=top - 0.01) != []
+
+
+def test_predict_is_unchanged_without_abstention_arguments(base_model):
+    clf = trained(base_model)
+    assert clf.predict("z") == clf.predict("z", abstain_below=None, abstain_ood=False)
+    assert clf.predict("z") != []
+
+
+def test_removing_examples_shrinks_the_class_radius(base_model):
+    clf = trained(base_model)
+    before = clf.memory.class_radius("A")
+    farthest = max(clf.memory.examples["A"],
+                   key=lambda ex: float(torch.norm(ex.embedding - clf.memory.prototypes["A"])))
+    clf.remove_examples([farthest.text], label="A", retrain=False)
+    assert clf.memory.class_radius("A") < before
+
+
+def test_class_radius_survives_save_and_reload(base_model, tmp_path):
+    """A reloaded memory holds few examples, so the spread must be persisted."""
+    clf = trained(base_model, n=20)
+    original = clf.memory.class_radius("A")
+    clf.save(str(tmp_path), include_onnx=False)
+
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    loaded._get_embeddings = Geometry()
+    assert len(loaded.memory.examples["A"]) < 20
+    assert loaded.memory.class_radius("A") == pytest.approx(original, abs=1e-5)
+    assert loaded.ood_score("a3") == pytest.approx(clf.ood_score("a3"), rel=1e-3, abs=1e-3)
+    assert loaded.is_ood("z")
+
+
+def test_ood_settings_round_trip_in_the_saved_config(base_model, tmp_path):
+    clf = trained(base_model, ood_threshold=2.0, ood_min_radius=0.2)
+    clf.save(str(tmp_path), include_onnx=False)
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    assert loaded.config.ood_threshold == 2.0
+    assert loaded.config.ood_min_radius == 0.2
+
+
+def test_old_saves_without_radii_still_load(base_model, tmp_path):
+    import json
+    clf = trained(base_model)
+    clf.save(str(tmp_path), include_onnx=False)
+    path = tmp_path / "config.json"
+    data = json.loads(path.read_text(encoding="utf-8"))
+    del data["ood_radii"]
+    del data["config"]["ood_threshold"]
+    path.write_text(json.dumps(data), encoding="utf-8")
+
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    loaded._get_embeddings = Geometry()
+    assert loaded.config.ood_threshold == 1.25
+    assert loaded.ood_score("a3") < float("inf")
diff --git a/tests/test_portable_export.py b/tests/test_portable_export.py
index edfaecc..f7b4d19 100644
--- a/tests/test_portable_export.py
+++ b/tests/test_portable_export.py
@@ -19,30 +19,12 @@
 
 from adaptive_classifier import AdaptiveClassifier
 
-WORDS = ("great terrible okay product service love hate awful fine good bad "
-         "slow fast cheap broke works the a is it very").split()
 TEXTS = ["great product love it", "works very good fast", "terrible awful service",
          "hate it broke bad", "okay fine the product", "it is okay fine"] * 3
 LABELS = ["pos", "pos", "neg", "neg", "neu", "neu"] * 3
 QUERIES = ["great fast product", "awful slow broke", "okay it is fine", "love the service"]
 
 
-@pytest.fixture(scope="module")
-def base_model(tmp_path_factory):
-    from transformers import BertConfig, BertModel, BertTokenizerFast
-
-    path = tmp_path_factory.mktemp("tiny_bert")
-    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + WORDS
-    (path / "vocab.txt").write_text("\n".join(vocab), encoding="utf-8")
-    torch.manual_seed(0)
-    BertModel(BertConfig(
-        vocab_size=len(vocab), hidden_size=32, num_hidden_layers=2,
-        num_attention_heads=2, intermediate_size=64,
-    )).save_pretrained(path)
-    BertTokenizerFast(str(path / "vocab.txt"), do_lower_case=True).save_pretrained(path)
-    return str(path)
-
-
 def _trained(base_model, **config):
     clf = AdaptiveClassifier(base_model, config=config, use_onnx=False, device="cpu")
     clf.add_examples(TEXTS, LABELS)

From 753ab30871560f98da7eb088a3dfc957f2c2c9c3 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 08:01:49 +0000
Subject: [PATCH 04/16] CI: run every test, cache models, add timeouts and a
 weekly run

- The integration job now runs `pytest tests/ -m integration` instead of a
  single named file, so integration-marked tests in any file are executed
  (together with the matrix job this covers the whole suite)
- Cache the Hugging Face model directory in both jobs
- Add a 45 minute timeout to the matrix job
- Add a weekly schedule and workflow_dispatch so Hub or dependency breakage
  is caught without waiting for a PR
- Cancel superseded runs on the same PR

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 .github/workflows/test.yml | 34 ++++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 1deaf4a..7592fc0 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -5,10 +5,21 @@ on:
     branches: [ main ]
   pull_request:
     branches: [ main ]
+  workflow_dispatch:
+  schedule:
+    # Weekly, so breakage on the Hugging Face Hub or in new transformers/torch
+    # releases shows up without waiting for the next PR.
+    - cron: '17 5 * * 1'
+
+# A newer push to the same PR supersedes the run in progress.
+concurrency:
+  group: tests-${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 jobs:
   unit-tests:
     runs-on: ubuntu-latest
+    timeout-minutes: 45
     strategy:
       fail-fast: false
       matrix:
@@ -23,12 +34,23 @@ jobs:
         python-version: ${{ matrix.python-version }}
         cache: pip
 
+    - name: Cache Hugging Face models
+      uses: actions/cache@v4
+      with:
+        path: ~/.cache/huggingface
+        key: hf-${{ runner.os }}-${{ hashFiles('tests/**/*.py') }}
+        restore-keys: |
+          hf-${{ runner.os }}-
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install -e ".[test]"
 
-    - name: Run unit tests
+    # Everything except tests marked `integration`. This includes the tests that
+    # download small models from the Hub; they need network access, which GitHub
+    # runners have. Tests marked `integration` run in the job below.
+    - name: Run tests
       env:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
       run: |
@@ -56,6 +78,14 @@ jobs:
         python-version: '3.12'
         cache: pip
 
+    - name: Cache Hugging Face models
+      uses: actions/cache@v4
+      with:
+        path: ~/.cache/huggingface
+        key: hf-${{ runner.os }}-${{ hashFiles('tests/**/*.py') }}
+        restore-keys: |
+          hf-${{ runner.os }}-
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
@@ -65,7 +95,7 @@ jobs:
       env:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
       run: |
-        pytest tests/test_enterprise_classifiers_integration.py -v -m "integration" --tb=short
+        pytest tests/ -v -m "integration" --tb=short
 
     - name: Integration test summary
       if: always()

From 62a31dd7874e2e4907e8c1b80b99b14d32e1688b Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 08:06:04 +0000
Subject: [PATCH 05/16] Add scikit-learn compatible SklearnAdaptiveClassifier

- New SklearnAdaptiveClassifier (fit, partial_fit, predict, predict_proba,
  predict_log_proba, score, classes_) wrapping AdaptiveClassifier, so it
  works in Pipeline, cross_val_score and GridSearchCV
- Strings or integers as labels, returned as given; input validation for
  X and y; partial_fit accepts new classes at any call
- The fitted AdaptiveClassifier is exposed as classifier_ for save,
  forget, remove_examples and OOD detection
- tests/test_sklearn_api.py (23 offline tests) including clone, CV,
  Pipeline and GridSearchCV
- Document in the README and docs/API.md

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 README.md                              |  22 +++
 docs/API.md                            |  23 +++
 src/adaptive_classifier/__init__.py    |   2 +
 src/adaptive_classifier/sklearn_api.py | 186 +++++++++++++++++++++
 tests/test_sklearn_api.py              | 222 +++++++++++++++++++++++++
 5 files changed, 455 insertions(+)
 create mode 100644 src/adaptive_classifier/sklearn_api.py
 create mode 100644 tests/test_sklearn_api.py

diff --git a/README.md b/README.md
index e03c5c0..653cad6 100644
--- a/README.md
+++ b/README.md
@@ -318,6 +318,28 @@ Any multilingual encoder works as the base model, and a saved classifier can be
 
 ## Advanced Usage
 
+### scikit-learn Interface
+
+`SklearnAdaptiveClassifier` works with `Pipeline`, `cross_val_score` and `GridSearchCV`. `X` is a list of strings.
+
+```python
+from adaptive_classifier import SklearnAdaptiveClassifier
+from sklearn.model_selection import cross_val_score
+
+clf = SklearnAdaptiveClassifier("sentence-transformers/all-MiniLM-L6-v2")
+clf.fit(texts, labels)
+clf.predict(["Where is my refund?"])          # class labels
+clf.predict_proba(["Where is my refund?"])    # columns follow clf.classes_
+cross_val_score(clf, texts, labels, cv=5)
+
+# The library's strength: keep learning, including brand-new classes
+clf.partial_fit(["App crashes on login"], ["bug"])
+
+clf.classifier_.save("./model")               # the underlying AdaptiveClassifier
+```
+
+`fit` starts over each time; `partial_fit` adds to what is already learned. A fitted estimator holds a FAISS index and can't be pickled, so save `classifier_` instead.
+
 ### Fixing Mistakes and Knowing When to Abstain
 
 ```python
diff --git a/docs/API.md b/docs/API.md
index e82b9c7..69199f7 100644
--- a/docs/API.md
+++ b/docs/API.md
@@ -353,3 +353,26 @@ classifier2 = AdaptiveClassifier("bert-base-uncased")
 # ... train both classifiers ...
 classifier1.merge_classifiers(classifier2)
 ```
+
+## SklearnAdaptiveClassifier
+
+```python
+SklearnAdaptiveClassifier(model_name="sentence-transformers/all-MiniLM-L6-v2",
+                          config=None, device=None, use_onnx="auto",
+                          trust_remote_code=False)
+```
+
+scikit-learn estimator wrapping `AdaptiveClassifier`. `X` is a 1-D sequence of strings (a single column is also accepted); `y` holds string or integer labels, which `predict` returns unchanged.
+
+| Method / attribute | Description |
+|---|---|
+| `fit(X, y)` | Fit from scratch; returns `self` |
+| `partial_fit(X, y, classes=None)` | Add examples and any new classes to what is already learned. If `classes` is given, `y` must stay within it |
+| `predict(X)` | Most likely class per text |
+| `predict_proba(X)` | Probabilities, columns ordered as `classes_` |
+| `predict_log_proba(X)` | Log of `predict_proba` |
+| `score(X, y)` | Accuracy |
+| `classes_` | Sorted labels seen so far |
+| `classifier_` | The underlying fitted `AdaptiveClassifier` (`save`, `push_to_hub`, `forget`, `remove_examples`, `is_ood`, ...) |
+
+Hyperparameters such as `prototype_weight` go inside `config`, so grid searches take `{"config": [{...}, {...}]}`. Fitted estimators cannot be pickled (they hold a FAISS index); persist `classifier_` with `save`.
diff --git a/src/adaptive_classifier/__init__.py b/src/adaptive_classifier/__init__.py
index 4189f43..a7e5738 100644
--- a/src/adaptive_classifier/__init__.py
+++ b/src/adaptive_classifier/__init__.py
@@ -2,6 +2,7 @@
 from .models import Example, AdaptiveHead, ModelConfig
 from .memory import PrototypeMemory
 from .multilabel import MultiLabelAdaptiveClassifier, MultiLabelAdaptiveHead
+from .sklearn_api import SklearnAdaptiveClassifier
 from huggingface_hub import ModelHubMixin
 
 __version__ = "0.2.0"
@@ -10,6 +11,7 @@
     "AdaptiveClassifier",
     "MultiLabelAdaptiveClassifier",
     "MultiLabelAdaptiveHead",
+    "SklearnAdaptiveClassifier",
     "Example",
     "AdaptiveHead",
     "ModelConfig",
diff --git a/src/adaptive_classifier/sklearn_api.py b/src/adaptive_classifier/sklearn_api.py
new file mode 100644
index 0000000..7d1c166
--- /dev/null
+++ b/src/adaptive_classifier/sklearn_api.py
@@ -0,0 +1,186 @@
+"""scikit-learn compatible interface to AdaptiveClassifier.
+
+`SklearnAdaptiveClassifier` takes a sequence of strings as `X`, so it drops into
+`Pipeline`, `cross_val_score` and `GridSearchCV`, and its `partial_fit` maps onto
+the library's strength: adding examples and new classes without retraining.
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+from sklearn.base import BaseEstimator, ClassifierMixin
+from sklearn.utils.validation import check_is_fitted
+
+from .classifier import AdaptiveClassifier
+
+
+class SklearnAdaptiveClassifier(ClassifierMixin, BaseEstimator):
+    """Text classifier with the scikit-learn estimator API.
+
+    Parameters
+    ----------
+    model_name : str
+        Hugging Face model (or local directory) used as the encoder.
+    config : dict, optional
+        Passed to `AdaptiveClassifier` (pooling, prototype_weight, ...).
+    device : str, optional
+        "cpu" or "cuda"; chosen automatically when None.
+    use_onnx : bool or "auto", default="auto"
+        Whether to use ONNX Runtime for the encoder.
+    trust_remote_code : bool, default=False
+        Forwarded to the Hugging Face loaders.
+
+    Attributes
+    ----------
+    classes_ : ndarray
+        Sorted class labels seen so far. Labels may be strings or integers and
+        are returned as given.
+    classifier_ : AdaptiveClassifier
+        The fitted underlying classifier. Use it for features that have no
+        scikit-learn equivalent: `save`, `push_to_hub`, `forget`,
+        `remove_examples`, `is_ood`.
+
+    Notes
+    -----
+    `fit` starts from scratch each time (it reloads the encoder, so
+    cross-validation pays that cost per fold). `partial_fit` adds to what is
+    already learned. A fitted estimator holds a FAISS index and cannot be
+    pickled; persist `classifier_` with `classifier_.save(...)` instead.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+        config: Optional[Dict[str, Any]] = None,
+        device: Optional[str] = None,
+        use_onnx: Union[bool, str] = "auto",
+        trust_remote_code: bool = False,
+    ):
+        self.model_name = model_name
+        self.config = config
+        self.device = device
+        self.use_onnx = use_onnx
+        self.trust_remote_code = trust_remote_code
+
+    # --- scikit-learn plumbing -------------------------------------------------
+
+    def __sklearn_tags__(self):
+        tags = super().__sklearn_tags__()
+        tags.input_tags.string = True
+        return tags
+
+    def _more_tags(self):  # scikit-learn < 1.6
+        return {"X_types": ["string"]}
+
+    # --- input handling ------------------------------------------------------------
+
+    @staticmethod
+    def _as_texts(X) -> List[str]:
+        array = np.asarray(X, dtype=object)
+        if array.ndim == 2 and array.shape[1] == 1:
+            array = array[:, 0]
+        if array.ndim != 1:
+            raise ValueError(
+                "X must be a 1-D sequence of strings (or a single column); "
+                f"got an array of shape {array.shape}"
+            )
+        if array.shape[0] == 0:
+            raise ValueError("X is empty")
+        texts = array.tolist()
+        bad = next((t for t in texts if not isinstance(t, str)), None)
+        if bad is not None:
+            raise ValueError(f"X must contain only strings; found {type(bad).__name__}")
+        return texts
+
+    def _as_labels(self, y, n_samples: int) -> List[Any]:
+        labels = np.asarray(y, dtype=object)
+        if labels.ndim == 2 and labels.shape[1] == 1:
+            labels = labels[:, 0]
+        if labels.ndim != 1:
+            raise ValueError(f"y must be 1-D; got shape {labels.shape}")
+        if labels.shape[0] != n_samples:
+            raise ValueError(
+                f"X and y have inconsistent lengths: {n_samples} != {labels.shape[0]}"
+            )
+        return [label.item() if isinstance(label, np.generic) else label
+                for label in labels.tolist()]
+
+    def _new_classifier(self) -> AdaptiveClassifier:
+        return AdaptiveClassifier(
+            self.model_name,
+            device=self.device,
+            config=self.config,
+            use_onnx=self.use_onnx,
+            trust_remote_code=self.trust_remote_code,
+        )
+
+    def _refresh_classes(self):
+        # The wrapped classifier keys everything by str(label); keep the
+        # caller's original label objects so predict returns what was passed in.
+        originals = list(self._original_labels.values())
+        try:
+            ordered = sorted(originals)
+        except TypeError:  # mixed label types
+            ordered = sorted(originals, key=str)
+        self.classes_ = np.array(ordered)
+
+    # --- fitting -------------------------------------------------------------------
+
+    def fit(self, X, y):
+        """Fit from scratch on texts `X` and labels `y`."""
+        texts = self._as_texts(X)
+        labels = self._as_labels(y, len(texts))
+        self.classifier_ = self._new_classifier()
+        self._original_labels = {}
+        self._add(texts, labels)
+        return self
+
+    def partial_fit(self, X, y, classes=None):
+        """Add examples (and any new classes) to what is already learned.
+
+        Unlike most scikit-learn estimators, new classes may appear in any
+        call. If `classes` is given, every label in `y` must be one of them.
+        """
+        texts = self._as_texts(X)
+        labels = self._as_labels(y, len(texts))
+        if classes is not None:
+            unexpected = set(labels) - set(np.asarray(classes, dtype=object).tolist())
+            if unexpected:
+                raise ValueError(f"y contains labels not in classes: {sorted(map(str, unexpected))}")
+        if not hasattr(self, "classifier_"):
+            self.classifier_ = self._new_classifier()
+            self._original_labels = {}
+        self._add(texts, labels)
+        return self
+
+    def _add(self, texts: List[str], labels: List[Any]):
+        for label in labels:
+            self._original_labels.setdefault(str(label), label)
+        self.classifier_.add_examples(texts, [str(label) for label in labels])
+        self._refresh_classes()
+
+    # --- prediction ----------------------------------------------------------------
+
+    def predict_proba(self, X):
+        """Class probabilities, columns ordered as `classes_`."""
+        check_is_fitted(self, "classifier_")
+        texts = self._as_texts(X)
+        keys = [str(label) for label in self.classes_]
+        column = {key: i for i, key in enumerate(keys)}
+        proba = np.zeros((len(texts), len(keys)))
+        batches = self.classifier_.predict_batch(texts, k=len(keys))
+        for row, scores in enumerate(batches):
+            for key, score in scores:
+                if key in column:
+                    proba[row, column[key]] = score
+        totals = proba.sum(axis=1, keepdims=True)
+        return np.divide(proba, totals, out=np.full_like(proba, 1.0 / len(keys)), where=totals > 0)
+
+    def predict_log_proba(self, X):
+        with np.errstate(divide="ignore"):
+            return np.log(self.predict_proba(X))
+
+    def predict(self, X):
+        """Most likely class for each text."""
+        proba = self.predict_proba(X)
+        return self.classes_[np.argmax(proba, axis=1)]
diff --git a/tests/test_sklearn_api.py b/tests/test_sklearn_api.py
new file mode 100644
index 0000000..75bdf0e
--- /dev/null
+++ b/tests/test_sklearn_api.py
@@ -0,0 +1,222 @@
+"""scikit-learn interface (offline, tiny local model).
+
+A randomly initialised encoder carries no meaning, so tests that check
+accuracy swap in hand-made embeddings: text "a3" sits near axis 0, "b2" near
+axis 1, and so on. Classes are therefore perfectly separable.
+"""
+
+import numpy as np
+import pytest
+import torch
+import torch.nn.functional as F
+from sklearn.base import clone
+from sklearn.exceptions import NotFittedError
+from sklearn.model_selection import GridSearchCV, cross_val_score
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import FunctionTransformer
+
+from adaptive_classifier import AdaptiveClassifier, SklearnAdaptiveClassifier
+
+DIM = 32
+PROTO_ONLY = {"new_class_example_threshold": 0}   # always use prototype_weight / neural_weight
+AXES = {"a": 0, "b": 1, "c": 2}
+
+
+def geometry(self, texts):
+    out = []
+    for text in texts:
+        v = torch.zeros(DIM)
+        v[AXES[text[0]]] = 1.0
+        g = torch.Generator().manual_seed(int(text[1:]))
+        out.append(F.normalize(v + 0.1 * torch.randn(DIM, generator=g), dim=0))
+    return out
+
+
+@pytest.fixture
+def separable(monkeypatch):
+    monkeypatch.setattr(AdaptiveClassifier, "_get_embeddings", geometry)
+
+
+def data(n=12, classes=("a", "b", "c")):
+    X = [f"{k}{i}" for k in classes for i in range(1, n + 1)]
+    y = [k.upper() for k in classes for _ in range(n)]
+    return X, y
+
+
+def est(base_model, **kw):
+    return SklearnAdaptiveClassifier(base_model, use_onnx=False, device="cpu", **kw)
+
+
+# --- estimator contract ----------------------------------------------------------
+
+def test_params_round_trip_and_clone(base_model):
+    e = est(base_model, config={"pooling": "mean"})
+    assert e.get_params()["model_name"] == base_model
+    assert e.get_params()["config"] == {"pooling": "mean"}
+    e.set_params(config={"pooling": "cls"}, use_onnx=False)
+    copy = clone(e)
+    assert copy.get_params() == e.get_params()
+    assert copy is not e
+
+
+def test_unfitted_estimator_raises(base_model):
+    with pytest.raises(NotFittedError):
+        est(base_model).predict(["a1"])
+    with pytest.raises(NotFittedError):
+        est(base_model).predict_proba(["a1"])
+
+
+def test_fit_returns_self_and_sets_classes(base_model, separable):
+    e = est(base_model)
+    X, y = data()
+    assert e.fit(X, y) is e
+    assert list(e.classes_) == ["A", "B", "C"]
+    assert isinstance(e.classifier_, AdaptiveClassifier)
+
+
+def test_predict_and_proba_agree(base_model, separable):
+    e = est(base_model).fit(*data())
+    X = ["a1", "b2", "c3", "a4"]
+    proba = e.predict_proba(X)
+    assert proba.shape == (4, 3)
+    assert np.allclose(proba.sum(axis=1), 1.0)
+    assert (proba >= 0).all()
+    assert list(e.classes_[proba.argmax(axis=1)]) == list(e.predict(X))
+    assert list(e.predict(X)) == ["A", "B", "C", "A"]
+    assert np.allclose(np.exp(e.predict_log_proba(X)), proba)
+
+
+def test_proba_columns_follow_classes_order_not_insertion_order(base_model, separable):
+    e = est(base_model).fit(*data(classes=("c", "a", "b")))   # inserted c, a, b
+    assert list(e.classes_) == ["A", "B", "C"]
+    assert e.predict_proba(["a1"])[0].argmax() == 0
+    assert e.predict_proba(["c1"])[0].argmax() == 2
+
+
+def test_score_is_accuracy(base_model, separable):
+    e = est(base_model).fit(*data())
+    assert e.score(["a1", "b1", "c1", "a2"], ["A", "B", "C", "B"]) == pytest.approx(0.75)
+
+
+def test_fit_starts_from_scratch(base_model, separable):
+    e = est(base_model).fit(*data())
+    e.fit(*data(classes=("a", "b")))
+    assert list(e.classes_) == ["A", "B"]
+
+
+def test_integer_labels_come_back_as_integers(base_model, separable):
+    X, y = data()
+    codes = {"A": 10, "B": 2, "C": 7}
+    e = est(base_model).fit(X, [codes[label] for label in y])
+    assert list(e.classes_) == [2, 7, 10]
+    predicted = e.predict(["a1", "b1", "c1"])
+    assert list(predicted) == [10, 2, 7]
+    assert np.issubdtype(predicted.dtype, np.integer)
+
+
+def test_numpy_labels_are_accepted(base_model, separable):
+    X, y = data()
+    e = est(base_model).fit(np.array(X), np.array(y))
+    assert list(e.predict(np.array(["a1", "b1"]))) == ["A", "B"]
+
+
+# --- partial_fit ---------------------------------------------------------------------
+
+def test_partial_fit_accumulates_and_adds_classes(base_model, separable):
+    e = est(base_model)
+    e.partial_fit(["a1", "a2", "b1", "b2"], ["A", "A", "B", "B"])
+    assert list(e.classes_) == ["A", "B"]
+    assert e.classifier_.training_history["A"] == 2
+
+    e.partial_fit(["a3", "c1", "c2"], ["A", "C", "C"])
+    assert list(e.classes_) == ["A", "B", "C"]
+    assert e.classifier_.training_history == {"A": 3, "B": 2, "C": 2}
+    assert e.predict(["c3"])[0] == "C"
+    assert e.predict_proba(["c3"]).shape == (1, 3)
+
+
+def test_partial_fit_validates_against_classes(base_model, separable):
+    e = est(base_model)
+    with pytest.raises(ValueError, match="not in classes"):
+        e.partial_fit(["a1", "d1"], ["A", "D"], classes=["A", "B"])
+    e.partial_fit(["a1", "b1"], ["A", "B"], classes=["A", "B"])
+    assert list(e.classes_) == ["A", "B"]
+
+
+# --- input validation ------------------------------------------------------------------
+
+def test_single_column_input_is_accepted(base_model, separable):
+    X, y = data(n=12)
+    e = est(base_model).fit(np.array(X).reshape(-1, 1), y)
+    assert list(e.predict(np.array(["a1"]).reshape(-1, 1))) == ["A"]
+
+
+@pytest.mark.parametrize("X, message", [
+    ([], "empty"),
+    ([["a1", "a2"], ["b1", "b2"]], "1-D"),
+    (["a1", 5], "strings"),
+    (np.zeros((3, 2)), "1-D"),
+])
+def test_bad_x_is_rejected(base_model, X, message):
+    with pytest.raises(ValueError, match=message):
+        est(base_model).fit(X, ["A"] * len(X))
+
+
+def test_mismatched_lengths_are_rejected(base_model):
+    with pytest.raises(ValueError, match="inconsistent"):
+        est(base_model).fit(["a1", "a2"], ["A"])
+
+
+def test_predict_rejects_non_text(base_model, separable):
+    e = est(base_model).fit(*data(n=12))
+    with pytest.raises(ValueError, match="strings"):
+        e.predict([1, 2])
+
+
+# --- the point: it works with scikit-learn tooling ---------------------------------------
+
+def test_cross_val_score(base_model, separable):
+    X, y = data(n=12)
+    # Folds see ~8 examples per class, under the default new-class threshold of 10;
+    # that regime leans on a barely-trained head, which is not what this test is about.
+    scores = cross_val_score(est(base_model, config=PROTO_ONLY), X, y, cv=3)
+    assert scores.shape == (3,)
+    assert scores.mean() > 0.9
+
+
+def test_pipeline_with_a_text_step(base_model, separable):
+    pipeline = Pipeline([
+        ("clean", FunctionTransformer(lambda texts: [t.strip() for t in texts])),
+        ("clf", est(base_model)),
+    ])
+    X, y = data(n=12)
+    pipeline.fit([f"  {t} " for t in X], y)
+    assert list(pipeline.predict(["  a1", "b1  "])) == ["A", "B"]
+    assert pipeline.score([" c1 "], ["C"]) == 1.0
+
+
+def test_grid_search_over_config(base_model, separable):
+    X, y = data(n=12)
+    search = GridSearchCV(
+        est(base_model),
+        {"config": [{**PROTO_ONLY, "prototype_weight": 0.9, "neural_weight": 0.1},
+                    {**PROTO_ONLY, "prototype_weight": 0.6, "neural_weight": 0.4}]},
+        cv=2,
+    )
+    search.fit(X, y)
+    assert search.best_score_ > 0.9
+    assert search.best_params_["config"]["prototype_weight"] in (0.9, 0.6)
+    assert list(search.predict(["b1"])) == ["B"]
+
+
+def test_underlying_classifier_features_stay_reachable(base_model, separable, tmp_path):
+    e = est(base_model).fit(*data())
+    e.classifier_.forget("C")
+    assert list(e.classifier_.label_to_id) == ["A", "B"]
+    e.classifier_.save(str(tmp_path), include_onnx=False)
+    assert (tmp_path / "config.json").exists()
+
+
+def test_string_input_tag_is_declared(base_model):
+    tags = est(base_model).__sklearn_tags__()
+    assert tags.input_tags.string

From 47007c06617dd06586759b0973e445e200eed5ff Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 08:46:53 +0000
Subject: [PATCH 06/16] Bump to 0.3.0, make CI offline and stable, add
 robustness and failure-mode tests

CI
- New `models` job downloads the test models once (with retries) via
  scripts/prefetch_test_models.py and caches them; the test jobs restore
  the cache and run with HF_HUB_OFFLINE=1. The previous run hit the Hub's
  1000-requests-per-5-minutes limit (HTTP 429) because four matrix legs
  shared one quota; every failure, including the ONNX ones, was a 429
- Matrix reduced to the ends of the supported range, Python 3.10 and 3.13
- tests/test_ci_models.py fails if a test names a model the prefetch list
  does not cover

Library fixes found by the new tests
- Non-string labels (e.g. ints) are rejected: they were accepted, then
  reloaded from disk as a duplicated mix of int and str classes
- None as a text or label, and mixed label types, raise a clear ValueError
  before any state changes
- Negative or non-integer k raises ValueError (k=-1 silently dropped the
  last prediction)
- Loading a directory without config.json says so instead of reporting an
  invalid Hub repo id

Tests
- tests/test_robustness.py: accuracy on separable synthetic data at several
  sizes, no catastrophic forgetting, class imbalance, label noise,
  determinism, order independence, and property-based checks (hypothesis)
  that predictions are always a valid ordered distribution
- tests/test_failure_modes.py: hostile and malformed input, damaged or
  missing files, Hub failures falling back to PyTorch, and randomized
  sequences of add/forget/remove checked for internal consistency
- A strict expected failure records a known weakness: with fewer than 10
  examples per class the default blend lets an undertrained neural head
  outvote correct prototypes (~75-92% training accuracy on separable data)
- Mutation-checked: 12 deliberate faults in the library, all caught

Also: version 0.3.0, CHANGELOG.md, hypothesis added to the test extra.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 .github/workflows/test.yml            |  56 +++-
 CHANGELOG.md                          |  24 ++
 pyproject.toml                        |   1 +
 scripts/prefetch_test_models.py       |  65 +++++
 src/adaptive_classifier/__init__.py   |   2 +-
 src/adaptive_classifier/classifier.py |  34 ++-
 tests/conftest.py                     |  10 +
 tests/synthetic.py                    |  56 ++++
 tests/test_ci_models.py               |  45 ++++
 tests/test_failure_modes.py           | 368 ++++++++++++++++++++++++++
 tests/test_robustness.py              | 237 +++++++++++++++++
 11 files changed, 882 insertions(+), 16 deletions(-)
 create mode 100644 CHANGELOG.md
 create mode 100644 scripts/prefetch_test_models.py
 create mode 100644 tests/synthetic.py
 create mode 100644 tests/test_ci_models.py
 create mode 100644 tests/test_failure_modes.py
 create mode 100644 tests/test_robustness.py

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 7592fc0..7bbdd2b 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -17,13 +17,43 @@ concurrency:
   cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 jobs:
+  models:
+    # Download the Hugging Face models once, with retries. The test jobs run
+    # offline from the cache, so a Hub rate limit (HTTP 429) or outage fails
+    # this one step instead of dozens of tests across the matrix.
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+
+    - name: Cache Hugging Face models
+      id: hf-cache
+      uses: actions/cache@v4
+      with:
+        path: ~/.cache/huggingface
+        key: hf-models-${{ hashFiles('scripts/prefetch_test_models.py') }}
+
+    - name: Download models
+      if: steps.hf-cache.outputs.cache-hit != 'true'
+      env:
+        HF_TOKEN: ${{ secrets.HF_TOKEN }}
+      run: |
+        pip install huggingface_hub
+        python scripts/prefetch_test_models.py
+
   unit-tests:
+    needs: models
     runs-on: ubuntu-latest
     timeout-minutes: 45
     strategy:
       fail-fast: false
       matrix:
-        python-version: ['3.10', '3.11', '3.12', '3.13']
+        # The two ends of the supported range (requires-python >=3.10).
+        python-version: ['3.10', '3.13']
 
     steps:
     - uses: actions/checkout@v4
@@ -34,25 +64,24 @@ jobs:
         python-version: ${{ matrix.python-version }}
         cache: pip
 
-    - name: Cache Hugging Face models
-      uses: actions/cache@v4
+    - name: Restore Hugging Face models
+      uses: actions/cache/restore@v4
       with:
         path: ~/.cache/huggingface
-        key: hf-${{ runner.os }}-${{ hashFiles('tests/**/*.py') }}
-        restore-keys: |
-          hf-${{ runner.os }}-
+        key: hf-models-${{ hashFiles('scripts/prefetch_test_models.py') }}
+        fail-on-cache-miss: true
 
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install -e ".[test]"
 
-    # Everything except tests marked `integration`. This includes the tests that
-    # download small models from the Hub; they need network access, which GitHub
-    # runners have. Tests marked `integration` run in the job below.
+    # Everything except tests marked `integration`, run offline against the
+    # models prefetched above. Tests marked `integration` run in the job below.
     - name: Run tests
       env:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        HF_HUB_OFFLINE: '1'
       run: |
         pytest tests/ -v --cov=adaptive_classifier --cov-report=xml --cov-report=term -m "not integration"
 
@@ -78,13 +107,12 @@ jobs:
         python-version: '3.12'
         cache: pip
 
-    - name: Cache Hugging Face models
-      uses: actions/cache@v4
+    - name: Restore Hugging Face models
+      uses: actions/cache/restore@v4
       with:
         path: ~/.cache/huggingface
-        key: hf-${{ runner.os }}-${{ hashFiles('tests/**/*.py') }}
-        restore-keys: |
-          hf-${{ runner.os }}-
+        key: hf-models-${{ hashFiles('scripts/prefetch_test_models.py') }}
+        fail-on-cache-miss: true
 
     - name: Install dependencies
       run: |
diff --git a/CHANGELOG.md b/CHANGELOG.md
new file mode 100644
index 0000000..3eb535d
--- /dev/null
+++ b/CHANGELOG.md
@@ -0,0 +1,24 @@
+# Changelog
+
+## 0.3.0
+
+### Added
+- `SklearnAdaptiveClassifier`: scikit-learn estimator (`fit`, `partial_fit`, `predict`, `predict_proba`, `score`) usable in `Pipeline`, `cross_val_score` and `GridSearchCV`.
+- `forget(labels)` and `remove_examples(texts, label=None, retrain=True)` to delete classes or fix wrong labels.
+- `ood_score()` / `is_ood()` out-of-distribution detection, and `predict(..., abstain_below=, abstain_ood=)`. New config keys `ood_threshold` and `ood_min_radius`.
+- `docs/deployment.md` and `examples/portable_inference.py`: run a saved classifier without PyTorch or FAISS (multilingual models, .NET, JavaScript).
+- `save()` now writes the tokenizer files and the resolved pooling mode, so a saved directory is self-contained.
+
+### Changed
+- **Python 3.10 or newer is required** (3.8 and 3.9 are end-of-life).
+- Packaging moved to a `[project]` table in `pyproject.toml`; `setup.py` and `requirements.txt` are gone. Install extras: `[test]`, `[scripts]`. `tqdm` is no longer a core dependency.
+- `clear_memory(labels=[...])` now removes those classes completely (it used to leave them in the label map and neural head, so they could still be predicted).
+
+### Fixed
+- Integer (or other non-string) labels corrupted a classifier on save/reload, duplicating classes. `add_examples` now rejects non-string labels with a clear error; convert with `str(label)`.
+- `predict(k=-1)` returned a silently truncated list; negative or non-integer `k` now raises `ValueError`.
+- `None` as a text or label, and mixed label types, now raise a clear `ValueError` before anything is changed.
+- Loading a directory without `config.json` now says so instead of reporting an invalid Hub repo id.
+
+### Known issues
+- With fewer than 10 examples per class the default blend gives the neural head 70% of the vote. A head trained on so few examples can outvote correct prototypes, so accuracy on small training sets is lower than prototype-only prediction. Tracked by a strict expected-failure test in `tests/test_robustness.py`.
diff --git a/pyproject.toml b/pyproject.toml
index 96cd5f9..01127b9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,6 +42,7 @@ test = [
     "pytest-cov>=4.0.0",
     "pytest-randomly>=3.12.0",
     "psutil",
+    "hypothesis>=6.80.0",
 ]
 # Only needed to run the scripts/ benchmarks and data generators.
 scripts = [
diff --git a/scripts/prefetch_test_models.py b/scripts/prefetch_test_models.py
new file mode 100644
index 0000000..812b4ae
--- /dev/null
+++ b/scripts/prefetch_test_models.py
@@ -0,0 +1,65 @@
+"""Download every Hugging Face model the test suite uses into the local cache.
+
+CI runs this once, with retries, and the test jobs then run with
+HF_HUB_OFFLINE=1. That way a Hub outage or rate limit (HTTP 429) fails one
+clearly-labelled step instead of dozens of unrelated tests, and the matrix
+legs do not each hit the Hub.
+
+    python scripts/prefetch_test_models.py
+"""
+
+import sys
+import time
+
+from huggingface_hub import snapshot_download
+
+# Keep in sync with the model names used in tests/ (tests/test_ci_models.py
+# fails when a test uses a model that is not listed here).
+MODELS = [
+    "prajjwal1/bert-tiny",
+    "bert-base-uncased",
+    "distilbert-base-uncased",
+    "distilbert/distilbert-base-cased",
+    "google-bert/bert-large-cased",
+    "answerdotai/ModernBERT-base",
+]
+
+ATTEMPTS = 6
+WEIGHTS = ["*.safetensors"]
+FALLBACK_WEIGHTS = ["*.bin"]
+SMALL_FILES = ["*.json", "*.txt", "*.model", "*.py"]
+
+
+def fetch(model: str):
+    """Download config, tokenizer and weights (safetensors, else .bin)."""
+    path = snapshot_download(model, allow_patterns=SMALL_FILES + WEIGHTS)
+    import glob
+    import os
+    if not glob.glob(os.path.join(path, "*.safetensors")):
+        snapshot_download(model, allow_patterns=SMALL_FILES + FALLBACK_WEIGHTS)
+
+
+def main() -> int:
+    failed = []
+    for model in MODELS:
+        for attempt in range(1, ATTEMPTS + 1):
+            try:
+                fetch(model)
+                print(f"ok    {model}")
+                break
+            except Exception as error:  # network, 429, 5xx
+                wait = min(120, 5 * 2 ** attempt)
+                print(f"retry {model} ({attempt}/{ATTEMPTS}): {type(error).__name__}: "
+                      f"{str(error)[:200]} - waiting {wait}s", flush=True)
+                if attempt == ATTEMPTS:
+                    failed.append(model)
+                else:
+                    time.sleep(wait)
+    if failed:
+        print(f"Could not download: {failed}", file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/src/adaptive_classifier/__init__.py b/src/adaptive_classifier/__init__.py
index a7e5738..cd2df02 100644
--- a/src/adaptive_classifier/__init__.py
+++ b/src/adaptive_classifier/__init__.py
@@ -5,7 +5,7 @@
 from .sklearn_api import SklearnAdaptiveClassifier
 from huggingface_hub import ModelHubMixin
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 __all__ = [
     "AdaptiveClassifier",
diff --git a/src/adaptive_classifier/classifier.py b/src/adaptive_classifier/classifier.py
index 16c2e3a..75e082e 100644
--- a/src/adaptive_classifier/classifier.py
+++ b/src/adaptive_classifier/classifier.py
@@ -6,6 +6,7 @@
 from typing import List, Dict, Optional, Tuple, Any, Set, Union
 import logging
 import copy
+import numbers
 from pathlib import Path
 from safetensors.torch import save_file, load_file
 import json
@@ -131,12 +132,36 @@ def _should_use_onnx(self, use_onnx: Union[bool, str]) -> bool:
             logger.warning(f"Invalid use_onnx value: {use_onnx}. Using auto-detection.")
             return self.device == "cpu"
 
+    @staticmethod
+    def _validate_training_inputs(texts: List[str], labels: List[str]):
+        """Reject inputs that would corrupt the classifier or fail obscurely later."""
+        bad_text = next((t for t in texts if not isinstance(t, str)), None)
+        if bad_text is not None or any(t is None for t in texts):
+            raise ValueError(
+                f"texts must all be strings; got {type(bad_text).__name__}"
+            )
+        bad_label = next((l for l in labels if not isinstance(l, str)), None)
+        if bad_label is not None or any(l is None for l in labels):
+            # Saving writes labels into JSON keys and tensor names, which turns
+            # other types into strings: a saved integer-labelled classifier
+            # reloads with duplicated, inconsistent classes.
+            raise ValueError(
+                f"labels must be strings; got {type(bad_label).__name__}. "
+                "Convert with str(label)."
+            )
+
+    @staticmethod
+    def _validate_k(k: int):
+        if not isinstance(k, numbers.Integral) or isinstance(k, bool) or k < 0:
+            raise ValueError(f"k must be a non-negative integer; got {k!r}")
+
     def add_examples(self, texts: List[str], labels: List[str]):
         """Add new examples with special handling for new classes."""
         if not texts or not labels:
             raise ValueError("Empty input lists")
         if len(texts) != len(labels):
             raise ValueError("Mismatched text and label lists")
+        self._validate_training_inputs(texts, labels)
 
         # Check if classifier has any existing classes (before updating mappings)
         has_existing_classes = len(self.label_to_id) > 0
@@ -417,6 +442,7 @@ def predict(
         """
         if not text:
             raise ValueError("Empty input text")
+        self._validate_k(k)
 
         if abstain_ood and self.is_ood(text):
             return []
@@ -823,6 +849,11 @@ def _from_pretrained(
        
         # Check if model_id is a local directory
         model_path = Path(model_id)
+        if model_path.is_dir() and not (model_path / "config.json").exists():
+            raise FileNotFoundError(
+                f"{model_id} is a directory but not a saved AdaptiveClassifier: "
+                "config.json is missing"
+            )
         try:
             if model_path.is_dir() and (model_path / "config.json").exists():
                 # Local directory with required files
@@ -1532,7 +1563,8 @@ def predict_batch(
         """Predict labels for a batch of texts with improved batching."""
         if not texts:
             raise ValueError("Empty input batch")
-        
+        self._validate_k(k)
+
         all_predictions = []
         
         # Process in batches
diff --git a/tests/conftest.py b/tests/conftest.py
index 7bbde3a..4559ac6 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -26,3 +26,13 @@ def base_model(tmp_path_factory):
     )).save_pretrained(path)
     BertTokenizerFast(str(path / "vocab.txt"), do_lower_case=True).save_pretrained(path)
     return str(path)
+
+
+@pytest.fixture
+def new_classifier(base_model):
+    """Factory for an untrained classifier on the tiny offline model."""
+    from adaptive_classifier import AdaptiveClassifier
+
+    def make(**config):
+        return AdaptiveClassifier(base_model, config=config, use_onnx=False, device="cpu")
+    return make
diff --git a/tests/synthetic.py b/tests/synthetic.py
new file mode 100644
index 0000000..39f60df
--- /dev/null
+++ b/tests/synthetic.py
@@ -0,0 +1,56 @@
+"""Synthetic, perfectly separable text data with hand-made embeddings.
+
+A randomly initialised encoder puts every text in nearly the same place, so
+accuracy tests replace `_get_embeddings` with this. Each class sits around its
+own axis; `noise` controls how spread out a class is. With these geometries a
+good classifier should be close to perfect, so any shortfall is the
+classifier's, not the data's.
+"""
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from adaptive_classifier import AdaptiveClassifier
+
+DIM = 32
+
+
+def point(axis: int, noise: float, seed: int) -> torch.Tensor:
+    v = torch.zeros(DIM)
+    v[axis % DIM] = 1.0
+    if noise:
+        g = torch.Generator().manual_seed(seed)
+        v = v + noise * torch.randn(DIM, generator=g)
+    return F.normalize(v, dim=0)
+
+
+class Dataset:
+    """Texts mapped to embeddings; install with `patch(monkeypatch)`."""
+
+    def __init__(self):
+        self.store = {}
+
+    def add(self, name: str, axis: int, noise: float, seed: int) -> str:
+        self.store[name] = point(axis, noise, seed)
+        return name
+
+    def classes(self, n_per: int, n_classes: int = 3, noise: float = 0.1,
+                seed: int = 0, prefix: str = "tr"):
+        """Return (texts, labels): `n_per` examples around each class axis."""
+        texts, labels = [], []
+        for c in range(n_classes):
+            for i in range(n_per):
+                texts.append(self.add(f"{prefix}{c}_{i}", c, noise, seed * 100003 + c * 1009 + i))
+                labels.append(f"C{c}")
+        return texts, labels
+
+    def patch(self, monkeypatch):
+        store = self.store
+        monkeypatch.setattr(AdaptiveClassifier, "_get_embeddings",
+                            lambda self, texts: [store[t] for t in texts])
+        return self
+
+
+def accuracy(clf: AdaptiveClassifier, texts, labels) -> float:
+    return float(np.mean([clf.predict(t)[0][0] == label for t, label in zip(texts, labels)]))
diff --git a/tests/test_ci_models.py b/tests/test_ci_models.py
new file mode 100644
index 0000000..551ccec
--- /dev/null
+++ b/tests/test_ci_models.py
@@ -0,0 +1,45 @@
+"""Guard the CI model prefetch list against drift.
+
+CI downloads models once (scripts/prefetch_test_models.py) and runs the suite
+offline. A test that names a Hub model missing from that list would pass
+locally and fail in CI with an obscure offline error, so catch it here.
+"""
+
+import importlib.util
+import re
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+HUB_NAME = re.compile(r"""["']((?:[\w.-]+/)?[\w.-]*(?:bert|BERT|MiniLM)[\w.-]*)["']""")
+# Local fixtures and names that are not Hub repos.
+IGNORED = {"tiny_bert"}
+
+
+def _prefetch_models():
+    spec = importlib.util.spec_from_file_location(
+        "prefetch_test_models", ROOT / "scripts" / "prefetch_test_models.py")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return set(module.MODELS)
+
+
+def _models_used_by_tests():
+    used = set()
+    for path in (ROOT / "tests").glob("test_*.py"):
+        if path.name == Path(__file__).name:
+            continue
+        used |= {m for m in HUB_NAME.findall(path.read_text(encoding="utf-8"))}
+    return used - IGNORED
+
+
+def test_every_model_named_in_tests_is_prefetched():
+    missing = _models_used_by_tests() - _prefetch_models()
+    assert not missing, (
+        f"Add {sorted(missing)} to MODELS in scripts/prefetch_test_models.py, "
+        "or CI (which runs offline) cannot load them"
+    )
+
+
+def test_prefetch_list_has_no_unused_models():
+    unused = _prefetch_models() - _models_used_by_tests()
+    assert not unused, f"Remove {sorted(unused)} from scripts/prefetch_test_models.py"
diff --git a/tests/test_failure_modes.py b/tests/test_failure_modes.py
new file mode 100644
index 0000000..416c3a7
--- /dev/null
+++ b/tests/test_failure_modes.py
@@ -0,0 +1,368 @@
+"""What happens when things go wrong.
+
+Positive-path tests show the library works when used as intended. These check
+the other half: bad input is rejected with a clear error (not silently
+accepted, not a cryptic crash), damaged files fail loudly, and a Hub outage
+degrades to a working classifier instead of a broken one.
+"""
+
+import json
+import logging
+import random
+
+import numpy as np
+import pytest
+import torch
+
+from adaptive_classifier import AdaptiveClassifier
+
+TEXTS = ["great love", "good fast", "terrible hate", "awful bad"] * 3
+LABELS = ["p", "p", "n", "n"] * 3
+
+
+@pytest.fixture
+def trained(new_classifier):
+    clf = new_classifier()
+    clf.add_examples(TEXTS, LABELS)
+    return clf
+
+
+# --- bad prediction input -------------------------------------------------------------
+
+@pytest.mark.parametrize("bad", ["", None])
+def test_empty_prediction_text_is_rejected(trained, bad):
+    with pytest.raises(ValueError, match="Empty"):
+        trained.predict(bad)
+
+
+@pytest.mark.parametrize("bad", [123, 4.5, ["list"], {"a": 1}, b"bytes"])
+def test_non_string_prediction_text_is_rejected(trained, bad):
+    with pytest.raises((ValueError, TypeError)):
+        trained.predict(bad)
+
+
+@pytest.mark.parametrize("k", [-1, -100, 1.5, "3", None])
+def test_invalid_k_is_rejected(trained, k):
+    with pytest.raises(ValueError, match="k must be"):
+        trained.predict("great", k=k)
+    with pytest.raises(ValueError, match="k must be"):
+        trained.predict_batch(["great"], k=k)
+
+
+def test_k_zero_returns_nothing_and_huge_k_returns_every_class(trained):
+    assert trained.predict("great", k=0) == []
+    assert len(trained.predict("great", k=10_000)) == 2
+
+
+def test_numpy_integers_are_accepted_for_k(trained):
+    assert len(trained.predict("great", k=np.int64(1))) == 1
+
+
+def test_empty_batch_is_rejected(trained):
+    with pytest.raises(ValueError, match="Empty"):
+        trained.predict_batch([])
+
+
+def test_untrained_classifier_predicts_nothing_instead_of_crashing(new_classifier):
+    clf = new_classifier()
+    assert clf.predict("great") == []
+    assert clf.predict_batch(["great", "bad"]) == [[], []]
+
+
+# --- hostile but legal text -----------------------------------------------------------------
+
+@pytest.mark.parametrize("text", [
+    "   ",
+    "\n\t",
+    "😀🎉 emoji only",
+    "مرحبا بالعالم",
+    "你好，世界",
+    "mixed ‮ right-to-left ​ zero-width",
+    "null\x00byte",
+    "a" * 50_000,
+    "great " * 20_000,
+    "<script>alert(1)</script> '; DROP TABLE x;--",
+])
+def test_unusual_text_gives_a_finite_valid_prediction(trained, text):
+    preds = trained.predict(text)
+    assert preds and {l for l, _ in preds} <= {"p", "n"}
+    assert all(np.isfinite(s) and 0 <= s <= 1 for _, s in preds)
+    assert sum(s for _, s in preds) == pytest.approx(1.0, abs=1e-5)
+
+
+def test_text_beyond_max_length_matches_its_truncation(new_classifier):
+    clf = new_classifier(max_length=16, pooling="mean")
+    clf.add_examples(TEXTS, LABELS)
+    long = "great love " * 200
+    head = clf.tokenizer.decode(clf.tokenizer(long, max_length=16, truncation=True)["input_ids"],
+                                skip_special_tokens=True)
+    a, b = dict(clf.predict(long)), dict(clf.predict(head))
+    assert all(a[l] == pytest.approx(b[l], abs=1e-4) for l in a)
+
+
+# --- bad training input ----------------------------------------------------------------------
+
+@pytest.mark.parametrize("texts, labels, message", [
+    ([], [], "Empty"),
+    (["a", "b"], ["x"], "Mismatched"),
+    ([None], ["x"], "strings"),
+    ([5], ["x"], "strings"),
+    (["a"], [None], "strings"),
+    (["a", "b"], [1, 2], "strings"),
+    (["a", "b"], ["x", 2], "strings"),
+    (["a"], [("x",)], "strings"),
+])
+def test_bad_training_input_is_rejected_before_anything_changes(new_classifier, texts, labels, message):
+    clf = new_classifier()
+    with pytest.raises(ValueError, match=message):
+        clf.add_examples(texts, labels)
+    assert clf.label_to_id == {} and clf.memory.prototypes == {}
+
+
+def test_rejected_input_leaves_a_trained_classifier_untouched(trained):
+    before = (dict(trained.label_to_id), dict(trained.training_history),
+              {l: len(v) for l, v in trained.memory.examples.items()})
+    with pytest.raises(ValueError):
+        trained.add_examples(["fine text", "another"], ["p", None])
+    after = (dict(trained.label_to_id), dict(trained.training_history),
+             {l: len(v) for l, v in trained.memory.examples.items()})
+    assert before == after
+
+
+def test_integer_labels_cannot_corrupt_a_saved_model(new_classifier, tmp_path):
+    """Integer labels used to be accepted, then reloaded as a duplicated set of
+    string and integer classes (JSON turns the keys into strings)."""
+    clf = new_classifier()
+    with pytest.raises(ValueError, match="strings"):
+        clf.add_examples(TEXTS, [1, 1, 2, 2] * 3)
+
+    clf.add_examples(TEXTS, [str(l) for l in [1, 1, 2, 2] * 3])
+    clf.save(str(tmp_path), include_onnx=False)
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    assert {l for l, _ in loaded.predict("great", k=10)} == {"1", "2"}
+
+
+def test_empty_string_text_is_allowed_and_embeds(new_classifier):
+    clf = new_classifier()
+    clf.add_examples(["", "good"], ["blank", "ok"])
+    assert {l for l, _ in clf.predict("good", k=5)} == {"blank", "ok"}
+
+
+def test_identical_text_with_conflicting_labels_stays_consistent(new_classifier):
+    clf = new_classifier()
+    clf.add_examples(["great"] * 4, ["p", "n", "p", "n"])
+    preds = clf.predict("great")
+    assert {l for l, _ in preds} == {"p", "n"}
+    assert abs(preds[0][1] - preds[1][1]) < 0.3
+
+
+def test_single_class_classifier_is_certain(new_classifier):
+    clf = new_classifier()
+    clf.add_examples(["great", "good"], ["only", "only"])
+    assert clf.predict("anything at all") == [("only", pytest.approx(1.0))]
+
+
+def test_memory_stays_bounded(new_classifier):
+    clf = new_classifier(max_examples_per_class=5)
+    clf.add_examples(["great"] * 40 + ["bad"] * 40, ["p"] * 40 + ["n"] * 40)
+    assert all(len(v) <= 5 for v in clf.memory.examples.values())
+    assert {l for l, _ in clf.predict("great")} == {"p", "n"}
+
+
+def test_forget_and_remove_on_unknown_things_fail_clearly(trained, new_classifier):
+    with pytest.raises(ValueError, match="Unknown"):
+        trained.forget("never-existed")
+    with pytest.raises(ValueError, match="Unknown"):
+        trained.remove_examples(["x"], label="never-existed")
+    assert new_classifier().remove_examples(["x"]) == 0     # nothing to remove is not an error
+
+
+# --- damaged or missing files ------------------------------------------------------------------
+
+@pytest.fixture
+def saved(trained, tmp_path):
+    trained.save(str(tmp_path), include_onnx=False)
+    return tmp_path
+
+
+def load(path):
+    return AdaptiveClassifier.load(str(path), use_onnx=False, device="cpu")
+
+
+def test_a_clean_save_loads_and_predicts_the_same(trained, saved):
+    assert load(saved).predict("great") == pytest.approx(trained.predict("great"), abs=1e-5) \
+        or [l for l, _ in load(saved).predict("great")] == [l for l, _ in trained.predict("great")]
+
+
+def test_directory_without_config_says_so(saved):
+    (saved / "config.json").unlink()
+    with pytest.raises(FileNotFoundError, match="config.json is missing"):
+        load(saved)
+
+
+def test_empty_directory_says_so(tmp_path):
+    with pytest.raises(FileNotFoundError, match="config.json is missing"):
+        load(tmp_path)
+
+
+def test_nonexistent_path_fails_loudly():
+    with pytest.raises(ValueError, match="Error loading model"):
+        load("/definitely/not/a/real/dir")
+
+
+@pytest.mark.parametrize("name", ["examples.json", "model.safetensors"])
+def test_missing_state_file_is_an_error_not_a_silent_empty_model(saved, name):
+    (saved / name).unlink()
+    with pytest.raises(FileNotFoundError):
+        load(saved)
+
+
+def test_truncated_weights_fail_loudly(saved):
+    data = (saved / "model.safetensors").read_bytes()
+    (saved / "model.safetensors").write_bytes(data[: len(data) // 2])
+    with pytest.raises(Exception, match="(?i)header|metadata|deserial|covered|safetensor"):
+        load(saved)
+
+
+def test_corrupt_config_json_fails_loudly(saved):
+    (saved / "config.json").write_text("{not json", encoding="utf-8")
+    with pytest.raises(json.JSONDecodeError):
+        load(saved)
+
+
+def test_config_that_disagrees_with_the_weights_fails_loudly(saved):
+    cfg = json.loads((saved / "config.json").read_text(encoding="utf-8"))
+    cfg["label_to_id"]["ghost"] = 2
+    cfg["id_to_label"]["2"] = "ghost"
+    (saved / "config.json").write_text(json.dumps(cfg), encoding="utf-8")
+    with pytest.raises(RuntimeError, match="size mismatch|loading state_dict"):
+        load(saved)
+
+
+# --- the Hub misbehaving -------------------------------------------------------------------------
+
+@pytest.mark.parametrize("error", [
+    OSError("429 Client Error: Too Many Requests"),
+    ConnectionError("Tunnel connection failed: 403 Forbidden"),
+    RuntimeError("export failed"),
+])
+def test_onnx_failure_falls_back_to_a_working_pytorch_model(base_model, monkeypatch, caplog, error):
+    """This is what happened in CI when the Hub rate-limited: the classifier must
+    keep working on PyTorch and say why, not fail or pretend ONNX loaded."""
+    pytest.importorskip("optimum.onnxruntime")
+    from optimum.onnxruntime import ORTModelForFeatureExtraction
+
+    def broken(*args, **kwargs):
+        raise error
+
+    monkeypatch.setattr(ORTModelForFeatureExtraction, "from_pretrained", broken)
+    with caplog.at_level(logging.WARNING):
+        clf = AdaptiveClassifier(base_model, use_onnx=True, device="cpu")
+
+    assert clf.use_onnx is False
+    assert "Falling back to PyTorch" in caplog.text
+    clf.add_examples(TEXTS, LABELS)
+    assert clf.predict("great")
+
+
+def test_missing_optimum_falls_back_with_install_hint(base_model, monkeypatch, caplog):
+    import builtins
+    real_import = builtins.__import__
+
+    def no_optimum(name, *args, **kwargs):
+        if name.startswith("optimum"):
+            raise ImportError("No module named 'optimum'")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", no_optimum)
+    with caplog.at_level(logging.WARNING):
+        clf = AdaptiveClassifier(base_model, use_onnx=True, device="cpu")
+    assert clf.use_onnx is False
+    assert "pip install optimum" in caplog.text
+
+
+def test_pooling_lookup_failure_falls_back_to_mean(trained, monkeypatch):
+    import huggingface_hub
+
+    def offline(*args, **kwargs):
+        raise OSError("Hub unreachable")
+
+    monkeypatch.setattr(huggingface_hub, "hf_hub_download", offline)
+    trained._resolved_pooling = None
+    assert trained._resolve_pooling() == "mean"
+
+
+# --- randomized sequences of operations ------------------------------------------------------------
+
+def check_consistent(clf):
+    """Every structure that tracks classes must agree after any operation."""
+    labels = set(clf.label_to_id)
+    n = len(labels)
+    assert sorted(clf.label_to_id.values()) == list(range(n)), "ids not contiguous"
+    assert {clf.id_to_label[i] for i in range(n)} == labels
+    assert all(clf.id_to_label[i] == l for l, i in clf.label_to_id.items())
+    assert set(clf.memory.prototypes) == labels
+    assert set(k for k, v in clf.memory.examples.items() if v) == labels
+    assert clf.memory.index.ntotal == n
+    assert set(clf.training_history) >= labels or n == 0
+    if clf.adaptive_head is not None:
+        assert clf.adaptive_head.model[-1].out_features == n
+    for text in ("great", "bad"):
+        preds = clf.predict(text, k=50)
+        assert {l for l, _ in preds} <= labels
+        assert all(np.isfinite(s) for _, s in preds)
+        if preds:
+            assert sum(s for _, s in preds) == pytest.approx(1.0, abs=1e-4)
+
+
+@pytest.mark.parametrize("seed", range(6))
+def test_random_operation_sequences_keep_every_structure_consistent(new_classifier, seed):
+    rng = random.Random(seed)
+    pool = ["great love", "good fast", "terrible hate", "awful bad", "okay fine", "cheap works",
+            "slow broke", "very good"]
+    names = ["a", "b", "c", "d", "e"]
+    clf = new_classifier()
+    added = []                                  # (text, label) pairs currently believed present
+
+    for _ in range(25):
+        op = rng.choice(["add", "add", "add", "forget", "remove", "remove_all_of_class", "predict"])
+        known = sorted(clf.label_to_id)
+        if op == "add" or not known:
+            texts = [rng.choice(pool) + f" {rng.randint(0, 3)}" for _ in range(rng.randint(1, 6))]
+            labels = [rng.choice(names) for _ in texts]
+            clf.add_examples(texts, labels)
+            added += list(zip(texts, labels))
+        elif op == "forget":
+            victim = rng.choice(known)
+            clf.forget(victim)
+            added = [(t, l) for t, l in added if l != victim]
+        elif op == "remove" and added:
+            text, label = rng.choice(added)
+            clf.remove_examples([text], label=label, retrain=rng.random() < 0.5)
+            added = [(t, l) for t, l in added if not (t == text and l == label)]
+        elif op == "remove_all_of_class":
+            victim = rng.choice(known)
+            doomed = sorted({t for t, l in added if l == victim})
+            if doomed:
+                clf.remove_examples(doomed, label=victim, retrain=False)
+                added = [(t, l) for t, l in added if l != victim]
+        else:
+            clf.predict(rng.choice(pool))
+        check_consistent(clf)
+
+
+def test_state_after_random_operations_survives_save_and_reload(new_classifier, tmp_path):
+    rng = random.Random(99)
+    clf = new_classifier(pooling="mean")
+    clf.add_examples(TEXTS, LABELS)
+    clf.add_examples(["okay fine", "cheap works"] * 6, ["c", "c"] * 6)
+    clf.forget("p")
+    clf.add_examples(["great love"] * 5, ["z"] * 5)
+    clf.remove_examples(["okay fine"], label="c", retrain=False)
+
+    clf.save(str(tmp_path), include_onnx=False)
+    loaded = load(tmp_path)
+    check_consistent(loaded)
+    assert set(loaded.label_to_id) == set(clf.label_to_id)
+    for text in ("great love", "terrible hate", "okay fine"):
+        assert [l for l, _ in loaded.predict(text, k=9)] == [l for l, _ in clf.predict(text, k=9)]
diff --git a/tests/test_robustness.py b/tests/test_robustness.py
new file mode 100644
index 0000000..3570263
--- /dev/null
+++ b/tests/test_robustness.py
@@ -0,0 +1,237 @@
+"""End-to-end quality and invariant checks.
+
+These do not need a benchmark run: the data is synthetic and perfectly
+separable (see tests/synthetic.py), so a shortfall is the classifier's fault.
+They are the tests that would have caught the weaknesses found while building
+the scikit-learn wrapper, and they check what a user relies on:
+
+  * it learns separable classes, at small and large sizes
+  * it keeps old classes when new ones arrive (no catastrophic forgetting)
+  * it survives class imbalance and label noise
+  * its output is always a valid, ordered probability distribution
+  * it is deterministic and independent of example order where documented
+
+Known weaknesses are recorded as strict expected failures, so they stay
+visible and the test fails (forcing the marker off) once someone fixes them.
+"""
+
+import numpy as np
+import pytest
+from hypothesis import HealthCheck, given, settings, strategies as st
+
+from adaptive_classifier import AdaptiveClassifier
+
+from .synthetic import Dataset, accuracy
+
+PROTOTYPE_ONLY = {"new_class_example_threshold": 0, "prototype_weight": 1.0, "neural_weight": 0.0}
+SEEDS = range(4)
+
+
+@pytest.fixture
+def data(monkeypatch):
+    return Dataset().patch(monkeypatch)
+
+
+def fit(new_classifier, texts, labels, **config):
+    clf = new_classifier(**config)
+    clf.add_examples(texts, labels)
+    return clf
+
+
+# --- it learns separable classes ------------------------------------------------------
+
+@pytest.mark.parametrize("n_per", [12, 30])
+@pytest.mark.parametrize("noise", [0.1, 0.2])
+def test_learns_separable_classes_with_enough_examples(new_classifier, data, n_per, noise):
+    train_acc, test_acc = [], []
+    for seed in SEEDS:
+        X, y = data.classes(n_per, noise=noise, seed=seed)
+        clf = fit(new_classifier, X, y)
+        Xt, yt = data.classes(15, noise=noise, seed=seed + 50, prefix="te")
+        train_acc.append(accuracy(clf, X, y))
+        test_acc.append(accuracy(clf, Xt, yt))
+    assert np.mean(train_acc) >= 0.98
+    assert np.mean(test_acc) >= 0.92
+
+
+@pytest.mark.parametrize("n_per", [1, 2, 3, 5, 8, 12])
+def test_prototype_only_is_accurate_at_every_size(new_classifier, data, n_per):
+    """The prototype path alone separates this data, even from a single example."""
+    for seed in SEEDS:
+        X, y = data.classes(n_per, noise=0.15, seed=seed)
+        clf = fit(new_classifier, X, y, **PROTOTYPE_ONLY)
+        Xt, yt = data.classes(15, noise=0.15, seed=seed + 50, prefix="te")
+        assert accuracy(clf, X, y) == 1.0
+        assert accuracy(clf, Xt, yt) >= 0.95
+
+
+@pytest.mark.xfail(
+    strict=True,
+    reason="Known weakness: with fewer than new_class_example_threshold (10) examples per "
+           "class the default blend gives the neural head 70% of the vote, and a head trained "
+           "on a handful of examples outvotes correct prototypes. Observed on separable data: "
+           "~75-92% training accuracy, ~66-83% held out, versus 100% from prototypes alone. "
+           "Remove this marker when fixed.",
+)
+@pytest.mark.parametrize("n_per", [2, 3, 5, 8])
+def test_default_config_fits_its_own_training_data_with_few_examples(new_classifier, data, n_per):
+    accs = []
+    for seed in SEEDS:
+        X, y = data.classes(n_per, noise=0.1, seed=seed)
+        accs.append(accuracy(fit(new_classifier, X, y), X, y))
+    assert np.mean(accs) >= 0.97, f"training accuracy {np.mean(accs):.2f} at {n_per} per class"
+
+
+# --- no catastrophic forgetting ---------------------------------------------------------
+
+def test_new_class_does_not_degrade_old_classes(new_classifier, data):
+    for seed in SEEDS:
+        X, y = data.classes(15, n_classes=2, noise=0.1, seed=seed)
+        Xt, yt = data.classes(15, n_classes=2, noise=0.1, seed=seed + 50, prefix="te")
+        clf = fit(new_classifier, X, y)
+        before = accuracy(clf, Xt, yt)
+
+        X3, y3 = data.classes(15, n_classes=3, noise=0.1, seed=seed + 7, prefix="new")
+        clf.add_examples(
+            [t for t, l in zip(X3, y3) if l == "C2"], [l for l in y3 if l == "C2"]
+        )
+        after = accuracy(clf, Xt, yt)
+        assert before >= 0.95
+        assert after >= before - 0.05, f"seed {seed}: {before:.2f} -> {after:.2f}"
+        assert "C2" in clf.label_to_id
+
+
+def test_new_class_is_itself_learned(new_classifier, data):
+    X, y = data.classes(15, n_classes=2, noise=0.1)
+    clf = fit(new_classifier, X, y)
+    X3, y3 = data.classes(15, n_classes=3, noise=0.1, seed=3, prefix="new")
+    new_texts = [t for t, l in zip(X3, y3) if l == "C2"]
+    clf.add_examples(new_texts, ["C2"] * len(new_texts))
+    Xt, yt = data.classes(15, n_classes=3, noise=0.1, seed=60, prefix="te")
+    new_only = [(t, l) for t, l in zip(Xt, yt) if l == "C2"]
+    recall = np.mean([clf.predict(t)[0][0] == "C2" for t, _ in new_only])
+    assert recall >= 0.9
+
+
+# --- imbalance and noise -------------------------------------------------------------------
+
+def test_minority_class_is_not_swallowed_by_a_large_one(new_classifier, data):
+    big, _ = data.classes(120, n_classes=1, noise=0.1, seed=1)
+    small_texts = [data.add(f"min_{i}", 1, 0.1, 900 + i) for i in range(12)]
+    clf = fit(new_classifier, big + small_texts, ["C0"] * len(big) + ["C1"] * len(small_texts))
+
+    held_out = [data.add(f"min_te_{i}", 1, 0.1, 5000 + i) for i in range(20)]
+    recall = np.mean([clf.predict(t)[0][0] == "C1" for t in held_out])
+    assert recall >= 0.9, f"minority recall {recall:.2f}"
+
+
+def test_a_few_wrong_labels_do_not_wreck_the_model(new_classifier, data):
+    rng = np.random.default_rng(0)
+    X, y = data.classes(30, noise=0.1, seed=2)
+    noisy = list(y)
+    for i in rng.choice(len(y), size=len(y) // 10, replace=False):    # flip 10%
+        noisy[i] = f"C{(int(y[i][1:]) + 1) % 3}"
+    clf = fit(new_classifier, X, noisy)
+    Xt, yt = data.classes(15, noise=0.1, seed=70, prefix="te")
+    assert accuracy(clf, Xt, yt) >= 0.85
+
+
+def test_wildly_overlapping_classes_degrade_gracefully(new_classifier, data):
+    """Indistinguishable classes must give near-uniform confidence, not a confident guess."""
+    X, y = data.classes(20, n_classes=2, noise=0.0, seed=0)        # identical points per class
+    same = [data.add(f"same_{i}", 0, 0.0, 0) for i in range(len(X))]
+    clf = fit(new_classifier, same, [("A", "B")[i % 2] for i in range(len(same))])
+    top = clf.predict(same[0])
+    assert len(top) == 2
+    assert abs(top[0][1] - top[1][1]) < 0.2
+
+
+# --- determinism and order --------------------------------------------------------------------
+
+def test_training_is_deterministic(new_classifier, data):
+    X, y = data.classes(12, noise=0.15, seed=4)
+    first = fit(new_classifier, X, y)
+    second = fit(new_classifier, X, y)
+    probe = data.add("probe", 0, 0.2, 77)
+    assert first.predict(probe) == second.predict(probe)
+
+
+def test_prototype_only_predictions_ignore_example_order(new_classifier, data):
+    X, y = data.classes(10, noise=0.15, seed=5)
+    order = np.random.default_rng(1).permutation(len(X))
+    a = fit(new_classifier, X, y, **PROTOTYPE_ONLY)
+    b = fit(new_classifier, [X[i] for i in order], [y[i] for i in order], **PROTOTYPE_ONLY)
+    for q in [data.add(f"q{i}", i % 3, 0.2, 800 + i) for i in range(10)]:
+        pa, pb = dict(a.predict(q)), dict(b.predict(q))
+        assert pa.keys() == pb.keys()
+        assert all(pa[k] == pytest.approx(pb[k], abs=1e-5) for k in pa)
+
+
+def test_adding_in_one_batch_or_many_gives_the_same_prototypes(new_classifier, data):
+    X, y = data.classes(10, noise=0.15, seed=6)
+    whole = fit(new_classifier, X, y, **PROTOTYPE_ONLY)
+    parts = new_classifier(**PROTOTYPE_ONLY)
+    for i in range(0, len(X), 7):
+        parts.add_examples(X[i:i + 7], y[i:i + 7])
+    for label in whole.memory.prototypes:
+        assert np.allclose(whole.memory.prototypes[label].numpy(),
+                           parts.memory.prototypes[label].numpy(), atol=1e-6)
+
+
+def test_predict_and_predict_batch_agree(new_classifier, data):
+    X, y = data.classes(12, noise=0.15, seed=8)
+    clf = fit(new_classifier, X, y)
+    queries = [data.add(f"b{i}", i % 3, 0.25, 300 + i) for i in range(12)]
+    singles = [clf.predict(q, k=3) for q in queries]
+    batch = clf.predict_batch(queries, k=3)
+    for one, many in zip(singles, batch):
+        assert [l for l, _ in one] == [l for l, _ in many]
+        assert all(a == pytest.approx(b, abs=1e-5) for (_, a), (_, b) in zip(one, many))
+
+
+# --- output invariants (property-based) --------------------------------------------------------
+
+WORDS = "great terrible okay product service love hate awful fine good bad slow fast cheap broke works".split()
+texts_st = st.lists(st.sampled_from(WORDS), min_size=1, max_size=8).map(" ".join)
+labels_st = st.sampled_from(["pos", "neg", "neu", "été", "类别", "a b", "x-1"])
+
+
+def assert_valid_distribution(preds, known, k):
+    assert len(preds) <= k
+    labels = [l for l, _ in preds]
+    scores = [s for _, s in preds]
+    assert len(set(labels)) == len(labels), "a label appeared twice"
+    assert set(labels) <= set(known)
+    assert all(np.isfinite(scores)) and all(0.0 <= s <= 1.0 + 1e-6 for s in scores)
+    assert scores == sorted(scores, reverse=True)
+
+
+@settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture])
+@given(
+    train=st.lists(st.tuples(texts_st, labels_st), min_size=1, max_size=14),
+    queries=st.lists(texts_st, min_size=1, max_size=4),
+    k=st.integers(min_value=0, max_value=9),
+)
+def test_predictions_are_always_a_valid_ordered_distribution(new_classifier, train, queries, k):
+    clf = new_classifier()
+    texts, labels = zip(*train)
+    clf.add_examples(list(texts), list(labels))
+    known = set(labels)
+
+    for q in queries:
+        preds = clf.predict(q, k=k)
+        assert_valid_distribution(preds, known, k)
+        if k >= len(known) and preds:
+            assert sum(s for _, s in preds) == pytest.approx(1.0, abs=1e-5)
+    for preds in clf.predict_batch(queries, k=max(k, 1)):
+        assert_valid_distribution(preds, known, max(k, 1))
+
+
+@settings(max_examples=15, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture])
+@given(train=st.lists(st.tuples(texts_st, labels_st), min_size=2, max_size=12), q=texts_st)
+def test_scores_and_ood_are_finite_for_any_training_set(new_classifier, train, q):
+    clf = new_classifier()
+    texts, labels = zip(*train)
+    clf.add_examples(list(texts), list(labels))
+    score = clf.ood_score(q)
+    assert np.isfinite(score) and score >= 0.0

From 6f0872710b343f8e63bf40ddb7ec782519df5e12 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 08:57:56 +0000
Subject: [PATCH 07/16] CI: prefetch models under canonical ids, fail fast on
 permanent errors

The models job failed on bert-base-uncased and distilbert-base-uncased: the
Hub's download-token endpoint returns 404 for the short alias names, on
every retry. Canonical ids (distilbert/distilbert-base-cased,
google-bert/bert-large-cased) downloaded fine.

- Tests and the prefetch list use google-bert/bert-base-uncased and
  distilbert/distilbert-base-uncased (same models)
- tests/test_ci_models.py rejects alias ids so they cannot return
- The prefetch script no longer retries HTTP 401/403/404, which are
  permanent, so such a failure shows up in seconds instead of minutes

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 scripts/prefetch_test_models.py               | 14 +++++++++++---
 tests/test_ci_models.py                       |  7 +++++++
 tests/test_classifier.py                      |  4 ++--
 tests/test_ewc.py                             |  4 ++--
 tests/test_new_class_accuracy_preservation.py |  8 ++++----
 5 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/scripts/prefetch_test_models.py b/scripts/prefetch_test_models.py
index 812b4ae..1c89038 100644
--- a/scripts/prefetch_test_models.py
+++ b/scripts/prefetch_test_models.py
@@ -14,11 +14,13 @@
 from huggingface_hub import snapshot_download
 
 # Keep in sync with the model names used in tests/ (tests/test_ci_models.py
-# fails when a test uses a model that is not listed here).
+# fails when a test uses a model that is not listed here). Use canonical
+# "owner/name" ids: the old short aliases ("bert-base-uncased") make the Hub's
+# download-token endpoint return 404, so they cannot be prefetched.
 MODELS = [
     "prajjwal1/bert-tiny",
-    "bert-base-uncased",
-    "distilbert-base-uncased",
+    "google-bert/bert-base-uncased",
+    "distilbert/distilbert-base-uncased",
     "distilbert/distilbert-base-cased",
     "google-bert/bert-large-cased",
     "answerdotai/ModernBERT-base",
@@ -48,6 +50,12 @@ def main() -> int:
                 print(f"ok    {model}")
                 break
             except Exception as error:  # network, 429, 5xx
+                status = getattr(getattr(error, "response", None), "status_code", None)
+                if status in (401, 403, 404):
+                    # Permanent (missing repo, bad id, no access); retrying only delays the failure.
+                    print(f"fail  {model}: HTTP {status} - {str(error)[:200]}", flush=True)
+                    failed.append(model)
+                    break
                 wait = min(120, 5 * 2 ** attempt)
                 print(f"retry {model} ({attempt}/{ATTEMPTS}): {type(error).__name__}: "
                       f"{str(error)[:200]} - waiting {wait}s", flush=True)
diff --git a/tests/test_ci_models.py b/tests/test_ci_models.py
index 551ccec..a6c30b2 100644
--- a/tests/test_ci_models.py
+++ b/tests/test_ci_models.py
@@ -43,3 +43,10 @@ def test_every_model_named_in_tests_is_prefetched():
 def test_prefetch_list_has_no_unused_models():
     unused = _prefetch_models() - _models_used_by_tests()
     assert not unused, f"Remove {sorted(unused)} from scripts/prefetch_test_models.py"
+
+
+def test_tests_use_canonical_model_ids():
+    """Short aliases such as "bert-base-uncased" cannot be prefetched: the Hub's
+    download-token endpoint returns 404 for them."""
+    aliases = {m for m in _models_used_by_tests() if "/" not in m}
+    assert not aliases, f"Use the canonical owner/name id instead of {sorted(aliases)}"
diff --git a/tests/test_classifier.py b/tests/test_classifier.py
index f3106b1..3735e35 100644
--- a/tests/test_classifier.py
+++ b/tests/test_classifier.py
@@ -8,7 +8,7 @@
 
 @pytest.fixture
 def base_classifier():
-    return AdaptiveClassifier("bert-base-uncased")
+    return AdaptiveClassifier("google-bert/bert-base-uncased")
 
 @pytest.fixture
 def sample_data():
@@ -206,7 +206,7 @@ def test_num_representative_examples(sample_data):
     config = {
         'num_representative_examples': 2  # Set to keep only 2 example per class
     }
-    classifier = AdaptiveClassifier("bert-base-uncased", config=config)
+    classifier = AdaptiveClassifier("google-bert/bert-base-uncased", config=config)
 
     # Add more examples than num_representative_examples
     texts, labels = sample_data
diff --git a/tests/test_ewc.py b/tests/test_ewc.py
index f237147..08b9f3e 100644
--- a/tests/test_ewc.py
+++ b/tests/test_ewc.py
@@ -91,7 +91,7 @@ def test_adaptive_classifier_with_many_classes():
     torch.manual_seed(42)
     
     # Create classifier
-    classifier = AdaptiveClassifier('distilbert-base-uncased', device='cpu')
+    classifier = AdaptiveClassifier('distilbert/distilbert-base-uncased', device='cpu')
     
     # Simulate many classes with few examples each
     num_classes = 20
@@ -155,7 +155,7 @@ def test_ewc_loss_computation(simple_model, small_dataset):
 
 def test_progressive_class_addition():
     """Test adding classes progressively (triggers EWC multiple times)."""
-    classifier = AdaptiveClassifier('distilbert-base-uncased', device='cpu')
+    classifier = AdaptiveClassifier('distilbert/distilbert-base-uncased', device='cpu')
     
     # Phase 1: Add initial classes
     phase1_texts = ["Good product", "Bad service", "Average quality"]
diff --git a/tests/test_new_class_accuracy_preservation.py b/tests/test_new_class_accuracy_preservation.py
index 78516c0..3197822 100644
--- a/tests/test_new_class_accuracy_preservation.py
+++ b/tests/test_new_class_accuracy_preservation.py
@@ -78,7 +78,7 @@ def many_class_data():
 @pytest.fixture
 def large_classifier():
     """Create classifier that will be used for many-class testing."""
-    return AdaptiveClassifier("bert-base-uncased", seed=42)
+    return AdaptiveClassifier("google-bert/bert-base-uncased", seed=42)
 
 
 def test_accuracy_preservation_after_adding_new_classes(large_classifier, many_class_data):
@@ -278,7 +278,7 @@ def test_weight_structure_preservation():
     """Test that head structure expands correctly and uses update_num_classes instead of reinitialization."""
     torch.manual_seed(42)
 
-    classifier = AdaptiveClassifier("bert-base-uncased", seed=42)
+    classifier = AdaptiveClassifier("google-bert/bert-base-uncased", seed=42)
 
     # Add initial classes
     initial_texts = ["Text about cats", "Text about dogs"]
@@ -304,7 +304,7 @@ def test_class_expansion_behavior():
     """Test that adding new classes expands the head instead of reinitializing."""
     torch.manual_seed(42)
 
-    classifier = AdaptiveClassifier("bert-base-uncased", seed=42)
+    classifier = AdaptiveClassifier("google-bert/bert-base-uncased", seed=42)
 
     # Add initial classes
     initial_texts = ["Text about cats", "Text about dogs"]
@@ -340,7 +340,7 @@ def test_improved_accuracy_preservation():
     torch.manual_seed(42)
     np.random.seed(42)
 
-    classifier = AdaptiveClassifier("bert-base-uncased", seed=42)
+    classifier = AdaptiveClassifier("google-bert/bert-base-uncased", seed=42)
 
     # Create a more controlled test with less training
     initial_texts = [

From fb8b3347aee0563ded864f535e8d81fb4cb04eef Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 10:35:56 +0000
Subject: [PATCH 08/16] Fix small-data accuracy: train the head properly,
 sharpen prototype scores, ramp the head in

The default blend scored 0.55-0.73 on perfectly separable synthetic data
with fewer than 10 examples per class, where prototypes alone scored
0.96-1.0. Experiments on synthetic Gaussian, XOR and incremental-class
data traced it to three causes, all fixed (this is a default-behaviour
change, hence the 0.3.0 minor bump):

- The neural head was badly undertrained: ~10 optimiser steps and an early
  stop at the first loss plateau. It could not learn a 2-D XOR from 120
  examples (0.49-0.70). It now trains for at least `head_steps` (300) steps
  with a cosine-decayed `head_learning_rate` (0.003) and solves XOR (1.00)
- Prototype scores were softmax(exp(-d)), which barely separated the
  nearest class from the rest (~0.45 vs 0.21), so a confident head always
  outvoted them. They are now softmax(-d / prototype_temperature), 0.25
- Classes below new_class_example_threshold used a fixed 0.3 prototype /
  0.7 head split. The head's weight now ramps up from zero with the number
  of examples, so there is no accuracy cliff at the threshold. Setting both
  new_class_* weights restores a fixed split

Backward compatibility: classifiers saved by an earlier version keep the
legacy scoring when loaded (prototype_temperature absent -> None).

Also updates the portable reference implementation and the deployment guide
to the new scoring math, the README section on predictions, the changelog,
and the real-model tests that encoded the old behaviour (blend weights, and
a 0.85-0.95 confidence window that pinned the old blend).

Tests: the known-weakness expected failure is now a regression test, plus
XOR (only the head can solve it, and prototypes alone cannot), head-only
and head_steps tests, ramp continuity and configuration, scoring sharpness,
legacy-save compatibility, and portable-reference parity for every scoring
branch. Mutation-checked: ten deliberate reversions, all caught.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                           |   9 +-
 README.md                              |  24 ++-
 docs/deployment.md                     |   7 +-
 examples/portable_inference.py         |  34 ++--
 src/adaptive_classifier/classifier.py  | 145 ++++++++--------
 src/adaptive_classifier/memory.py      |  41 ++---
 src/adaptive_classifier/models.py      |  34 +++-
 tests/test_failure_modes.py            |   2 +-
 tests/test_pooling.py                  |  23 ++-
 tests/test_portable_export.py          |  25 +++
 tests/test_reported_confidence_drop.py |  19 ++-
 tests/test_robustness.py               | 218 +++++++++++++++++++++++--
 12 files changed, 420 insertions(+), 161 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3eb535d..93d16b8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,16 +9,19 @@
 - `docs/deployment.md` and `examples/portable_inference.py`: run a saved classifier without PyTorch or FAISS (multilingual models, .NET, JavaScript).
 - `save()` now writes the tokenizer files and the resolved pooling mode, so a saved directory is self-contained.
 
-### Changed
+### Changed (default behaviour: predictions change)
+- **The neural head is now actually trained.** It used to get about ten optimiser steps and stop at the first loss plateau, so it could not learn even a 2-D XOR from 120 examples. It now trains for at least `head_steps` (default 300) steps with a cosine-decayed `head_learning_rate` (default 0.003). Small memories take roughly 1-2 s per `add_examples` call; lower `head_steps` if you add examples very frequently.
+- **Prototype scores are sharper.** They are now `softmax(-distance² / prototype_temperature)` (default 0.25) instead of `softmax(exp(-distance))`, which barely separated the nearest class from the rest (about 0.45 vs 0.21), so a confident head could always outvote it. `prototype_temperature=None` restores the old scoring.
+- **Classes with few examples lean on their prototypes.** Below `new_class_example_threshold` examples the head's weight ramps up from zero, instead of the old fixed 0.3 prototype / 0.7 head split. On separable synthetic data the old default scored 0.55-0.73 with fewer than 10 examples per class where prototypes alone scored 0.96-1.0; accuracy no longer drops off a cliff at the threshold. Trade-off: a class added with a single example is recognised less often than before, because the old split biased the head toward new classes at the expense of existing ones. Set both `new_class_prototype_weight` and `new_class_neural_weight` to restore a fixed split.
+- Classifiers saved by an earlier version keep their old scoring when loaded, so their predictions do not change.
 - **Python 3.10 or newer is required** (3.8 and 3.9 are end-of-life).
 - Packaging moved to a `[project]` table in `pyproject.toml`; `setup.py` and `requirements.txt` are gone. Install extras: `[test]`, `[scripts]`. `tqdm` is no longer a core dependency.
 - `clear_memory(labels=[...])` now removes those classes completely (it used to leave them in the label map and neural head, so they could still be predicted).
 
 ### Fixed
+- Accuracy on small training sets (see above): the default configuration got only 75-92% of its own training data right on perfectly separable data with fewer than 10 examples per class.
 - Integer (or other non-string) labels corrupted a classifier on save/reload, duplicating classes. `add_examples` now rejects non-string labels with a clear error; convert with `str(label)`.
 - `predict(k=-1)` returned a silently truncated list; negative or non-integer `k` now raises `ValueError`.
 - `None` as a text or label, and mixed label types, now raise a clear `ValueError` before anything is changed.
 - Loading a directory without `config.json` now says so instead of reporting an invalid Hub repo id.
 
-### Known issues
-- With fewer than 10 examples per class the default blend gives the neural head 70% of the vote. A head trained on so few examples can outvote correct prototypes, so accuracy on small training sets is lower than prototype-only prediction. Tracked by a strict expected-failure test in `tests/test_robustness.py`.
diff --git a/README.md b/README.md
index 653cad6..778f4e4 100644
--- a/README.md
+++ b/README.md
@@ -144,16 +144,9 @@ Classifiers saved before 0.2.0 keep CLS pooling when reloaded, so upgrading does
 not change their predictions. Retrain, or pass `pooling` explicitly, to pick up
 the new behaviour.
 
-## Prototype and neural weights
+## How a prediction is made
 
-A prediction mixes two signals: similarity to the class prototypes, and the
-trained neural head. `prototype_weight` and `neural_weight` control the mix.
-
-Both keys existed before 0.2.0 but were ignored. Every prediction path used a
-hardcoded 0.7 / 0.3, so `predict` and `predict_batch` could also disagree once a
-caller changed the config. They work now, and the best split depends on how well
-your encoder separates your classes: an encoder that bunches everything together
-produces weak prototypes and is better off leaning on the head.
+A prediction mixes two signals: the **prototype score** (how close the text is to each class's mean embedding) and the **neural head** (a small trained network on the same embeddings). `prototype_weight` and `neural_weight` (default 0.7 / 0.3) control the mix.
 
 ```python
 classifier = AdaptiveClassifier(
@@ -162,10 +155,15 @@ classifier = AdaptiveClassifier(
 )
 ```
 
-A class with few examples has an unreliable prototype, so it leans on the head
-until it is established. That is configurable too, via
-`new_class_example_threshold` (default 10), `new_class_prototype_weight` (0.3)
-and `new_class_neural_weight` (0.7).
+The best split depends on how well your encoder separates your classes: an encoder that bunches everything together produces weak prototypes and is better off leaning on the head.
+
+**Classes with few examples.** Until a class has `new_class_example_threshold` examples (default 10), the head's weight ramps up linearly from zero to `neural_weight` and the prototype takes the rest, so a brand-new class is judged almost entirely by its prototype. (Before 0.3.0 such classes used a fixed 0.3 / 0.7 split in favour of the head, and a head fitted to a handful of points could confidently outvote a correct prototype. To get a fixed split back, set both `new_class_prototype_weight` and `new_class_neural_weight`.)
+
+**Prototype sharpness.** Prototype scores are `softmax(-distance² / prototype_temperature)` with a default temperature of 0.25. Lower values make the nearest class stand out more. Set `prototype_temperature` to `None` for the scoring used before 0.3.0, which barely separated the nearest class from the rest.
+
+**Head training.** The head is trained for at least `head_steps` optimiser steps (default 300) with a cosine-decayed `head_learning_rate` (default 0.003). That takes roughly 1-2 seconds on a small memory; lower `head_steps` (for example to 100) if you add examples very frequently. Before 0.3.0 the head got about ten steps, too few to learn even a simple XOR pattern, which prototypes alone cannot separate.
+
+Classifiers saved before 0.3.0 keep their old scoring when reloaded, so upgrading does not change their predictions. Retrain, or set the keys above, to adopt the new behaviour.
 
 ## ⚡ Quick Start
 
diff --git a/docs/deployment.md b/docs/deployment.md
index b9b7afd..943b9f5 100644
--- a/docs/deployment.md
+++ b/docs/deployment.md
@@ -73,9 +73,12 @@ Given the files above, `predict(text)` is:
    - `mean`: sum of token vectors weighted by `attention_mask`, divided by the number of real tokens.
    - `cls`: the first token's vector.
 4. **L2-normalise** the vector.
-5. **Prototype scores.** For each class prototype `p`, compute `exp(-||p - v||²)` (squared Euclidean distance), then take a softmax across classes.
+5. **Prototype scores.** For each class prototype `p`, compute the squared Euclidean distance `d = ||p - v||²`, then take a softmax of `-d / prototype_temperature` across classes (the default temperature is 0.25). If `prototype_temperature` is `null` in `config.json`, the classifier predates 0.3.0 and uses the legacy score instead: a softmax of `exp(-d)`.
 6. **Neural head scores** (skip if `model.safetensors` has no head). Compute `h1 = ReLU(W0·v + b0)`, `h2 = ReLU(W3·h1 + b3)`, `logits = W6·h2 + b6`, then softmax. Dropout is inactive at inference. Weights are stored as `(out, in)`.
-7. **Blend per class.** `score = proto_weight · proto_score + neural_weight · head_score`. A class with fewer than `new_class_example_threshold` (default 10) entries in `training_history` uses the `new_class_*` weights (default 0.3 / 0.7) instead of `prototype_weight` / `neural_weight` (default 0.7 / 0.3).
+7. **Blend per class.** `score = proto_weight · proto_score + neural_weight · head_score`, where the weights depend on how many examples the class was trained on (`training_history` in `config.json`):
+   - With at least `new_class_example_threshold` examples (default 10): `prototype_weight` and `neural_weight` (defaults 0.7 and 0.3).
+   - With fewer, and both `new_class_prototype_weight` and `new_class_neural_weight` set: those two fixed values (this is how classifiers saved before 0.3.0 behave: 0.3 and 0.7).
+   - With fewer, and those two unset (`null`, the default since 0.3.0): the head's weight ramps up from zero. With `share = examples / threshold`, the prototype weight is `(1 - share) + share · prototype_weight` and the neural weight is `share · neural_weight`.
 8. **Normalise** the blended scores to sum to 1 and sort descending.
 
 This applies to the default (non-strategic) mode. Strategic mode adds extra steps that are not covered here.
diff --git a/examples/portable_inference.py b/examples/portable_inference.py
index 7c0be20..70f5d55 100644
--- a/examples/portable_inference.py
+++ b/examples/portable_inference.py
@@ -43,8 +43,12 @@ def __init__(self, directory, pooling=None, prefer_quantized=True):
         self.prototype_weight = settings.get("prototype_weight", 0.7)
         self.neural_weight = settings.get("neural_weight", 0.3)
         self.new_threshold = settings.get("new_class_example_threshold", 10)
-        self.new_prototype_weight = settings.get("new_class_prototype_weight", 0.3)
-        self.new_neural_weight = settings.get("new_class_neural_weight", 0.7)
+        # None (the default since 0.3.0) means the head's weight ramps up with
+        # the number of examples; both set means a fixed split below the threshold.
+        self.new_prototype_weight = settings.get("new_class_prototype_weight")
+        self.new_neural_weight = settings.get("new_class_neural_weight")
+        # None selects the legacy scoring used before 0.3.0.
+        self.temperature = settings.get("prototype_temperature")
 
         # Current versions save the resolved mode. Older directories say 'auto'
         # or have no pooling key; pass `pooling` yourself for those: use what
@@ -110,24 +114,34 @@ def _head_probs(self, emb):
                 x = np.maximum(x, 0.0)
         return softmax(x)
 
+    def _weights(self, label):
+        """Prototype and neural weight for one class (see AdaptiveClassifier._blend_weights)."""
+        trained = self.cfg.get("training_history", {}).get(label, 0)
+        if trained >= self.new_threshold:
+            return self.prototype_weight, self.neural_weight
+        if self.new_prototype_weight is not None and self.new_neural_weight is not None:
+            return self.new_prototype_weight, self.new_neural_weight
+        share = trained / self.new_threshold
+        return (1 - share) + share * self.prototype_weight, share * self.neural_weight
+
     def predict(self, text, k=5):
         emb = self.embed(text)
 
-        # Prototype score: exp(-squared L2 distance), softmaxed over classes.
+        # Prototype score: softmax over classes of -squared_L2_distance / temperature
+        # (legacy, when temperature is null: softmax of exp(-squared_L2_distance)).
         labels = list(self.prototypes)
-        sims = np.array([np.exp(-np.sum((p - emb) ** 2)) for p in self.prototypes.values()])
-        proto_scores = dict(zip(labels, softmax(sims)))
+        dist = np.array([np.sum((p - emb) ** 2) for p in self.prototypes.values()])
+        if self.temperature:
+            proto_scores = dict(zip(labels, softmax(-dist / self.temperature)))
+        else:
+            proto_scores = dict(zip(labels, softmax(np.exp(-dist))))
         head_scores = (
             dict(zip(self.labels, self._head_probs(emb))) if self.head is not None else {}
         )
 
-        history = self.cfg.get("training_history", {})
         combined = {}
         for label in self.labels:
-            if history.get(label, 0) < self.new_threshold:
-                wp, wn = self.new_prototype_weight, self.new_neural_weight
-            else:
-                wp, wn = self.prototype_weight, self.neural_weight
+            wp, wn = self._weights(label)
             combined[label] = proto_scores.get(label, 0.0) * wp + head_scores.get(label, 0.0) * wn
 
         total = sum(combined.values())
diff --git a/src/adaptive_classifier/classifier.py b/src/adaptive_classifier/classifier.py
index 75e082e..7197938 100644
--- a/src/adaptive_classifier/classifier.py
+++ b/src/adaptive_classifier/classifier.py
@@ -6,6 +6,7 @@
 from typing import List, Dict, Optional, Tuple, Any, Set, Union
 import logging
 import copy
+import math
 import numbers
 from pathlib import Path
 from safetensors.torch import save_file, load_file
@@ -949,6 +950,15 @@ def _from_pretrained(
         if isinstance(saved_config, dict) and 'pooling' not in saved_config:
             config_dict['config'] = {**saved_config, 'pooling': 'cls'}
 
+        # Likewise, classifiers saved before 0.3.0 were built (head included)
+        # with the legacy prototype scoring and carry explicit new-class weights
+        # in their stored config. Keep that scoring so their predictions do not
+        # change on upgrade; retrain, or set `prototype_temperature`, to adopt
+        # the new behaviour.
+        saved_config = config_dict.get('config')
+        if isinstance(saved_config, dict) and 'prototype_temperature' not in saved_config:
+            config_dict['config'] = {**saved_config, 'prototype_temperature': None}
+
         # Load examples
         with open(model_path / "examples.json", "r", encoding="utf-8") as f:
             saved_examples = json.load(f)
@@ -1434,16 +1444,27 @@ def _initialize_adaptive_head(self):
     def _blend_weights(self, label: str) -> Tuple[float, float]:
         """Prototype and neural weights for one label.
 
-        A class the model has barely seen has an unreliable prototype, so the
-        neural head carries more of the decision until the class is
-        established. Both regimes are configurable; before 0.2.0 these were
-        hardcoded and `prototype_weight` did nothing.
+        A class with fewer than `new_class_example_threshold` examples has too
+        little evidence behind the neural head, so the head's weight ramps up
+        linearly from zero and the prototype carries the rest. With a fixed
+        `new_class_prototype_weight` and `new_class_neural_weight` both set,
+        such classes use that split instead (the behaviour before 0.3.0, whose
+        defaults were 0.3 / 0.7).
         """
+        config = self.config
+        established = (config.prototype_weight, config.neural_weight)
         trained = self.training_history.get(label, 0)
-        if trained < getattr(self.config, 'new_class_example_threshold', 10):
-            return (getattr(self.config, 'new_class_prototype_weight', 0.3),
-                    getattr(self.config, 'new_class_neural_weight', 0.7))
-        return (self.config.prototype_weight, self.config.neural_weight)
+        threshold = getattr(config, 'new_class_example_threshold', 10)
+        if trained >= threshold:
+            return established
+
+        fixed_proto = getattr(config, 'new_class_prototype_weight', None)
+        fixed_neural = getattr(config, 'new_class_neural_weight', None)
+        if fixed_proto is not None and fixed_neural is not None:
+            return (fixed_proto, fixed_neural)
+
+        share = trained / threshold          # 0 for a brand-new class, 1 at the threshold
+        return ((1 - share) + share * established[0], share * established[1])
 
     def _resolve_pooling(self) -> str:
         """Decide how to pool token embeddings into one vector per text.
@@ -1677,101 +1698,63 @@ def merge_classifiers(self, other: 'AdaptiveClassifier') -> 'AdaptiveClassifier'
         return self
     
     def _train_adaptive_head(self, epochs: int = 10):
-        """Train the adaptive head with improved stability."""
+        """Train the adaptive head on every stored example.
+
+        Runs `max(epochs * batches_per_epoch, head_steps)` optimiser steps with
+        a cosine-decayed learning rate and no early stopping. The step floor
+        matters on small memories: ten epochs over a few dozen examples is only
+        about ten steps, which is far too few for the head to learn anything.
+        """
         if not self.memory.examples:
             return
-            
-        # Prepare training data
+
         all_embeddings = []
         all_labels = []
-        
+
         # Sort examples for deterministic order
         for label in sorted(self.memory.examples.keys()):
             examples = sorted(self.memory.examples[label], key=lambda x: x.text)
             for example in examples:
                 all_embeddings.append(example.embedding)
-                # Convert string labels to numeric indices
                 all_labels.append(self.label_to_id[example.label])
-        
-        all_embeddings = torch.stack(all_embeddings)
-        # Ensure labels are Long tensor
-        all_labels = torch.tensor(all_labels, dtype=torch.long, device=self.device)
-        
-        # Normalize embeddings for stable training
-        all_embeddings = F.normalize(all_embeddings, p=2, dim=1)
-        
-        # Create deterministic data loader
-        dataset = torch.utils.data.TensorDataset(all_embeddings, all_labels)
-        loader = torch.utils.data.DataLoader(
-            dataset,
-            batch_size=min(32, len(all_embeddings)),
-            shuffle=True,
-            generator=torch.Generator().manual_seed(42)
-        )
-        
-        # Training setup
+
+        embeddings = F.normalize(torch.stack(all_embeddings), p=2, dim=1).to(self.device)
+        labels = torch.tensor(all_labels, dtype=torch.long, device=self.device)
+
+        batch_size = min(32, len(embeddings))
+        batches_per_epoch = math.ceil(len(embeddings) / batch_size)
+        total_steps = max(epochs * batches_per_epoch,
+                          int(getattr(self.config, 'head_steps', 0) or 0))
+
         self.adaptive_head.train()
         criterion = nn.CrossEntropyLoss()
         optimizer = torch.optim.AdamW(
             self.adaptive_head.parameters(),
-            lr=0.001,
+            lr=getattr(self.config, 'head_learning_rate', 0.001),
             weight_decay=0.01,
             betas=(0.9, 0.999)
         )
-        
-        # Learning rate scheduler for stability
-        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-            optimizer,
-            mode='min',
-            factor=0.5,
-            patience=2,
-        )
-        
-        best_loss = float('inf')
-        patience_counter = 0
-        patience = 3
-        
-        for epoch in range(epochs):
-            total_loss = 0
-            for batch_embeddings, batch_labels in loader:
-                batch_embeddings = batch_embeddings.to(self.device)
-                batch_labels = batch_labels.to(self.device)
-                
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps)
+        generator = torch.Generator().manual_seed(42)
+
+        step = 0
+        while step < total_steps:
+            order = torch.randperm(len(embeddings), generator=generator)
+            for start in range(0, len(order), batch_size):
+                index = order[start:start + batch_size]
                 optimizer.zero_grad()
-                outputs = self.adaptive_head(batch_embeddings)
-                
-                # Add shape debugging
-                if epoch == 0 and total_loss == 0:  # Only for first batch of first epoch
-                    logger.debug(f"outputs shape: {outputs.shape}")
-                    logger.debug(f"batch_labels shape: {batch_labels.shape}")
-                    logger.debug(f"batch_labels content: {batch_labels}")
-                
-                loss = criterion(outputs, batch_labels)
+                loss = criterion(self.adaptive_head(embeddings[index]), labels[index])
                 loss.backward()
-                
-                torch.nn.utils.clip_grad_norm_(
-                    self.adaptive_head.parameters(),
-                    max_norm=1.0
-                )
+                torch.nn.utils.clip_grad_norm_(self.adaptive_head.parameters(), max_norm=1.0)
                 optimizer.step()
-                
-                total_loss += loss.item()
-            
-            avg_loss = total_loss / len(loader)
-            scheduler.step(avg_loss)
-            
-            # Early stopping check
-            if avg_loss < best_loss:
-                best_loss = avg_loss
-                patience_counter = 0
-            else:
-                patience_counter += 1
-                if patience_counter >= patience:
-                    logger.debug(f"Early stopping at epoch {epoch + 1}")
+                scheduler.step()
+                step += 1
+                if step >= total_steps:
                     break
-        
+
+        self.adaptive_head.eval()
         self.train_steps += 1
-    
+
     def _update_adaptive_head(self):
         """Update adaptive head for new classes."""
         num_classes = len(self.label_to_id)
diff --git a/src/adaptive_classifier/memory.py b/src/adaptive_classifier/memory.py
index ddcb74d..6c69d55 100644
--- a/src/adaptive_classifier/memory.py
+++ b/src/adaptive_classifier/memory.py
@@ -116,26 +116,27 @@ def get_nearest_prototypes(
             k = min(k, self.index.ntotal)
             distances, indices = self.index.search(query_np, k)
             
-            # Convert distances to similarities using exponential scaling
-            similarities = np.exp(-distances[0])  # Apply to first (only) query result
-            
-            # Convert to labels and scores
-            results = []
-            for idx, similarity in zip(indices[0], similarities):
-                if idx >= 0:  # Valid index
-                    label = self.index_to_label[int(idx)]
-                    score = float(similarity)
-                    results.append((label, score))
-            
-            # Normalize scores with softmax
-            if results:
-                scores = torch.tensor([score for _, score in results])
-                normalized_scores = torch.nn.functional.softmax(scores, dim=0)
-                results = [
-                    (label, float(score)) 
-                    for (label, _), score in zip(results, normalized_scores)
-                ]
-            
+            temperature = getattr(self.config, 'prototype_temperature', None)
+            valid = [(int(i), float(d)) for i, d in zip(indices[0], distances[0]) if i >= 0]
+            if not valid:
+                return []
+
+            if temperature:
+                # Softmax over negative squared distance: the nearest class gets
+                # most of the mass when it is clearly nearer than the others.
+                scores = torch.softmax(
+                    torch.tensor([-d / temperature for _, d in valid]), dim=0
+                )
+            else:
+                # Legacy scoring (before 0.3.0): exp(-distance), then softmax.
+                scores = torch.softmax(
+                    torch.tensor([float(np.exp(-d)) for _, d in valid]), dim=0
+                )
+
+            results = [
+                (self.index_to_label[i], float(score))
+                for (i, _), score in zip(valid, scores)
+            ]
             return results
     
     def _update_prototype(self, label: str):
diff --git a/src/adaptive_classifier/models.py b/src/adaptive_classifier/models.py
index 14b44d1..f4e7b78 100644
--- a/src/adaptive_classifier/models.py
+++ b/src/adaptive_classifier/models.py
@@ -161,11 +161,34 @@ def __init__(self, config: Optional[Dict[str, Any]] = None):
         # to be a knob that works.
         self.prototype_weight = self.config.get('prototype_weight', 0.7)
         self.neural_weight = self.config.get('neural_weight', 0.3)
-        # A class with few examples has an unreliable prototype, so the neural
-        # head is trusted more until the class is established.
+        # A class with few examples has little evidence behind either signal, and
+        # a neural head fitted to a handful of points is the less trustworthy of
+        # the two (it can confidently outvote a correct prototype). Until a class
+        # has `new_class_example_threshold` examples the head's weight therefore
+        # ramps up linearly from zero to `neural_weight`, with the prototype
+        # taking the rest. Before 0.3.0 the opposite held: such classes used a
+        # fixed 0.3 prototype / 0.7 head split, which scored 0.55-0.73 on
+        # separable synthetic data that prototypes alone got 100% on.
+        # Setting both new_class_* weights restores a fixed split for such classes.
         self.new_class_example_threshold = self.config.get('new_class_example_threshold', 10)
-        self.new_class_prototype_weight = self.config.get('new_class_prototype_weight', 0.3)
-        self.new_class_neural_weight = self.config.get('new_class_neural_weight', 0.7)
+        self.new_class_prototype_weight = self.config.get('new_class_prototype_weight', None)
+        self.new_class_neural_weight = self.config.get('new_class_neural_weight', None)
+
+        # Prototype scores are softmax(-squared_distance / temperature) over
+        # classes. A small temperature sharpens them; before 0.3.0 the score was
+        # softmax(exp(-distance)), whose output barely distinguished the nearest
+        # class from the rest (about 0.45 vs 0.21), so a confident head could
+        # always outvote it. None restores that legacy scoring, which
+        # classifiers saved before 0.3.0 keep when they are loaded.
+        self.prototype_temperature = self.config.get('prototype_temperature', 0.25)
+
+        # Neural head training. The head is trained for at least `head_steps`
+        # optimiser steps with a cosine-decayed `head_learning_rate`. Before
+        # 0.3.0 it got about 10 steps and stopped at the first loss plateau,
+        # which left it unable to learn even a 2-D XOR from 120 examples.
+        # Lower `head_steps` (e.g. 100) if you add examples very frequently.
+        self.head_steps = self.config.get('head_steps', 300)
+        self.head_learning_rate = self.config.get('head_learning_rate', 0.003)
         self.min_confidence = self.config.get('min_confidence', 0.1)
 
         # Out-of-distribution detection. A query is scored by how far it sits
@@ -226,6 +249,9 @@ def to_dict(self) -> Dict[str, Any]:
             'new_class_example_threshold': self.new_class_example_threshold,
             'new_class_prototype_weight': self.new_class_prototype_weight,
             'new_class_neural_weight': self.new_class_neural_weight,
+            'prototype_temperature': self.prototype_temperature,
+            'head_steps': self.head_steps,
+            'head_learning_rate': self.head_learning_rate,
             'min_confidence': self.min_confidence,
             'ood_threshold': self.ood_threshold,
             'ood_min_radius': self.ood_min_radius,
diff --git a/tests/test_failure_modes.py b/tests/test_failure_modes.py
index 416c3a7..305e27e 100644
--- a/tests/test_failure_modes.py
+++ b/tests/test_failure_modes.py
@@ -321,7 +321,7 @@ def test_random_operation_sequences_keep_every_structure_consistent(new_classifi
     pool = ["great love", "good fast", "terrible hate", "awful bad", "okay fine", "cheap works",
             "slow broke", "very good"]
     names = ["a", "b", "c", "d", "e"]
-    clf = new_classifier()
+    clf = new_classifier(head_steps=20)
     added = []                                  # (text, label) pairs currently believed present
 
     for _ in range(25):
diff --git a/tests/test_pooling.py b/tests/test_pooling.py
index c0efc32..671d0e3 100644
--- a/tests/test_pooling.py
+++ b/tests/test_pooling.py
@@ -132,6 +132,8 @@ def test_blend_weight_defaults():
     assert config.prototype_weight == 0.7
     assert config.neural_weight == 0.3
     assert config.new_class_example_threshold == 10
+    assert config.new_class_prototype_weight is None      # None: ramp the head in with evidence
+    assert config.new_class_neural_weight is None
 
 
 def test_blend_weights_are_configurable():
@@ -141,11 +143,23 @@ def test_blend_weights_are_configurable():
     assert clf._blend_weights("a") == (0.2, 0.8)
 
 
-def test_new_classes_lean_on_the_neural_head():
-    """A class with few examples has an unreliable prototype."""
+def test_new_classes_lean_on_their_prototypes():
+    """A class with few examples has no evidence behind the neural head, so the head's
+    weight ramps up from zero as the class accumulates examples (0.3.0; before that
+    such classes got a fixed 0.3 prototype / 0.7 head split)."""
     clf = AdaptiveClassifier(MODEL, device="cpu")
-    clf.training_history = {"established": 50, "fresh": 2}
+    clf.training_history = {"established": 50, "fresh": 2, "unseen": 0}
     assert clf._blend_weights("established") == (0.7, 0.3)
+    assert clf._blend_weights("unseen") == (1.0, 0.0)
+    prototype, neural = clf._blend_weights("fresh")
+    assert neural == pytest.approx(0.06) and prototype == pytest.approx(0.94)
+
+
+def test_fixed_new_class_weights_can_still_be_set():
+    clf = AdaptiveClassifier(MODEL, device="cpu",
+                             config={"new_class_prototype_weight": 0.3,
+                                     "new_class_neural_weight": 0.7})
+    clf.training_history = {"fresh": 2}
     assert clf._blend_weights("fresh") == (0.3, 0.7)
 
 
@@ -153,7 +167,8 @@ def test_new_class_threshold_is_configurable():
     clf = AdaptiveClassifier(MODEL, device="cpu",
                              config={"new_class_example_threshold": 100})
     clf.training_history = {"a": 50}
-    assert clf._blend_weights("a") == (0.3, 0.7)      # still "new" at 50
+    prototype, neural = clf._blend_weights("a")        # halfway up the ramp at 50 of 100
+    assert neural == pytest.approx(0.15) and prototype == pytest.approx(0.85)
 
 
 def test_blend_weight_changes_the_prediction():
diff --git a/tests/test_portable_export.py b/tests/test_portable_export.py
index f7b4d19..5b884ed 100644
--- a/tests/test_portable_export.py
+++ b/tests/test_portable_export.py
@@ -193,3 +193,28 @@ def test_portable_reference_needs_a_pooling_mode_for_legacy_configs(base_model,
     actual = dict(portable.predict(QUERIES[0], k=3))
     for label in expected:
         assert float(actual[label]) == pytest.approx(expected[label], abs=1e-4)
+
+
+# --- the reference tracks every branch of the scoring math ---------------------------------
+
+@pytest.mark.skipif(not _onnx_available(), reason="optimum[onnxruntime] not installed")
+@pytest.mark.parametrize("config", [
+    {"prototype_temperature": None},                                   # legacy scoring
+    {"prototype_temperature": 0.1},
+    {"new_class_prototype_weight": 0.3, "new_class_neural_weight": 0.7},  # fixed split below threshold
+    {"new_class_example_threshold": 4},                                # some classes established
+    {"prototype_weight": 0.5, "neural_weight": 0.5},
+], ids=["legacy-scoring", "temperature-0.1", "fixed-new-class", "mixed-established", "even-blend"])
+def test_portable_reference_matches_predict_for_every_scoring_branch(base_model, tmp_path, config):
+    pytest.importorskip("tokenizers")
+    clf = _trained(base_model, pooling="mean", **config)
+    clf.add_examples(["great product love it"] * 6, ["pos"] * 6)      # 'pos' now has more examples than the rest
+    clf.save(str(tmp_path), include_onnx=True, quantize_onnx=False)
+
+    portable = _load_reference()(tmp_path, prefer_quantized=False)
+    for query in QUERIES:
+        expected = dict(clf.predict(query, k=3))
+        actual = dict(portable.predict(query, k=3))
+        assert actual.keys() == expected.keys()
+        for label in expected:
+            assert float(actual[label]) == pytest.approx(expected[label], abs=1e-4)
diff --git a/tests/test_reported_confidence_drop.py b/tests/test_reported_confidence_drop.py
index 6469a9e..78e8d3f 100644
--- a/tests/test_reported_confidence_drop.py
+++ b/tests/test_reported_confidence_drop.py
@@ -56,14 +56,19 @@ def test_reported_confidence_values():
         assert abs(cat_conf_before - cat_conf_after) < 0.01, \
             f"Cat confidence dropped from {cat_conf_before:.4f} to {cat_conf_after:.4f}"
         
-        # Both before and after should be around 0.8997 (blended prediction)
-        # not 0.9997 (pure neural prediction)
-        assert 0.85 < fish_conf_before < 0.95, \
-            f"Before save confidence should be blended, got {fish_conf_before:.4f}"
-        
-        assert 0.85 < fish_conf_after < 0.95, \
-            f"After load confidence should be blended, got {fish_conf_after:.4f}"
+        # The reported bug was a pure-neural score (0.9997) before saving versus a
+        # blended one (0.8997) afterwards. Both sides must now be the same blend:
+        # the right label, confident but not the saturated head-only value. (The
+        # exact blended value moved in 0.3.0, when prototype scores were sharpened
+        # and the head's weight started ramping up with the number of examples, so
+        # this no longer pins a 0.85-0.95 window.)
+        assert result_fish_before[0][0] == "foo" and result_fish_after[0][0] == "foo"
+        assert result_cat_before[0][0] == "bar" and result_cat_after[0][0] == "bar"
+        assert 0.5 < fish_conf_before < 0.9995, \
+            f"Before save confidence should be a blend, got {fish_conf_before:.4f}"
 
+        assert 0.5 < fish_conf_after < 0.9995, \
+            f"After load confidence should be a blend, got {fish_conf_after:.4f}"
 
 if __name__ == "__main__":
     test_reported_confidence_values()
\ No newline at end of file
diff --git a/tests/test_robustness.py b/tests/test_robustness.py
index 3570263..5d90a65 100644
--- a/tests/test_robustness.py
+++ b/tests/test_robustness.py
@@ -11,19 +11,23 @@
   * its output is always a valid, ordered probability distribution
   * it is deterministic and independent of example order where documented
 
-Known weaknesses are recorded as strict expected failures, so they stay
-visible and the test fails (forcing the marker off) once someone fixes them.
+A weakness that is known but not yet fixed should be recorded here as a strict
+expected failure (`xfail(strict=True)`), so it stays visible and the test fails,
+forcing the marker off, once someone fixes it.
 """
 
 import numpy as np
 import pytest
+import torch
+import torch.nn.functional as F
 from hypothesis import HealthCheck, given, settings, strategies as st
 
 from adaptive_classifier import AdaptiveClassifier
 
-from .synthetic import Dataset, accuracy
+from .synthetic import DIM, Dataset, accuracy
 
-PROTOTYPE_ONLY = {"new_class_example_threshold": 0, "prototype_weight": 1.0, "neural_weight": 0.0}
+PROTOTYPE_ONLY = {"new_class_example_threshold": 0, "prototype_weight": 1.0, "neural_weight": 0.0,
+                  "head_steps": 10}
 SEEDS = range(4)
 
 
@@ -65,21 +69,203 @@ def test_prototype_only_is_accurate_at_every_size(new_classifier, data, n_per):
         assert accuracy(clf, Xt, yt) >= 0.95
 
 
-@pytest.mark.xfail(
-    strict=True,
-    reason="Known weakness: with fewer than new_class_example_threshold (10) examples per "
-           "class the default blend gives the neural head 70% of the vote, and a head trained "
-           "on a handful of examples outvotes correct prototypes. Observed on separable data: "
-           "~75-92% training accuracy, ~66-83% held out, versus 100% from prototypes alone. "
-           "Remove this marker when fixed.",
-)
 @pytest.mark.parametrize("n_per", [2, 3, 5, 8])
 def test_default_config_fits_its_own_training_data_with_few_examples(new_classifier, data, n_per):
-    accs = []
+    """Regression: below 10 examples per class the neural head (trained on a handful of
+    points, and given 70% of the vote) used to outvote correct prototypes, so the model
+    got only ~75-92% of its own training data right on perfectly separable data."""
+    accs, held = [], []
     for seed in SEEDS:
         X, y = data.classes(n_per, noise=0.1, seed=seed)
-        accs.append(accuracy(fit(new_classifier, X, y), X, y))
+        clf = fit(new_classifier, X, y)
+        Xt, yt = data.classes(15, noise=0.1, seed=seed + 50, prefix="te")
+        accs.append(accuracy(clf, X, y))
+        held.append(accuracy(clf, Xt, yt))
     assert np.mean(accs) >= 0.97, f"training accuracy {np.mean(accs):.2f} at {n_per} per class"
+    assert np.mean(held) >= 0.95, f"held-out accuracy {np.mean(held):.2f} at {n_per} per class"
+
+
+def test_accuracy_does_not_fall_off_a_cliff_at_the_new_class_threshold(new_classifier, data):
+    """Before 0.3.0 accuracy jumped from ~0.7 to 1.0 between 8 and 12 examples per class."""
+    def mean_heldout(n_per):
+        scores = []
+        for seed in SEEDS:
+            X, y = data.classes(n_per, noise=0.3, seed=seed)
+            clf = fit(new_classifier, X, y)
+            Xt, yt = data.classes(20, noise=0.3, seed=seed + 50, prefix="te")
+            scores.append(accuracy(clf, Xt, yt))
+        return np.mean(scores)
+
+    below, above = mean_heldout(8), mean_heldout(12)
+    assert below >= above - 0.08, f"8 per class: {below:.2f}, 12 per class: {above:.2f}"
+    assert below >= 0.8
+
+
+# --- the neural head is a real component, not decoration ------------------------------------------
+
+def make_xor(data, n, seed, noise=0.12, tag="tr"):
+    """Class A = (+,+) and (-,-) blobs, class B = (+,-) and (-,+): both class means are ~0,
+    so prototypes carry no information and only a trained head can separate them."""
+    g = torch.Generator().manual_seed(seed)
+    texts, labels = [], []
+    for i in range(n):
+        for sx, sy, label in [(1, 1, "A"), (-1, -1, "A"), (1, -1, "B"), (-1, 1, "B")]:
+            v = torch.zeros(DIM)
+            v[0], v[1] = sx, sy
+            name = f"xor_{tag}_{seed}_{i}_{sx}_{sy}"
+            data.store[name] = F.normalize(v + noise * torch.randn(DIM, generator=g), dim=0)
+            texts.append(name)
+            labels.append(label)
+    return texts, labels
+
+
+@pytest.mark.parametrize("n_per", [8, 12, 30])
+def test_default_blend_solves_a_problem_only_the_head_can(new_classifier, data, n_per):
+    """Prototypes cannot separate XOR; the default blend must still succeed via the head.
+    Before 0.3.0 the head was too undertrained to learn it (0.49-0.70)."""
+    accs = []
+    for seed in range(3):
+        X, y = make_xor(data, n_per, seed)
+        clf = fit(new_classifier, X, y)
+        Xt, yt = make_xor(data, 15, seed + 50, tag="te")
+        accs.append(accuracy(clf, Xt, yt))
+    assert np.mean(accs) >= 0.95, f"XOR accuracy {np.mean(accs):.2f}"
+
+
+def test_prototypes_alone_cannot_solve_xor(new_classifier, data):
+    """Guards the test above: if this started passing, the XOR data would no longer
+    prove that the head is doing the work."""
+    X, y = make_xor(data, 12, 0)
+    clf = fit(new_classifier, X, y, **PROTOTYPE_ONLY)
+    Xt, yt = make_xor(data, 15, 50, tag="te")
+    assert accuracy(clf, Xt, yt) <= 0.65
+
+
+def test_head_alone_learns_separable_data(new_classifier, data):
+    head_only = {"prototype_weight": 0.0, "neural_weight": 1.0, "new_class_example_threshold": 0}
+    accs = []
+    for seed in range(3):
+        X, y = data.classes(30, noise=0.3, seed=seed)
+        clf = fit(new_classifier, X, y, **head_only)
+        Xt, yt = data.classes(20, noise=0.3, seed=seed + 50, prefix="te")
+        accs.append(accuracy(clf, Xt, yt))
+    assert np.mean(accs) >= 0.85, f"head-only accuracy {np.mean(accs):.2f}"
+
+
+def test_head_steps_is_a_working_knob(new_classifier, data):
+    X, y = data.classes(20, noise=0.3, seed=1)
+    little = fit(new_classifier, X, y, head_steps=1, prototype_weight=0.0, neural_weight=1.0,
+                 new_class_example_threshold=0)
+    plenty = fit(new_classifier, X, y, head_steps=400, prototype_weight=0.0, neural_weight=1.0,
+                 new_class_example_threshold=0)
+    Xt, yt = data.classes(20, noise=0.3, seed=51, prefix="te")
+    assert accuracy(plenty, Xt, yt) > accuracy(little, Xt, yt) + 0.1
+
+
+# --- prototype scoring ------------------------------------------------------------------------------
+
+def test_prototype_scores_clearly_favour_the_nearest_class(new_classifier, data):
+    X, y = data.classes(10, noise=0.1, seed=3)
+    sharp = fit(new_classifier, X, y, **PROTOTYPE_ONLY)
+    legacy = fit(new_classifier, X, y, prototype_temperature=None, **PROTOTYPE_ONLY)
+    query = data.add("center", 0, 0.0, 0)
+    top_sharp = sharp.predict(query)[0][1]
+    top_legacy = legacy.predict(query)[0][1]
+    assert top_sharp > 0.7, f"sharp scoring gave the nearest class only {top_sharp:.2f}"
+    assert top_legacy < 0.55, "legacy scoring is expected to be flat"
+
+
+def test_temperature_changes_confidence_but_not_the_ranking_of_an_easy_case(new_classifier, data):
+    X, y = data.classes(10, noise=0.1, seed=4)
+    q = data.add("q", 1, 0.1, 5)
+    for tau in (None, 0.05, 0.25, 1.0):
+        clf = fit(new_classifier, X, y, prototype_temperature=tau, **PROTOTYPE_ONLY)
+        assert clf.predict(q)[0][0] == "C1"
+
+
+def test_prototype_scores_form_a_distribution_for_any_temperature(new_classifier, data):
+    X, y = data.classes(6, noise=0.2, seed=6)
+    for tau in (None, 0.01, 0.25, 5.0):
+        clf = fit(new_classifier, X, y, prototype_temperature=tau, **PROTOTYPE_ONLY)
+        scores = [s for _, s in clf.predict(data.add("p", 2, 0.3, 9))]
+        assert sum(scores) == pytest.approx(1.0, abs=1e-5) and all(np.isfinite(scores))
+
+
+# --- how the blend depends on how much a class has seen ----------------------------------------------
+
+def test_blend_weights_ramp_up_without_a_jump(new_classifier):
+    clf = new_classifier()
+    weights = []
+    for n in range(0, 14):
+        clf.training_history = {"c": n}
+        weights.append(clf._blend_weights("c"))
+    neural = [w[1] for w in weights]
+    assert neural[0] == 0.0 and weights[0][0] == 1.0          # unseen class: prototype only
+    assert neural == sorted(neural)                           # never decreases
+    assert max(b - a for a, b in zip(neural, neural[1:])) <= 0.31 / 10 + 1e-9   # smooth steps
+    assert weights[10] == (0.7, 0.3) == weights[13]           # established at the threshold
+    assert all(sum(w) == pytest.approx(1.0) for w in weights)
+
+
+def test_fixed_new_class_weights_still_work(new_classifier):
+    clf = new_classifier(new_class_prototype_weight=0.3, new_class_neural_weight=0.7)
+    clf.training_history = {"fresh": 2, "established": 50}
+    assert clf._blend_weights("fresh") == (0.3, 0.7)
+    assert clf._blend_weights("established") == (0.7, 0.3)
+
+
+def test_ramp_scales_with_the_configured_threshold_and_weights(new_classifier):
+    clf = new_classifier(new_class_example_threshold=4, prototype_weight=0.5, neural_weight=0.5)
+    clf.training_history = {"c": 2}
+    assert clf._blend_weights("c") == pytest.approx((0.75, 0.25))     # halfway from (1, 0) to (.5, .5)
+    clf.training_history = {"c": 4}
+    assert clf._blend_weights("c") == (0.5, 0.5)
+
+
+def test_zero_threshold_disables_the_ramp(new_classifier):
+    clf = new_classifier(new_class_example_threshold=0)
+    clf.training_history = {"c": 0}
+    assert clf._blend_weights("c") == (0.7, 0.3)
+
+
+# --- models saved before 0.3.0 must not change ------------------------------------------------------------
+
+def test_a_pre_0_3_0_save_keeps_its_scoring(new_classifier, base_model, data, tmp_path):
+    """Old configs have no prototype_temperature and explicit new-class weights; loading one must
+    reproduce its old predictions exactly rather than silently switching to the new scoring."""
+    import json
+    from adaptive_classifier import AdaptiveClassifier
+
+    X, y = data.classes(6, noise=0.2, seed=7)
+    legacy = fit(new_classifier, X, y, prototype_temperature=None,
+                 new_class_prototype_weight=0.3, new_class_neural_weight=0.7)
+    expected = {t: legacy.predict(t, k=5) for t in X[:6]}
+    legacy.save(str(tmp_path), include_onnx=False)
+
+    path = tmp_path / "config.json"
+    data_json = json.loads(path.read_text(encoding="utf-8"))
+    del data_json["config"]["prototype_temperature"]               # what a 0.2.x save looks like
+    for key in ("head_steps", "head_learning_rate"):
+        data_json["config"].pop(key, None)
+    path.write_text(json.dumps(data_json), encoding="utf-8")
+
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    assert loaded.config.prototype_temperature is None
+    for text, preds in expected.items():
+        got = loaded.predict(text, k=5)
+        assert [l for l, _ in got] == [l for l, _ in preds]
+        assert all(a == pytest.approx(b, abs=1e-4) for (_, a), (_, b) in zip(got, preds))
+
+
+def test_new_saves_record_the_new_settings(new_classifier, data, tmp_path):
+    import json
+    X, y = data.classes(4, seed=8)
+    fit(new_classifier, X, y).save(str(tmp_path), include_onnx=False)
+    config = json.loads((tmp_path / "config.json").read_text(encoding="utf-8"))["config"]
+    assert config["prototype_temperature"] == 0.25
+    assert config["head_steps"] == 300
+    assert config["new_class_prototype_weight"] is None
+    assert config["new_class_neural_weight"] is None
 
 
 # --- no catastrophic forgetting ---------------------------------------------------------
@@ -213,7 +399,7 @@ def assert_valid_distribution(preds, known, k):
     k=st.integers(min_value=0, max_value=9),
 )
 def test_predictions_are_always_a_valid_ordered_distribution(new_classifier, train, queries, k):
-    clf = new_classifier()
+    clf = new_classifier(head_steps=20)
     texts, labels = zip(*train)
     clf.add_examples(list(texts), list(labels))
     known = set(labels)
@@ -230,7 +416,7 @@ def test_predictions_are_always_a_valid_ordered_distribution(new_classifier, tra
 @settings(max_examples=15, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture])
 @given(train=st.lists(st.tuples(texts_st, labels_st), min_size=2, max_size=12), q=texts_st)
 def test_scores_and_ood_are_finite_for_any_training_set(new_classifier, train, q):
-    clf = new_classifier()
+    clf = new_classifier(head_steps=20)
     texts, labels = zip(*train)
     clf.add_examples(list(texts), list(labels))
     score = clf.ood_score(q)

From 8e50c99ce5509d74d92725d8df79b0cc6810baa7 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 10:46:37 +0000
Subject: [PATCH 09/16] Add confidence calibration and conformal prediction
 sets

- calibrate(texts, labels) fits one temperature on held-out labelled data
  (minimising negative log-likelihood), so a reported 0.9 means about 90%.
  predict, predict_batch and abstain_below then use calibrated
  probabilities; the winning class never changes. On synthetic overlapping
  classes held-out ECE drops from 0.20-0.34 to 0.02-0.03
- predict_set(text, alpha) returns a split-conformal set that contains the
  true label with probability >= 1 - alpha (never empty)
- calibration_report() gives accuracy, ECE and NLL for checking on a
  separate test set
- Calibration is saved with the model and discarded when the class set
  changes (a new class, forget)
- Fix: predict_batch scored only k classes before blending, so with more
  than k classes its scores differed from predict(); it now scores all
  classes and trims to k
- Pure maths lives in adaptive_classifier.calibration

tests/test_calibration.py (45 tests): the maths against known answers
(recovering a known miscalibration, conformal order statistics, ECE), and
the classifier on separate train / calibration / test splits. Mutation-
checked with 14 deliberate faults, all caught.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                           |   4 +
 README.md                              |  18 ++
 docs/API.md                            |  24 ++
 src/adaptive_classifier/calibration.py |  98 +++++++
 src/adaptive_classifier/classifier.py  | 168 +++++++++--
 tests/test_calibration.py              | 371 +++++++++++++++++++++++++
 6 files changed, 665 insertions(+), 18 deletions(-)
 create mode 100644 src/adaptive_classifier/calibration.py
 create mode 100644 tests/test_calibration.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 93d16b8..b1dde8a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,9 @@
 - `docs/deployment.md` and `examples/portable_inference.py`: run a saved classifier without PyTorch or FAISS (multilingual models, .NET, JavaScript).
 - `save()` now writes the tokenizer files and the resolved pooling mode, so a saved directory is self-contained.
 
+- `calibrate(texts, labels)` fits one temperature on held-out labelled data so confidences match how often the model is right, and `calibration_report(...)` measures accuracy, expected calibration error and NLL. `predict`, `predict_batch` and `abstain_below` then use calibrated probabilities. Calibration is saved with the model and discarded when the set of classes changes.
+- `predict_set(text, alpha=0.1)`: split-conformal prediction sets that contain the true label with probability at least `1 - alpha`.
+
 ### Changed (default behaviour: predictions change)
 - **The neural head is now actually trained.** It used to get about ten optimiser steps and stop at the first loss plateau, so it could not learn even a 2-D XOR from 120 examples. It now trains for at least `head_steps` (default 300) steps with a cosine-decayed `head_learning_rate` (default 0.003). Small memories take roughly 1-2 s per `add_examples` call; lower `head_steps` if you add examples very frequently.
 - **Prototype scores are sharper.** They are now `softmax(-distance² / prototype_temperature)` (default 0.25) instead of `softmax(exp(-distance))`, which barely separated the nearest class from the rest (about 0.45 vs 0.21), so a confident head could always outvote it. `prototype_temperature=None` restores the old scoring.
@@ -19,6 +22,7 @@
 - `clear_memory(labels=[...])` now removes those classes completely (it used to leave them in the label map and neural head, so they could still be predicted).
 
 ### Fixed
+- `predict_batch(k=...)` scored only `k` classes before blending, so with more than `k` classes its scores differed from `predict` for the same text. It now scores every class and trims to `k`.
 - Accuracy on small training sets (see above): the default configuration got only 75-92% of its own training data right on perfectly separable data with fewer than 10 examples per class.
 - Integer (or other non-string) labels corrupted a classifier on save/reload, duplicating classes. `add_examples` now rejects non-string labels with a clear error; convert with `str(label)`.
 - `predict(k=-1)` returned a silently truncated list; negative or non-integer `k` now raises `ValueError`.
diff --git a/README.md b/README.md
index 778f4e4..cfee60a 100644
--- a/README.md
+++ b/README.md
@@ -338,6 +338,24 @@ clf.classifier_.save("./model")               # the underlying AdaptiveClassifie
 
 `fit` starts over each time; `partial_fit` adds to what is already learned. A fitted estimator holds a FAISS index and can't be pickled, so save `classifier_` instead.
 
+### Calibrated Confidence and Prediction Sets
+
+Raw scores are not probabilities: a model can say 0.9 and be right 70% of the time. Fit a calibration on labelled examples the model was *not* trained on:
+
+```python
+info = classifier.calibrate(held_out_texts, held_out_labels)
+info["ece_before"], info["ece_after"]       # expected calibration error, lower is better
+
+classifier.predict(text)                    # confidences now match how often the model is right
+classifier.predict(text, abstain_below=0.8) # so thresholds like this mean what they say
+
+# A set of labels that contains the true one at least 90% of the time
+classifier.predict_set(text, alpha=0.1)     # [("billing", 0.55), ("refunds", 0.31)]
+classifier.calibration_report(test_texts, test_labels)   # check it on separate data
+```
+
+One easy input gives a one-label set; an ambiguous one gives several. Calibration is saved with the model and discarded if you add or forget a class; re-run it after substantial new data.
+
 ### Fixing Mistakes and Knowing When to Abstain
 
 ```python
diff --git a/docs/API.md b/docs/API.md
index 69199f7..0ee6de1 100644
--- a/docs/API.md
+++ b/docs/API.md
@@ -130,6 +130,30 @@ Clear stored examples and prototypes.
 Parameters:
 - `labels`: Optional list of labels to clear. If given, this is the same as `forget(labels)` (unknown labels are ignored). If None, clears all stored examples and prototypes
 
+#### calibrate
+
+```python
+def calibrate(texts: List[str], labels: List[str]) -> Dict[str, float]
+```
+
+Fit confidence calibration (temperature scaling) on labelled examples the model was not trained on. Afterwards `predict`, `predict_batch` and `abstain_below` use calibrated probabilities; the winning class never changes. Returns `temperature` (above 1 means the model was over-confident), `n`, `accuracy`, and `ece_before` / `ece_after` / `nll_before` / `nll_after`. Adding a new class or forgetting one discards the calibration.
+
+#### calibration_report
+
+```python
+def calibration_report(texts, labels, calibrated=True, bins=10) -> Dict[str, float]
+```
+
+Accuracy, mean confidence, expected calibration error (`ece`) and `nll` on labelled data. Run it on a test set separate from the calibration set.
+
+#### predict_set
+
+```python
+def predict_set(text: str, alpha: float = 0.1) -> List[Tuple[str, float]]
+```
+
+Split-conformal prediction set: contains the true label with probability at least `1 - alpha` for inputs that look like the calibration data (an average guarantee, not per input). Never empty. Requires `calibrate` first.
+
 #### forget
 
 ```python
diff --git a/src/adaptive_classifier/calibration.py b/src/adaptive_classifier/calibration.py
new file mode 100644
index 0000000..1b8d421
--- /dev/null
+++ b/src/adaptive_classifier/calibration.py
@@ -0,0 +1,98 @@
+"""Confidence calibration and conformal prediction sets.
+
+A classifier's scores are only useful as probabilities if "0.9" means right
+about 90% of the time. Two tools here, both fitted on labelled examples the
+model was not trained on:
+
+* temperature scaling rescales the whole score distribution with one number
+  `T` (`p_i ** (1/T)`, renormalised): T > 1 softens an over-confident model,
+  T < 1 sharpens an under-confident one. It never changes which class wins.
+* split conformal prediction turns the calibrated scores into a *set* of
+  labels that contains the true one with probability at least `1 - alpha`
+  (for data exchangeable with the calibration set).
+
+Everything here is a pure function of arrays so it can be tested on its own.
+"""
+
+import math
+from typing import Dict, List, Sequence
+
+import numpy as np
+from scipy.optimize import minimize_scalar
+
+EPS = 1e-12
+
+
+def apply_temperature(probs: np.ndarray, temperature: float) -> np.ndarray:
+    """Rescale a probability vector (or rows of a matrix) by `p ** (1/T)`, renormalised."""
+    probs = np.clip(np.asarray(probs, dtype=float), EPS, None)
+    scaled = np.exp(np.log(probs) / temperature)
+    return scaled / scaled.sum(axis=-1, keepdims=True)
+
+
+def negative_log_likelihood(probs: np.ndarray, targets: np.ndarray) -> float:
+    probs = np.clip(np.asarray(probs, dtype=float), EPS, 1.0)
+    return float(-np.mean(np.log(probs[np.arange(len(targets)), targets])))
+
+
+def fit_temperature(probs: np.ndarray, targets: Sequence[int],
+                    bounds=(0.05, 20.0)) -> float:
+    """The temperature minimising the negative log-likelihood of `targets`."""
+    probs = np.asarray(probs, dtype=float)
+    targets = np.asarray(targets, dtype=int)
+    if probs.ndim != 2 or len(probs) != len(targets):
+        raise ValueError("probs must be (n_samples, n_classes) and match targets")
+    if len(targets) < 2:
+        raise ValueError("Need at least 2 calibration examples")
+
+    def loss(log_t: float) -> float:
+        return negative_log_likelihood(apply_temperature(probs, math.exp(log_t)), targets)
+
+    result = minimize_scalar(loss, bounds=(math.log(bounds[0]), math.log(bounds[1])),
+                             method="bounded", options={"xatol": 1e-4})
+    return float(math.exp(result.x))
+
+
+def nonconformity_scores(probs: np.ndarray, targets: Sequence[int]) -> np.ndarray:
+    """1 minus the probability given to the true class (higher = more surprising)."""
+    probs = np.asarray(probs, dtype=float)
+    return 1.0 - probs[np.arange(len(targets)), np.asarray(targets, dtype=int)]
+
+
+def conformal_threshold(scores: Sequence[float], alpha: float) -> float:
+    """The score quantile that gives `1 - alpha` coverage; infinity when there are too few scores."""
+    if not 0.0 < alpha < 1.0:
+        raise ValueError(f"alpha must be between 0 and 1; got {alpha}")
+    scores = np.sort(np.asarray(scores, dtype=float))
+    n = len(scores)
+    rank = math.ceil((n + 1) * (1.0 - alpha))        # 1-based rank of the quantile
+    if n == 0 or rank > n:
+        return float("inf")
+    return float(scores[rank - 1])
+
+
+def expected_calibration_error(probs: np.ndarray, targets: Sequence[int], bins: int = 10) -> float:
+    """Average gap between confidence and accuracy, weighted by how many predictions fall in each bin."""
+    probs = np.asarray(probs, dtype=float)
+    targets = np.asarray(targets, dtype=int)
+    confidence = probs.max(axis=1)
+    correct = (probs.argmax(axis=1) == targets).astype(float)
+    edges = np.linspace(0.0, 1.0, bins + 1)
+    ece = 0.0
+    for low, high in zip(edges[:-1], edges[1:]):
+        in_bin = (confidence > low) & (confidence <= high)
+        if in_bin.any():
+            ece += in_bin.mean() * abs(correct[in_bin].mean() - confidence[in_bin].mean())
+    return float(ece)
+
+
+def summarise(probs: np.ndarray, targets: Sequence[int], bins: int = 10) -> Dict[str, float]:
+    probs = np.asarray(probs, dtype=float)
+    targets = np.asarray(targets, dtype=int)
+    return {
+        "n": int(len(targets)),
+        "accuracy": float((probs.argmax(axis=1) == targets).mean()),
+        "mean_confidence": float(probs.max(axis=1).mean()),
+        "ece": expected_calibration_error(probs, targets, bins),
+        "nll": negative_log_likelihood(probs, targets),
+    }
diff --git a/src/adaptive_classifier/classifier.py b/src/adaptive_classifier/classifier.py
index 7197938..35f22d7 100644
--- a/src/adaptive_classifier/classifier.py
+++ b/src/adaptive_classifier/classifier.py
@@ -19,6 +19,7 @@
 from .models import Example, AdaptiveHead, ModelConfig
 from .memory import PrototypeMemory
 from .ewc import EWC
+from . import calibration as calibration_math
 from .strategic import (
     StrategicCostFunction, CostFunctionFactory, StrategicOptimizer, StrategicEvaluator
 )
@@ -28,7 +29,11 @@
 
 class AdaptiveClassifier(ModelHubMixin):
     """A flexible classifier that can adapt to new classes and examples."""
-    
+
+    # Set by `calibrate`. A class-level default keeps instances built without
+    # __init__ (the loaders do that) working.
+    calibration: Optional[Dict[str, Any]] = None
+
     def __init__(
         self,
         model_name: str,
@@ -171,6 +176,10 @@ def add_examples(self, texts: List[str], labels: List[str]):
         new_classes = set(labels) - set(self.label_to_id.keys())
         is_adding_new_classes = len(new_classes) > 0
 
+        if is_adding_new_classes and self.calibration:
+            logger.warning("Discarding confidence calibration: the set of classes changed")
+            self.calibration = None
+
         # Update label mappings - sort new classes alphabetically for consistent IDs
         for label in sorted(new_classes):
             idx = len(self.label_to_id)
@@ -449,11 +458,15 @@ def predict(
             return []
 
         # If strategic mode is not enabled, use regular prediction
+        # Calibration rescales the whole distribution, so it needs every class.
+        internal_k = len(self.id_to_label) if self.calibration else k
         if not self.strategic_mode:
-            predictions = self._predict_regular(text, k)
+            predictions = self._predict_regular(text, internal_k)
         else:
             # Dual prediction system: blend strategic and regular predictions
-            predictions = self._predict_dual(text, k)
+            predictions = self._predict_dual(text, internal_k)
+        if self.calibration:
+            predictions = self._calibrate(predictions)[:k]
 
         if abstain_below is not None and (not predictions or predictions[0][1] < abstain_below):
             return []
@@ -485,6 +498,120 @@ def is_ood(self, text: str, threshold: Optional[float] = None) -> bool:
         limit = self.config.ood_threshold if threshold is None else threshold
         return self.ood_score(text) > limit
 
+    # --- calibration ---------------------------------------------------------------
+
+    def _calibrate(self, predictions: List[Tuple[str, float]]) -> List[Tuple[str, float]]:
+        """Apply the fitted temperature to a full (all-class) prediction list."""
+        if not self.calibration or not predictions:
+            return predictions
+        probs = calibration_math.apply_temperature(
+            np.array([score for _, score in predictions]), self.calibration['temperature']
+        )
+        calibrated = [(label, float(p)) for (label, _), p in zip(predictions, probs)]
+        return sorted(calibrated, key=lambda item: item[1], reverse=True)
+
+    def _distribution_matrix(self, texts: List[str]):
+        """Uncalibrated scores for each text as a matrix with columns in label-id order."""
+        saved, self.calibration = self.calibration, None
+        try:
+            n = len(self.id_to_label)
+            matrix = np.zeros((len(texts), n))
+            for row, text in enumerate(texts):
+                for label, score in self.predict(text, k=n):
+                    matrix[row, self.label_to_id[label]] = score
+        finally:
+            self.calibration = saved
+        return matrix
+
+    def _calibration_inputs(self, texts: List[str], labels: List[str]):
+        if not texts:
+            raise ValueError("Empty input lists")
+        if len(texts) != len(labels):
+            raise ValueError("Mismatched text and label lists")
+        self._validate_training_inputs(texts, labels)
+        unknown = sorted(set(labels) - set(self.label_to_id))
+        if unknown:
+            raise ValueError(f"Unknown labels: {unknown}")
+        return np.array([self.label_to_id[label] for label in labels])
+
+    def calibrate(self, texts: List[str], labels: List[str]) -> Dict[str, float]:
+        """Fit confidence calibration on labelled examples the model was *not* trained on.
+
+        Learns one temperature so that confidences match how often the model is
+        actually right (a score of 0.9 should mean about 90%), and records the
+        scores needed by `predict_set`. After this, `predict`, `predict_batch`
+        and `abstain_below` use the calibrated probabilities; the predicted
+        class never changes. Calibrating on data the model was trained on
+        makes it look better than it is, so hold some data out.
+
+        Re-run it after substantial additions of new data. Adding a new class,
+        or forgetting one, discards the calibration because the class set it
+        described no longer exists.
+
+        Returns:
+            The fitted `temperature` (above 1: the model was over-confident),
+            and `ece` / `nll` before and after (lower is better).
+        """
+        targets = self._calibration_inputs(texts, labels)
+        if len(texts) < 2 or len(self.label_to_id) < 2:
+            raise ValueError("Calibration needs at least 2 examples and 2 classes")
+        if len(texts) < 20:
+            logger.warning(f"Calibrating on only {len(texts)} examples; the estimate will be noisy")
+
+        probs = self._distribution_matrix(texts)
+        temperature = calibration_math.fit_temperature(probs, targets)
+        calibrated = calibration_math.apply_temperature(probs, temperature)
+        scores = calibration_math.nonconformity_scores(calibrated, targets)
+        self.calibration = {
+            'temperature': temperature,
+            'scores': sorted(float(score) for score in scores),
+            'n': int(len(targets)),
+        }
+        before = calibration_math.summarise(probs, targets)
+        after = calibration_math.summarise(calibrated, targets)
+        return {
+            'temperature': temperature,
+            'n': before['n'],
+            'ece_before': before['ece'], 'ece_after': after['ece'],
+            'nll_before': before['nll'], 'nll_after': after['nll'],
+            'accuracy': before['accuracy'],
+        }
+
+    def calibration_report(self, texts: List[str], labels: List[str],
+                           calibrated: bool = True, bins: int = 10) -> Dict[str, float]:
+        """Accuracy, mean confidence, expected calibration error (ECE) and NLL on labelled data.
+
+        Use it on a *separate* test set to check that calibration generalises.
+        With `calibrated=False` the fitted temperature is ignored.
+        """
+        targets = self._calibration_inputs(texts, labels)
+        probs = self._distribution_matrix(texts)
+        if calibrated and self.calibration:
+            probs = calibration_math.apply_temperature(probs, self.calibration['temperature'])
+        return calibration_math.summarise(probs, targets, bins)
+
+    def predict_set(self, text: str, alpha: float = 0.1) -> List[Tuple[str, float]]:
+        """Labels that contain the true one with probability at least `1 - alpha`.
+
+        A split-conformal prediction set: when the model is sure the set has
+        one label, and it grows with uncertainty. The guarantee holds for
+        inputs that look like the calibration data (exchangeable with it) and
+        is an average over inputs, not a promise for each one. The top label is
+        always included, so the set is never empty (which can only make
+        coverage higher). Requires `calibrate` first.
+
+        Returns:
+            (label, calibrated probability) pairs, most likely first.
+        """
+        if not self.calibration:
+            raise ValueError("predict_set needs calibration; call calibrate(texts, labels) first")
+        if not text:
+            raise ValueError("Empty input text")
+        threshold = calibration_math.conformal_threshold(self.calibration['scores'], alpha)
+        full = self.predict(text, k=len(self.id_to_label))
+        chosen = [(label, p) for label, p in full if 1.0 - p <= threshold]
+        return chosen or full[:1]
+
     def forget(self, labels: Union[str, List[str]]):
         """Remove one or more classes entirely.
 
@@ -566,6 +693,9 @@ def _drop_classes(self, drop: Set[str]):
         """Remove classes from every structure that tracks them."""
         ordered = sorted(self.label_to_id.items(), key=lambda item: item[1])
         keep = [(label, idx) for label, idx in ordered if label not in drop]
+        if self.calibration:
+            logger.warning("Discarding confidence calibration: the set of classes changed")
+            self.calibration = None
 
         for label in drop:
             self.memory.remove_class(label)
@@ -716,6 +846,7 @@ def _save_pretrained(
             'train_steps': self.train_steps,
             'training_history': self.training_history,  # Save cumulative training counts
             'ood_radii': {label: self.memory.class_radius(label) for label in self.memory.prototypes},
+            'calibration': self.calibration,
             'config': saved_settings,
             'library_name': 'adaptive-classifier'  # Tell HuggingFace Hub this requires the adaptive-classifier library
         }
@@ -1068,6 +1199,7 @@ def _from_pretrained(
         
         # Restore training history with backward compatibility
         classifier.training_history = config_dict.get('training_history', {})
+        classifier.calibration = config_dict.get('calibration')
 
         # Load tensors
         tensors = load_file(model_path / "model.safetensors")
@@ -1597,13 +1729,16 @@ def predict_batch(
             
             # Get predictions for each embedding
             batch_predictions = []
+            num_classes = len(self.id_to_label)
             for embedding in batch_embeddings:
-                # Get prototype predictions
+                # Score every class, as `predict` does, and trim to k at the end.
+                # Restricting the lookup to k classes made the scores depend on k
+                # whenever there were more than k classes.
                 proto_preds = self.memory.get_nearest_prototypes(
                     embedding,
-                    k=k
+                    k=num_classes or k
                 )
-                
+
                 # Get neural predictions if available
                 if self.adaptive_head is not None:
                     self.adaptive_head.eval()
@@ -1614,18 +1749,15 @@ def predict_batch(
                         # Squeeze batch dimension
                         logits = logits.squeeze(0)
                         probs = F.softmax(logits, dim=0)
-                        
-                        values, indices = torch.topk(
-                            probs,
-                            min(k, len(self.id_to_label))
-                        )
+
+                        values, indices = torch.topk(probs, num_classes)
                         head_preds = [
                             (self.id_to_label[idx.item()], val.item())
                             for val, idx in zip(values, indices)
                         ]
                 else:
                     head_preds = []
-                
+
                 # Combine predictions with the configured weights. This path
                 # used to hardcode 0.7/0.3, so `predict` and `predict_batch`
                 # could disagree once a caller changed the config.
@@ -1637,22 +1769,22 @@ def predict_batch(
                     combined_scores[label] = (
                         combined_scores.get(label, 0) + score * self._blend_weights(label)[1]
                     )
-                
+
                 # Sort and normalize predictions
                 predictions = sorted(
                     combined_scores.items(),
                     key=lambda x: x[1],
                     reverse=True
                 )
-                
-                # Normalize scores
                 total = sum(score for _, score in predictions)
                 if total > 0:
-                    predictions = [(label, score/total) 
+                    predictions = [(label, score/total)
                                 for label, score in predictions]
-                
+
+                if self.calibration:
+                    predictions = self._calibrate(predictions)
                 batch_predictions.append(predictions[:k])
-            
+
             all_predictions.extend(batch_predictions)
         
         return all_predictions
diff --git a/tests/test_calibration.py b/tests/test_calibration.py
new file mode 100644
index 0000000..b8a86e9
--- /dev/null
+++ b/tests/test_calibration.py
@@ -0,0 +1,371 @@
+"""Confidence calibration and conformal prediction sets.
+
+Two layers: the pure maths in `adaptive_classifier.calibration` is checked
+against cases with known answers, and the classifier integration is checked on
+synthetic overlapping classes (so the model really is uncertain, which is what
+makes calibration matter) with separate train / calibration / test splits.
+"""
+
+import json
+import logging
+
+import numpy as np
+import pytest
+
+from adaptive_classifier import AdaptiveClassifier
+from adaptive_classifier import calibration as cal
+
+from .synthetic import Dataset
+
+NOISE = 0.45          # classes overlap, so confidence has to be earned
+
+
+# --- the maths ------------------------------------------------------------------------------
+
+def test_temperature_one_is_the_identity_and_large_temperature_is_uniform():
+    p = np.array([[0.7, 0.2, 0.1], [0.4, 0.4, 0.2]])
+    assert cal.apply_temperature(p, 1.0) == pytest.approx(p)
+    assert cal.apply_temperature(p, 1e6) == pytest.approx(np.full_like(p, 1 / 3), abs=1e-4)
+
+
+def test_temperature_never_changes_the_ranking_and_rows_still_sum_to_one():
+    rng = np.random.default_rng(0)
+    p = rng.dirichlet(np.ones(5), size=50)
+    for t in (0.2, 0.7, 1.5, 6.0):
+        out = cal.apply_temperature(p, t)
+        assert out.sum(axis=1) == pytest.approx(np.ones(50))
+        assert (out.argmax(axis=1) == p.argmax(axis=1)).all()
+        assert (np.argsort(-out, axis=1) == np.argsort(-p, axis=1)).all()
+
+
+def test_low_temperature_sharpens_and_high_temperature_softens():
+    p = np.array([[0.6, 0.3, 0.1]])
+    assert cal.apply_temperature(p, 0.5).max() > 0.6 > cal.apply_temperature(p, 2.0).max()
+
+
+def test_zero_probabilities_do_not_produce_nan():
+    out = cal.apply_temperature(np.array([[1.0, 0.0, 0.0]]), 0.5)
+    assert np.isfinite(out).all() and out.sum() == pytest.approx(1.0)
+
+
+@pytest.mark.parametrize("true_t", [0.5, 2.0, 4.0])
+def test_fit_temperature_recovers_a_known_miscalibration(true_t):
+    """Labels drawn from the true probabilities; the model reports them rescaled by true_t."""
+    rng = np.random.default_rng(1)
+    truth = rng.dirichlet(np.ones(4) * 0.6, size=4000)
+    targets = np.array([rng.choice(4, p=row) for row in truth])
+    reported = cal.apply_temperature(truth, 1.0 / true_t)         # model is wrong by factor true_t
+    assert cal.fit_temperature(reported, targets) == pytest.approx(true_t, rel=0.12)
+
+
+def test_fit_temperature_leaves_a_calibrated_model_alone():
+    rng = np.random.default_rng(2)
+    truth = rng.dirichlet(np.ones(3), size=4000)
+    targets = np.array([rng.choice(3, p=row) for row in truth])
+    assert cal.fit_temperature(truth, targets) == pytest.approx(1.0, abs=0.1)
+
+
+def test_fit_temperature_validates_input():
+    with pytest.raises(ValueError):
+        cal.fit_temperature(np.array([[0.5, 0.5]]), [0])           # one example
+    with pytest.raises(ValueError):
+        cal.fit_temperature(np.array([[0.5, 0.5], [0.5, 0.5]]), [0])   # length mismatch
+    with pytest.raises(ValueError):
+        cal.fit_temperature(np.array([0.5, 0.5, 0.5]), [0, 1, 0])      # not a matrix
+
+
+def test_conformal_threshold_picks_the_right_order_statistic():
+    scores = np.arange(1, 10) / 10                                 # 0.1 ... 0.9, n = 9
+    assert cal.conformal_threshold(scores, 0.1) == pytest.approx(0.9)    # ceil(10 * 0.9) = 9th
+    assert cal.conformal_threshold(scores, 0.5) == pytest.approx(0.5)    # ceil(10 * 0.5) = 5th
+    assert cal.conformal_threshold(scores, 0.05) == float("inf")        # too few points for 95%
+    assert cal.conformal_threshold([], 0.1) == float("inf")
+
+
+@pytest.mark.parametrize("alpha", [0.0, 1.0, -0.1, 1.5])
+def test_conformal_threshold_rejects_an_invalid_alpha(alpha):
+    with pytest.raises(ValueError, match="alpha"):
+        cal.conformal_threshold([0.1, 0.2], alpha)
+
+
+def test_expected_calibration_error_known_cases():
+    perfect = np.array([[1.0, 0.0]] * 10)
+    assert cal.expected_calibration_error(perfect, [0] * 10) == pytest.approx(0.0)
+    # Always 100% sure, right half the time -> off by 0.5
+    sure = np.array([[1.0, 0.0]] * 10)
+    assert cal.expected_calibration_error(sure, [0] * 5 + [1] * 5) == pytest.approx(0.5)
+    # Says 0.6 and is right 60% of the time -> calibrated
+    hedged = np.array([[0.6, 0.4]] * 10)
+    assert cal.expected_calibration_error(hedged, [0] * 6 + [1] * 4) == pytest.approx(0.0, abs=1e-9)
+
+
+# --- the classifier -----------------------------------------------------------------------------
+
+@pytest.fixture
+def data(monkeypatch):
+    return Dataset().patch(monkeypatch)
+
+
+@pytest.fixture
+def splits(data):
+    """Train / calibration / test, all drawn from the same overlapping classes."""
+    train = data.classes(30, n_classes=4, noise=NOISE, seed=1)
+    calib = data.classes(60, n_classes=4, noise=NOISE, seed=2, prefix="cal")
+    test = data.classes(150, n_classes=4, noise=NOISE, seed=3, prefix="te")
+    return train, calib, test
+
+
+@pytest.fixture
+def fitted(new_classifier, splits):
+    (X, y), _, _ = splits
+    clf = new_classifier()
+    clf.add_examples(X, y)
+    return clf
+
+
+def test_calibration_makes_confidence_honest_on_unseen_data(fitted, splits):
+    _, (Xc, yc), (Xt, yt) = splits
+    before = fitted.calibration_report(Xt, yt, calibrated=False)
+    fitted.calibrate(Xc, yc)
+    after = fitted.calibration_report(Xt, yt)
+
+    assert before["ece"] > 0.12, "the test needs a model that starts out miscalibrated"
+    assert after["ece"] < before["ece"] / 2
+    assert after["ece"] < 0.08
+    assert abs(after["mean_confidence"] - after["accuracy"]) < 0.06
+
+
+def test_calibration_never_changes_which_class_wins(fitted, splits):
+    (Xtr, _), (Xc, yc), (Xt, _) = splits
+    top_before = [fitted.predict(t)[0][0] for t in Xt[:60]]
+    fitted.calibrate(Xc, yc)
+    assert [fitted.predict(t)[0][0] for t in Xt[:60]] == top_before
+
+
+def test_calibrate_returns_a_summary_that_shows_the_improvement(fitted, splits):
+    _, (Xc, yc), _ = splits
+    info = fitted.calibrate(Xc, yc)
+    assert set(info) >= {"temperature", "n", "ece_before", "ece_after", "nll_before", "nll_after", "accuracy"}
+    assert info["n"] == len(Xc)
+    assert info["ece_after"] < info["ece_before"] and info["nll_after"] < info["nll_before"]
+    assert fitted.calibration["temperature"] == pytest.approx(info["temperature"])
+
+
+def test_calibrated_predictions_are_still_a_valid_distribution(fitted, splits):
+    _, (Xc, yc), (Xt, _) = splits
+    fitted.calibrate(Xc, yc)
+    for t in Xt[:20]:
+        preds = fitted.predict(t, k=10)
+        scores = [s for _, s in preds]
+        assert sum(scores) == pytest.approx(1.0, abs=1e-5)
+        assert scores == sorted(scores, reverse=True)
+        assert all(0 <= s <= 1 for s in scores)
+
+
+def test_k_truncation_does_not_change_the_calibrated_scores(fitted, splits):
+    _, (Xc, yc), (Xt, _) = splits
+    fitted.calibrate(Xc, yc)
+    full = dict(fitted.predict(Xt[0], k=4))
+    top2 = fitted.predict(Xt[0], k=2)
+    assert all(full[label] == pytest.approx(score, abs=1e-6) for label, score in top2)
+
+
+def test_abstain_below_applies_to_the_calibrated_confidence(fitted, splits):
+    _, (Xc, yc), (Xt, _) = splits
+    fitted.calibrate(Xc, yc)
+    top = fitted.predict(Xt[0])[0][1]
+    assert fitted.predict(Xt[0], abstain_below=top + 0.001) == []
+    assert fitted.predict(Xt[0], abstain_below=top - 0.001) != []
+
+
+def test_predict_batch_matches_predict_after_calibration(fitted, splits):
+    _, (Xc, yc), (Xt, _) = splits
+    fitted.calibrate(Xc, yc)
+    batch = fitted.predict_batch(Xt[:15], k=2)
+    for text, many in zip(Xt[:15], batch):
+        one = fitted.predict(text, k=2)
+        assert [l for l, _ in one] == [l for l, _ in many]
+        assert all(a == pytest.approx(b, abs=1e-5) for (_, a), (_, b) in zip(one, many))
+
+
+def test_predict_batch_scores_do_not_depend_on_k(new_classifier, data):
+    """Regression: predict_batch used to look up only k prototypes, so with more than k
+    classes its scores differed from predict() for the same text."""
+    X, y = data.classes(12, n_classes=5, noise=0.3, seed=4)
+    clf = new_classifier()
+    clf.add_examples(X, y)
+    q = data.add("q", 1, 0.3, 77)
+    single = dict(clf.predict(q, k=2))
+    batched = dict(clf.predict_batch([q], k=2)[0])
+    assert single.keys() == batched.keys()
+    for label in single:
+        assert batched[label] == pytest.approx(single[label], abs=1e-5)
+
+
+# --- conformal sets -------------------------------------------------------------------------------
+
+@pytest.mark.parametrize("alpha", [0.2, 0.1])
+def test_prediction_sets_cover_the_truth_at_roughly_the_requested_rate(fitted, splits, alpha):
+    _, (Xc, yc), (Xt, yt) = splits
+    fitted.calibrate(Xc, yc)
+    covered = [label in {l for l, _ in fitted.predict_set(t, alpha=alpha)} for t, label in zip(Xt, yt)]
+    # Guarantee is 1 - alpha on average; allow for sampling noise on 600 test points.
+    assert np.mean(covered) >= 1 - alpha - 0.05
+
+
+def test_stricter_alpha_gives_bigger_sets(fitted, splits):
+    _, (Xc, yc), (Xt, _) = splits
+    fitted.calibrate(Xc, yc)
+    size = lambda alpha: np.mean([len(fitted.predict_set(t, alpha=alpha)) for t in Xt[:120]])
+    assert size(0.02) >= size(0.1) >= size(0.4)
+    assert size(0.02) > size(0.4)
+
+
+def test_sets_are_small_for_easy_inputs_and_larger_for_ambiguous_ones(fitted, splits, data):
+    """Queries sitting between two class axes are genuinely ambiguous; those on an axis are not."""
+    import torch
+    from .synthetic import point
+
+    _, (Xc, yc), _ = splits
+    fitted.calibrate(Xc, yc)
+    easy = [data.add(f"easy{i}", i % 4, 0.1, 100 + i) for i in range(40)]
+    ambiguous = []
+    for i in range(40):
+        name = f"between{i}"
+        data.store[name] = torch.nn.functional.normalize(
+            point(i % 3, 0.1, 200 + i) + point(i % 3 + 1, 0.1, 300 + i), dim=0)
+        ambiguous.append(name)
+
+    mean_size = lambda texts: np.mean([len(fitted.predict_set(t, alpha=0.1)) for t in texts])
+    assert mean_size(easy) < 1.1
+    assert mean_size(ambiguous) > mean_size(easy) + 0.3
+
+
+def test_a_prediction_set_is_never_empty_and_starts_with_the_top_label(fitted, splits):
+    _, (Xc, yc), (Xt, _) = splits
+    fitted.calibrate(Xc, yc)
+    for t in Xt[:60]:
+        chosen = fitted.predict_set(t, alpha=0.5)
+        assert chosen and chosen[0][0] == fitted.predict(t)[0][0]
+
+
+def test_too_little_calibration_data_gives_the_whole_label_set(new_classifier, data):
+    X, y = data.classes(10, n_classes=3, noise=0.3, seed=5)
+    clf = new_classifier()
+    clf.add_examples(X, y)
+    Xc, yc = data.classes(3, n_classes=3, noise=0.3, seed=6, prefix="cal")     # 9 points
+    clf.calibrate(Xc, yc)
+    assert len(clf.predict_set(Xc[0], alpha=0.01)) == 3                         # cannot promise 99% with 9 points
+
+
+# --- validation and state ----------------------------------------------------------------------------
+
+def test_predict_set_requires_calibration(fitted, splits):
+    with pytest.raises(ValueError, match="calibrate"):
+        fitted.predict_set(splits[2][0][0])
+
+
+def test_predict_set_validates_alpha_and_text(fitted, splits):
+    _, (Xc, yc), (Xt, _) = splits
+    fitted.calibrate(Xc, yc)
+    with pytest.raises(ValueError, match="alpha"):
+        fitted.predict_set(Xt[0], alpha=0)
+    with pytest.raises(ValueError, match="Empty"):
+        fitted.predict_set("")
+
+
+@pytest.mark.parametrize("texts, labels, message", [
+    ([], [], "Empty"),
+    (["a"], ["C0", "C1"], "Mismatched"),
+    (["x"], ["not-a-class"], "Unknown"),
+    ([None], ["C0"], "strings"),
+    (["x"], [3], "strings"),
+])
+def test_calibrate_rejects_bad_input(fitted, texts, labels, message):
+    with pytest.raises(ValueError, match=message):
+        fitted.calibrate(texts, labels)
+    assert fitted.calibration is None
+
+
+def test_calibration_needs_two_classes(new_classifier, data):
+    X, y = data.classes(10, n_classes=1, noise=0.3, seed=7)
+    clf = new_classifier()
+    clf.add_examples(X, y)
+    with pytest.raises(ValueError, match="2 classes"):
+        clf.calibrate(X[:5], y[:5])
+
+
+def test_a_small_calibration_set_warns(fitted, splits, caplog):
+    _, (Xc, yc), _ = splits
+    with caplog.at_level(logging.WARNING):
+        fitted.calibrate(Xc[:10], yc[:10])
+    assert "noisy" in caplog.text
+
+
+def test_uncalibrated_predictions_are_unchanged(new_classifier, splits):
+    (X, y), _, (Xt, _) = splits
+    a, b = new_classifier(), new_classifier()
+    a.add_examples(X, y)
+    b.add_examples(X, y)
+    assert a.calibration is None
+    assert a.predict(Xt[0]) == b.predict(Xt[0])
+
+
+def test_adding_a_new_class_discards_the_calibration(fitted, splits, data, caplog):
+    _, (Xc, yc), _ = splits
+    fitted.calibrate(Xc, yc)
+    new_text = data.add("novel", 7, 0.1, 5)
+    with caplog.at_level(logging.WARNING):
+        fitted.add_examples([new_text], ["C-new"])
+    assert fitted.calibration is None and "Discarding" in caplog.text
+
+
+def test_adding_examples_to_known_classes_keeps_the_calibration(fitted, splits, data):
+    _, (Xc, yc), _ = splits
+    fitted.calibrate(Xc, yc)
+    fitted.add_examples([data.add("more", 0, 0.3, 11)], ["C0"])
+    assert fitted.calibration is not None
+
+
+def test_forgetting_a_class_discards_the_calibration(fitted, splits):
+    _, (Xc, yc), _ = splits
+    fitted.calibrate(Xc, yc)
+    fitted.forget("C3")
+    assert fitted.calibration is None
+    with pytest.raises(ValueError, match="calibrate"):
+        fitted.predict_set("tr0_0")
+
+
+def test_calibration_survives_save_and_reload(fitted, splits, tmp_path, data):
+    _, (Xc, yc), (Xt, yt) = splits
+    fitted.calibrate(Xc, yc)
+    fitted.save(str(tmp_path), include_onnx=False)
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+
+    assert loaded.calibration["temperature"] == pytest.approx(fitted.calibration["temperature"])
+    assert loaded.calibration["scores"] == pytest.approx(fitted.calibration["scores"])
+    for t in Xt[:10]:
+        a, b = fitted.predict(t, k=4), loaded.predict(t, k=4)
+        assert [l for l, _ in a] == [l for l, _ in b]
+        assert all(x == pytest.approx(y, abs=1e-4) for (_, x), (_, y) in zip(a, b))
+        assert {l for l, _ in fitted.predict_set(t)} == {l for l, _ in loaded.predict_set(t)}
+
+
+def test_saves_without_calibration_load_as_uncalibrated(fitted, tmp_path):
+    fitted.save(str(tmp_path), include_onnx=False)
+    path = tmp_path / "config.json"
+    config = json.loads(path.read_text(encoding="utf-8"))
+    config.pop("calibration", None)                                   # an older save
+    path.write_text(json.dumps(config), encoding="utf-8")
+    assert AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu").calibration is None
+
+
+def test_sklearn_wrapper_uses_calibrated_probabilities(base_model, splits, data):
+    from adaptive_classifier import SklearnAdaptiveClassifier
+    (X, y), (Xc, yc), (Xt, yt) = splits
+    est = SklearnAdaptiveClassifier(base_model, use_onnx=False, device="cpu").fit(X, y)
+    before = est.predict_proba(Xt[:40]).max(axis=1).mean()
+    est.classifier_.calibrate(Xc, yc)
+    after = est.predict_proba(Xt[:40]).max(axis=1).mean()
+    assert after != pytest.approx(before, abs=1e-3)
+    assert est.predict_proba(Xt[:5]).sum(axis=1) == pytest.approx(np.ones(5), abs=1e-5)

From a725359ba812c558d61ada472c38d03860b7edab Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 10:53:34 +0000
Subject: [PATCH 10/16] Add active learning (suggest_labels) and drift
 detection (drift_report)

- suggest_labels(texts, n, strategy, diverse) ranks unlabeled texts by how
  useful a label would be: margin, entropy, least_confidence, or ood (to
  find classes the model has not seen). diverse=True discounts texts
  similar to earlier picks so near-duplicates are not all chosen
- drift_report(texts) scores a window of incoming texts with ood_score and
  runs a one-sided binomial test of the out-of-distribution share against
  an expected rate; returns the rate, mean score, p-value and a verdict

tests/test_active_learning.py (38 tests): each strategy's ranking is
checked against an independent recomputation from predict(), boundary
texts outrank clear ones, the ood strategy finds a hidden class, diverse
selection avoids twins, and drift is flagged for shifted streams but not
for in-distribution ones. Mutation-checked with 12 deliberate faults; two
initially survived and exposed missing tests, now added.

Honest note kept in the docs: on heavily overlapping synthetic classes
uncertainty sampling did not beat random labelling (0.67 vs 0.69), so no
efficiency claim is made.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 CHANGELOG.md                          |   3 +
 README.md                             |  16 ++
 docs/API.md                           |  16 ++
 src/adaptive_classifier/classifier.py | 131 ++++++++++++++
 tests/test_active_learning.py         | 239 ++++++++++++++++++++++++++
 5 files changed, 405 insertions(+)
 create mode 100644 tests/test_active_learning.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b1dde8a..c82a51f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,9 @@
 - `calibrate(texts, labels)` fits one temperature on held-out labelled data so confidences match how often the model is right, and `calibration_report(...)` measures accuracy, expected calibration error and NLL. `predict`, `predict_batch` and `abstain_below` then use calibrated probabilities. Calibration is saved with the model and discarded when the set of classes changes.
 - `predict_set(text, alpha=0.1)`: split-conformal prediction sets that contain the true label with probability at least `1 - alpha`.
 
+- `suggest_labels(texts, n, strategy, diverse)`: active learning. Ranks unlabeled texts by how much a label would help (`margin`, `entropy`, `least_confidence`, or `ood` to discover unseen classes), optionally spreading the picks so near-duplicates are not all chosen.
+- `drift_report(texts)`: tests whether a window of incoming texts has moved away from the known classes (share of out-of-distribution texts against an expected rate, one-sided binomial test).
+
 ### Changed (default behaviour: predictions change)
 - **The neural head is now actually trained.** It used to get about ten optimiser steps and stop at the first loss plateau, so it could not learn even a 2-D XOR from 120 examples. It now trains for at least `head_steps` (default 300) steps with a cosine-decayed `head_learning_rate` (default 0.003). Small memories take roughly 1-2 s per `add_examples` call; lower `head_steps` if you add examples very frequently.
 - **Prototype scores are sharper.** They are now `softmax(-distance² / prototype_temperature)` (default 0.25) instead of `softmax(exp(-distance))`, which barely separated the nearest class from the rest (about 0.45 vs 0.21), so a confident head could always outvote it. `prototype_temperature=None` restores the old scoring.
diff --git a/README.md b/README.md
index cfee60a..eb0b69a 100644
--- a/README.md
+++ b/README.md
@@ -356,6 +356,22 @@ classifier.calibration_report(test_texts, test_labels)   # check it on separate
 
 One easy input gives a one-label set; an ambiguous one gives several. Calibration is saved with the model and discarded if you add or forget a class; re-run it after substantial new data.
 
+### Active Learning and Drift
+
+```python
+# Which of these 5,000 unlabeled texts should a person label next?
+for index, score in classifier.suggest_labels(unlabeled, n=20, strategy="margin", diverse=True):
+    print(unlabeled[index])
+
+classifier.suggest_labels(unlabeled, n=20, strategy="ood")   # hunt for classes you have not defined yet
+
+# Has this week's traffic moved away from what the model knows?
+report = classifier.drift_report(recent_texts)
+report["drifted"], report["ood_rate"], report["p_value"]
+```
+
+`margin` (the default) finds texts on a decision boundary; `ood` finds texts unlike any known class. Uncertainty sampling is not guaranteed to beat labeling at random: when classes overlap so much that the confusion is irreducible noise, labeling the confusing texts teaches the model little (it did not in a synthetic test with heavily overlapping classes). It earns its keep when ambiguity comes from missing data, and `ood` is the reliable way to find classes you have not defined yet.
+
 ### Fixing Mistakes and Knowing When to Abstain
 
 ```python
diff --git a/docs/API.md b/docs/API.md
index 0ee6de1..0416958 100644
--- a/docs/API.md
+++ b/docs/API.md
@@ -154,6 +154,22 @@ def predict_set(text: str, alpha: float = 0.1) -> List[Tuple[str, float]]
 
 Split-conformal prediction set: contains the true label with probability at least `1 - alpha` for inputs that look like the calibration data (an average guarantee, not per input). Never empty. Requires `calibrate` first.
 
+#### suggest_labels
+
+```python
+def suggest_labels(texts, n=10, strategy="margin", diverse=False) -> List[Tuple[int, float]]
+```
+
+Rank unlabeled texts by how useful a label would be. Strategies: `margin` (top two classes close together), `entropy`, `least_confidence`, and `ood` (far from every known class, for discovering new classes). With `diverse=True` texts similar to an earlier pick are discounted so the picks are spread out. Returns up to `n` `(index into texts, score)` pairs, best first; uses calibrated probabilities if `calibrate` was run.
+
+#### drift_report
+
+```python
+def drift_report(texts, expected_ood_rate=0.05, alpha=0.01, threshold=None) -> Dict[str, Any]
+```
+
+Tests whether the share of out-of-distribution texts in a window of incoming data is higher than `expected_ood_rate` (one-sided binomial test). Returns `n`, `ood_rate`, `mean_ood_score`, `p_value` and `drifted` (`p_value < alpha`). A window of a few dozen texts cannot detect small shifts; use a stricter `alpha` when checking repeatedly.
+
 #### forget
 
 ```python
diff --git a/src/adaptive_classifier/classifier.py b/src/adaptive_classifier/classifier.py
index 35f22d7..78b57f1 100644
--- a/src/adaptive_classifier/classifier.py
+++ b/src/adaptive_classifier/classifier.py
@@ -612,6 +612,137 @@ def predict_set(self, text: str, alpha: float = 0.1) -> List[Tuple[str, float]]:
         chosen = [(label, p) for label, p in full if 1.0 - p <= threshold]
         return chosen or full[:1]
 
+    # --- active learning and drift -----------------------------------------------------
+
+    _SUGGEST_STRATEGIES = ("margin", "entropy", "least_confidence", "ood")
+
+    def suggest_labels(
+        self,
+        texts: List[str],
+        n: int = 10,
+        strategy: str = "margin",
+        diverse: bool = False,
+    ) -> List[Tuple[int, float]]:
+        """Pick the unlabeled texts that would teach the model the most if labelled.
+
+        Strategies (higher score = more worth labelling):
+
+        * ``"margin"``: how close the top two classes are (``1 - (p1 - p2)``).
+          A good default: it finds texts sitting on a decision boundary.
+        * ``"entropy"``: uncertainty spread over all classes, scaled to 0-1.
+        * ``"least_confidence"``: ``1 - p1``.
+        * ``"ood"``: how far the text is from every known class (see
+          `ood_score`). Use it to discover classes the model has never seen.
+
+        With ``diverse=True`` the picks are spread out: after each choice,
+        texts similar to it are discounted, so you do not label ten near
+        duplicates of the same confusing case.
+
+        Probabilities are calibrated if `calibrate` was run.
+
+        Returns:
+            Up to `n` ``(index into texts, score)`` pairs, best first.
+        """
+        if not texts:
+            raise ValueError("Empty input list")
+        if strategy not in self._SUGGEST_STRATEGIES:
+            raise ValueError(f"strategy must be one of {self._SUGGEST_STRATEGIES}; got {strategy!r}")
+        if not isinstance(n, numbers.Integral) or isinstance(n, bool) or n < 1:
+            raise ValueError(f"n must be a positive integer; got {n!r}")
+        if not self.label_to_id:
+            raise ValueError("suggest_labels needs a trained classifier")
+        if any(not isinstance(t, str) or not t for t in texts):
+            raise ValueError("texts must all be non-empty strings")
+
+        scores = self._informativeness(texts, strategy)
+        count = min(int(n), len(texts))
+
+        if not diverse:
+            order = sorted(range(len(texts)), key=lambda i: scores[i], reverse=True)
+            return [(i, float(scores[i])) for i in order[:count]]
+
+        embeddings = torch.stack(self._get_embeddings(texts))
+        chosen: List[int] = []
+        discount = np.ones(len(texts))
+        for _ in range(count):
+            adjusted = np.where(np.isin(np.arange(len(texts)), chosen), -np.inf, scores * discount)
+            best = int(np.argmax(adjusted))
+            chosen.append(best)
+            # Texts similar to the one just chosen are worth less from now on.
+            similarity = (embeddings @ embeddings[best]).clamp(min=0).numpy()
+            discount = np.minimum(discount, 1.0 - similarity)
+        return [(i, float(scores[i])) for i in chosen]
+
+    def _informativeness(self, texts: List[str], strategy: str) -> np.ndarray:
+        if strategy == "ood":
+            return np.array(self._ood_scores(texts))
+        classes = len(self.id_to_label)
+        batches = self.predict_batch(texts, k=classes)
+        scores = np.zeros(len(texts))
+        for row, preds in enumerate(batches):
+            probs = np.array(sorted((p for _, p in preds), reverse=True))
+            if probs.size == 0:
+                continue
+            if strategy == "least_confidence":
+                scores[row] = 1.0 - probs[0]
+            elif strategy == "margin":
+                scores[row] = 1.0 - (probs[0] - (probs[1] if probs.size > 1 else 0.0))
+            else:                                                    # entropy
+                nonzero = probs[probs > 0]
+                entropy = float(-(nonzero * np.log(nonzero)).sum())
+                scores[row] = entropy / math.log(classes) if classes > 1 else 0.0
+        return scores
+
+    def _ood_scores(self, texts: List[str]) -> List[float]:
+        with torch.no_grad():
+            embeddings = self._get_embeddings(texts)
+        return [self.memory.ood_score(embedding)[0] for embedding in embeddings]
+
+    def drift_report(
+        self,
+        texts: List[str],
+        expected_ood_rate: float = 0.05,
+        alpha: float = 0.01,
+        threshold: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """Has incoming data moved away from what the classifier knows?
+
+        Scores each text with `ood_score` and tests whether the share of
+        out-of-distribution texts is larger than `expected_ood_rate` (what you
+        see on healthy traffic; 5% by default) with a one-sided binomial test.
+        Run it on a recent window of production inputs. A single window of a
+        few dozen texts cannot detect small shifts, and repeated checks need a
+        stricter `alpha`.
+
+        Returns:
+            ``n``, ``ood_rate``, ``mean_ood_score``, ``p_value`` and
+            ``drifted`` (True when ``p_value < alpha``).
+        """
+        if not texts:
+            raise ValueError("Empty input list")
+        if not self.label_to_id:
+            raise ValueError("drift_report needs a trained classifier")
+        if not 0.0 <= expected_ood_rate < 1.0:
+            raise ValueError(f"expected_ood_rate must be in [0, 1); got {expected_ood_rate}")
+        if not 0.0 < alpha < 1.0:
+            raise ValueError(f"alpha must be between 0 and 1; got {alpha}")
+        if any(not isinstance(t, str) or not t for t in texts):
+            raise ValueError("texts must all be non-empty strings")
+
+        from scipy.stats import binomtest
+
+        limit = self.config.ood_threshold if threshold is None else threshold
+        scores = np.array(self._ood_scores(texts))
+        flagged = int((scores > limit).sum())
+        p_value = float(binomtest(flagged, len(texts), expected_ood_rate, alternative="greater").pvalue)
+        return {
+            'n': len(texts),
+            'ood_rate': flagged / len(texts),
+            'mean_ood_score': float(scores.mean()),
+            'p_value': p_value,
+            'drifted': p_value < alpha,
+        }
+
     def forget(self, labels: Union[str, List[str]]):
         """Remove one or more classes entirely.
 
diff --git a/tests/test_active_learning.py b/tests/test_active_learning.py
new file mode 100644
index 0000000..4fb748e
--- /dev/null
+++ b/tests/test_active_learning.py
@@ -0,0 +1,239 @@
+"""suggest_labels (what to label next) and drift_report (has the data moved?).
+
+What is checked is the contract: the ranking is the one each strategy defines,
+it is deterministic, and it does the job it exists for (surface the confusing
+and the unfamiliar). It is deliberately NOT asserted that uncertainty sampling
+beats random labelling on every dataset: on heavily overlapping classes, where
+the confusion is irreducible noise, a synthetic experiment showed it does not.
+"""
+
+import numpy as np
+import pytest
+import torch
+import torch.nn.functional as F
+
+from .synthetic import Dataset, point
+
+STRATEGIES = ["margin", "entropy", "least_confidence", "ood"]
+
+
+@pytest.fixture
+def data(monkeypatch):
+    return Dataset().patch(monkeypatch)
+
+
+@pytest.fixture
+def clf(new_classifier, data):
+    X, y = data.classes(20, n_classes=3, noise=0.15, seed=1)
+    c = new_classifier(head_steps=60)
+    c.add_examples(X, y)
+    return c
+
+
+def add_between(data, name, a, b, noise=0.05, seed=0):
+    """A point halfway between the axes of classes a and b."""
+    data.store[name] = F.normalize(point(a, noise, seed) + point(b, noise, seed + 1), dim=0)
+    return name
+
+
+# --- ranking ---------------------------------------------------------------------------------
+
+@pytest.mark.parametrize("strategy", ["margin", "entropy", "least_confidence"])
+def test_boundary_texts_outrank_clear_ones(clf, data, strategy):
+    clear = [data.add(f"clear{i}", i % 3, 0.05, 500 + i) for i in range(6)]
+    boundary = [add_between(data, f"edge{i}", 0, 1, seed=600 + 2 * i) for i in range(3)]
+    pool = clear + boundary
+    ranked = [i for i, _ in clf.suggest_labels(pool, n=3, strategy=strategy)]
+    assert set(ranked) == {6, 7, 8}, f"{strategy} picked {ranked}"
+
+
+def test_ood_strategy_surfaces_a_class_the_model_has_never_seen(clf, data):
+    known = [data.add(f"k{i}", i % 3, 0.15, 700 + i) for i in range(18)]
+    unseen = [data.add(f"u{i}", 9, 0.15, 800 + i) for i in range(2)]        # axis 9: no such class
+    pool = known + unseen
+    picks = {i for i, _ in clf.suggest_labels(pool, n=2, strategy="ood")}
+    assert picks == {18, 19}
+
+
+def test_most_of_a_hidden_class_is_found_within_a_small_budget(clf, data):
+    """10% of the pool belongs to a class the model has never seen; random labelling would
+    find ~1 in 10, the ood strategy should find almost all of them first."""
+    known = [data.add(f"k{i}", i % 3, 0.15, 900 + i) for i in range(45)]
+    hidden = [data.add(f"h{i}", 7, 0.15, 950 + i) for i in range(5)]
+    picks = [i for i, _ in clf.suggest_labels(known + hidden, n=5, strategy="ood")]
+    assert sum(i >= 45 for i in picks) >= 4
+
+
+def reference_scores(clf, texts, strategy):
+    """Recompute each strategy's score independently from predict()."""
+    out = []
+    for text in texts:
+        p = sorted((s for _, s in clf.predict(text, k=len(clf.id_to_label))), reverse=True)
+        if strategy == "least_confidence":
+            out.append(1 - p[0])
+        elif strategy == "margin":
+            out.append(1 - (p[0] - p[1]))
+        else:
+            q = np.array([x for x in p if x > 0])
+            out.append(float(-(q * np.log(q)).sum() / np.log(len(p))))
+    return out
+
+
+@pytest.mark.parametrize("strategy", ["margin", "entropy", "least_confidence"])
+def test_scores_are_what_the_strategy_defines(clf, data, strategy):
+    pool = [data.add(f"s{i}", i % 3, 0.4, 1000 + i) for i in range(12)]
+    got = dict(clf.suggest_labels(pool, n=12, strategy=strategy))
+    expected = reference_scores(clf, pool, strategy)
+    for i, score in got.items():
+        assert score == pytest.approx(expected[i], abs=1e-4)
+
+
+@pytest.mark.parametrize("strategy", STRATEGIES)
+def test_results_are_sorted_unique_and_in_range(clf, data, strategy):
+    pool = [data.add(f"r{i}", i % 3, 0.3, 1100 + i) for i in range(15)]
+    out = clf.suggest_labels(pool, n=8, strategy=strategy)
+    indices = [i for i, _ in out]
+    scores = [s for _, s in out]
+    assert len(out) == 8 and len(set(indices)) == 8
+    assert all(0 <= i < 15 for i in indices)
+    assert scores == sorted(scores, reverse=True)
+    assert all(np.isfinite(scores))
+
+
+def test_n_larger_than_the_pool_returns_everything(clf, data):
+    pool = [data.add(f"p{i}", i % 3, 0.3, 1200 + i) for i in range(4)]
+    assert sorted(i for i, _ in clf.suggest_labels(pool, n=50)) == [0, 1, 2, 3]
+
+
+def test_suggestions_are_deterministic(clf, data):
+    pool = [data.add(f"d{i}", i % 3, 0.35, 1300 + i) for i in range(10)]
+    assert clf.suggest_labels(pool, n=5) == clf.suggest_labels(pool, n=5)
+
+
+def test_calibration_changes_scores_but_not_the_most_ambiguous_pick(clf, data):
+    pool = [data.add(f"c{i}", i % 3, 0.05, 1400 + i) for i in range(6)] + [add_between(data, "mid", 1, 2)]
+    before = clf.suggest_labels(pool, n=1)[0]
+    Xc, yc = data.classes(30, n_classes=3, noise=0.45, seed=9, prefix="cal")
+    clf.calibrate(Xc, yc)
+    after = clf.suggest_labels(pool, n=1)[0]
+    assert before[0] == after[0] == 6
+
+
+# --- diversity ----------------------------------------------------------------------------------
+
+def test_diverse_selection_avoids_near_duplicates(clf, data):
+    # Candidates between several class pairs; the most ambiguous one is then duplicated five times.
+    candidates = [add_between(data, f"cand{i}", i % 3, (i + 1) % 3, noise=0.15, seed=1500 + 2 * i)
+                  for i in range(8)]
+    top = clf.suggest_labels(candidates, n=1, strategy="margin")[0][0]
+    for i in range(5):
+        data.store[f"twin{i}"] = data.store[candidates[top]]
+    twins = [f"twin{i}" for i in range(5)]
+    others = [c for j, c in enumerate(candidates) if j != top]
+    pool = twins + others
+
+    plain = [i for i, _ in clf.suggest_labels(pool, n=3, strategy="margin")]
+    spread = [i for i, _ in clf.suggest_labels(pool, n=3, strategy="margin", diverse=True)]
+    assert sum(i < 5 for i in plain) == 3, f"the twins should fill the plain top 3: {plain}"
+    assert sum(i < 5 for i in spread) == 1, f"diverse selection should take one twin: {spread}"
+    assert len(set(spread)) == 3
+
+
+def test_diverse_selection_with_only_duplicates_still_returns_distinct_indices(clf, data):
+    base = add_between(data, "dup_base", 0, 1)
+    pool = []
+    for i in range(4):
+        data.store[f"dup{i}"] = data.store[base]
+        pool.append(f"dup{i}")
+    picks = [i for i, _ in clf.suggest_labels(pool, n=3, strategy="margin", diverse=True)]
+    assert len(picks) == 3 and len(set(picks)) == 3
+
+
+@pytest.mark.parametrize("diverse", [False, True])
+def test_n_larger_than_the_pool_returns_each_text_once(clf, data, diverse):
+    pool = [data.add(f"big{i}", i % 3, 0.3, 1700 + i) for i in range(4)]
+    picks = [i for i, _ in clf.suggest_labels(pool, n=10, diverse=diverse)]
+    assert sorted(picks) == [0, 1, 2, 3]
+
+
+# --- validation -----------------------------------------------------------------------------------
+
+@pytest.mark.parametrize("kwargs, message", [
+    ({"texts": []}, "Empty"),
+    ({"texts": ["tr0_0"], "n": 0}, "positive integer"),
+    ({"texts": ["tr0_0"], "n": -3}, "positive integer"),
+    ({"texts": ["tr0_0"], "n": 1.5}, "positive integer"),
+    ({"texts": ["tr0_0"], "strategy": "vibes"}, "strategy"),
+    ({"texts": ["tr0_0", ""]}, "non-empty"),
+    ({"texts": ["tr0_0", None]}, "non-empty"),
+])
+def test_bad_arguments_are_rejected(clf, kwargs, message):
+    with pytest.raises(ValueError, match=message):
+        clf.suggest_labels(**kwargs)
+
+
+def test_needs_a_trained_classifier(new_classifier):
+    with pytest.raises(ValueError, match="trained"):
+        new_classifier().suggest_labels(["a"])
+    with pytest.raises(ValueError, match="trained"):
+        new_classifier().drift_report(["a"])
+
+
+# --- drift ------------------------------------------------------------------------------------------
+
+def test_data_like_the_training_set_is_not_drift(clf, data):
+    fresh, _ = data.classes(40, n_classes=3, noise=0.15, seed=21, prefix="live")
+    report = clf.drift_report(fresh)
+    assert report["drifted"] is False and report["p_value"] > 0.01
+    assert report["n"] == 120 and report["ood_rate"] < 0.1
+
+
+def test_a_shifted_stream_is_flagged(clf, data):
+    live = [data.add(f"ok{i}", i % 3, 0.15, 2000 + i) for i in range(30)]
+    live += [data.add(f"new{i}", 9, 0.15, 2100 + i) for i in range(30)]          # half the traffic is unfamiliar
+    report = clf.drift_report(live)
+    assert report["drifted"] is True and report["p_value"] < 1e-6
+    assert report["ood_rate"] > 0.4
+
+
+def test_mean_score_grows_with_the_amount_of_shift(clf, data):
+    def stream(shift_fraction, tag):
+        total = 60
+        moved = int(total * shift_fraction)
+        return ([data.add(f"{tag}ok{i}", i % 3, 0.15, 2200 + i) for i in range(total - moved)]
+                + [data.add(f"{tag}mv{i}", 9, 0.15, 2300 + i) for i in range(moved)])
+
+    scores = [clf.drift_report(stream(f, f"f{int(f * 10)}"))["mean_ood_score"] for f in (0.0, 0.25, 0.5, 1.0)]
+    assert scores == sorted(scores) and scores[-1] > 1.5 * scores[0]
+
+
+def test_a_tiny_window_cannot_prove_drift(clf, data):
+    report = clf.drift_report([data.add("lone", 9, 0.15, 2400)])
+    assert report["ood_rate"] == 1.0 and report["drifted"] is False       # one sample: p = 0.05, above alpha
+
+
+def test_expected_rate_and_alpha_control_the_verdict(clf, data):
+    live = [data.add(f"x{i}", i % 3, 0.15, 2500 + i) for i in range(45)]
+    live += [data.add(f"y{i}", 9, 0.15, 2600 + i) for i in range(5)]              # 10% unfamiliar
+    assert clf.drift_report(live, expected_ood_rate=0.01)["drifted"] is True
+    assert clf.drift_report(live, expected_ood_rate=0.3)["drifted"] is False
+    assert clf.drift_report(live, expected_ood_rate=0.01, alpha=1e-9)["drifted"] is False
+
+
+def test_threshold_override_changes_what_counts_as_unfamiliar(clf, data):
+    live = [data.add(f"t{i}", i % 3, 0.3, 2700 + i) for i in range(30)]
+    strict = clf.drift_report(live, threshold=0.2)
+    lenient = clf.drift_report(live, threshold=50.0)
+    assert strict["ood_rate"] > lenient["ood_rate"] == 0.0
+
+
+@pytest.mark.parametrize("kwargs, message", [
+    ({"texts": []}, "Empty"),
+    ({"texts": ["tr0_0"], "expected_ood_rate": 1.0}, "expected_ood_rate"),
+    ({"texts": ["tr0_0"], "expected_ood_rate": -0.1}, "expected_ood_rate"),
+    ({"texts": ["tr0_0"], "alpha": 0}, "alpha"),
+    ({"texts": ["tr0_0", ""]}, "non-empty"),
+])
+def test_drift_report_rejects_bad_arguments(clf, kwargs, message):
+    with pytest.raises(ValueError, match=message):
+        clf.drift_report(**kwargs)

From 9013ac45c9ad8b89fc00aa58771b1be657921af3 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 11:08:08 +0000
Subject: [PATCH 11/16] Add HTTP serving, a thread-safe classifier, async
 methods and a Dockerfile

Thread safety (a bug fix that serving needs)
- Predicting while another thread added or forgot classes crashed with
  "selected index k out of range", and concurrent add_examples calls
  corrupted each other's gradient computation ("modified by an inplace
  operation"). Both reproduced reliably. Every public operation now holds a
  per-instance reentrant lock, so one classifier can be shared between
  threads. The stress tests fail 5 of 5 runs when the lock is removed

Serving
- adaptive_classifier.serving: create_app(classifier) builds a FastAPI app
  (/health, /info, /predict, /predict_batch, /predict_set, /ood, and
  /examples, /forget, /remove_examples when updates are allowed);
  python -m adaptive_classifier.serving ./model runs it
- Updates are off by default; --allow-updates enables them, an API key
  (bearer token, constant-time compare) protects them, --save-dir persists
  them. Request size limits, validation (422) and classifier errors (400)
- apredict / apredict_batch / aadd_examples run in a worker thread so they
  do not block an event loop
- New serve extra; fastapi and httpx added to the test extra
- docker/Dockerfile (CPU torch, non-root, health check) and docs/serving.md.
  The Dockerfile was not built: there was no Docker daemon available

Self-contained saves
- Loading a saved directory used to fetch the tokenizer from the base
  model's repository even though save() writes it alongside the model, so
  an ONNX-saved classifier could not start without Hub access. The saved
  tokenizer is now used (and fetched from the Hub repo with the other
  files), for both the ONNX and PyTorch paths

tests/test_concurrency.py (11) and tests/test_serving.py (43), plus tokenizer
loading tests. Mutation-checked: 15 deliberate faults, all caught.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 .dockerignore                         |  11 +
 CHANGELOG.md                          |   5 +
 README.md                             |  10 +
 docker/Dockerfile                     |  40 ++++
 docs/serving.md                       |  71 +++++++
 pyproject.toml                        |   7 +
 src/adaptive_classifier/classifier.py | 104 +++++++++-
 src/adaptive_classifier/serving.py    | 238 +++++++++++++++++++++
 tests/test_concurrency.py             | 179 ++++++++++++++++
 tests/test_portable_export.py         |  40 ++++
 tests/test_serving.py                 | 285 ++++++++++++++++++++++++++
 11 files changed, 989 insertions(+), 1 deletion(-)
 create mode 100644 .dockerignore
 create mode 100644 docker/Dockerfile
 create mode 100644 docs/serving.md
 create mode 100644 src/adaptive_classifier/serving.py
 create mode 100644 tests/test_concurrency.py
 create mode 100644 tests/test_serving.py

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 0000000..b46b346
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,11 @@
+.git
+.github
+tests
+docs
+scripts
+examples
+build
+dist
+**/__pycache__
+**/*.egg-info
+*.webp
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c82a51f..387ddf1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,10 @@
 - `suggest_labels(texts, n, strategy, diverse)`: active learning. Ranks unlabeled texts by how much a label would help (`margin`, `entropy`, `least_confidence`, or `ood` to discover unseen classes), optionally spreading the picks so near-duplicates are not all chosen.
 - `drift_report(texts)`: tests whether a window of incoming texts has moved away from the known classes (share of out-of-distribution texts against an expected rate, one-sided binomial test).
 
+- **Serving.** `python -m adaptive_classifier.serving ./model` runs a FastAPI server (`/predict`, `/predict_batch`, `/predict_set`, `/ood`, and optional authenticated `/examples`, `/forget`, `/remove_examples`); `create_app(classifier)` returns the app for embedding. Install with `pip install "adaptive-classifier[serve]"`. See `docs/serving.md`. `docker/Dockerfile` is provided but has not been built in the environment it was written in.
+- Async methods `apredict`, `apredict_batch` and `aadd_examples` that run in a worker thread.
+- A saved directory is self-contained for loading: the tokenizer saved next to the model is used instead of fetching the base model's repository, so an ONNX-saved classifier starts without Hub access.
+
 ### Changed (default behaviour: predictions change)
 - **The neural head is now actually trained.** It used to get about ten optimiser steps and stop at the first loss plateau, so it could not learn even a 2-D XOR from 120 examples. It now trains for at least `head_steps` (default 300) steps with a cosine-decayed `head_learning_rate` (default 0.003). Small memories take roughly 1-2 s per `add_examples` call; lower `head_steps` if you add examples very frequently.
 - **Prototype scores are sharper.** They are now `softmax(-distance² / prototype_temperature)` (default 0.25) instead of `softmax(exp(-distance))`, which barely separated the nearest class from the rest (about 0.45 vs 0.21), so a confident head could always outvote it. `prototype_temperature=None` restores the old scoring.
@@ -25,6 +29,7 @@
 - `clear_memory(labels=[...])` now removes those classes completely (it used to leave them in the label map and neural head, so they could still be predicted).
 
 ### Fixed
+- **Thread safety.** Predicting while another thread added or forgot classes crashed with errors such as "selected index k out of range", and concurrent `add_examples` calls corrupted each other. Every public operation now holds a per-instance lock, so one classifier can be shared between threads. Throughput scales with processes, not threads.
 - `predict_batch(k=...)` scored only `k` classes before blending, so with more than `k` classes its scores differed from `predict` for the same text. It now scores every class and trims to `k`.
 - Accuracy on small training sets (see above): the default configuration got only 75-92% of its own training data right on perfectly separable data with fewer than 10 examples per class.
 - Integer (or other non-string) labels corrupted a classifier on save/reload, duplicating classes. `add_examples` now rejects non-string labels with a clear error; convert with `str(label)`.
diff --git a/README.md b/README.md
index eb0b69a..b1a5e5f 100644
--- a/README.md
+++ b/README.md
@@ -372,6 +372,16 @@ report["drifted"], report["ood_rate"], report["p_value"]
 
 `margin` (the default) finds texts on a decision boundary; `ood` finds texts unlike any known class. Uncertainty sampling is not guaranteed to beat labeling at random: when classes overlap so much that the confusion is irreducible noise, labeling the confusing texts teaches the model little (it did not in a synthetic test with heavily overlapping classes). It earns its keep when ambiguity comes from missing data, and `ood` is the reliable way to find classes you have not defined yet.
 
+### Serving
+
+```bash
+pip install "adaptive-classifier[serve]"
+python -m adaptive_classifier.serving ./my_classifier --port 8000
+curl -s localhost:8000/predict -H 'content-type: application/json' -d '{"text": "where is my refund?"}'
+```
+
+Predictions, batches, conformal sets and out-of-distribution checks over HTTP, plus optional API-key-protected endpoints for adding examples and classes while it runs. One classifier is safe to share between threads, and `await classifier.apredict(...)` keeps an event loop responsive. See [docs/serving.md](docs/serving.md) and `docker/Dockerfile`.
+
 ### Fixing Mistakes and Knowing When to Abstain
 
 ```python
diff --git a/docker/Dockerfile b/docker/Dockerfile
new file mode 100644
index 0000000..8481f14
--- /dev/null
+++ b/docker/Dockerfile
@@ -0,0 +1,40 @@
+# Serve a saved AdaptiveClassifier over HTTP.
+#
+#   docker build -f docker/Dockerfile -t adaptive-classifier .
+#   docker run --rm -p 8000:8000 -v "$PWD/my_classifier:/model:ro" adaptive-classifier
+#   curl -s localhost:8000/predict -H 'content-type: application/json' -d '{"text": "where is my refund?"}'
+#
+# Updates are off by default. To allow them (and persist what is learned), mount the
+# model read-write and pass flags, for example:
+#   docker run --rm -p 8000:8000 -v "$PWD/my_classifier:/model" \
+#     -e ADAPTIVE_CLASSIFIER_API_KEY=change-me adaptive-classifier \
+#     /model --host 0.0.0.0 --allow-updates --save-dir /model
+#
+# A classifier saved with the default `save()` includes an ONNX copy of the encoder and the
+# tokenizer, so it starts without reaching the Hugging Face Hub. One saved with
+# `include_onnx=False` fetches its base model from the Hub on startup; mount a populated
+# Hugging Face cache at /cache/huggingface (or bake the model in) for offline use.
+FROM python:3.12-slim
+
+ENV PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    HF_HOME=/cache/huggingface
+
+WORKDIR /app
+
+# CPU-only PyTorch keeps the image a few GB smaller than the default CUDA build.
+RUN pip install --index-url https://download.pytorch.org/whl/cpu torch
+
+COPY pyproject.toml README.md LICENSE ./
+COPY src ./src
+RUN pip install ".[serve]"
+
+RUN useradd --create-home --uid 1000 app && mkdir -p /cache/huggingface && chown -R app /cache
+USER app
+
+EXPOSE 8000
+HEALTHCHECK --interval=30s --timeout=5s --start-period=60s --retries=3 \
+    CMD python -c "import urllib.request,sys; sys.exit(0 if urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=4).status == 200 else 1)"
+
+ENTRYPOINT ["python", "-m", "adaptive_classifier.serving"]
+CMD ["/model", "--host", "0.0.0.0", "--port", "8000"]
diff --git a/docs/serving.md b/docs/serving.md
new file mode 100644
index 0000000..5c39328
--- /dev/null
+++ b/docs/serving.md
@@ -0,0 +1,71 @@
+# Serving a classifier over HTTP
+
+```bash
+pip install "adaptive-classifier[serve]"
+python -m adaptive_classifier.serving ./my_classifier --port 8000
+```
+
+`./my_classifier` is a directory written by `classifier.save(...)`. It is self-contained: it holds the
+encoder (ONNX), the tokenizer and the learned state, so the server starts without reaching the Hugging
+Face Hub. (Directories saved with `include_onnx=False`, or by an older version, fetch the base model
+from the Hub when they start.)
+
+## Endpoints
+
+| Method and path | Body | Returns |
+|---|---|---|
+| `GET /health` | | `{"status": "ok", "classes": 3}` |
+| `GET /info` | | classes, examples per class, whether calibrated, whether updates are allowed |
+| `POST /predict` | `{"text", "k"?, "abstain_below"?, "abstain_ood"?}` | `{"predictions": [{"label", "score"}], "abstained": false}` |
+| `POST /predict_batch` | `{"texts": [...], "k"?}` | `{"predictions": [[...], ...]}` |
+| `POST /predict_set` | `{"text", "alpha"?}` | conformal label set (needs `calibrate`) |
+| `POST /ood` | `{"text", "threshold"?}` | `{"score", "threshold", "ood"}` |
+| `POST /examples` | `{"texts": [...], "labels": [...]}` | add examples or new classes (updates only) |
+| `POST /forget` | `{"labels": [...]}` | remove classes (updates only) |
+| `POST /remove_examples` | `{"texts": [...], "label"?}` | remove examples (updates only) |
+
+```bash
+curl -s localhost:8000/predict -H 'content-type: application/json' \
+     -d '{"text": "where is my refund?", "k": 3}'
+```
+
+Bad input gets a 4xx with a message: 422 for a malformed request (empty text, wrong types, batch too
+large) and 400 when the classifier rejects it (unknown label, calibration missing, ...). The interactive
+API docs are at `/docs`.
+
+## Updates
+
+Learning at runtime is the point of the library, but it is off by default on a server:
+
+```bash
+export ADAPTIVE_CLASSIFIER_API_KEY=change-me
+python -m adaptive_classifier.serving ./my_classifier --allow-updates --save-dir ./my_classifier
+```
+
+Update endpoints then require `Authorization: Bearer change-me`; prediction endpoints stay open.
+`--save-dir` saves the classifier after every update, so a restart keeps what was learned.
+Without an API key anyone who can reach the server can change the model, and the server warns about it.
+
+## Concurrency and scaling
+
+Every classifier operation takes a lock, so one process can safely serve concurrent requests while
+learning. That means requests are processed one at a time inside a process. For more throughput run
+several **processes** (`--workers`, or several containers behind a load balancer); with a read-only
+server this scales cleanly. Each process has its own copy of the model, so updates only reach the
+worker that received them: run a single process if you enable updates.
+
+From Python, the same pieces are available without HTTP: `await classifier.apredict(...)`,
+`apredict_batch` and `aadd_examples` run in a worker thread so they do not block an event loop, and
+`create_app(classifier)` returns an ASGI app you can mount in an existing FastAPI project.
+
+## Docker
+
+```bash
+docker build -f docker/Dockerfile -t adaptive-classifier .
+docker run --rm -p 8000:8000 -v "$PWD/my_classifier:/model:ro" adaptive-classifier
+```
+
+The image runs as a non-root user, has a health check on `/health`, and uses CPU-only PyTorch. The
+Dockerfile has not been built or run in the environment this was written in (no Docker daemon was
+available); the application it runs is covered by the server tests, but treat the image build as
+unverified until you have run it once.
diff --git a/pyproject.toml b/pyproject.toml
index 01127b9..927bd75 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -43,6 +43,13 @@ test = [
     "pytest-randomly>=3.12.0",
     "psutil",
     "hypothesis>=6.80.0",
+    "fastapi>=0.100.0",
+    "httpx>=0.24.0",
+]
+# HTTP server: python -m adaptive_classifier.serving ./model
+serve = [
+    "fastapi>=0.100.0",
+    "uvicorn>=0.23.0",
 ]
 # Only needed to run the scripts/ benchmarks and data generators.
 scripts = [
diff --git a/src/adaptive_classifier/classifier.py b/src/adaptive_classifier/classifier.py
index 78b57f1..f7e1cb7 100644
--- a/src/adaptive_classifier/classifier.py
+++ b/src/adaptive_classifier/classifier.py
@@ -5,9 +5,12 @@
 from transformers import AutoModel, AutoTokenizer
 from typing import List, Dict, Optional, Tuple, Any, Set, Union
 import logging
+import asyncio
 import copy
+import functools
 import math
 import numbers
+import threading
 from pathlib import Path
 from safetensors.torch import save_file, load_file
 import json
@@ -27,6 +30,37 @@
 
 logger = logging.getLogger(__name__)
 
+_LOCK_CREATION = threading.Lock()
+
+# Files that make a directory loadable by AutoTokenizer.
+_TOKENIZER_FILES = (
+    "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
+    "vocab.txt", "vocab.json", "merges.txt", "added_tokens.json",
+    "spiece.model", "sentencepiece.bpe.model",
+)
+
+
+def _has_local_tokenizer(directory: Path) -> bool:
+    return (directory / "tokenizer.json").exists() or (directory / "tokenizer_config.json").exists()
+
+
+def _synchronized(method):
+    """Run a method while holding the instance lock.
+
+    Predictions read several structures (label maps, prototypes, the neural
+    head) that `add_examples`, `forget` and friends rewrite together; a reader
+    that sees them half-updated fails with errors like "selected index k out of
+    range". One reentrant lock per classifier makes every public operation
+    atomic, so a classifier can be shared between threads (for example behind a
+    web server). Throughput scales with processes, not threads.
+    """
+    @functools.wraps(method)
+    def wrapper(self, *args, **kwargs):
+        with self._lock:
+            return method(self, *args, **kwargs)
+    return wrapper
+
+
 class AdaptiveClassifier(ModelHubMixin):
     """A flexible classifier that can adapt to new classes and examples."""
 
@@ -34,6 +68,15 @@ class AdaptiveClassifier(ModelHubMixin):
     # __init__ (the loaders do that) working.
     calibration: Optional[Dict[str, Any]] = None
 
+    @property
+    def _lock(self) -> "threading.RLock":
+        # Created lazily because the loaders build instances without __init__.
+        lock = self.__dict__.get('_instance_lock')
+        if lock is None:
+            with _LOCK_CREATION:
+                lock = self.__dict__.setdefault('_instance_lock', threading.RLock())
+        return lock
+
     def __init__(
         self,
         model_name: str,
@@ -161,6 +204,7 @@ def _validate_k(k: int):
         if not isinstance(k, numbers.Integral) or isinstance(k, bool) or k < 0:
             raise ValueError(f"k must be a non-negative integer; got {k!r}")
 
+    @_synchronized
     def add_examples(self, texts: List[str], labels: List[str]):
         """Add new examples with special handling for new classes."""
         if not texts or not labels:
@@ -426,6 +470,7 @@ def _perform_strategic_training(self):
             
             logger.debug("Performed strategic training step")
     
+    @_synchronized
     def predict(
         self,
         text: str,
@@ -472,6 +517,7 @@ def predict(
             return []
         return predictions
 
+    @_synchronized
     def ood_score(self, text: str) -> float:
         """How far `text` sits outside the known classes.
 
@@ -488,6 +534,7 @@ def ood_score(self, text: str) -> float:
         score, _ = self.memory.ood_score(embedding)
         return score
 
+    @_synchronized
     def is_ood(self, text: str, threshold: Optional[float] = None) -> bool:
         """True when `text` scores above the out-of-distribution threshold.
 
@@ -498,6 +545,20 @@ def is_ood(self, text: str, threshold: Optional[float] = None) -> bool:
         limit = self.config.ood_threshold if threshold is None else threshold
         return self.ood_score(text) > limit
 
+    # --- async API ---------------------------------------------------------------------
+
+    async def apredict(self, text: str, k: int = 5, **kwargs) -> List[Tuple[str, float]]:
+        """`predict` without blocking the event loop (runs in a worker thread)."""
+        return await asyncio.to_thread(self.predict, text, k, **kwargs)
+
+    async def apredict_batch(self, texts: List[str], k: int = 5, **kwargs) -> List[List[Tuple[str, float]]]:
+        """`predict_batch` without blocking the event loop."""
+        return await asyncio.to_thread(self.predict_batch, texts, k, **kwargs)
+
+    async def aadd_examples(self, texts: List[str], labels: List[str]):
+        """`add_examples` without blocking the event loop."""
+        return await asyncio.to_thread(self.add_examples, texts, labels)
+
     # --- calibration ---------------------------------------------------------------
 
     def _calibrate(self, predictions: List[Tuple[str, float]]) -> List[Tuple[str, float]]:
@@ -534,6 +595,7 @@ def _calibration_inputs(self, texts: List[str], labels: List[str]):
             raise ValueError(f"Unknown labels: {unknown}")
         return np.array([self.label_to_id[label] for label in labels])
 
+    @_synchronized
     def calibrate(self, texts: List[str], labels: List[str]) -> Dict[str, float]:
         """Fit confidence calibration on labelled examples the model was *not* trained on.
 
@@ -577,6 +639,7 @@ class never changes. Calibrating on data the model was trained on
             'accuracy': before['accuracy'],
         }
 
+    @_synchronized
     def calibration_report(self, texts: List[str], labels: List[str],
                            calibrated: bool = True, bins: int = 10) -> Dict[str, float]:
         """Accuracy, mean confidence, expected calibration error (ECE) and NLL on labelled data.
@@ -590,6 +653,7 @@ def calibration_report(self, texts: List[str], labels: List[str],
             probs = calibration_math.apply_temperature(probs, self.calibration['temperature'])
         return calibration_math.summarise(probs, targets, bins)
 
+    @_synchronized
     def predict_set(self, text: str, alpha: float = 0.1) -> List[Tuple[str, float]]:
         """Labels that contain the true one with probability at least `1 - alpha`.
 
@@ -616,6 +680,7 @@ def predict_set(self, text: str, alpha: float = 0.1) -> List[Tuple[str, float]]:
 
     _SUGGEST_STRATEGIES = ("margin", "entropy", "least_confidence", "ood")
 
+    @_synchronized
     def suggest_labels(
         self,
         texts: List[str],
@@ -698,6 +763,7 @@ def _ood_scores(self, texts: List[str]) -> List[float]:
             embeddings = self._get_embeddings(texts)
         return [self.memory.ood_score(embedding)[0] for embedding in embeddings]
 
+    @_synchronized
     def drift_report(
         self,
         texts: List[str],
@@ -743,6 +809,7 @@ def drift_report(
             'drifted': p_value < alpha,
         }
 
+    @_synchronized
     def forget(self, labels: Union[str, List[str]]):
         """Remove one or more classes entirely.
 
@@ -760,6 +827,7 @@ def forget(self, labels: Union[str, List[str]]):
             raise ValueError(f"Unknown labels: {unknown}")
         self._drop_classes(set(labels))
 
+    @_synchronized
     def remove_examples(
         self,
         texts: List[str],
@@ -939,6 +1007,7 @@ def _predict_dual(self, text: str, k: int = 5) -> List[Tuple[str, float]]:
         
         return blended_predictions[:k]
     
+    @_synchronized
     def _save_pretrained(
         self,
         save_directory: Union[str, Path],
@@ -1162,6 +1231,24 @@ def _from_pretrained(
                     local_files_only=local_files_only,
                 )
 
+                # Tokenizer files (written by save() since 0.3.0), so loading does not
+                # depend on the base model's repository. Older repos do not have them.
+                for tokenizer_file in _TOKENIZER_FILES:
+                    try:
+                        hf_hub_download(
+                            repo_id=model_id,
+                            filename=tokenizer_file,
+                            revision=revision,
+                            cache_dir=cache_dir,
+                            force_download=force_download,
+                            proxies=proxies,
+                            resume_download=resume_download,
+                            token=token,
+                            local_files_only=local_files_only,
+                        )
+                    except Exception:
+                        pass
+
                 # Try to download ONNX files if they exist
                 try:
                     # Download quantized ONNX model (primary)
@@ -1281,7 +1368,8 @@ def _from_pretrained(
                 file_name=onnx_file,
                 trust_remote_code=trust_remote_code
             )
-            classifier.tokenizer = AutoTokenizer.from_pretrained(config_dict['model_name'], trust_remote_code=trust_remote_code)
+            tokenizer_source = str(model_path) if _has_local_tokenizer(model_path) else config_dict['model_name']
+            classifier.tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, trust_remote_code=trust_remote_code)
 
             # Initialize memory and other components
             classifier.embedding_dim = classifier.model.config.hidden_size
@@ -1321,6 +1409,12 @@ def _from_pretrained(
                 trust_remote_code=trust_remote_code
             )
 
+        # A tokenizer saved with the classifier wins over the one the constructor fetched,
+        # so what runs matches what was saved (and needs no network for the tokenizer).
+        if _has_local_tokenizer(model_path):
+            classifier.tokenizer = AutoTokenizer.from_pretrained(
+                str(model_path), trust_remote_code=trust_remote_code)
+
         # Restore label mappings
         classifier.label_to_id = config_dict['label_to_id']
         classifier.id_to_label = {
@@ -1486,6 +1580,7 @@ def _format_class_distribution(self, stats: Dict[str, Any]) -> str:
             
         return "\n".join(lines)
 
+    @_synchronized
     def export_onnx(
         self,
         save_directory: Union[str, Path],
@@ -1685,6 +1780,7 @@ def to(self, device: str) -> 'AdaptiveClassifier':
             self.adaptive_head = self.adaptive_head.to(device)
         return self
     
+    @_synchronized
     def get_memory_stats(self) -> Dict[str, Any]:
         """Get memory statistics.
         
@@ -1814,6 +1910,7 @@ def _get_embeddings(self, texts: List[str]) -> List[torch.Tensor]:
         # Return embeddings as list
         return [emb.cpu() for emb in embeddings]
 
+    @_synchronized
     def get_example_statistics(self) -> Dict[str, Any]:
         """Get statistics about stored examples and model state."""
         stats = {
@@ -1838,6 +1935,7 @@ def get_example_statistics(self) -> Dict[str, Any]:
         
         return stats
 
+    @_synchronized
     def predict_batch(
         self,
         texts: List[str],
@@ -1920,6 +2018,7 @@ def predict_batch(
         
         return all_predictions
 
+    @_synchronized
     def clear_memory(self, labels: Optional[List[str]] = None):
         """Clear memory for specified labels or all if none specified.
 
@@ -1934,6 +2033,7 @@ def clear_memory(self, labels: Optional[List[str]] = None):
         if known:
             self.forget(known)
 
+    @_synchronized
     def merge_classifiers(self, other: 'AdaptiveClassifier') -> 'AdaptiveClassifier':
         """Merge another classifier into this one."""
         # Verify compatibility
@@ -2143,6 +2243,7 @@ def _strategic_training_step(self, all_embeddings: torch.Tensor, all_labels: tor
         
         logger.debug("Completed strategic training step")
     
+    @_synchronized
     def predict_strategic(self, text: str, k: int = 5) -> List[Tuple[str, float]]:
         """Predict assuming the input might be strategically modified.
         
@@ -2190,6 +2291,7 @@ def classifier_func(x):
             logger.warning(f"Strategic prediction failed: {e}. Falling back to regular prediction.")
             return self._predict_regular(text, k)
     
+    @_synchronized
     def predict_robust(self, text: str, k: int = 5) -> List[Tuple[str, float]]:
         """Predict assuming input has already been strategically modified.
         
diff --git a/src/adaptive_classifier/serving.py b/src/adaptive_classifier/serving.py
new file mode 100644
index 0000000..db2bf1b
--- /dev/null
+++ b/src/adaptive_classifier/serving.py
@@ -0,0 +1,238 @@
+"""Serve an AdaptiveClassifier over HTTP.
+
+    pip install "adaptive-classifier[serve]"
+    python -m adaptive_classifier.serving ./my_classifier --port 8000
+
+or, from Python::
+
+    from adaptive_classifier.serving import create_app
+    app = create_app(classifier)             # any ASGI server can run it
+
+Endpoints: ``GET /health``, ``GET /info``, ``POST /predict``,
+``POST /predict_batch``, ``POST /predict_set``, ``POST /ood``, and, only when
+updates are allowed, ``POST /examples``, ``POST /forget`` and
+``POST /remove_examples``.
+
+The classifier takes a lock around every operation, so one process is safe to
+share across requests; scale throughput with ``--workers`` (separate
+processes, each with its own copy and its own updates) rather than threads.
+Updates therefore only reach the worker that handled them. Run a single
+worker if you enable them, and pass ``--save-dir`` to persist what was learned.
+"""
+
+import argparse
+import asyncio
+import hmac
+import logging
+from typing import Any, Dict, List, Optional
+
+from .classifier import AdaptiveClassifier
+
+logger = logging.getLogger(__name__)
+
+
+def create_app(
+    classifier: AdaptiveClassifier,
+    *,
+    allow_updates: bool = False,
+    api_key: Optional[str] = None,
+    max_batch: int = 256,
+    max_text_chars: int = 20_000,
+    save_dir: Optional[str] = None,
+):
+    """Build a FastAPI app around `classifier`.
+
+    Args:
+        classifier: A trained (or loaded) classifier.
+        allow_updates: Enable ``/examples``, ``/forget`` and ``/remove_examples``.
+        api_key: If set, update endpoints require ``Authorization: Bearer <key>``.
+            Strongly recommended with ``allow_updates``.
+        max_batch: Largest accepted batch (texts per request).
+        max_text_chars: Longest accepted text, in characters.
+        save_dir: If set, the classifier is saved there after every update.
+    """
+    try:
+        from fastapi import Depends, FastAPI, HTTPException, Request
+        from fastapi.responses import JSONResponse
+        from pydantic import BaseModel, Field
+    except ImportError as error:          # pragma: no cover - depends on the environment
+        raise ImportError(
+            'Serving needs FastAPI: pip install "adaptive-classifier[serve]"'
+        ) from error
+
+    if allow_updates and not api_key:
+        logger.warning("Updates are enabled without an api_key: anyone who can reach this "
+                       "server can change the classifier")
+
+    class Prediction(BaseModel):
+        label: str
+        score: float
+
+    class PredictRequest(BaseModel):
+        text: str = Field(min_length=1, max_length=max_text_chars)
+        k: int = Field(default=5, ge=0, le=1000)
+        abstain_below: Optional[float] = Field(default=None, ge=0.0, le=1.0)
+        abstain_ood: bool = False
+
+    class PredictResponse(BaseModel):
+        predictions: List[Prediction]
+        abstained: bool
+
+    class BatchRequest(BaseModel):
+        texts: List[str] = Field(min_length=1, max_length=max_batch)
+        k: int = Field(default=5, ge=0, le=1000)
+
+    class BatchResponse(BaseModel):
+        predictions: List[List[Prediction]]
+
+    class SetRequest(BaseModel):
+        text: str = Field(min_length=1, max_length=max_text_chars)
+        alpha: float = Field(default=0.1, gt=0.0, lt=1.0)
+
+    class OodRequest(BaseModel):
+        text: str = Field(min_length=1, max_length=max_text_chars)
+        threshold: Optional[float] = Field(default=None, gt=0.0)
+
+    class ExamplesRequest(BaseModel):
+        texts: List[str] = Field(min_length=1, max_length=max_batch)
+        labels: List[str] = Field(min_length=1, max_length=max_batch)
+
+    class ForgetRequest(BaseModel):
+        labels: List[str] = Field(min_length=1, max_length=max_batch)
+
+    class RemoveRequest(BaseModel):
+        texts: List[str] = Field(min_length=1, max_length=max_batch)
+        label: Optional[str] = None
+
+    app = FastAPI(title="Adaptive Classifier", version=_version())
+
+    @app.exception_handler(ValueError)
+    async def bad_request(_: Request, error: ValueError):
+        # The classifier raises ValueError for unusable input (unknown label, empty text, ...).
+        return JSONResponse(status_code=400, content={"detail": str(error)})
+
+    def pairs(items) -> List[Dict[str, Any]]:
+        return [{"label": label, "score": float(score)} for label, score in items]
+
+    def require_updates(request: Request):
+        if not allow_updates:
+            raise HTTPException(status_code=403, detail="Updates are disabled on this server")
+        if api_key:
+            header = request.headers.get("authorization", "")
+            supplied = header[7:] if header.lower().startswith("bearer ") else ""
+            if not hmac.compare_digest(supplied.encode(), api_key.encode()):
+                raise HTTPException(status_code=401, detail="Missing or invalid API key",
+                                    headers={"WWW-Authenticate": "Bearer"})
+
+    def persist():
+        if save_dir:
+            classifier.save(save_dir)
+
+    @app.get("/health")
+    async def health():
+        return {"status": "ok", "classes": len(classifier.label_to_id)}
+
+    @app.get("/info")
+    async def info():
+        stats = classifier.get_memory_stats()
+        return {
+            "version": _version(),
+            "model": getattr(classifier, "_model_name", None),
+            "classes": sorted(classifier.label_to_id),
+            "examples_per_class": stats.get("examples_per_class", {}),
+            "calibrated": bool(classifier.calibration),
+            "updates_allowed": allow_updates,
+        }
+
+    @app.post("/predict", response_model=PredictResponse)
+    async def predict(body: PredictRequest):
+        result = await classifier.apredict(
+            body.text, body.k, abstain_below=body.abstain_below, abstain_ood=body.abstain_ood
+        )
+        asked_to_abstain = body.abstain_below is not None or body.abstain_ood
+        abstained = asked_to_abstain and not result and body.k > 0
+        return {"predictions": pairs(result), "abstained": abstained}
+
+    @app.post("/predict_batch", response_model=BatchResponse)
+    async def predict_batch(body: BatchRequest):
+        if any(not t or len(t) > max_text_chars for t in body.texts):
+            raise HTTPException(status_code=422, detail=f"Each text must be 1-{max_text_chars} characters")
+        result = await classifier.apredict_batch(body.texts, body.k)
+        return {"predictions": [pairs(r) for r in result]}
+
+    @app.post("/predict_set")
+    async def predict_set(body: SetRequest):
+        result = await asyncio.to_thread(classifier.predict_set, body.text, body.alpha)
+        return {"labels": pairs(result)}
+
+    @app.post("/ood")
+    async def ood(body: OodRequest):
+        score = await asyncio.to_thread(classifier.ood_score, body.text)
+        limit = classifier.config.ood_threshold if body.threshold is None else body.threshold
+        return {"score": None if score == float("inf") else score, "threshold": limit,
+                "ood": score > limit}
+
+    @app.post("/examples", dependencies=[Depends(require_updates)])
+    async def add_examples(body: ExamplesRequest):
+        if any(not t or len(t) > max_text_chars for t in body.texts):
+            raise HTTPException(status_code=422, detail=f"Each text must be 1-{max_text_chars} characters")
+        await classifier.aadd_examples(body.texts, body.labels)
+        persist()
+        return {"added": len(body.texts), "classes": sorted(classifier.label_to_id)}
+
+    @app.post("/forget", dependencies=[Depends(require_updates)])
+    async def forget(body: ForgetRequest):
+        await asyncio.to_thread(classifier.forget, body.labels)
+        persist()
+        return {"classes": sorted(classifier.label_to_id)}
+
+    @app.post("/remove_examples", dependencies=[Depends(require_updates)])
+    async def remove_examples(body: RemoveRequest):
+        removed = await asyncio.to_thread(classifier.remove_examples, body.texts, body.label)
+        persist()
+        return {"removed": removed, "classes": sorted(classifier.label_to_id)}
+
+    return app
+
+
+def _version() -> str:
+    from . import __version__
+    return __version__
+
+
+def main(argv: Optional[List[str]] = None) -> None:
+    parser = argparse.ArgumentParser(
+        prog="python -m adaptive_classifier.serving",
+        description="Serve a saved AdaptiveClassifier over HTTP.",
+    )
+    parser.add_argument("model", help="Directory from classifier.save(), or a Hub id")
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--allow-updates", action="store_true",
+                        help="Enable /examples, /forget and /remove_examples")
+    parser.add_argument("--api-key", default=None,
+                        help="Bearer token required for updates (or set ADAPTIVE_CLASSIFIER_API_KEY)")
+    parser.add_argument("--save-dir", default=None, help="Save the classifier here after every update")
+    parser.add_argument("--max-batch", type=int, default=256)
+    parser.add_argument("--no-onnx", action="store_true", help="Use PyTorch even on CPU")
+    args = parser.parse_args(argv)
+
+    import os
+    try:
+        import uvicorn
+    except ImportError as error:          # pragma: no cover
+        raise SystemExit('uvicorn is required: pip install "adaptive-classifier[serve]"') from error
+
+    classifier = AdaptiveClassifier.load(args.model, use_onnx=False if args.no_onnx else "auto")
+    app = create_app(
+        classifier,
+        allow_updates=args.allow_updates,
+        api_key=args.api_key or os.environ.get("ADAPTIVE_CLASSIFIER_API_KEY"),
+        max_batch=args.max_batch,
+        save_dir=args.save_dir,
+    )
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py
new file mode 100644
index 0000000..efb9483
--- /dev/null
+++ b/tests/test_concurrency.py
@@ -0,0 +1,179 @@
+"""A classifier shared between threads (a web server) must stay consistent.
+
+Predictions read the label maps, the prototypes and the neural head, which
+`add_examples`, `forget` and friends rewrite together. Without a lock a reader
+could see them half-updated and crash ("selected index k out of range"); the
+stress test below reproduces that reliably when the lock is removed.
+"""
+
+import asyncio
+import threading
+import time
+
+import pytest
+
+from adaptive_classifier import AdaptiveClassifier
+
+from .synthetic import Dataset
+
+
+@pytest.fixture
+def data(monkeypatch):
+    return Dataset().patch(monkeypatch)
+
+
+@pytest.fixture
+def clf(new_classifier, data):
+    X, y = data.classes(12, n_classes=3, noise=0.2, seed=1)
+    c = new_classifier(head_steps=30)
+    c.add_examples(X, y)
+    return c
+
+
+def run_threads(targets):
+    errors = []
+
+    def guarded(fn):
+        def run():
+            try:
+                fn()
+            except Exception as error:               # noqa: BLE001 - collecting for the assertion
+                errors.append(f"{type(error).__name__}: {str(error)[:100]}")
+        return run
+
+    threads = [threading.Thread(target=guarded(fn)) for fn in targets]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout=120)
+    assert not any(t.is_alive() for t in threads), "a thread is stuck (deadlock?)"
+    return errors
+
+
+def test_predictions_survive_concurrent_class_changes(clf, data):
+    stop = threading.Event()
+
+    def reader():
+        while not stop.is_set():
+            preds = clf.predict("tr0_0", k=10)
+            assert all(score == score for _, score in preds)
+            clf.predict_batch(["tr1_0", "tr2_1"], k=10)
+
+    def writer():
+        try:
+            for round_ in range(8):
+                extra, labels = data.classes(3, n_classes=3, noise=0.2, seed=60 + round_, prefix=f"w{round_}")
+                clf.add_examples(extra, [f"N{round_}" if label == "C2" else label for label in labels])
+                if round_ % 2:
+                    clf.forget(f"N{round_}")
+        finally:
+            stop.set()
+
+    errors = run_threads([reader, reader, reader, reader, writer])
+    assert errors == []
+    # The final state is internally consistent.
+    n = len(clf.label_to_id)
+    assert sorted(clf.label_to_id.values()) == list(range(n))
+    assert clf.adaptive_head.model[-1].out_features == n
+    assert set(clf.memory.prototypes) == set(clf.label_to_id)
+
+
+def test_concurrent_adds_lose_no_examples(new_classifier, data):
+    clf = new_classifier(head_steps=10)
+    batches = []
+    for t in range(6):
+        X, y = data.classes(4, n_classes=2, noise=0.2, seed=200 + t, prefix=f"t{t}")
+        batches.append((X, y))
+    errors = run_threads([lambda b=b: clf.add_examples(*b) for b in batches])
+    assert errors == []
+    assert clf.training_history == {"C0": 24, "C1": 24}
+    assert sum(len(v) for v in clf.memory.examples.values()) == 48
+
+
+def test_the_lock_is_reentrant_so_methods_can_call_each_other(clf, data):
+    Xc, yc = data.classes(20, n_classes=3, noise=0.3, seed=7, prefix="cal")
+    clf.calibrate(Xc, yc)                          # calibrate -> predict, all under the lock
+    assert clf.predict_set("tr0_0")                # predict_set -> predict
+
+
+def test_each_classifier_has_its_own_lock(clf, new_classifier):
+    other = new_classifier()
+    assert clf._lock is clf._lock and clf._lock is not other._lock
+
+
+def test_a_loaded_classifier_has_a_working_lock(clf, tmp_path):
+    clf.save(str(tmp_path), include_onnx=False)
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    with loaded._lock:                             # built without __init__
+        pass
+    assert loaded.predict("tr0_0")
+
+
+def test_a_failing_operation_releases_the_lock(clf):
+    with pytest.raises(ValueError):
+        clf.add_examples(["x"], [None])
+    result = []
+    t = threading.Thread(target=lambda: result.append(clf.predict("tr0_0")))
+    t.start()
+    t.join(timeout=10)
+    assert result, "the lock was left held after an exception"
+
+
+# --- async API ----------------------------------------------------------------------------------
+
+def test_async_methods_match_the_sync_ones(clf, data):
+    async def go():
+        return (
+            await clf.apredict("tr0_0", k=3),
+            await clf.apredict_batch(["tr1_0", "tr2_0"], k=3),
+        )
+    single, batch = asyncio.run(go())
+    assert single == clf.predict("tr0_0", k=3)
+    assert batch == clf.predict_batch(["tr1_0", "tr2_0"], k=3)
+
+
+def test_async_methods_pass_keyword_arguments_through(clf):
+    top = clf.predict("tr0_0")[0][1]
+    assert asyncio.run(clf.apredict("tr0_0", abstain_below=top + 0.01)) == []
+    assert asyncio.run(clf.apredict("tr0_0", abstain_below=top - 0.01)) != []
+
+
+def test_aadd_examples_adds_and_validates(clf, data):
+    extra, labels = data.classes(3, n_classes=1, noise=0.2, seed=300, prefix="ax")
+    asyncio.run(clf.aadd_examples(extra, ["C0"] * 3))
+    assert clf.training_history["C0"] == 15
+    with pytest.raises(ValueError, match="strings"):
+        asyncio.run(clf.aadd_examples(["x"], [3]))
+
+
+def test_async_calls_do_not_block_the_event_loop(clf, monkeypatch):
+    original = clf.predict
+
+    def slow(*args, **kwargs):
+        time.sleep(0.4)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(clf, "predict", slow)
+
+    async def go():
+        ticks = 0
+
+        async def ticker():
+            nonlocal ticks
+            while True:
+                await asyncio.sleep(0.02)
+                ticks += 1
+
+        task = asyncio.create_task(ticker())
+        await clf.apredict("tr0_0")
+        task.cancel()
+        return ticks
+
+    assert asyncio.run(go()) >= 8          # the loop kept running during the 0.4s call
+
+
+def test_many_concurrent_async_predictions_all_complete(clf):
+    async def go():
+        return await asyncio.gather(*[clf.apredict(f"tr{i % 3}_{i % 5}") for i in range(40)])
+    results = asyncio.run(go())
+    assert len(results) == 40 and all(r for r in results)
diff --git a/tests/test_portable_export.py b/tests/test_portable_export.py
index 5b884ed..1e245bd 100644
--- a/tests/test_portable_export.py
+++ b/tests/test_portable_export.py
@@ -218,3 +218,43 @@ def test_portable_reference_matches_predict_for_every_scoring_branch(base_model,
         assert actual.keys() == expected.keys()
         for label in expected:
             assert float(actual[label]) == pytest.approx(expected[label], abs=1e-4)
+
+
+# --- a saved directory carries its own tokenizer ----------------------------------------------
+
+@pytest.mark.skipif(not _onnx_available(), reason="optimum[onnxruntime] not installed")
+def test_onnx_load_uses_the_saved_tokenizer_not_the_base_model_repo(base_model, tmp_path):
+    """The stored model name pointed at a Hub repo; loading used to fetch the tokenizer from it,
+    so a saved classifier could not start without network access to that repo."""
+    clf = _trained(base_model, pooling="mean")
+    clf.save(str(tmp_path), include_onnx=True, quantize_onnx=False)
+    config_path = tmp_path / "config.json"
+    data = json.loads(config_path.read_text(encoding="utf-8"))
+    data["model_name"] = "nobody/this-repo-does-not-exist"
+    config_path.write_text(json.dumps(data), encoding="utf-8")
+
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=True, prefer_quantized=False, device="cpu")
+    assert loaded.use_onnx
+    assert loaded.tokenizer.name_or_path == str(tmp_path)
+    for query in QUERIES:
+        a, b = dict(clf.predict(query, k=3)), dict(loaded.predict(query, k=3))
+        assert all(a[l] == pytest.approx(b[l], abs=1e-4) for l in a)
+
+
+def test_loading_without_saved_tokenizer_files_still_works(base_model, tmp_path):
+    """Directories saved before 0.3.0 have no tokenizer files and fall back to the model name."""
+    clf = _trained(base_model, pooling="mean")
+    clf.save(str(tmp_path), include_onnx=False)
+    for name in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.txt"):
+        (tmp_path / name).unlink(missing_ok=True)
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    assert loaded.predict(QUERIES[0])
+
+
+def test_pytorch_load_also_uses_the_saved_tokenizer(base_model, tmp_path):
+    clf = _trained(base_model, pooling="mean")
+    clf.save(str(tmp_path), include_onnx=False)
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    assert loaded.tokenizer.name_or_path == str(tmp_path)
+    for query in QUERIES:
+        assert [l for l, _ in loaded.predict(query)] == [l for l, _ in clf.predict(query)]
diff --git a/tests/test_serving.py b/tests/test_serving.py
new file mode 100644
index 0000000..3de5209
--- /dev/null
+++ b/tests/test_serving.py
@@ -0,0 +1,285 @@
+"""The HTTP server (adaptive_classifier.serving), exercised through FastAPI's TestClient."""
+
+import threading
+
+import pytest
+
+pytest.importorskip("fastapi")
+pytest.importorskip("httpx")
+
+from fastapi.testclient import TestClient
+
+from adaptive_classifier import AdaptiveClassifier
+from adaptive_classifier import serving
+from adaptive_classifier.serving import create_app
+
+from .synthetic import Dataset
+
+
+@pytest.fixture
+def data(monkeypatch):
+    return Dataset().patch(monkeypatch)
+
+
+@pytest.fixture
+def clf(new_classifier, data):
+    X, y = data.classes(12, n_classes=3, noise=0.2, seed=1)
+    c = new_classifier(head_steps=30)
+    c.add_examples(X, y)
+    return c
+
+
+@pytest.fixture
+def client(clf):
+    return TestClient(create_app(clf))
+
+
+def updatable(clf, **kwargs):
+    return TestClient(create_app(clf, allow_updates=True, **kwargs))
+
+
+# --- reading ---------------------------------------------------------------------------------
+
+def test_health_and_info(client, clf):
+    assert client.get("/health").json() == {"status": "ok", "classes": 3}
+    info = client.get("/info").json()
+    assert info["classes"] == ["C0", "C1", "C2"]
+    assert info["examples_per_class"] == {"C0": 12, "C1": 12, "C2": 12}
+    assert info["calibrated"] is False and info["updates_allowed"] is False
+
+
+def test_predict_returns_what_the_classifier_returns(client, clf):
+    body = client.post("/predict", json={"text": "tr1_3", "k": 3}).json()
+    expected = clf.predict("tr1_3", k=3)
+    assert body["abstained"] is False
+    assert [(p["label"], pytest.approx(p["score"])) for p in body["predictions"]] == expected
+
+
+def test_predict_respects_k(client):
+    assert len(client.post("/predict", json={"text": "tr0_0", "k": 1}).json()["predictions"]) == 1
+    assert client.post("/predict", json={"text": "tr0_0", "k": 0}).json()["predictions"] == []
+
+
+def test_abstention_is_reported(client, clf):
+    top = clf.predict("tr0_0")[0][1]
+    r = client.post("/predict", json={"text": "tr0_0", "abstain_below": min(top + 0.01, 1.0)}).json()
+    assert r == {"predictions": [], "abstained": True}
+    ok = client.post("/predict", json={"text": "tr0_0", "abstain_below": top - 0.01}).json()
+    assert ok["abstained"] is False and ok["predictions"]
+
+
+def test_abstain_ood_flags_foreign_text(client, data):
+    foreign = data.add("far", 9, 0.1, 5)
+    assert client.post("/predict", json={"text": foreign, "abstain_ood": True}).json()["abstained"] is True
+    assert client.post("/predict", json={"text": "tr0_0", "abstain_ood": True}).json()["abstained"] is False
+
+
+def test_predict_batch_matches_the_classifier(client, clf):
+    texts = ["tr0_0", "tr1_1", "tr2_2"]
+    body = client.post("/predict_batch", json={"texts": texts, "k": 2}).json()["predictions"]
+    expected = clf.predict_batch(texts, k=2)
+    for got, want in zip(body, expected):
+        assert [(p["label"], pytest.approx(p["score"])) for p in got] == want
+
+
+def test_ood_endpoint(client, data):
+    near = client.post("/ood", json={"text": "tr0_0"}).json()
+    far = client.post("/ood", json={"text": data.add("far2", 9, 0.1, 6)}).json()
+    assert near["ood"] is False and far["ood"] is True and far["score"] > near["score"]
+    strict = client.post("/ood", json={"text": "tr0_0", "threshold": 0.001}).json()
+    assert strict["ood"] is True and strict["threshold"] == 0.001
+
+
+def test_predict_set_needs_calibration_then_works(client, clf, data):
+    r = client.post("/predict_set", json={"text": "tr0_0"})
+    assert r.status_code == 400 and "calibrate" in r.json()["detail"]
+
+    Xc, yc = data.classes(30, n_classes=3, noise=0.4, seed=8, prefix="cal")
+    clf.calibrate(Xc, yc)
+    r = client.post("/predict_set", json={"text": "tr0_0", "alpha": 0.2})
+    assert r.status_code == 200 and r.json()["labels"][0]["label"] == "C0"
+    assert client.get("/info").json()["calibrated"] is True
+
+
+# --- input validation ------------------------------------------------------------------------------
+
+@pytest.mark.parametrize("path, body", [
+    ("/predict", {"text": ""}),
+    ("/predict", {}),
+    ("/predict", {"text": "x", "k": -1}),
+    ("/predict", {"text": "x", "k": "many"}),
+    ("/predict", {"text": "x", "abstain_below": 2}),
+    ("/predict", {"text": "x" * 20_001}),
+    ("/predict_batch", {"texts": []}),
+    ("/predict_batch", {"texts": ["a", ""]}),
+    ("/predict_set", {"text": "x", "alpha": 0}),
+    ("/predict_set", {"text": "x", "alpha": 1}),
+    ("/ood", {"text": ""}),
+])
+def test_invalid_requests_are_rejected(client, path, body):
+    assert client.post(path, json=body).status_code == 422
+
+
+def test_batch_size_is_capped(clf):
+    small = TestClient(create_app(clf, max_batch=3))
+    assert small.post("/predict_batch", json={"texts": ["tr0_0"] * 3}).status_code == 200
+    assert small.post("/predict_batch", json={"texts": ["tr0_0"] * 4}).status_code == 422
+
+
+def test_text_length_is_capped(clf):
+    small = TestClient(create_app(clf, max_text_chars=5))
+    assert small.post("/predict", json={"text": "tr0_0"}).status_code == 200      # exactly 5 characters
+    assert small.post("/predict", json={"text": "tr0_10"}).status_code == 422     # 6 characters
+    assert small.post("/predict_batch", json={"texts": ["tr0_10"]}).status_code == 422
+
+
+def test_malformed_json_is_a_client_error(client):
+    r = client.post("/predict", content="{not json", headers={"content-type": "application/json"})
+    assert r.status_code == 422
+
+
+# --- updates ------------------------------------------------------------------------------------------
+
+@pytest.mark.parametrize("path, body", [
+    ("/examples", {"texts": ["a"], "labels": ["C0"]}),
+    ("/forget", {"labels": ["C0"]}),
+    ("/remove_examples", {"texts": ["tr0_0"]}),
+])
+def test_updates_are_off_by_default(client, clf, path, body):
+    r = client.post(path, json=body)
+    assert r.status_code == 403
+    assert sorted(clf.label_to_id) == ["C0", "C1", "C2"]
+
+
+def test_adding_examples_and_a_new_class(clf, data):
+    http = updatable(clf)
+    novel = [data.add(f"nv{i}", 5, 0.1, 40 + i) for i in range(6)]
+    r = http.post("/examples", json={"texts": novel, "labels": ["C5"] * 6})
+    assert r.status_code == 200
+    assert r.json() == {"added": 6, "classes": ["C0", "C1", "C2", "C5"]}
+    assert http.post("/predict", json={"text": novel[0]}).json()["predictions"][0]["label"] == "C5"
+
+
+def test_forgetting_and_removing(clf):
+    http = updatable(clf)
+    r = http.post("/remove_examples", json={"texts": ["tr0_0", "tr0_1"], "label": "C0"})
+    assert r.json()["removed"] == 2 and clf.training_history["C0"] == 10
+    r = http.post("/forget", json={"labels": ["C2"]})
+    assert r.json()["classes"] == ["C0", "C1"]
+    assert "C2" not in {p["label"] for p in http.post("/predict", json={"text": "tr2_0"}).json()["predictions"]}
+
+
+def test_classifier_errors_become_400_with_the_message(clf):
+    http = updatable(clf)
+    r = http.post("/forget", json={"labels": ["nope"]})
+    assert r.status_code == 400 and "Unknown labels" in r.json()["detail"]
+    r = http.post("/examples", json={"texts": ["a", "b"], "labels": ["C0"]})
+    assert r.status_code == 400 and "Mismatched" in r.json()["detail"]
+    assert sorted(clf.label_to_id) == ["C0", "C1", "C2"]
+
+
+def test_non_string_labels_are_rejected(clf):
+    r = updatable(clf).post("/examples", json={"texts": ["a"], "labels": [3]})
+    assert r.status_code == 422
+
+
+@pytest.mark.parametrize("headers, status", [
+    ({}, 401),
+    ({"Authorization": "Bearer wrong"}, 401),
+    ({"Authorization": "wrong-scheme secret"}, 401),
+    ({"Authorization": "Bearer "}, 401),
+    ({"Authorization": "Bearer secret"}, 200),
+    ({"authorization": "bearer secret"}, 200),
+])
+def test_api_key_protects_updates(clf, headers, status):
+    http = updatable(clf, api_key="secret")
+    r = http.post("/remove_examples", json={"texts": ["nothing"]}, headers=headers)
+    assert r.status_code == status
+    if status == 401:
+        assert r.headers["www-authenticate"] == "Bearer"
+
+
+def test_api_key_is_not_needed_to_read(clf):
+    http = updatable(clf, api_key="secret")
+    assert http.post("/predict", json={"text": "tr0_0"}).status_code == 200
+    assert http.get("/health").status_code == 200
+
+
+def test_updates_without_a_key_log_a_warning(clf, caplog):
+    with caplog.at_level("WARNING"):
+        create_app(clf, allow_updates=True)
+    assert "without an api_key" in caplog.text
+
+
+def test_updates_are_saved_when_a_save_dir_is_given(clf, data, tmp_path):
+    http = updatable(clf, save_dir=str(tmp_path))
+    novel = [data.add(f"sv{i}", 6, 0.1, 80 + i) for i in range(5)]
+    assert http.post("/examples", json={"texts": novel, "labels": ["C6"] * 5}).status_code == 200
+    reloaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    assert "C6" in reloaded.label_to_id
+
+
+def test_requests_in_parallel_do_not_corrupt_the_classifier(clf, data):
+    http = updatable(clf)
+    statuses = []
+
+    def read():
+        for _ in range(15):
+            statuses.append(http.post("/predict", json={"text": "tr0_0", "k": 10}).status_code)
+
+    def write():
+        for i in range(4):
+            texts = [data.add(f"par{i}_{j}", 4 + i, 0.1, 500 + 10 * i + j) for j in range(3)]
+            statuses.append(http.post("/examples", json={"texts": texts, "labels": [f"P{i}"] * 3}).status_code)
+
+    threads = [threading.Thread(target=read) for _ in range(3)] + [threading.Thread(target=write)]
+    [t.start() for t in threads]
+    [t.join(timeout=120) for t in threads]
+    assert statuses and set(statuses) == {200}
+    assert sorted(clf.label_to_id.values()) == list(range(len(clf.label_to_id)))
+
+
+# --- command line -------------------------------------------------------------------------------------
+
+@pytest.fixture
+def saved(clf, tmp_path):
+    clf.save(str(tmp_path), include_onnx=False)
+    return str(tmp_path)
+
+
+@pytest.fixture
+def launched(monkeypatch):
+    calls = {}
+    import uvicorn
+    monkeypatch.setattr(uvicorn, "run", lambda app, host, port: calls.update(app=app, host=host, port=port))
+    return calls
+
+
+def test_cli_loads_a_saved_classifier_and_serves_it(saved, launched):
+    serving.main([saved, "--port", "9001", "--host", "0.0.0.0", "--no-onnx"])
+    assert (launched["host"], launched["port"]) == ("0.0.0.0", 9001)
+    client = TestClient(launched["app"])
+    assert client.get("/health").json()["classes"] == 3
+    assert client.post("/forget", json={"labels": ["C0"]}).status_code == 403
+
+
+def test_cli_flags_enable_updates_and_a_key(saved, launched):
+    serving.main([saved, "--allow-updates", "--api-key", "k1", "--no-onnx"])
+    client = TestClient(launched["app"])
+    assert client.post("/remove_examples", json={"texts": ["z"]}).status_code == 401
+    assert client.post("/remove_examples", json={"texts": ["z"]},
+                       headers={"Authorization": "Bearer k1"}).status_code == 200
+
+
+def test_cli_reads_the_key_from_the_environment(saved, launched, monkeypatch):
+    monkeypatch.setenv("ADAPTIVE_CLASSIFIER_API_KEY", "from-env")
+    serving.main([saved, "--allow-updates", "--no-onnx"])
+    client = TestClient(launched["app"])
+    assert client.post("/forget", json={"labels": ["C0"]}).status_code == 401
+    assert client.post("/forget", json={"labels": ["C0"]},
+                       headers={"Authorization": "Bearer from-env"}).status_code == 200
+
+
+def test_cli_fails_clearly_for_a_missing_model(launched, tmp_path):
+    with pytest.raises(FileNotFoundError, match="config.json"):
+        serving.main([str(tmp_path), "--no-onnx"])

From 76c7c1bae324f5c42bbc6331fc0e106c37d20824 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 11:18:01 +0000
Subject: [PATCH 12/16] Add uvicorn to the test extra

The serving CLI tests monkeypatch uvicorn.run, so they need uvicorn
installed. CI failed with four "No module named 'uvicorn'" errors (406
other tests passed); it worked locally only because uvicorn had been
installed by hand. An audit of every third-party import in tests/ confirms
nothing else is missing from a clean install of the test extra.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index 927bd75..dd955c3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -44,6 +44,7 @@ test = [
     "psutil",
     "hypothesis>=6.80.0",
     "fastapi>=0.100.0",
+    "uvicorn>=0.23.0",
     "httpx>=0.24.0",
 ]
 # HTTP server: python -m adaptive_classifier.serving ./model

From 2132d33bc103b5a648b30d2aaac85af5476525c3 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 12:16:00 +0000
Subject: [PATCH 13/16] Add a few-shot benchmark against SetFit and logistic
 regression

scripts/benchmark_fewshot.py compares adaptive-classifier (default and
prototypes-only), logistic regression on the same frozen embeddings, and
SetFit, all on one shared encoder with the same stratified few-shot sample and
test subset. Per run it records accuracy, macro-F1, calibration error, train
time, test-set prediction time, median single-query latency, and the cost of
adding one more labelled example (SetFit has no incremental update, so its
training time is reported and marked as a full retrain). Samples are nested
across shot counts so results are comparable.

Commands: run, report (Markdown, best per column in bold), check (streams one
example from each split to verify dataset ids and columns before a long run) and
plan (validates workflow inputs so they are never interpolated into a shell).

.github/workflows/benchmark.yml runs it by hand: a plan job validates inputs
and checks the datasets, one parallel job per dataset runs the benchmark, and a
summary job merges the results into the run summary. docs/benchmark.md explains
the methods and, deliberately, the limits of what the numbers show.

Found while verifying against the real Hub: PolyAI/banking77 is a loading
script that current `datasets` releases refuse to run, so the benchmark uses
the parquet copy mteb/banking77. SetFit also wrote ~260 MB of checkpoints into
the working directory; they now go to a temporary directory (and checkpoints/
is git-ignored).

tests/test_benchmark_harness.py (69 tests, all offline): sampling, metrics,
aggregation, the report, input validation including shell-injection attempts,
dataset verification against a stubbed `datasets`, every method run end to end
on a tiny local model, and the command line. SetFit runs end to end when it is
installed. The workflow's own shell steps were executed from the YAML with a
simulated Actions environment.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 .github/workflows/benchmark.yml | 167 +++++++++
 .gitignore                      |   3 +
 CHANGELOG.md                    |   2 +
 README.md                       |   4 +
 docs/benchmark.md               |  96 +++++
 pyproject.toml                  |   5 +
 scripts/benchmark_fewshot.py    | 636 ++++++++++++++++++++++++++++++++
 tests/conftest.py               |  22 ++
 tests/test_benchmark_harness.py | 482 ++++++++++++++++++++++++
 tests/test_failure_modes.py     |   6 +-
 10 files changed, 1421 insertions(+), 2 deletions(-)
 create mode 100644 .github/workflows/benchmark.yml
 create mode 100644 docs/benchmark.md
 create mode 100644 scripts/benchmark_fewshot.py
 create mode 100644 tests/test_benchmark_harness.py

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
new file mode 100644
index 0000000..1eab524
--- /dev/null
+++ b/.github/workflows/benchmark.yml
@@ -0,0 +1,167 @@
+name: Few-shot benchmark
+
+# Compares adaptive-classifier with SetFit and logistic regression on public datasets.
+# Run it by hand: Actions -> Few-shot benchmark -> Run workflow. It is not part of the test
+# suite because it downloads datasets, trains SetFit on a CPU, and takes a long time.
+# Results appear in the run summary and as downloadable artifacts. See docs/benchmark.md.
+
+on:
+  workflow_dispatch:
+    inputs:
+      datasets:
+        description: "Datasets, comma separated (ag_news, emotion, banking77)"
+        default: "ag_news,emotion"
+      shots:
+        description: "Training examples per class, space separated"
+        default: "4 16"
+      seeds:
+        description: "Random seeds (different few-shot samples), 1-10"
+        default: "3"
+      methods:
+        description: "Methods to run (adaptive adaptive-proto logreg setfit)"
+        default: "adaptive adaptive-proto logreg setfit"
+      model:
+        description: "Sentence encoder shared by every method"
+        default: "sentence-transformers/all-MiniLM-L6-v2"
+      max_test:
+        description: "Test examples per dataset (a fixed random subset)"
+        default: "1000"
+
+permissions:
+  contents: read
+
+jobs:
+  plan:
+    # Validates the inputs and that each dataset loads, in seconds, before the long jobs start.
+    # Inputs reach the shell only as environment variables, never interpolated into a command.
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    outputs:
+      datasets: ${{ steps.plan.outputs.datasets }}
+      shots: ${{ steps.plan.outputs.shots }}
+      seeds: ${{ steps.plan.outputs.seeds }}
+      methods: ${{ steps.plan.outputs.methods }}
+      model: ${{ steps.plan.outputs.model }}
+      max_test: ${{ steps.plan.outputs.max_test }}
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+
+    - name: Validate inputs
+      id: plan
+      env:
+        IN_DATASETS: ${{ inputs.datasets }}
+        IN_SHOTS: ${{ inputs.shots }}
+        IN_SEEDS: ${{ inputs.seeds }}
+        IN_METHODS: ${{ inputs.methods }}
+        IN_MODEL: ${{ inputs.model }}
+        IN_MAX_TEST: ${{ inputs.max_test }}
+      run: |
+        pip install numpy
+        python scripts/benchmark_fewshot.py plan \
+          --datasets "$IN_DATASETS" --shots "$IN_SHOTS" --seeds "$IN_SEEDS" \
+          --methods "$IN_METHODS" --model "$IN_MODEL" --max-test "$IN_MAX_TEST" > plan.json
+        cat plan.json
+        python - <<'PY' >> "$GITHUB_OUTPUT"
+        import json
+        plan = json.load(open("plan.json"))
+        print("datasets=" + json.dumps(plan["datasets"]))
+        for key in ("shots", "seeds", "methods", "model", "max_test"):
+            print(f"{key}={plan[key]}")
+        PY
+
+    - name: Check that every dataset loads
+      env:
+        HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        DATASETS: ${{ steps.plan.outputs.datasets }}
+      run: |
+        pip install datasets
+        python scripts/benchmark_fewshot.py check $(python -c "import json,os; print(' '.join(json.loads(os.environ['DATASETS'])))")
+
+  benchmark:
+    needs: plan
+    runs-on: ubuntu-latest
+    timeout-minutes: 350
+    strategy:
+      fail-fast: false
+      matrix:
+        dataset: ${{ fromJSON(needs.plan.outputs.datasets) }}
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+        cache: pip
+
+    - name: Cache Hugging Face downloads
+      uses: actions/cache@v4
+      with:
+        path: ~/.cache/huggingface
+        key: hf-benchmark-${{ needs.plan.outputs.model }}-${{ matrix.dataset }}
+
+    - name: Install
+      run: |
+        python -m pip install --upgrade pip
+        pip install torch --index-url https://download.pytorch.org/whl/cpu
+        pip install -e ".[benchmark]"
+        pip list 2>/dev/null | grep -iE "^(torch|transformers|sentence-transformers|setfit|datasets|scikit-learn|adaptive-classifier) "
+
+    - name: Run benchmark
+      env:
+        HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        DATASET: ${{ matrix.dataset }}
+        SHOTS: ${{ needs.plan.outputs.shots }}
+        SEEDS: ${{ needs.plan.outputs.seeds }}
+        METHODS: ${{ needs.plan.outputs.methods }}
+        MODEL: ${{ needs.plan.outputs.model }}
+        MAX_TEST: ${{ needs.plan.outputs.max_test }}
+      run: |
+        # SHOTS and METHODS are space-separated lists validated by the plan job.
+        python scripts/benchmark_fewshot.py run \
+          --dataset "$DATASET" --shots $SHOTS --seeds "$SEEDS" --methods $METHODS \
+          --model "$MODEL" --max-test "$MAX_TEST" --out "results/$DATASET.json"
+
+    - name: Upload results
+      if: always()
+      uses: actions/upload-artifact@v4
+      with:
+        name: results-${{ matrix.dataset }}
+        path: results/*.json
+        if-no-files-found: warn
+
+  summary:
+    needs: [plan, benchmark]
+    if: always() && needs.plan.result == 'success'
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+
+    - uses: actions/download-artifact@v4
+      with:
+        pattern: results-*
+        path: results
+        merge-multiple: true
+
+    - name: Build the report
+      run: |
+        pip install numpy
+        if ls results/*.json >/dev/null 2>&1; then
+          python scripts/benchmark_fewshot.py report results/*.json --out benchmark-report.md > /dev/null
+          cat benchmark-report.md >> "$GITHUB_STEP_SUMMARY"
+        else
+          echo "No result files were produced; see the benchmark jobs." | tee benchmark-report.md >> "$GITHUB_STEP_SUMMARY"
+        fi
+
+    - uses: actions/upload-artifact@v4
+      with:
+        name: benchmark-report
+        path: benchmark-report.md
diff --git a/.gitignore b/.gitignore
index 3d255ad..d27df0a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -173,3 +173,6 @@ cython_debug/
 **/.DS_store
 demo_classifier/
 scripts/benchmark_results/
+
+# SetFit training checkpoints
+checkpoints/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 387ddf1..d3d940e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,8 @@
 - Async methods `apredict`, `apredict_batch` and `aadd_examples` that run in a worker thread.
 - A saved directory is self-contained for loading: the tokenizer saved next to the model is used instead of fetching the base model's repository, so an ONNX-saved classifier starts without Hub access.
 
+- **Few-shot benchmark** against SetFit and logistic regression: `scripts/benchmark_fewshot.py` and a manually triggered `Few-shot benchmark` workflow (accuracy, macro-F1, calibration error, training time, latency, and the cost of adding one example). Install with `pip install -e ".[benchmark]"`. See `docs/benchmark.md`.
+
 ### Changed (default behaviour: predictions change)
 - **The neural head is now actually trained.** It used to get about ten optimiser steps and stop at the first loss plateau, so it could not learn even a 2-D XOR from 120 examples. It now trains for at least `head_steps` (default 300) steps with a cosine-decayed `head_learning_rate` (default 0.003). Small memories take roughly 1-2 s per `add_examples` call; lower `head_steps` if you add examples very frequently.
 - **Prototype scores are sharper.** They are now `softmax(-distance² / prototype_temperature)` (default 0.25) instead of `softmax(exp(-distance))`, which barely separated the nearest class from the rest (about 0.45 vs 0.21), so a confident head could always outvote it. `prototype_temperature=None` restores the old scoring.
diff --git a/README.md b/README.md
index b1a5e5f..01b78fb 100644
--- a/README.md
+++ b/README.md
@@ -382,6 +382,10 @@ curl -s localhost:8000/predict -H 'content-type: application/json' -d '{"text":
 
 Predictions, batches, conformal sets and out-of-distribution checks over HTTP, plus optional API-key-protected endpoints for adding examples and classes while it runs. One classifier is safe to share between threads, and `await classifier.apredict(...)` keeps an event loop responsive. See [docs/serving.md](docs/serving.md) and `docker/Dockerfile`.
 
+### How does it compare?
+
+A benchmark against [SetFit](https://github.com/huggingface/setfit) and logistic regression on the same encoder (accuracy, macro-F1, calibration, training time, latency, and the cost of adding one example) is in `scripts/benchmark_fewshot.py`, with a manually triggered GitHub workflow. See [docs/benchmark.md](docs/benchmark.md) for how it works and what it does and does not show.
+
 ### Fixing Mistakes and Knowing When to Abstain
 
 ```python
diff --git a/docs/benchmark.md b/docs/benchmark.md
new file mode 100644
index 0000000..8c657d2
--- /dev/null
+++ b/docs/benchmark.md
@@ -0,0 +1,96 @@
+# Few-shot benchmark: adaptive-classifier vs SetFit vs logistic regression
+
+`scripts/benchmark_fewshot.py` measures how well each method does when it only has a handful
+of labelled examples per class, which is the situation this library is built for. It runs as a
+manually triggered GitHub Actions workflow, or on your own machine.
+
+## What is compared
+
+All methods use the **same sentence encoder** (default `sentence-transformers/all-MiniLM-L6-v2`),
+the **same stratified training sample**, and the **same test examples**.
+
+| method | what it is |
+|---|---|
+| `adaptive` | this library with default settings |
+| `adaptive-proto` | this library with the neural head switched off (prototypes only) |
+| `logreg` | logistic regression on the frozen, L2-normalised embeddings |
+| `setfit` | [SetFit](https://github.com/huggingface/setfit): contrastive fine-tuning of the encoder, then a logistic-regression head |
+
+`adaptive-proto` and `logreg` are included so that a difference between `adaptive` and `setfit`
+can be traced: is it the encoder fine-tuning, or the classifier on top?
+
+## Metrics
+
+| metric | meaning |
+|---|---|
+| accuracy, macro-F1 | on a fixed random subset of the test split (`--max-test`, default 1000) |
+| ECE | expected calibration error of the reported confidences: do 80% confidences come true 80% of the time? Lower is better. Reported **without** any calibration step |
+| train (s) | time to train from raw text, including embedding the training examples (model loading is not counted) |
+| predict test set (s), latency p50 (ms) | time for the whole test set from raw text, and the median time for a single query |
+| add one example (s) | the cost of learning one more labelled example. `adaptive` updates in place. `logreg` refits on cached features. **SetFit has no incremental update, so its full training time is reported** and marked "(full retrain)" |
+
+Each (dataset, shots, seed) draws a stratified sample; a smaller sample for a seed is always a
+subset of a larger one, so results across shot counts are comparable. Results are mean ± standard
+deviation over seeds. The best value in each column is bold.
+
+## Running it on GitHub
+
+Actions → **Few-shot benchmark** → **Run workflow**. The inputs are validated before any long job
+starts, and each dataset is checked to load, so a typo fails in seconds. Each dataset runs as its
+own job (in parallel); the combined Markdown report appears in the run summary and as the
+`benchmark-report` artifact, and the raw numbers as `results-<dataset>` artifacts.
+
+Defaults (`ag_news,emotion`, 4 and 16 shots, 3 seeds) take a while on GitHub's CPU runners, mostly
+because SetFit fine-tunes the encoder. `banking77` (77 classes) is much slower: use few shots and
+one seed first.
+
+## Running it locally
+
+```bash
+pip install -e ".[benchmark]"
+python scripts/benchmark_fewshot.py check ag_news emotion banking77        # do the datasets load?
+python scripts/benchmark_fewshot.py run --dataset ag_news --shots 4 16 --seeds 3 --out results/ag_news.json
+python scripts/benchmark_fewshot.py report results/*.json --out results/summary.md
+```
+
+`--methods` selects a subset (for example `adaptive logreg` to skip SetFit), `--model` changes the
+encoder, and `--dataset synthetic` runs a tiny generated dataset with no network access, only to
+check that the harness works.
+
+## Datasets
+
+| name | source | classes |
+|---|---|---|
+| `ag_news` | `fancyzhx/ag_news` | 4 |
+| `emotion` | `dair-ai/emotion` | 6 |
+| `banking77` | `mteb/banking77` | 77 |
+
+`PolyAI/banking77` is not used: it is a Hub loading script, which current `datasets` releases
+refuse to run. The `mteb` copy has the same data as parquet. To add a dataset, add an entry to
+`DATASETS` in the script (Hub id, text column, label column, test split name) and run `check`.
+
+## Limits of what this shows
+
+Read the numbers with these in mind:
+
+- **One encoder.** Everything shares MiniLM-L6. A different encoder can change the ranking, and
+  SetFit is usually run with larger ones.
+- **Untuned settings.** SetFit uses its common defaults (1 epoch, 20 pair iterations, batch 16);
+  logistic regression uses `C=10`; adaptive-classifier uses its defaults. None was tuned per dataset.
+- **Few seeds, subsampled test set.** Differences of a point or two are within noise. With 3 seeds
+  the standard deviations are themselves rough.
+- **Three classification datasets of short English text.** Nothing here speaks to long documents,
+  other languages, or very many classes beyond `banking77`.
+- **CPU timings are indicative.** They depend on the machine and on whether a GPU is present
+  (the report records the device).
+- **No LLM baseline.** Few-shot prompting of a language model is a real alternative and is not
+  included.
+- **The first full runs have not been analysed yet.** Check the report before quoting any number
+  from it, and look at the run's `error` fields for failed cells (shown as `failed`).
+
+## Checking the harness itself
+
+The parts that can be tested offline (sampling, metrics, aggregation, the report, the command line,
+input validation, and a smoke run of every method on a tiny local model) are covered by
+`tests/test_benchmark_harness.py`, so a surprising number is more likely to be about the methods than
+about the harness.
diff --git a/pyproject.toml b/pyproject.toml
index dd955c3..15ad957 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,6 +52,11 @@ serve = [
     "fastapi>=0.100.0",
     "uvicorn>=0.23.0",
 ]
+# Few-shot benchmark against SetFit: scripts/benchmark_fewshot.py (see docs/benchmark.md).
+benchmark = [
+    "setfit>=1.1.0",
+    "datasets>=2.14.0",
+]
 # Only needed to run the scripts/ benchmarks and data generators.
 scripts = [
     "tqdm>=4.65.0",
diff --git a/scripts/benchmark_fewshot.py b/scripts/benchmark_fewshot.py
new file mode 100644
index 0000000..13cb208
--- /dev/null
+++ b/scripts/benchmark_fewshot.py
@@ -0,0 +1,636 @@
+"""Few-shot text classification benchmark: adaptive-classifier vs SetFit vs logistic regression.
+
+Everything uses the same sentence encoder (``--model``). For each dataset, number of
+examples per class ("shots") and random seed, every method trains on the same
+stratified sample from the training split and is scored on the same test subset.
+
+Methods
+    adaptive        this library with default settings
+    adaptive-proto  this library, prototypes only (neural head switched off)
+    logreg          logistic regression on the frozen, L2-normalised embeddings
+    setfit          SetFit (contrastive fine-tuning of the encoder, then a logistic head)
+
+Measured per run: accuracy, macro-F1, expected calibration error (ECE) of the reported
+confidences, training time, whole-test-set prediction time, median single-query
+latency, and the cost of adding one more labelled example (adaptive-classifier updates in
+place; logistic regression refits on cached features; SetFit has to retrain from scratch,
+so its training time is reported as the update cost).
+
+    pip install -e ".[benchmark]"
+    python scripts/benchmark_fewshot.py check ag_news emotion banking77   # quick schema check
+    python scripts/benchmark_fewshot.py run --dataset ag_news --shots 4 16 --seeds 3 \\
+        --out results/ag_news.json
+    python scripts/benchmark_fewshot.py report results/*.json --out results/summary.md
+
+``--dataset synthetic`` generates a tiny dataset with no network access; it exists to
+smoke-test the harness, and its numbers mean nothing.
+"""
+
+import argparse
+import json
+import platform
+import statistics
+import subprocess
+import sys
+import tempfile
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence
+
+import numpy as np
+
+METHODS = ["adaptive", "adaptive-proto", "logreg", "setfit"]
+
+# Hub id, text column, label column. Add entries here to benchmark another dataset.
+DATASETS = {
+    "ag_news": {"hub": "fancyzhx/ag_news", "text": "text", "label": "label", "test": "test"},
+    "emotion": {"hub": "dair-ai/emotion", "text": "text", "label": "label", "test": "test"},
+    # PolyAI/banking77 is a loading script, which current `datasets` versions refuse to run;
+    # this parquet copy has the same data. label_text holds the intent names.
+    "banking77": {"hub": "mteb/banking77", "text": "text", "label": "label_text", "test": "test"},
+}
+
+
+# --------------------------------------------------------------------------------------
+# Pure helpers (no model, no network)
+# --------------------------------------------------------------------------------------
+
+def sample_few_shot(labels: Sequence[str], shots: int, seed: int) -> List[int]:
+    """Indices of a stratified sample with up to `shots` examples per class.
+
+    Each class is shuffled once per seed and its first `shots` examples are taken, so a
+    smaller sample is always a subset of a larger one for the same seed. That keeps
+    results across shot counts comparable, with less noise from which examples were drawn.
+    """
+    if shots < 1:
+        raise ValueError(f"shots must be at least 1; got {shots}")
+    rng = np.random.default_rng(seed)
+    by_class: Dict[str, List[int]] = {}
+    for index, label in enumerate(labels):
+        by_class.setdefault(label, []).append(index)
+    chosen: List[int] = []
+    for label in sorted(by_class):
+        members = by_class[label]
+        order = rng.permutation(len(members))
+        chosen.extend(members[i] for i in order[:shots])
+    return sorted(chosen)
+
+
+def subsample(texts: Sequence[str], labels: Sequence[str], limit: Optional[int], seed: int = 0):
+    """A fixed random subset of a labelled split (all of it when `limit` is None or large)."""
+    if limit is None or limit >= len(texts):
+        return list(texts), list(labels)
+    order = np.random.default_rng(seed).permutation(len(texts))[:limit]
+    return [texts[i] for i in order], [labels[i] for i in order]
+
+
+def expected_calibration_error(probs: np.ndarray, targets: np.ndarray, bins: int = 10) -> float:
+    confidence = probs.max(axis=1)
+    correct = (probs.argmax(axis=1) == targets).astype(float)
+    edges = np.linspace(0.0, 1.0, bins + 1)
+    ece = 0.0
+    for low, high in zip(edges[:-1], edges[1:]):
+        in_bin = (confidence > low) & (confidence <= high)
+        if in_bin.any():
+            ece += in_bin.mean() * abs(correct[in_bin].mean() - confidence[in_bin].mean())
+    return float(ece)
+
+
+def score_predictions(probs: np.ndarray, targets: np.ndarray) -> Dict[str, float]:
+    """Accuracy, macro-F1 and ECE for a (n_samples, n_classes) probability matrix."""
+    from sklearn.metrics import f1_score
+
+    predicted = probs.argmax(axis=1)
+    return {
+        "accuracy": float((predicted == targets).mean()),
+        "macro_f1": float(f1_score(targets, predicted, average="macro", zero_division=0)),
+        "ece": expected_calibration_error(probs, targets),
+    }
+
+
+def _mean_std(values: List[float]):
+    if not values:
+        return None, None
+    return statistics.fmean(values), (statistics.stdev(values) if len(values) > 1 else None)
+
+
+def aggregate(results: List[Dict[str, Any]]) -> Dict[Any, Dict[str, Any]]:
+    """Mean and standard deviation over seeds, keyed by (dataset, method, shots)."""
+    groups: Dict[Any, List[Dict[str, Any]]] = {}
+    for row in results:
+        groups.setdefault((row["dataset"], row["method"], row["shots"]), []).append(row)
+    summary = {}
+    for key, rows in groups.items():
+        ok = [r for r in rows if not r.get("error")]
+        entry: Dict[str, Any] = {"runs": len(rows), "failed": len(rows) - len(ok)}
+        for metric in ("accuracy", "macro_f1", "ece", "train_s", "predict_s",
+                       "latency_ms_p50", "update_s"):
+            mean, std = _mean_std([r[metric] for r in ok if r.get(metric) is not None])
+            entry[metric] = (mean, std)
+        entry["update_is_retrain"] = any(r.get("update_is_retrain") for r in ok)
+        summary[key] = entry
+    return summary
+
+
+def _cell(pair, digits=3, scale=1.0, bold=False):
+    mean, std = pair
+    if mean is None:
+        return "failed"
+    text = f"{mean * scale:.{digits}f}"
+    if std is not None:
+        text += f" ± {std * scale:.{digits}f}"
+    return f"**{text}**" if bold else text
+
+
+def _table(summary, dataset, methods, shots, metric, higher_is_better=True, digits=3, scale=1.0):
+    header = "| method | " + " | ".join(f"{s}-shot" for s in shots) + " |"
+    lines = [header, "|---|" + "---|" * len(shots)]
+    best = {}
+    for s in shots:
+        values = [(m, summary[(dataset, m, s)][metric][0]) for m in methods
+                  if (dataset, m, s) in summary and summary[(dataset, m, s)][metric][0] is not None]
+        if values:
+            pick = max if higher_is_better else min
+            best[s] = pick(v for _, v in values)
+    for m in methods:
+        cells = []
+        for s in shots:
+            entry = summary.get((dataset, m, s))
+            if entry is None:
+                cells.append("–")
+                continue
+            pair = entry[metric]
+            is_best = pair[0] is not None and s in best and abs(pair[0] - best[s]) < 1e-12
+            cells.append(_cell(pair, digits, scale, bold=is_best))
+        lines.append(f"| {m} | " + " | ".join(cells) + " |")
+    return "\n".join(lines)
+
+
+def render_markdown(results: List[Dict[str, Any]], meta: Optional[Dict[str, Any]] = None) -> str:
+    """A readable report: one section per dataset; the best value in each column is bold."""
+    meta = meta or {}
+    summary = aggregate(results)
+    datasets = sorted({r["dataset"] for r in results})
+    out = ["# Few-shot benchmark", ""]
+    if meta:
+        facts = [f"**{k}**: {v}" for k, v in meta.items() if v not in (None, "")]
+        out += ["  \n".join(facts), ""]
+    errors = [r for r in results if r.get("error")]
+    if errors:
+        out += [f"> **{len(errors)} run(s) failed.** Cells marked `failed` have no result; "
+                "see the `error` field in the JSON.", ""]
+
+    for dataset in datasets:
+        methods = [m for m in METHODS if any(k[0] == dataset and k[1] == m for k in summary)]
+        methods += sorted({k[1] for k in summary if k[0] == dataset} - set(methods))
+        shots = sorted({k[2] for k in summary if k[0] == dataset})
+        sample = next(r for r in results if r["dataset"] == dataset)
+        out += [f"## {dataset}", "",
+                f"{sample.get('n_classes', '?')} classes, "
+                f"{sample.get('n_test', '?')} test examples; mean ± std over seeds. "
+                "Best per column in bold.", ""]
+        out += ["### Accuracy", "", _table(summary, dataset, methods, shots, "accuracy"), ""]
+        out += ["### Macro-F1", "", _table(summary, dataset, methods, shots, "macro_f1"), ""]
+        out += ["### Calibration error (ECE, lower is better)", "",
+                _table(summary, dataset, methods, shots, "ece", higher_is_better=False), ""]
+        out += ["### Cost", ""]
+        cost_shots = [shots[-1]]
+        out += [f"At {cost_shots[0]} examples per class.", "",
+                "| method | train (s) | predict test set (s) | latency p50 (ms) | add one example (s) |",
+                "|---|---|---|---|---|"]
+        for m in methods:
+            entry = summary.get((dataset, m, cost_shots[0]))
+            if entry is None:
+                continue
+            update = _cell(entry["update_s"], 2)
+            if entry["update_is_retrain"] and update != "failed":
+                update += " (full retrain)"
+            out.append(f"| {m} | {_cell(entry['train_s'], 1)} | {_cell(entry['predict_s'], 1)} | "
+                       f"{_cell(entry['latency_ms_p50'], 1)} | {update} |")
+        out.append("")
+    return "\n".join(out).rstrip() + "\n"
+
+
+# --------------------------------------------------------------------------------------
+# Data
+# --------------------------------------------------------------------------------------
+
+def synthetic_dataset(n_classes: int = 4, per_class_train: int = 40, n_test: int = 80, seed: int = 0):
+    """Texts made of class-specific words plus shared noise words. For smoke tests only."""
+    rng = np.random.default_rng(seed)
+    class_words = [[f"w{c}x{i}" for i in range(6)] for c in range(n_classes)]
+    shared = [f"s{i}" for i in range(10)]
+
+    def make(c):
+        words = list(rng.choice(class_words[c], 4)) + list(rng.choice(shared, 3))
+        rng.shuffle(words)
+        return " ".join(words)
+
+    train = [(make(c), f"class_{c}") for c in range(n_classes) for _ in range(per_class_train)]
+    test = [(make(c % n_classes), f"class_{c % n_classes}") for c in range(n_test)]
+    return {
+        "train_texts": [t for t, _ in train], "train_labels": [l for _, l in train],
+        "test_texts": [t for t, _ in test], "test_labels": [l for _, l in test],
+    }
+
+
+def load_dataset_splits(name: str) -> Dict[str, List[str]]:
+    if name == "synthetic":
+        return synthetic_dataset()
+    if name not in DATASETS:
+        raise SystemExit(f"Unknown dataset {name!r}; choose from {sorted(DATASETS) + ['synthetic']}")
+    from datasets import load_dataset
+
+    spec = DATASETS[name]
+    data = load_dataset(spec["hub"])
+    names = getattr(data["train"].features[spec["label"]], "names", None)
+
+    def split(part):
+        rows = data[part]
+        labels = rows[spec["label"]]
+        return list(rows[spec["text"]]), [names[i] if names else str(i) for i in labels]
+
+    train_texts, train_labels = split("train")
+    test_texts, test_labels = split(spec["test"])
+    return {"train_texts": train_texts, "train_labels": train_labels,
+            "test_texts": test_texts, "test_labels": test_labels}
+
+
+def verify_dataset(name: str) -> str:
+    """Check that a dataset can be loaded and has the expected columns, without downloading it.
+
+    Streams the first example of the train and test splits. Returns a one-line description,
+    or raises with a message naming what is wrong (unknown Hub id, missing column or split).
+    """
+    if name == "synthetic":
+        return "synthetic: generated locally"
+    if name not in DATASETS:
+        raise SystemExit(f"Unknown dataset {name!r}; choose from {sorted(DATASETS) + ['synthetic']}")
+    from datasets import load_dataset
+
+    spec = DATASETS[name]
+    problems = []
+    classes = None
+    for part in ("train", spec["test"]):
+        try:
+            stream = load_dataset(spec["hub"], split=part, streaming=True)
+            first = next(iter(stream))
+        except Exception as error:
+            problems.append(f"{spec['hub']} split {part!r}: {type(error).__name__}: {str(error)[:200]}")
+            continue
+        for column in (spec["text"], spec["label"]):
+            if column not in first:
+                problems.append(f"{spec['hub']} split {part!r} has no column {column!r} "
+                                f"(columns: {sorted(first)})")
+        if part == "train":
+            feature = (stream.features or {}).get(spec["label"])
+            classes = len(feature.names) if getattr(feature, "names", None) else None
+    if problems:
+        raise RuntimeError(f"{name}: " + "; ".join(problems))
+    return f"{name}: {spec['hub']} ok" + (f", {classes} classes" if classes else "")
+
+
+# --------------------------------------------------------------------------------------
+# Methods
+# --------------------------------------------------------------------------------------
+
+class Timer:
+    def __enter__(self):
+        self.start = time.perf_counter()
+        return self
+
+    def __exit__(self, *exc):
+        self.seconds = time.perf_counter() - self.start
+
+
+def _latency_p50_ms(predict_one, texts: Sequence[str], n: int = 30) -> float:
+    predict_one(texts[0])                                        # warm-up
+    times = []
+    for text in texts[:n]:
+        with Timer() as t:
+            predict_one(text)
+        times.append(t.seconds * 1000)
+    return float(statistics.median(times))
+
+
+class AdaptiveMethod:
+    """adaptive-classifier. `prototypes_only` switches the neural head off."""
+
+    def __init__(self, model: str, prototypes_only: bool = False, device: Optional[str] = None):
+        self.model, self.device = model, device
+        self.config = ({"prototype_weight": 1.0, "neural_weight": 0.0,
+                        "new_class_example_threshold": 0, "head_steps": 1}
+                       if prototypes_only else {})
+
+    def run(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        from adaptive_classifier import AdaptiveClassifier
+
+        classifier = AdaptiveClassifier(self.model, config=self.config, use_onnx=False,
+                                        device=self.device)                # load time not counted
+        with Timer() as train:
+            classifier.add_examples(data["train_texts"], data["train_labels"])
+        classes = sorted(set(data["train_labels"]))
+        column = {label: i for i, label in enumerate(classes)}
+
+        with Timer() as predict:
+            batches = classifier.predict_batch(data["test_texts"], k=len(classes))
+        probs = np.zeros((len(data["test_texts"]), len(classes)))
+        for row, scores in enumerate(batches):
+            for label, score in scores:
+                probs[row, column[label]] = score
+
+        latency = _latency_p50_ms(lambda t: classifier.predict(t, k=1), data["test_texts"])
+        with Timer() as update:
+            classifier.add_examples([data["extra_text"]], [data["extra_label"]])
+        return {"probs": probs, "classes": classes, "train_s": train.seconds,
+                "predict_s": predict.seconds, "latency_ms_p50": latency,
+                "update_s": update.seconds, "update_is_retrain": False}
+
+
+class LogRegMethod:
+    """Logistic regression on frozen embeddings from the same encoder."""
+
+    def __init__(self, model: str, device: Optional[str] = None):
+        self.model, self.device = model, device
+
+    def run(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        from sklearn.linear_model import LogisticRegression
+
+        from adaptive_classifier import AdaptiveClassifier
+
+        encoder = AdaptiveClassifier(self.model, use_onnx=False, device=self.device)
+
+        def embed(texts):
+            out = []
+            for i in range(0, len(texts), 64):
+                out += [e.numpy() for e in encoder._get_embeddings(list(texts[i:i + 64]))]
+            return np.stack(out)
+
+        with Timer() as train:
+            features = embed(data["train_texts"])
+            model = LogisticRegression(max_iter=1000, C=10.0).fit(features, data["train_labels"])
+        classes = list(model.classes_)
+
+        with Timer() as predict:
+            probs = model.predict_proba(embed(data["test_texts"]))
+        latency = _latency_p50_ms(
+            lambda t: model.predict_proba(embed([t])), data["test_texts"])
+
+        extra = embed([data["extra_text"]])
+        with Timer() as update:                       # refit on cached features plus the new example
+            LogisticRegression(max_iter=1000, C=10.0).fit(
+                np.vstack([features, extra]), data["train_labels"] + [data["extra_label"]])
+        return {"probs": probs, "classes": classes, "train_s": train.seconds,
+                "predict_s": predict.seconds, "latency_ms_p50": latency,
+                "update_s": update.seconds, "update_is_retrain": True}
+
+
+class SetFitMethod:
+    """SetFit: contrastive fine-tuning of the encoder followed by a logistic-regression head."""
+
+    def __init__(self, model: str, device: Optional[str] = None):
+        self.model = model
+
+    def run(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        from datasets import Dataset
+        from setfit import SetFitModel, Trainer, TrainingArguments
+
+        classes = sorted(set(data["train_labels"]))
+        index = {label: i for i, label in enumerate(classes)}
+        model = SetFitModel.from_pretrained(self.model)                      # load time not counted
+        # SetFit writes checkpoints; keep them out of the working directory.
+        scratch = tempfile.mkdtemp(prefix="setfit_checkpoints_")
+        args = TrainingArguments(output_dir=scratch, batch_size=16, num_epochs=1, num_iterations=20,
+                                 body_learning_rate=2e-5, seed=data["seed"],
+                                 show_progress_bar=False, report_to="none")
+        trainer = Trainer(
+            model=model, args=args,
+            train_dataset=Dataset.from_dict({
+                "text": data["train_texts"],
+                "label": [index[l] for l in data["train_labels"]]}))
+        with Timer() as train:
+            trainer.train()
+
+        with Timer() as predict:
+            probs = np.asarray(model.predict_proba(data["test_texts"]))
+        latency = _latency_p50_ms(lambda t: model.predict_proba([t]), data["test_texts"])
+        return {"probs": probs, "classes": classes, "train_s": train.seconds,
+                "predict_s": predict.seconds, "latency_ms_p50": latency,
+                "update_s": train.seconds, "update_is_retrain": True}
+
+
+def build_method(name: str, model: str, device: Optional[str]):
+    if name == "adaptive":
+        return AdaptiveMethod(model, device=device)
+    if name == "adaptive-proto":
+        return AdaptiveMethod(model, prototypes_only=True, device=device)
+    if name == "logreg":
+        return LogRegMethod(model, device=device)
+    if name == "setfit":
+        return SetFitMethod(model, device=device)
+    raise SystemExit(f"Unknown method {name!r}; choose from {METHODS}")
+
+
+def plan(datasets: str, shots: str, seeds: str, methods: str, model: str, max_test: str) -> Dict[str, Any]:
+    """Validate and normalise benchmark settings given as plain strings (workflow inputs).
+
+    Raises ValueError naming the first bad value, so a typo fails in seconds rather than
+    after the long jobs have started. Nothing here is ever interpolated into a shell command.
+    """
+    import re
+
+    def words(text):
+        return [w for w in re.split(r"[,\s]+", text.strip()) if w]
+
+    chosen = words(datasets)
+    unknown = [d for d in chosen if d not in DATASETS and d != "synthetic"]
+    if not chosen or unknown:
+        raise ValueError(f"datasets must be a non-empty list from {sorted(DATASETS) + ['synthetic']}; "
+                         f"got {datasets!r}")
+    try:
+        shot_list = [int(w) for w in words(shots)]
+    except ValueError:
+        raise ValueError(f"shots must be whole numbers; got {shots!r}") from None
+    if not shot_list or min(shot_list) < 1 or max(shot_list) > 512:
+        raise ValueError(f"shots must be between 1 and 512; got {shots!r}")
+    try:
+        seed_count = int(seeds)
+    except ValueError:
+        raise ValueError(f"seeds must be a whole number; got {seeds!r}") from None
+    if not 1 <= seed_count <= 10:
+        raise ValueError(f"seeds must be between 1 and 10; got {seeds!r}")
+    method_list = words(methods) or list(METHODS)
+    bad = [m for m in method_list if m not in METHODS]
+    if bad:
+        raise ValueError(f"unknown methods {bad}; choose from {METHODS}")
+    if not re.fullmatch(r"[A-Za-z0-9_.\-]+(/[A-Za-z0-9_.\-]+)?", model.strip()):
+        raise ValueError(f"model must look like 'name' or 'owner/name'; got {model!r}")
+    try:
+        test_limit = int(max_test)
+    except ValueError:
+        raise ValueError(f"max_test must be a whole number; got {max_test!r}") from None
+    if not 10 <= test_limit <= 100_000:
+        raise ValueError(f"max_test must be between 10 and 100000; got {max_test!r}")
+    return {
+        "datasets": list(dict.fromkeys(chosen)),
+        "shots": " ".join(str(k) for k in sorted(set(shot_list))),
+        "seeds": seed_count,
+        "methods": " ".join(dict.fromkeys(method_list)),
+        "model": model.strip(),
+        "max_test": test_limit,
+    }
+
+
+# --------------------------------------------------------------------------------------
+# Running
+# --------------------------------------------------------------------------------------
+
+def run_benchmark(dataset: str, shots: Sequence[int], seeds: int, methods: Sequence[str],
+                  model: str, max_test: Optional[int], device: Optional[str] = None,
+                  log=print) -> List[Dict[str, Any]]:
+    splits = load_dataset_splits(dataset)
+    test_texts, test_labels = subsample(splits["test_texts"], splits["test_labels"], max_test)
+    train_texts, train_labels = splits["train_texts"], splits["train_labels"]
+    classes = sorted(set(train_labels))
+    class_index = {label: i for i, label in enumerate(classes)}
+    targets = np.array([class_index[l] for l in test_labels])
+    results: List[Dict[str, Any]] = []
+
+    for seed in range(seeds):
+        # One extra labelled example (from a class we know) to time the cost of an update.
+        extra_rng = np.random.default_rng(10_000 + seed)
+        extra_i = int(extra_rng.integers(len(train_texts)))
+        for k in shots:
+            picked = sample_few_shot(train_labels, k, seed)
+            data = {
+                "train_texts": [train_texts[i] for i in picked],
+                "train_labels": [train_labels[i] for i in picked],
+                "test_texts": test_texts, "seed": seed,
+                "extra_text": train_texts[extra_i], "extra_label": train_labels[extra_i],
+            }
+            for name in methods:
+                row = {"dataset": dataset, "method": name, "shots": k, "seed": seed,
+                       "n_classes": len(classes), "n_train": len(picked), "n_test": len(test_texts)}
+                try:
+                    out = build_method(name, model, device).run(data)
+                    order = [out["classes"].index(c) for c in classes]       # align columns
+                    probs = np.asarray(out["probs"])[:, order]
+                    row.update(score_predictions(probs, targets))
+                    row.update({key: out[key] for key in
+                                ("train_s", "predict_s", "latency_ms_p50", "update_s",
+                                 "update_is_retrain")})
+                    log(f"{dataset} {name:15s} shots={k:<3d} seed={seed} "
+                        f"acc={row['accuracy']:.3f} f1={row['macro_f1']:.3f} "
+                        f"ece={row['ece']:.3f} train={row['train_s']:.1f}s")
+                except Exception as error:                                   # keep going; report at the end
+                    row["error"] = f"{type(error).__name__}: {str(error)[:300]}"
+                    log(f"{dataset} {name:15s} shots={k:<3d} seed={seed} FAILED {row['error']}")
+                results.append(row)
+    return results
+
+
+def environment_meta(model: str, max_test: Optional[int], seeds: int) -> Dict[str, Any]:
+    def version(module):
+        try:
+            return __import__(module).__version__
+        except Exception:
+            return None
+
+    try:
+        commit = subprocess.run(["git", "rev-parse", "--short", "HEAD"], capture_output=True,
+                                text=True, check=True).stdout.strip()
+    except Exception:
+        commit = None
+    try:
+        import torch
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    except Exception:
+        device = None
+    import adaptive_classifier
+    return {
+        "date": datetime.now(timezone.utc).strftime("%Y-%m-%d"),
+        "model": model, "seeds": seeds, "max test examples": max_test,
+        "device": device, "commit": commit,
+        "adaptive-classifier": adaptive_classifier.__version__,
+        "setfit": version("setfit"), "torch": version("torch"),
+        "python": platform.python_version(),
+    }
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    run = sub.add_parser("run", help="run the benchmark on one dataset")
+    run.add_argument("--dataset", required=True, help=f"{sorted(DATASETS)} or synthetic")
+    run.add_argument("--shots", type=int, nargs="+", default=[4, 16])
+    run.add_argument("--seeds", type=int, default=3)
+    run.add_argument("--methods", nargs="+", default=METHODS, choices=METHODS)
+    run.add_argument("--model", default="sentence-transformers/all-MiniLM-L6-v2")
+    run.add_argument("--max-test", type=int, default=1000)
+    run.add_argument("--device", default=None)
+    run.add_argument("--out", required=True, help="JSON file to write")
+
+    planner = sub.add_parser("plan", help="validate workflow inputs and print them as JSON")
+    planner.add_argument("--datasets", default="ag_news,emotion")
+    planner.add_argument("--shots", default="4 16")
+    planner.add_argument("--seeds", default="3")
+    planner.add_argument("--methods", default=" ".join(METHODS))
+    planner.add_argument("--model", default="sentence-transformers/all-MiniLM-L6-v2")
+    planner.add_argument("--max-test", default="1000")
+
+    check = sub.add_parser("check", help="verify that datasets load and have the expected columns")
+    check.add_argument("datasets", nargs="+", help=f"{sorted(DATASETS)} or synthetic")
+
+    report = sub.add_parser("report", help="turn result files into a Markdown report")
+    report.add_argument("files", nargs="+")
+    report.add_argument("--out", default=None, help="write here as well as printing")
+
+    args = parser.parse_args(argv)
+
+    if args.command == "plan":
+        try:
+            print(json.dumps(plan(args.datasets, args.shots, args.seeds, args.methods,
+                                  args.model, args.max_test)))
+        except ValueError as error:
+            print(f"Invalid benchmark settings: {error}", file=sys.stderr)
+            return 2
+        return 0
+
+    if args.command == "check":
+        failures = 0
+        for name in args.datasets:
+            try:
+                print(verify_dataset(name))
+            except RuntimeError as error:
+                failures += 1
+                print(f"FAIL {error}", file=sys.stderr)
+        return 1 if failures else 0
+
+    if args.command == "run":
+        results = run_benchmark(args.dataset, args.shots, args.seeds, args.methods, args.model,
+                                args.max_test, args.device)
+        meta = environment_meta(args.model, args.max_test, args.seeds)
+        path = Path(args.out)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps({"meta": meta, "results": results}, indent=2), encoding="utf-8")
+        failed = [r for r in results if r.get("error")]
+        if failed:
+            print(f"{len(failed)} of {len(results)} runs failed", file=sys.stderr)
+        return 1 if failed else 0
+
+    results, meta = [], {}
+    for name in args.files:
+        payload = json.loads(Path(name).read_text(encoding="utf-8"))
+        results += payload["results"]
+        meta = meta or payload.get("meta", {})
+    text = render_markdown(results, meta)
+    print(text)
+    if args.out:
+        Path(args.out).write_text(text, encoding="utf-8")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tests/conftest.py b/tests/conftest.py
index 4559ac6..e0d923b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -36,3 +36,25 @@ def new_classifier(base_model):
     def make(**config):
         return AdaptiveClassifier(base_model, config=config, use_onnx=False, device="cpu")
     return make
+
+
+@pytest.fixture(scope="session")
+def synthetic_text_model(tmp_path_factory):
+    """A tiny BERT whose vocabulary holds the words of the benchmark's synthetic dataset.
+
+    Randomly initialised, so it is not a good encoder, but each token has its own
+    embedding, which is enough for texts of class-specific words to be separable.
+    """
+    from transformers import BertConfig, BertModel, BertTokenizerFast
+
+    path = tmp_path_factory.mktemp("synthetic_bert")
+    words = [f"w{c}x{i}" for c in range(8) for i in range(6)] + [f"s{i}" for i in range(10)]
+    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + words
+    (path / "vocab.txt").write_text("\n".join(vocab), encoding="utf-8")
+    torch.manual_seed(0)
+    BertModel(BertConfig(
+        vocab_size=len(vocab), hidden_size=32, num_hidden_layers=2,
+        num_attention_heads=2, intermediate_size=64,
+    )).save_pretrained(path)
+    BertTokenizerFast(str(path / "vocab.txt"), do_lower_case=True).save_pretrained(path)
+    return str(path)
diff --git a/tests/test_benchmark_harness.py b/tests/test_benchmark_harness.py
new file mode 100644
index 0000000..593a274
--- /dev/null
+++ b/tests/test_benchmark_harness.py
@@ -0,0 +1,482 @@
+"""The few-shot benchmark harness (scripts/benchmark_fewshot.py).
+
+The benchmark itself needs network access and takes a long time, so it runs as a manual
+workflow. These tests check everything that can be checked offline, so that when it does
+run, a surprising number is the library's and not the harness's: sampling, metrics,
+aggregation, the report, the command line, and a smoke run of each method on a tiny model.
+"""
+
+import importlib.util
+import json
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from adaptive_classifier import calibration
+
+SCRIPT = Path(__file__).resolve().parent.parent / "scripts" / "benchmark_fewshot.py"
+
+
+@pytest.fixture(scope="module")
+def bench():
+    spec = importlib.util.spec_from_file_location("benchmark_fewshot", SCRIPT)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def labels_of(n_per_class, classes=("a", "b", "c")):
+    return [c for c in classes for _ in range(n_per_class)]
+
+
+# --- sampling ------------------------------------------------------------------------------
+
+def test_sample_takes_exactly_k_per_class(bench):
+    labels = labels_of(20)
+    picked = bench.sample_few_shot(labels, 5, seed=0)
+    assert len(picked) == 15
+    for c in "abc":
+        assert sum(labels[i] == c for i in picked) == 5
+
+
+def test_sample_is_sorted_unique_and_in_range(bench):
+    labels = labels_of(10)
+    picked = bench.sample_few_shot(labels, 4, seed=3)
+    assert picked == sorted(set(picked)) and all(0 <= i < len(labels) for i in picked)
+
+
+def test_sample_is_deterministic_and_seed_dependent(bench):
+    labels = labels_of(50)
+    assert bench.sample_few_shot(labels, 8, 1) == bench.sample_few_shot(labels, 8, 1)
+    assert bench.sample_few_shot(labels, 8, 1) != bench.sample_few_shot(labels, 8, 2)
+
+
+def test_smaller_samples_are_subsets_of_larger_ones(bench):
+    labels = labels_of(50)
+    four, sixteen = set(bench.sample_few_shot(labels, 4, 7)), set(bench.sample_few_shot(labels, 16, 7))
+    assert four < sixteen
+
+
+def test_a_class_with_fewer_examples_than_shots_contributes_all_of_them(bench):
+    labels = ["a"] * 10 + ["b"] * 2
+    picked = bench.sample_few_shot(labels, 5, 0)
+    assert sum(labels[i] == "b" for i in picked) == 2 and sum(labels[i] == "a" for i in picked) == 5
+
+
+def test_sample_does_not_depend_on_row_order_within_labels(bench):
+    labels = labels_of(12)
+    assert len(bench.sample_few_shot(list(reversed(labels)), 3, 0)) == 9
+
+
+@pytest.mark.parametrize("shots", [0, -1])
+def test_sample_rejects_nonpositive_shots(bench, shots):
+    with pytest.raises(ValueError, match="shots"):
+        bench.sample_few_shot(["a", "b"], shots, 0)
+
+
+def test_subsample(bench):
+    texts = [f"t{i}" for i in range(100)]
+    labels = [str(i % 3) for i in range(100)]
+    all_t, all_l = bench.subsample(texts, labels, None)
+    assert (all_t, all_l) == (texts, labels)
+    assert bench.subsample(texts, labels, 500)[0] == texts
+    sub_t, sub_l = bench.subsample(texts, labels, 20)
+    assert len(sub_t) == 20 and len(set(sub_t)) == 20
+    assert all(labels[texts.index(t)] == l for t, l in zip(sub_t, sub_l))        # still aligned
+    assert bench.subsample(texts, labels, 20) == (sub_t, sub_l)                  # fixed subset
+
+
+# --- metrics -------------------------------------------------------------------------------------
+
+def test_perfect_predictions_score_perfectly(bench):
+    probs = np.eye(3)[[0, 1, 2, 0, 1, 2]]
+    scores = bench.score_predictions(probs, np.array([0, 1, 2, 0, 1, 2]))
+    assert scores == {"accuracy": 1.0, "macro_f1": 1.0, "ece": pytest.approx(0.0)}
+
+
+def test_known_accuracy_and_f1(bench):
+    probs = np.array([[0.9, 0.1], [0.8, 0.2], [0.3, 0.7], [0.6, 0.4]])
+    targets = np.array([0, 0, 1, 1])                      # last one wrong
+    scores = bench.score_predictions(probs, targets)
+    assert scores["accuracy"] == pytest.approx(0.75)
+    assert scores["macro_f1"] == pytest.approx(0.7333, abs=1e-3)
+
+
+def test_ece_matches_the_library_implementation(bench):
+    rng = np.random.default_rng(0)
+    probs = rng.dirichlet(np.ones(4), size=200)
+    targets = rng.integers(0, 4, size=200)
+    assert bench.expected_calibration_error(probs, targets) == pytest.approx(
+        calibration.expected_calibration_error(probs, targets))
+
+
+def test_ece_of_an_overconfident_model(bench):
+    probs = np.array([[1.0, 0.0]] * 10)
+    assert bench.expected_calibration_error(probs, np.array([0] * 5 + [1] * 5)) == pytest.approx(0.5)
+
+
+# --- aggregation and report ------------------------------------------------------------------------
+
+def row(dataset="ds", method="adaptive", shots=4, seed=0, accuracy=0.5, **extra):
+    base = {"dataset": dataset, "method": method, "shots": shots, "seed": seed, "n_classes": 3,
+            "n_train": 12, "n_test": 100, "accuracy": accuracy, "macro_f1": accuracy - 0.05,
+            "ece": 0.1, "train_s": 2.0, "predict_s": 1.0, "latency_ms_p50": 5.0, "update_s": 0.5,
+            "update_is_retrain": False}
+    base.update(extra)
+    return base
+
+
+def test_aggregate_computes_mean_and_sample_std(bench):
+    rows = [row(seed=s, accuracy=a) for s, a in enumerate([0.6, 0.7, 0.8])]
+    summary = bench.aggregate(rows)[("ds", "adaptive", 4)]
+    mean, std = summary["accuracy"]
+    assert mean == pytest.approx(0.7) and std == pytest.approx(0.1)
+    assert summary["runs"] == 3 and summary["failed"] == 0
+
+
+def test_aggregate_single_seed_has_no_std(bench):
+    assert bench.aggregate([row()])[("ds", "adaptive", 4)]["accuracy"] == (0.5, None)
+
+
+def test_aggregate_excludes_failed_runs_but_counts_them(bench):
+    rows = [row(seed=0, accuracy=0.6), {"dataset": "ds", "method": "adaptive", "shots": 4,
+                                         "seed": 1, "error": "RuntimeError: boom"}]
+    summary = bench.aggregate(rows)[("ds", "adaptive", 4)]
+    assert summary["accuracy"][0] == pytest.approx(0.6)
+    assert (summary["runs"], summary["failed"]) == (2, 1)
+
+
+def test_aggregate_keeps_methods_datasets_and_shots_apart(bench):
+    rows = [row(method="a", accuracy=0.1), row(method="b", accuracy=0.9),
+            row(dataset="other", accuracy=0.4), row(shots=16, accuracy=0.7)]
+    summary = bench.aggregate(rows)
+    assert len(summary) == 4 and summary[("ds", "b", 4)]["accuracy"][0] == 0.9
+
+
+def test_report_has_a_section_per_dataset_and_bolds_the_best(bench):
+    rows = [row(dataset="alpha", method="adaptive", accuracy=0.80),
+            row(dataset="alpha", method="logreg", accuracy=0.70),
+            row(dataset="beta", method="adaptive", accuracy=0.60)]
+    text = bench.render_markdown(rows, {"model": "tiny", "seeds": 1})
+    assert "## alpha" in text and "## beta" in text
+    assert "**model**: tiny" in text
+    assert "**0.800**" in text and "0.700" in text and "**0.700**" not in text
+
+
+def test_report_flags_failures_and_marks_gaps(bench):
+    rows = [row(method="adaptive", shots=4), row(method="adaptive", shots=16),
+            row(method="setfit", shots=4),
+            {"dataset": "ds", "method": "setfit", "shots": 16, "seed": 0, "error": "OOM"}]
+    text = bench.render_markdown(rows)
+    assert "1 run(s) failed" in text and "failed" in text.split("### Accuracy")[1]
+    only_one = bench.render_markdown([row(method="adaptive", shots=4), row(method="logreg", shots=16)])
+    assert "–" in only_one
+
+
+def test_report_cost_table_marks_retraining(bench):
+    rows = [row(method="adaptive", shots=16), row(method="setfit", shots=16, update_is_retrain=True,
+                                                    update_s=90.0)]
+    text = bench.render_markdown(rows)
+    cost = text.split("### Cost")[1]
+    assert "90.00 (full retrain)" in cost and "0.50 (full retrain)" not in cost
+
+
+def test_report_with_std_shows_plus_minus(bench):
+    rows = [row(seed=0, accuracy=0.6), row(seed=1, accuracy=0.8)]
+    assert "±" in bench.render_markdown(rows)
+
+
+def test_report_orders_known_methods_first(bench):
+    rows = [row(method="zzz-custom"), row(method="setfit"), row(method="adaptive")]
+    text = bench.render_markdown(rows).split("### Accuracy")[1]
+    assert text.index("| adaptive") < text.index("| setfit") < text.index("| zzz-custom")
+
+
+# --- the synthetic dataset -----------------------------------------------------------------------------
+
+def test_synthetic_dataset_is_deterministic_and_balanced(bench):
+    a, b = bench.synthetic_dataset(), bench.synthetic_dataset()
+    assert a == b
+    assert sorted(set(a["train_labels"])) == [f"class_{c}" for c in range(4)]
+    assert {a["train_labels"].count(l) for l in set(a["train_labels"])} == {40}
+    assert set(a["test_labels"]) <= set(a["train_labels"])
+
+
+def test_unknown_dataset_and_method_are_rejected(bench):
+    with pytest.raises(SystemExit):
+        bench.load_dataset_splits("no-such-dataset")
+    with pytest.raises(SystemExit):
+        bench.build_method("no-such-method", "model", None)
+
+
+# --- running each method for real, on a tiny local model ---------------------------------------------------
+
+@pytest.fixture
+def smoke(bench, synthetic_text_model):
+    def go(methods, shots=(5,), seeds=1):
+        return bench.run_benchmark("synthetic", list(shots), seeds, methods, synthetic_text_model,
+                                   max_test=60, log=lambda *_: None)
+    return go
+
+
+def test_adaptive_methods_and_logreg_run_end_to_end(smoke):
+    results = smoke(["adaptive", "adaptive-proto", "logreg"])
+    assert len(results) == 3 and not [r for r in results if r.get("error")]
+    for r in results:
+        assert 0.25 < r["accuracy"] <= 1.0, f"{r['method']} is no better than chance: {r['accuracy']}"
+        assert 0.0 <= r["macro_f1"] <= 1.0 and 0.0 <= r["ece"] <= 1.0
+        assert r["train_s"] > 0 and r["predict_s"] > 0 and r["latency_ms_p50"] > 0 and r["update_s"] > 0
+        assert r["n_classes"] == 4 and r["n_test"] == 60 and r["n_train"] == 20
+
+
+def test_adding_an_example_is_an_update_for_adaptive_and_a_refit_for_logreg(smoke):
+    by_method = {r["method"]: r for r in smoke(["adaptive", "logreg"])}
+    assert by_method["adaptive"]["update_is_retrain"] is False
+    assert by_method["logreg"]["update_is_retrain"] is True
+
+
+def test_results_are_reproducible_across_runs(smoke):
+    first, second = smoke(["adaptive-proto", "logreg"]), smoke(["adaptive-proto", "logreg"])
+    for a, b in zip(first, second):
+        assert a["accuracy"] == b["accuracy"] and a["macro_f1"] == pytest.approx(b["macro_f1"])
+
+
+def test_every_seed_and_shot_count_gets_a_row_per_method(smoke):
+    results = smoke(["adaptive-proto", "logreg"], shots=(3, 6), seeds=2)
+    assert len(results) == 2 * 2 * 2
+    assert {(r["shots"], r["seed"], r["method"]) for r in results} == {
+        (k, s, m) for k in (3, 6) for s in (0, 1) for m in ("adaptive-proto", "logreg")}
+
+
+def test_a_failing_method_is_recorded_and_the_others_still_run(bench, synthetic_text_model, monkeypatch):
+    real = bench.build_method
+
+    def build(name, model, device):
+        if name == "logreg":
+            raise RuntimeError("boom")
+        return real(name, model, device)
+
+    monkeypatch.setattr(bench, "build_method", build)
+    results = bench.run_benchmark("synthetic", [4], 1, ["logreg", "adaptive-proto"],
+                                  synthetic_text_model, 40, log=lambda *_: None)
+    by = {r["method"]: r for r in results}
+    assert "boom" in by["logreg"]["error"] and "accuracy" not in by["logreg"]
+    assert by["adaptive-proto"]["accuracy"] > 0
+
+
+def test_setfit_does_not_write_checkpoints_into_the_working_directory(bench, synthetic_text_model,
+                                                                      tmp_path, monkeypatch):
+    try:
+        import datasets  # noqa: F401
+        import setfit  # noqa: F401
+    except ImportError as error:
+        pytest.skip(f"SetFit is not usable in this environment: {error}")
+    monkeypatch.chdir(tmp_path)
+    results = bench.run_benchmark("synthetic", [4], 1, ["setfit"], synthetic_text_model, 40,
+                                  log=lambda *_: None)
+    assert not results[0].get("error"), results[0].get("error")
+    assert list(tmp_path.iterdir()) == [], "SetFit left files in the working directory"
+
+
+def test_setfit_runs_end_to_end(bench, synthetic_text_model):
+    try:
+        import datasets  # noqa: F401
+        import setfit  # noqa: F401
+    except ImportError as error:        # absent, or installed against an incompatible stack
+        pytest.skip(f"SetFit is not usable in this environment: {error}")
+    results = bench.run_benchmark("synthetic", [4], 1, ["setfit"], synthetic_text_model, 40,
+                                  log=lambda *_: None)
+    assert not results[0].get("error"), results[0].get("error")
+    assert 0.0 < results[0]["accuracy"] <= 1.0
+    assert results[0]["update_is_retrain"] is True and results[0]["update_s"] == results[0]["train_s"]
+
+
+# --- the command line ----------------------------------------------------------------------------------------
+
+def test_run_writes_a_results_file_and_report_renders_it(bench, synthetic_text_model, tmp_path, capsys):
+    out = tmp_path / "nested" / "synthetic.json"
+    code = bench.main(["run", "--dataset", "synthetic", "--shots", "4", "--seeds", "2",
+                       "--methods", "adaptive-proto", "logreg", "--model", synthetic_text_model,
+                       "--max-test", "40", "--out", str(out)])
+    assert code == 0 and out.exists()
+    payload = json.loads(out.read_text())
+    assert payload["meta"]["model"] == synthetic_text_model and payload["meta"]["seeds"] == 2
+    assert len(payload["results"]) == 4
+
+    summary = tmp_path / "summary.md"
+    assert bench.main(["report", str(out), "--out", str(summary)]) == 0
+    text = summary.read_text()
+    assert "## synthetic" in text and "adaptive-proto" in text and "logreg" in text
+    assert "# Few-shot benchmark" in capsys.readouterr().out
+
+
+def test_run_exits_nonzero_but_still_writes_results_when_a_method_fails(bench, synthetic_text_model,
+                                                                         tmp_path, monkeypatch):
+    real = bench.build_method
+    monkeypatch.setattr(bench, "build_method",
+                        lambda n, m, d: (_ for _ in ()).throw(RuntimeError("down")) if n == "logreg"
+                        else real(n, m, d))
+    out = tmp_path / "r.json"
+    code = bench.main(["run", "--dataset", "synthetic", "--shots", "4", "--seeds", "1",
+                       "--methods", "adaptive-proto", "logreg", "--model", synthetic_text_model,
+                       "--max-test", "40", "--out", str(out)])
+    assert code == 1
+    errors = [r for r in json.loads(out.read_text())["results"] if r.get("error")]
+    assert len(errors) == 1 and "down" in errors[0]["error"]
+
+
+def test_report_merges_several_result_files(bench, tmp_path):
+    for name, dataset in (("a.json", "alpha"), ("b.json", "beta")):
+        (tmp_path / name).write_text(json.dumps({"meta": {"model": "m"}, "results": [row(dataset=dataset)]}))
+    out = tmp_path / "s.md"
+    assert bench.main(["report", str(tmp_path / "a.json"), str(tmp_path / "b.json"), "--out", str(out)]) == 0
+    assert "## alpha" in out.read_text() and "## beta" in out.read_text()
+
+
+def test_command_line_rejects_unknown_methods(bench, tmp_path):
+    with pytest.raises(SystemExit):
+        bench.main(["run", "--dataset", "synthetic", "--methods", "magic", "--out", str(tmp_path / "x")])
+
+
+# --- checking datasets before a long run ------------------------------------------------------------------------
+
+class FakeStream:
+    def __init__(self, rows, features=None):
+        self.rows, self.features = rows, features
+
+    def __iter__(self):
+        return iter(self.rows)
+
+
+class FakeFeature:
+    names = ["a", "b", "c"]
+
+
+def install_fake_datasets(monkeypatch, loader):
+    import sys
+    import types
+    fake = types.ModuleType("datasets")
+    fake.load_dataset = loader
+    monkeypatch.setitem(sys.modules, "datasets", fake)
+
+
+def test_verify_accepts_a_well_formed_dataset(bench, monkeypatch):
+    def loader(hub, split, streaming):
+        assert streaming is True, "verification must not download the whole dataset"
+        return FakeStream([{"text": "hello", "label": 1}], {"label": FakeFeature()})
+
+    install_fake_datasets(monkeypatch, loader)
+    assert bench.verify_dataset("ag_news") == "ag_news: fancyzhx/ag_news ok, 3 classes"
+
+
+def test_verify_names_the_hub_id_and_split_when_loading_fails(bench, monkeypatch):
+    def loader(hub, split, streaming):
+        raise FileNotFoundError(f"Dataset '{hub}' doesn't exist on the Hub")
+
+    install_fake_datasets(monkeypatch, loader)
+    with pytest.raises(RuntimeError) as info:
+        bench.verify_dataset("emotion")
+    assert "dair-ai/emotion" in str(info.value) and "'train'" in str(info.value)
+    assert "doesn't exist" in str(info.value)
+
+
+def test_verify_reports_a_missing_column(bench, monkeypatch):
+    install_fake_datasets(monkeypatch, lambda hub, split, streaming:
+                          FakeStream([{"sentence": "x", "label": 0}], {}))
+    with pytest.raises(RuntimeError, match="no column 'text'"):
+        bench.verify_dataset("banking77")
+
+
+def test_no_default_dataset_depends_on_a_loading_script(bench):
+    """Recent `datasets` releases refuse to run Hub loading scripts, so a dataset that only
+    has one fails at load time. PolyAI/banking77 was one; the parquet copy replaces it."""
+    assert all(spec["hub"] != "PolyAI/banking77" for spec in bench.DATASETS.values())
+
+
+def test_verify_reports_every_problem_not_just_the_first(bench, monkeypatch):
+    def loader(hub, split, streaming):
+        if split == "test":
+            raise ValueError("no test split")
+        return FakeStream([{"text": "x"}], {})
+
+    install_fake_datasets(monkeypatch, loader)
+    with pytest.raises(RuntimeError) as info:
+        bench.verify_dataset("ag_news")
+    assert "'test'" in str(info.value) and "no column 'label'" in str(info.value)
+
+
+def test_check_command_exit_codes(bench, monkeypatch, capsys):
+    install_fake_datasets(monkeypatch, lambda hub, split, streaming:
+                          FakeStream([{"text": "x", "label": 0}], {"label": FakeFeature()}))
+    assert bench.main(["check", "synthetic", "ag_news"]) == 0
+    assert "synthetic" in capsys.readouterr().out
+
+    install_fake_datasets(monkeypatch, lambda *a, **k: (_ for _ in ()).throw(OSError("offline")))
+    assert bench.main(["check", "ag_news", "emotion"]) == 1
+    assert capsys.readouterr().err.count("FAIL") == 2
+
+
+def test_loading_works_for_datasets_without_class_names(bench, monkeypatch):
+    class Split(dict):
+        features = {"label": object()}                      # a plain value: no .names
+
+    train = Split(text=["a", "b", "c"], label=[0, 1, 0])
+    test = Split(text=["d"], label=[1])
+    install_fake_datasets(monkeypatch, lambda hub: {"train": train, "test": test})
+    loaded = bench.load_dataset_splits("ag_news")
+    assert loaded["train_labels"] == ["0", "1", "0"] and loaded["test_labels"] == ["1"]
+
+
+def test_loading_maps_class_ids_to_names(bench, monkeypatch):
+    class Split(dict):
+        features = {"label": FakeFeature()}
+
+    train = Split(text=["a", "b"], label=[2, 0])
+    install_fake_datasets(monkeypatch, lambda hub: {"train": train, "test": Split(text=["c"], label=[1])})
+    loaded = bench.load_dataset_splits("emotion")
+    assert loaded["train_labels"] == ["c", "a"] and loaded["test_labels"] == ["b"]
+
+
+# --- validating workflow inputs -----------------------------------------------------------------------------------
+
+def good_plan(bench, **overrides):
+    args = dict(datasets="ag_news,emotion", shots="4 16", seeds="3", methods="adaptive logreg",
+                model="sentence-transformers/all-MiniLM-L6-v2", max_test="1000")
+    args.update(overrides)
+    return bench.plan(**args)
+
+
+def test_plan_normalises_valid_input(bench):
+    p = good_plan(bench, datasets=" ag_news , emotion,ag_news ", shots="16,4 4", methods="logreg  adaptive")
+    assert p == {"datasets": ["ag_news", "emotion"], "shots": "4 16", "seeds": 3,
+                 "methods": "logreg adaptive", "model": "sentence-transformers/all-MiniLM-L6-v2",
+                 "max_test": 1000}
+
+
+def test_plan_defaults_to_every_method(bench):
+    assert good_plan(bench, methods="")["methods"] == " ".join(bench.METHODS)
+
+
+@pytest.mark.parametrize("field, value", [
+    ("datasets", ""), ("datasets", "ag_news,imaginary"), ("datasets", "ag_news; rm -rf /"),
+    ("shots", ""), ("shots", "0"), ("shots", "4 many"), ("shots", "1000"), ("shots", "-3"),
+    ("seeds", "0"), ("seeds", "11"), ("seeds", "three"), ("seeds", ""),
+    ("methods", "adaptive magic"), ("methods", "adaptive;ls"),
+    ("model", ""), ("model", "a b"), ("model", "x; echo hi"), ("model", "$(whoami)"),
+    ("model", "../../etc/passwd"), ("model", "a/b/c"),
+    ("max_test", "5"), ("max_test", "lots"), ("max_test", "10000000"),
+])
+def test_plan_rejects_bad_values_and_names_the_setting(bench, field, value):
+    with pytest.raises(ValueError) as info:
+        good_plan(bench, **{field: value})
+    assert field.replace("_", " ") in str(info.value) or field in str(info.value)
+
+
+def test_plan_command_prints_json_or_fails_with_a_message(bench, capsys):
+    assert bench.main(["plan", "--datasets", "emotion", "--shots", "8", "--seeds", "2"]) == 0
+    printed = json.loads(capsys.readouterr().out)
+    assert printed["datasets"] == ["emotion"] and printed["shots"] == "8" and printed["seeds"] == 2
+
+    assert bench.main(["plan", "--datasets", "nope"]) == 2
+    assert "Invalid benchmark settings" in capsys.readouterr().err
diff --git a/tests/test_failure_modes.py b/tests/test_failure_modes.py
index 305e27e..a870907 100644
--- a/tests/test_failure_modes.py
+++ b/tests/test_failure_modes.py
@@ -249,8 +249,10 @@ def test_config_that_disagrees_with_the_weights_fails_loudly(saved):
 def test_onnx_failure_falls_back_to_a_working_pytorch_model(base_model, monkeypatch, caplog, error):
     """This is what happened in CI when the Hub rate-limited: the classifier must
     keep working on PyTorch and say why, not fail or pretend ONNX loaded."""
-    pytest.importorskip("optimum.onnxruntime")
-    from optimum.onnxruntime import ORTModelForFeatureExtraction
+    try:
+        from optimum.onnxruntime import ORTModelForFeatureExtraction
+    except ImportError:                     # not installed, or built for a different transformers
+        pytest.skip("optimum.onnxruntime is not importable in this environment")
 
     def broken(*args, **kwargs):
         raise error

From 83a799ea099506dca7f41d7f3057639c69687682 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 12:27:42 +0000
Subject: [PATCH 14/16] Use a placeholder encoder name in the benchmark harness
 tests

tests/test_ci_models.py failed in CI because the harness tests pass
"sentence-transformers/all-MiniLM-L6-v2" to plan() as input data, and the guard
reads any model-like string in tests/ as a model the suite needs offline. Nothing
downloads it. Use "acme/encoder-small" for the dummy value instead of adding a
model to the prefetch list for no reason. All other tests in the run passed
(476 passed on Python 3.13).

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 tests/test_benchmark_harness.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_benchmark_harness.py b/tests/test_benchmark_harness.py
index 593a274..eb34a22 100644
--- a/tests/test_benchmark_harness.py
+++ b/tests/test_benchmark_harness.py
@@ -442,7 +442,7 @@ class Split(dict):
 
 def good_plan(bench, **overrides):
     args = dict(datasets="ag_news,emotion", shots="4 16", seeds="3", methods="adaptive logreg",
-                model="sentence-transformers/all-MiniLM-L6-v2", max_test="1000")
+                model="acme/encoder-small", max_test="1000")
     args.update(overrides)
     return bench.plan(**args)
 
@@ -450,7 +450,7 @@ def good_plan(bench, **overrides):
 def test_plan_normalises_valid_input(bench):
     p = good_plan(bench, datasets=" ag_news , emotion,ag_news ", shots="16,4 4", methods="logreg  adaptive")
     assert p == {"datasets": ["ag_news", "emotion"], "shots": "4 16", "seeds": 3,
-                 "methods": "logreg adaptive", "model": "sentence-transformers/all-MiniLM-L6-v2",
+                 "methods": "logreg adaptive", "model": "acme/encoder-small",
                  "max_test": 1000}
 
 

From 9fff1f0714032e7617a8349e0d2cd42502ffc024 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 13:02:37 +0000
Subject: [PATCH 15/16] Remove the SetFit comparison from the README and
 changelog

The benchmark harness stays as internal tooling, but the library's public
documentation is about continual learning and does not need to advertise a
comparison with another method. Also neutralise two comments that named it.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 .gitignore     | 2 +-
 CHANGELOG.md   | 1 -
 README.md      | 4 ----
 pyproject.toml | 2 +-
 4 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/.gitignore b/.gitignore
index d27df0a..f0797e7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -174,5 +174,5 @@ cython_debug/
 demo_classifier/
 scripts/benchmark_results/
 
-# SetFit training checkpoints
+# Training checkpoints written by the benchmark harness
 checkpoints/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d3d940e..882f966 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,7 +19,6 @@
 - Async methods `apredict`, `apredict_batch` and `aadd_examples` that run in a worker thread.
 - A saved directory is self-contained for loading: the tokenizer saved next to the model is used instead of fetching the base model's repository, so an ONNX-saved classifier starts without Hub access.
 
-- **Few-shot benchmark** against SetFit and logistic regression: `scripts/benchmark_fewshot.py` and a manually triggered `Few-shot benchmark` workflow (accuracy, macro-F1, calibration error, training time, latency, and the cost of adding one example). Install with `pip install -e ".[benchmark]"`. See `docs/benchmark.md`.
 
 ### Changed (default behaviour: predictions change)
 - **The neural head is now actually trained.** It used to get about ten optimiser steps and stop at the first loss plateau, so it could not learn even a 2-D XOR from 120 examples. It now trains for at least `head_steps` (default 300) steps with a cosine-decayed `head_learning_rate` (default 0.003). Small memories take roughly 1-2 s per `add_examples` call; lower `head_steps` if you add examples very frequently.
diff --git a/README.md b/README.md
index 01b78fb..b1a5e5f 100644
--- a/README.md
+++ b/README.md
@@ -382,10 +382,6 @@ curl -s localhost:8000/predict -H 'content-type: application/json' -d '{"text":
 
 Predictions, batches, conformal sets and out-of-distribution checks over HTTP, plus optional API-key-protected endpoints for adding examples and classes while it runs. One classifier is safe to share between threads, and `await classifier.apredict(...)` keeps an event loop responsive. See [docs/serving.md](docs/serving.md) and `docker/Dockerfile`.
 
-### How does it compare?
-
-A benchmark against [SetFit](https://github.com/huggingface/setfit) and logistic regression on the same encoder (accuracy, macro-F1, calibration, training time, latency, and the cost of adding one example) is in `scripts/benchmark_fewshot.py`, with a manually triggered GitHub workflow. See [docs/benchmark.md](docs/benchmark.md) for how it works and what it does and does not show.
-
 ### Fixing Mistakes and Knowing When to Abstain
 
 ```python
diff --git a/pyproject.toml b/pyproject.toml
index 15ad957..90e91bd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,7 +52,7 @@ serve = [
     "fastapi>=0.100.0",
     "uvicorn>=0.23.0",
 ]
-# Few-shot benchmark against SetFit: scripts/benchmark_fewshot.py (see docs/benchmark.md).
+# Benchmark harness: scripts/benchmark_fewshot.py (see docs/benchmark.md).
 benchmark = [
     "setfit>=1.1.0",
     "datasets>=2.14.0",

From 4a60c492c2f631f3d9a74e682b835411c16832e2 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 8 Oct 2026 16:34:13 +0000
Subject: [PATCH 16/16] Keep the fixed new-class split when loading saves from
 before 0.2.0

Saves from before 0.2.0 store neither prototype_temperature nor the new-class
weights. The loader already pinned the temperature to the legacy scoring, but
a missing new_class_prototype_weight / new_class_neural_weight picked up the new
default (a ramp), so a class with fewer than 10 examples was blended
(0.85, 0.15) instead of the (0.3, 0.7) it had in 0.2.0 and predictions changed.
Found by loading one stripped-down save under both versions. Missing weights
are now pinned to the old fixed split; saves that stored their own weights
keep them.

The classifiers published in the adaptive-classifier Hugging Face org were not
affected (they all have 100 examples per class): 10 of them give identical
predictions under 0.2.0 and 0.3.0 (225 predictions, no changed labels, largest
score difference 3e-8).

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 src/adaptive_classifier/classifier.py | 11 ++++++-
 tests/test_robustness.py              | 43 +++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/src/adaptive_classifier/classifier.py b/src/adaptive_classifier/classifier.py
index f7e1cb7..24a2d91 100644
--- a/src/adaptive_classifier/classifier.py
+++ b/src/adaptive_classifier/classifier.py
@@ -1306,7 +1306,16 @@ def _from_pretrained(
         # the new behaviour.
         saved_config = config_dict.get('config')
         if isinstance(saved_config, dict) and 'prototype_temperature' not in saved_config:
-            config_dict['config'] = {**saved_config, 'prototype_temperature': None}
+            # Saves from before 0.2.0 do not store the new-class weights either, and
+            # the fixed 0.3 / 0.7 split was the behaviour then. Without this they would
+            # pick up the new default (a ramp) and score classes with few examples
+            # differently. Saves that did store them keep their own values.
+            config_dict['config'] = {
+                'new_class_prototype_weight': 0.3,
+                'new_class_neural_weight': 0.7,
+                **saved_config,
+                'prototype_temperature': None,
+            }
 
         # Load examples
         with open(model_path / "examples.json", "r", encoding="utf-8") as f:
diff --git a/tests/test_robustness.py b/tests/test_robustness.py
index 5d90a65..685b1b1 100644
--- a/tests/test_robustness.py
+++ b/tests/test_robustness.py
@@ -257,6 +257,49 @@ def test_a_pre_0_3_0_save_keeps_its_scoring(new_classifier, base_model, data, tm
         assert all(a == pytest.approx(b, abs=1e-4) for (_, a), (_, b) in zip(got, preds))
 
 
+def test_a_pre_0_2_0_save_with_small_classes_keeps_the_fixed_new_class_split(new_classifier, data, tmp_path):
+    """Saves from before 0.2.0 store neither the new-class weights nor a temperature. A class
+    with fewer than 10 examples must still be blended 0.3 / 0.7 on load, not pick up the ramp
+    that is the default for newly built classifiers (found by loading a stripped-down save
+    under both versions: the blend was (0.3, 0.7) in 0.2.0 and (0.85, 0.15) in 0.3.0)."""
+    import json
+    from adaptive_classifier import AdaptiveClassifier
+
+    X, y = data.classes(6, noise=0.2, seed=11)                      # 6 examples per class: under the threshold
+    legacy = fit(new_classifier, X, y, prototype_temperature=None,
+                 new_class_prototype_weight=0.3, new_class_neural_weight=0.7)
+    expected = {t: legacy.predict(t, k=5) for t in X[:6]}
+    legacy.save(str(tmp_path), include_onnx=False)
+
+    path = tmp_path / "config.json"
+    saved = json.loads(path.read_text(encoding="utf-8"))
+    for key in ("prototype_temperature", "new_class_prototype_weight", "new_class_neural_weight",
+                "head_steps", "head_learning_rate"):
+        saved["config"].pop(key, None)                                # a pre-0.2.0 config has none of these
+    path.write_text(json.dumps(saved), encoding="utf-8")
+
+    loaded = AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")
+    assert loaded._blend_weights("C0") == (0.3, 0.7)
+    for text, preds in expected.items():
+        got = loaded.predict(text, k=5)
+        assert [l for l, _ in got] == [l for l, _ in preds]
+        assert all(a == pytest.approx(b, abs=1e-4) for (_, a), (_, b) in zip(got, preds))
+
+
+def test_a_save_that_stored_its_own_new_class_weights_keeps_them(new_classifier, data, tmp_path):
+    import json
+    from adaptive_classifier import AdaptiveClassifier
+
+    X, y = data.classes(6, noise=0.2, seed=12)
+    fit(new_classifier, X, y, new_class_prototype_weight=0.6, new_class_neural_weight=0.4).save(
+        str(tmp_path), include_onnx=False)
+    path = tmp_path / "config.json"
+    saved = json.loads(path.read_text(encoding="utf-8"))
+    del saved["config"]["prototype_temperature"]                      # older than 0.3.0, but it stored the weights
+    path.write_text(json.dumps(saved), encoding="utf-8")
+    assert AdaptiveClassifier.load(str(tmp_path), use_onnx=False, device="cpu")._blend_weights("C0") == (0.6, 0.4)
+
+
 def test_new_saves_record_the_new_settings(new_classifier, data, tmp_path):
     import json
     X, y = data.classes(4, seed=8)