tinydistill 0.3.2__tar.gz → 0.3.3__tar.gz

{tinydistill-0.3.2 → tinydistill-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: tinydistill
-Version: 0.3.2
+Version: 0.3.3
 Summary: Distil an LLM classification prompt into a small, fast model that runs on a CPU.
 Project-URL: Homepage, https://github.com/Andreas-Hum/tinydistill
 Author: Andreas-Hum
@@ -28,10 +28,13 @@ Requires-Dist: bitsandbytes>=0.43; extra == '4bit'
 Provides-Extra: dev
 Requires-Dist: build>=1.2; extra == 'dev'
 Requires-Dist: datasets>=2.18; extra == 'dev'
+Requires-Dist: matplotlib>=3.7; extra == 'dev'
 Requires-Dist: pytest>=8; extra == 'dev'
 Requires-Dist: ruff>=0.5; extra == 'dev'
 Provides-Extra: hf
 Requires-Dist: datasets>=2.18; extra == 'hf'
+Provides-Extra: plot
+Requires-Dist: matplotlib>=3.7; extra == 'plot'
 Description-Content-Type: text/markdown
 
 # tinydistill
@@ -86,6 +89,7 @@ browser and shows what happens inside it.</i></p>
 pip install tinydistill               # label, train, export, predict
 pip install "tinydistill[hf]"         # + Hugging Face datasets as input (hf:<dataset>:<split>)
 pip install "tinydistill[hf,4bit]"    # + 4-bit teachers with bitsandbytes (CUDA)
+pip install "tinydistill[plot]"       # + charts of a run (matplotlib)
 ```
 
 Python 3.11 or newer, on Linux, macOS or Windows. Labelling with a large teacher wants a GPU; training and
@@ -94,6 +98,19 @@ inference are fine on a CPU. For a CPU-only PyTorch, install it first:
 
 ## Quickstart
 
+To try a distilled model before training your own, the Banking77 student from the [results](#results) is on
+the [Hugging Face Hub](https://huggingface.co/Shroomy/tinydistill-banking77-minilm):
+
+```python
+from tinydistill import Predictor  # pip install tinydistill
+
+clf = Predictor("hf:Shroomy/tinydistill-banking77-minilm")
+clf.describe(["I still haven't received my new card", "My card got stolen"])
+# [{'label': 'card_arrival', 'confidence': 0.4672}, {'label': 'lost_or_stolen_card', 'confidence': 0.9921}]
+```
+
+To distil your own:
+
 ```sh
 curl -LO https://raw.githubusercontent.com/Andreas-Hum/tinydistill/main/examples/banking77/task.yaml
 tinydistill run task.yaml --data hf:PolyAI/banking77:train --test hf:PolyAI/banking77:test \
@@ -110,8 +127,13 @@ clf = Predictor("runs/banking77/onnx")  # or a model on the Hub: Predictor("hf:u
 clf.predict(["I still haven't received my new card"])  # a list of label names
 ```
 
-Or try it without installing anything: the [Colab notebook](https://colab.research.google.com/github/Andreas-Hum/tinydistill/blob/main/notebooks/quickstart.ipynb) distils a sentiment classifier in a
-few minutes on a free GPU.
+Or try it without installing anything, on a free Colab GPU:
+
+| Notebook | What it shows |
+|---|---|
+| [Quickstart](https://colab.research.google.com/github/Andreas-Hum/tinydistill/blob/main/notebooks/quickstart.ipynb) | distil a sentiment classifier in about ten minutes, with charts of what the student learned |
+| [Your own task](https://colab.research.google.com/github/Andreas-Hum/tinydistill/blob/main/notebooks/your_task.ipynb) | write a task for your labels, check the teacher, distil on your texts, find the mistakes, download the model |
+| [Deploy](https://colab.research.google.com/github/Andreas-Hum/tinydistill/blob/main/notebooks/deploy.ipynb) | the model from Python, behind an HTTP server, in a web page, and on the Hugging Face Hub |
 
 Labelling the full dataset with the default teacher wants a GPU. For a quick look on a CPU, add
 `--teacher Qwen/Qwen2.5-0.5B-Instruct --limit 200 --test-limit 200`. `--limit` takes the first rows, and the
@@ -299,6 +321,30 @@ results show the mean and standard deviation.
   adds them to the training texts and trains, exports and evaluates a new student on the same test set, so
   the student learns the inputs it found hard in production.
 
+### Charts and metrics
+
+`tinydistill plot RUN` (with the `plot` extra) writes the charts of a run to `RUN/plots/`; in Python every
+chart is a function in `tinydistill.plots` that returns a matplotlib figure, so it shows inline in a notebook:
+
+| Chart | What it shows |
+|---|---|
+| `soft_labels` | for a few test texts, the teacher's probabilities (the soft labels) next to the student's |
+| `speed_accuracy` | accuracy against latency for the teacher and every form of the student, with the baseline and upper bound |
+| `training` | the loss per epoch on the training texts and, with `--val-share`, on the held-out texts and the kept epoch |
+| `per_label` | F1 of every label for teacher and student, the teacher's weakest first |
+| `confusions` | the confusion matrix over the labels involved in the most mistakes |
+| `cascade` | accuracy as the student passes more of its least confident texts to the teacher |
+
+<p align="center">
+  <img alt="Soft labels on three Banking77 test texts: the teacher's probabilities next to the student's" src="https://raw.githubusercontent.com/Andreas-Hum/tinydistill/main/docs/images/banking77-soft-labels.png" width="49%">
+  <img alt="Accuracy against latency on Banking77: the students on a CPU, the teacher on a GPU, the baseline and the upper bound" src="https://raw.githubusercontent.com/Andreas-Hum/tinydistill/main/docs/images/banking77-speed-accuracy.png" width="49%">
+</p>
+
+Every run's results report accuracy and macro-F1, agreement with the teacher, latency at batch 1 and 32,
+throughput and size; tasks with more than five labels also get top-3 accuracy, multi-label tasks micro-F1, and
+ordered labels the mean distance in label steps. `tinydistill analyse RUN` adds a page of per-label scores and
+example mistakes.
+
 ### From Python
 
 The exported model:

{tinydistill-0.3.2 → tinydistill-0.3.3}/README.md

@@ -50,6 +50,7 @@ browser and shows what happens inside it.</i></p>
 pip install tinydistill               # label, train, export, predict
 pip install "tinydistill[hf]"         # + Hugging Face datasets as input (hf:<dataset>:<split>)
 pip install "tinydistill[hf,4bit]"    # + 4-bit teachers with bitsandbytes (CUDA)
+pip install "tinydistill[plot]"       # + charts of a run (matplotlib)
 ```
 
 Python 3.11 or newer, on Linux, macOS or Windows. Labelling with a large teacher wants a GPU; training and
@@ -58,6 +59,19 @@ inference are fine on a CPU. For a CPU-only PyTorch, install it first:
 
 ## Quickstart
 
+To try a distilled model before training your own, the Banking77 student from the [results](#results) is on
+the [Hugging Face Hub](https://huggingface.co/Shroomy/tinydistill-banking77-minilm):
+
+```python
+from tinydistill import Predictor  # pip install tinydistill
+
+clf = Predictor("hf:Shroomy/tinydistill-banking77-minilm")
+clf.describe(["I still haven't received my new card", "My card got stolen"])
+# [{'label': 'card_arrival', 'confidence': 0.4672}, {'label': 'lost_or_stolen_card', 'confidence': 0.9921}]
+```
+
+To distil your own:
+
 ```sh
 curl -LO https://raw.githubusercontent.com/Andreas-Hum/tinydistill/main/examples/banking77/task.yaml
 tinydistill run task.yaml --data hf:PolyAI/banking77:train --test hf:PolyAI/banking77:test \
@@ -74,8 +88,13 @@ clf = Predictor("runs/banking77/onnx")  # or a model on the Hub: Predictor("hf:u
 clf.predict(["I still haven't received my new card"])  # a list of label names
 ```
 
-Or try it without installing anything: the [Colab notebook](https://colab.research.google.com/github/Andreas-Hum/tinydistill/blob/main/notebooks/quickstart.ipynb) distils a sentiment classifier in a
-few minutes on a free GPU.
+Or try it without installing anything, on a free Colab GPU:
+
+| Notebook | What it shows |
+|---|---|
+| [Quickstart](https://colab.research.google.com/github/Andreas-Hum/tinydistill/blob/main/notebooks/quickstart.ipynb) | distil a sentiment classifier in about ten minutes, with charts of what the student learned |
+| [Your own task](https://colab.research.google.com/github/Andreas-Hum/tinydistill/blob/main/notebooks/your_task.ipynb) | write a task for your labels, check the teacher, distil on your texts, find the mistakes, download the model |
+| [Deploy](https://colab.research.google.com/github/Andreas-Hum/tinydistill/blob/main/notebooks/deploy.ipynb) | the model from Python, behind an HTTP server, in a web page, and on the Hugging Face Hub |
 
 Labelling the full dataset with the default teacher wants a GPU. For a quick look on a CPU, add
 `--teacher Qwen/Qwen2.5-0.5B-Instruct --limit 200 --test-limit 200`. `--limit` takes the first rows, and the
@@ -263,6 +282,30 @@ results show the mean and standard deviation.
   adds them to the training texts and trains, exports and evaluates a new student on the same test set, so
   the student learns the inputs it found hard in production.
 
+### Charts and metrics
+
+`tinydistill plot RUN` (with the `plot` extra) writes the charts of a run to `RUN/plots/`; in Python every
+chart is a function in `tinydistill.plots` that returns a matplotlib figure, so it shows inline in a notebook:
+
+| Chart | What it shows |
+|---|---|
+| `soft_labels` | for a few test texts, the teacher's probabilities (the soft labels) next to the student's |
+| `speed_accuracy` | accuracy against latency for the teacher and every form of the student, with the baseline and upper bound |
+| `training` | the loss per epoch on the training texts and, with `--val-share`, on the held-out texts and the kept epoch |
+| `per_label` | F1 of every label for teacher and student, the teacher's weakest first |
+| `confusions` | the confusion matrix over the labels involved in the most mistakes |
+| `cascade` | accuracy as the student passes more of its least confident texts to the teacher |
+
+<p align="center">
+  <img alt="Soft labels on three Banking77 test texts: the teacher's probabilities next to the student's" src="https://raw.githubusercontent.com/Andreas-Hum/tinydistill/main/docs/images/banking77-soft-labels.png" width="49%">
+  <img alt="Accuracy against latency on Banking77: the students on a CPU, the teacher on a GPU, the baseline and the upper bound" src="https://raw.githubusercontent.com/Andreas-Hum/tinydistill/main/docs/images/banking77-speed-accuracy.png" width="49%">
+</p>
+
+Every run's results report accuracy and macro-F1, agreement with the teacher, latency at batch 1 and 32,
+throughput and size; tasks with more than five labels also get top-3 accuracy, multi-label tasks micro-F1, and
+ordered labels the mean distance in label steps. `tinydistill analyse RUN` adds a page of per-label scores and
+example mistakes.
+
 ### From Python
 
 The exported model:

{tinydistill-0.3.2 → tinydistill-0.3.3}/examples/banking77/run.sh

@@ -47,4 +47,11 @@ tinydistill compare runs/banking77-1.5b runs/banking77-1.5b-8shot runs/banking77
 mkdir -p examples/banking77/students
 for dir in runs/banking77-students/*/; do cp "$dir/results.json" "examples/banking77/students/$(basename "$dir").json"; done
 tinydistill compare runs/banking77-7b-8shot runs/banking77-students/* > examples/banking77/students.md
+python - <<'PY'  # the charts in the README
+from tinydistill import plots
+plots.speed_accuracy("runs/banking77-7b-8shot").savefig("docs/images/banking77-speed-accuracy.png", dpi=130)
+plots.soft_labels("runs/banking77-students/all-MiniLM-L6-v2-20-epochs").savefig(
+    "docs/images/banking77-soft-labels.png", dpi=130
+)
+PY
 python examples/update_readme.py

{tinydistill-0.3.2 → tinydistill-0.3.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "tinydistill"
-version = "0.3.2"
+version = "0.3.3"
 description = "Distil an LLM classification prompt into a small, fast model that runs on a CPU."
 readme = "README.md"
 license = "MIT"
@@ -34,7 +34,8 @@ dependencies = [
 [project.optional-dependencies]
 hf = ["datasets>=2.18"]
 4bit = ["bitsandbytes>=0.43", "accelerate>=0.30"]
-dev = ["pytest>=8", "ruff>=0.5", "datasets>=2.18", "build>=1.2"]
+plot = ["matplotlib>=3.7"]
+dev = ["pytest>=8", "ruff>=0.5", "datasets>=2.18", "build>=1.2", "matplotlib>=3.7"]
 
 [project.scripts]
 tinydistill = "tinydistill.cli:main"

{tinydistill-0.3.2 → tinydistill-0.3.3}/src/tinydistill/__init__.py

@@ -29,4 +29,4 @@ __all__ = [
     "read_examples",
     "train_student",
 ]
-__version__ = "0.3.2"
+__version__ = "0.3.3"

{tinydistill-0.3.2 → tinydistill-0.3.3}/src/tinydistill/analyse.py

@@ -53,30 +53,42 @@ CSS = """
 """
 
 
-def build(run: Path, top: int = 15) -> tuple[str, str]:
-    """(HTML page, short Markdown summary) for a single-label run with a gold-labelled test set."""
+def predictions(run: Path) -> dict:
+    """For the gold-labelled test texts of a single-label run: the task, texts, gold labels, and the teacher's
+    and the deployed student's predicted labels."""
     from .export import Predictor
     from .task import load_task
 
     task = load_task(run / "task.yaml")
     if task.multi_label:
-        raise SystemExit("analyse is for single-label tasks; see results.md for multi-label runs")
+        raise SystemExit("this view is for single-label tasks; see results.md for multi-label runs")
     rows = [
         r
         for r in map(json.loads, (run / "test.jsonl").read_text(encoding="utf-8").splitlines())
         if r.get("gold") is not None
     ]
     if not rows:
-        raise SystemExit(f"{run}/test.jsonl has no gold labels to analyse")
-    names = task.labels
+        raise SystemExit(f"{run}/test.jsonl has no gold labels")
     texts = [r["text"] for r in rows]
-    gold = [r["gold"] for r in rows]
-    teacher = [max(range(len(names)), key=r["scores"].__getitem__) for r in rows]
-    student = [int(i) for i in Predictor(run / "onnx").logits(texts).argmax(-1)]
+    n = len(task.labels)
+    return {
+        "task": task,
+        "texts": texts,
+        "gold": [r["gold"] for r in rows],
+        "teacher": [max(range(n), key=r["scores"].__getitem__) for r in rows],
+        "student": [int(i) for i in Predictor(run / "onnx").logits(texts).argmax(-1)],
+    }
+
+
+def build(run: Path, top: int = 15) -> tuple[str, str]:
+    """(HTML page, short Markdown summary) for a single-label run with a gold-labelled test set."""
+    p = predictions(run)
+    task, texts, gold, teacher, student = p["task"], p["texts"], p["gold"], p["teacher"], p["student"]
+    names = task.labels
 
     t_lab, s_lab = per_label(teacher, gold, len(names)), per_label(student, gold, len(names))
     order = sorted(range(len(names)), key=lambda c: t_lab[c]["f1"])  # the teacher's weakest labels first
-    disagree = [i for i in range(len(rows)) if teacher[i] != student[i]]
+    disagree = [i for i in range(len(texts)) if teacher[i] != student[i]]
     t_acc = sum(t == g for t, g in zip(teacher, gold, strict=True)) / len(gold)
     s_acc = sum(s == g for s, g in zip(student, gold, strict=True)) / len(gold)
     e = html.escape
@@ -115,7 +127,7 @@ def build(run: Path, top: int = 15) -> tuple[str, str]:
 <title>Error analysis: {e(task.name)}</title><style>{CSS}</style></head>
 <body><main>
 <h1>Error analysis: {e(task.name)}</h1>
-<p class="sub">{len(rows)} gold-labelled test texts. Teacher accuracy {t_acc:.3f}, student (deployed model)
+<p class="sub">{len(texts)} gold-labelled test texts. Teacher accuracy {t_acc:.3f}, student (deployed model)
 accuracy {s_acc:.3f}; they disagree on {len(disagree)} texts. Written by <code>tinydistill analyse</code>.</p>
 
 <h2>Per label, the teacher's weakest first</h2>
@@ -141,7 +153,7 @@ on the student side.</p>
     worst = [names[c] for c in order[:5]]
     pairs = Counter({(names[g], names[p]): n for g, p, n, _ in confusions(teacher, gold, 5)})
     summary = (
-        f"Teacher accuracy {t_acc:.3f}, student {s_acc:.3f}; they disagree on {len(disagree)} of {len(rows)} texts.\n"
+        f"Teacher accuracy {t_acc:.3f}, student {s_acc:.3f}; they disagree on {len(disagree)} of {len(texts)} texts.\n"
         f"The teacher's weakest labels: {', '.join(worst)}.\n"
         "Its most common mistakes: "
         + "; ".join(f"{g} -> {p} ({n})" for (g, p), n in pairs.most_common(5))

{tinydistill-0.3.2 → tinydistill-0.3.3}/src/tinydistill/cli.py

@@ -358,6 +358,7 @@ def cmd_eval(a: argparse.Namespace) -> None:
         ordinal_metrics,
         render,
         time_calls,
+        top_k_accuracy,
     )
     from .export import Predictor
     from .student import load_student, student_logits
@@ -386,7 +387,9 @@ def cmd_eval(a: argparse.Namespace) -> None:
         raise SystemExit("no texts to evaluate on: label a --test set (or training data) first")
     threads = torch.get_num_threads()
 
-    def quality(pred: list) -> dict:
+    def quality(pred: list, scores: list | None = None) -> dict:
+        """Metrics of predictions; with each text's label scores, also top-3 accuracy (tasks with more than
+        five labels)."""
         if not gold:
             return {}
         if multi:
@@ -394,6 +397,8 @@ def cmd_eval(a: argparse.Namespace) -> None:
                 "agreement": accuracy(pred, teacher_pred)
             }
         extra = ordinal_metrics(pred, gold) if task.ordinal else {}
+        if scores is not None and len(task.labels) > 5:
+            extra["top3_accuracy"] = top_k_accuracy(scores, gold)
         return extra | {
             "accuracy": accuracy(pred, gold),
             "macro_f1": macro_f1(pred, gold),
@@ -418,8 +423,8 @@ def cmd_eval(a: argparse.Namespace) -> None:
     examples = [(r["text"], r["gold"]) for r in _read_jsonl(run / "examples.jsonl")] or None
     members: list[dict] = []
     for ti, name in enumerate(teachers):
-        own = [decide(r["teacher_scores"][ti]) for r in labelled] if len(teachers) > 1 else teacher_pred
-        q = quality(own)
+        own_scores = [r["teacher_scores"][ti] if len(teachers) > 1 else r["scores"] for r in labelled]
+        q = quality([decide(s) for s in own_scores], own_scores)
         q.pop("agreement", None)
         for dev in devices:
             label = f"teacher, {prompt}" + (f": {name.split('/')[-1]}" if len(teachers) > 1 else "")
@@ -463,7 +468,7 @@ def cmd_eval(a: argparse.Namespace) -> None:
         models.extend(members)
     else:  # the ensemble is the actual teacher: it runs every member, so its latency is their sum
         models.extend(m | {"role": "teacher member"} for m in members)
-        q = quality(teacher_pred)
+        q = quality(teacher_pred, [r["scores"] for r in labelled])
         q.pop("agreement", None)
         for dev in devices:
             on_dev = [m for m in members if m["device"] == dev]
@@ -493,9 +498,8 @@ def cmd_eval(a: argparse.Namespace) -> None:
         per_seed = []
         for d in dirs:
             model, tok = load_student(d, "cpu")
-            per_seed.append(
-                quality([decide(z) for z in student_logits(model, tok, texts).tolist()] if texts else [])
-            )
+            z = student_logits(model, tok, texts).tolist() if texts else []
+            per_seed.append(quality([decide(x) for x in z], z))
             bar.update()
         row = {"role": role, "name": row_name, "device": "cpu", **aggregate(per_seed)}
         if role == "student":  # timings and size of the shipped (first-seed) model
@@ -517,15 +521,19 @@ def cmd_eval(a: argparse.Namespace) -> None:
             continue
         log(f"evaluating {f.name} (onnxruntime, CPU)")
         pr = Predictor(run / "onnx", int8=int8, threads=threads)
-        pred = [decide(z) for z in pr.logits(texts).tolist()] if texts else []
+        z = pr.logits(texts).tolist() if texts else []
+        pred = [decide(x) for x in z]
         row = {
             "role": "student",
             "name": f"student, ONNX {'int8' if int8 else 'fp32'}",
             "device": "cpu",
-            **quality(pred),
+            **quality(pred, z),
             "size_mb": dir_size_mb([f]),
         }
-        row.update(bench(lambda b, pr=pr: pr.logits(b), bench_texts, a.reps, max(3, a.reps // 5)))
+        # timed on real batches (batch 32 = one onnxruntime call with 32 texts)
+        row.update(
+            bench(lambda b, pr=pr: pr.logits(b, batch_size=len(b)), bench_texts, a.reps, max(3, a.reps // 5))
+        )
         models.append(row)
         bar.update()
 
@@ -680,6 +688,16 @@ def cmd_serve(a: argparse.Namespace) -> None:
         server.server_close()
 
 
+def cmd_plot(a: argparse.Namespace) -> None:
+    from .plots import save_all
+
+    files = save_all(Path(a.run), a.out)
+    if not files:
+        raise SystemExit(f"{a.run} has no data to plot: run `tinydistill eval` first")
+    for f in files:
+        print(f)
+
+
 def cmd_web(a: argparse.Namespace) -> None:
     from .export import model_path
     from .web import build_site
@@ -1044,6 +1062,13 @@ def build_parser() -> argparse.ArgumentParser:
     srv.add_argument("--threshold", type=float, default=0.5, help="with --cascade: see `predict --help`")
     srv.add_argument("--log", help="with --cascade: append the texts the teacher answered (for `refresh`)")
     srv.set_defaults(fn=cmd_serve)
+    plo = sub.add_parser(
+        "plot",
+        help="charts of a run as PNG files: training, speed and accuracy, per label, confusions, cascade",
+    )
+    plo.add_argument("run", help="run directory")
+    plo.add_argument("--out", help="folder for the PNG files (default: RUN/plots)")
+    plo.set_defaults(fn=cmd_plot)
     web = sub.add_parser("web", help="a static web page that runs the model in the browser")
     web.add_argument("model", help="directory written by `export` (run/onnx), or hf:user/name on the Hub")
     web.add_argument("--out", default="site", help="folder to write the page and model into")

{tinydistill-0.3.2 → tinydistill-0.3.3}/src/tinydistill/evaluate.py

@@ -87,6 +87,12 @@ def bench(fn: Callable[[list[str]], object], texts: list[str], reps1: int, reps3
     return out
 
 
+def top_k_accuracy(scores: list[list[float]], gold: list[int], k: int = 3) -> float:
+    """Share of texts whose gold label is among the k highest-scoring labels."""
+    hits = sum(g in sorted(range(len(s)), key=s.__getitem__)[-k:] for s, g in zip(scores, gold, strict=True))
+    return hits / len(gold) if gold else float("nan")
+
+
 def aggregate(per_seed: list[dict]) -> dict:
     """Mean of each metric over seeds; with several seeds also the sample standard deviation and values."""
     if not per_seed or not per_seed[0]:
@@ -141,7 +147,12 @@ def _pm(row: dict, key: str) -> str:
     return _fmt(row.get(key), ".3f") + (f" ± {std:.3f}" if std is not None else "")
 
 
-EXTRA_COLUMNS = (("micro_f1", "Micro-F1"), ("mae", "MAE (label steps)"), ("within_one", "Within one step"))
+EXTRA_COLUMNS = (
+    ("top3_accuracy", "Top-3 accuracy"),
+    ("micro_f1", "Micro-F1"),
+    ("mae", "MAE (label steps)"),
+    ("within_one", "Within one step"),
+)
 
 
 def markdown_table(results: dict) -> str:

{tinydistill-0.3.2 → tinydistill-0.3.3}/src/tinydistill/export.py

@@ -178,8 +178,13 @@ class Predictor:
         self.multi_label = config.get("problem_type") == "multi_label_classification"
         self.max_length = max_length
         self.file = file
+        self.int8 = int8
 
-    def logits(self, texts: list[str], batch_size: int = 32) -> np.ndarray:
+    def logits(self, texts: list[str], batch_size: int | None = None) -> np.ndarray:
+        """Label scores per text. The int8 model runs one text at a time by default: its dynamic quantisation
+        scales activations over the whole batch, so in a batch a text's scores would depend on the other texts
+        (on a CPU this costs little throughput). The fp32 model runs 32 texts at a time."""
+        batch_size = batch_size or (1 if self.int8 else 32)
         out = []
         for b in range(0, len(texts), batch_size):
             enc = self.tok(
@@ -196,7 +201,7 @@ class Predictor:
             out.append(self.session.run(["logits"], feed)[0])
         return np.concatenate(out) if out else np.zeros((0, len(self.labels)), dtype=np.float32)
 
-    def predict_proba(self, texts: list[str], batch_size: int = 32) -> np.ndarray:
+    def predict_proba(self, texts: list[str], batch_size: int | None = None) -> np.ndarray:
         """Label probabilities per text: a softmax, or for a multi-label model each label's sigmoid."""
         z = self.logits(texts, batch_size)
         if self.multi_label:
@@ -204,7 +209,7 @@ class Predictor:
         z = np.exp(z - z.max(-1, keepdims=True))
         return z / z.sum(-1, keepdims=True)
 
-    def predict(self, texts: list[str], batch_size: int = 32) -> list:
+    def predict(self, texts: list[str], batch_size: int | None = None) -> list:
         """The top label per text, or for a multi-label model the list of labels with probability >= 0.5."""
         z = self.logits(texts, batch_size)
         if self.multi_label:

{tinydistill-0.3.2 → tinydistill-0.3.3}/src/tinydistill/hub.py

@@ -25,6 +25,13 @@ def model_card(onnx_dir: Path, repo: str, license: str | None = None) -> str:
     ]
     if license:
         front.append(f"license: {license}")
+    student_id = str(meta.get("student", ""))
+    if "/" in student_id and not student_id.startswith("static"):  # links the card to the encoder's page
+        front.append(f"base_model: {student_id}")
+    data = [str(meta.get(k) or "") for k in ("data", "test")]
+    datasets = sorted({d.split(":")[1] for d in data if d.startswith("hf:") and len(d.split(":")) > 1})
+    if datasets:  # links the card to the datasets' pages
+        front += ["datasets:", *(f"- {d}" for d in datasets)]
     front.append("---")
     teacher, student = meta.get("teacher", "an LLM"), meta.get("student", "a small encoder")
     lines = [

tinydistill-0.3.3/src/tinydistill/plots.py

@@ -0,0 +1,326 @@
+"""Charts of a run, as matplotlib figures (``pip install "tinydistill[plot]"``). Each function takes a run
+directory and returns a figure, so in a notebook it shows inline; ``save_all`` writes every chart the run has
+data for to ``RUN/plots/*.png`` (that is what `tinydistill plot RUN` does).
+
+The student is blue and the teacher orange throughout, as in the web page."""
+
+from __future__ import annotations
+
+import json
+from collections import Counter
+from pathlib import Path
+
+STUDENT, TEACHER, REFERENCE = "#2a78d6", "#eb6834", "#8a8984"
+SURFACE, INK, MUTED, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#e4e3df"
+
+
+def _plt():
+    try:
+        import matplotlib
+
+        if not _in_notebook():
+            matplotlib.use("Agg")  # files only, no window
+        import matplotlib.pyplot as plt
+    except ImportError as e:
+        raise SystemExit('plots need matplotlib: pip install "tinydistill[plot]"') from e
+    plt.rcParams.update(
+        {
+            "figure.facecolor": SURFACE,
+            "axes.facecolor": SURFACE,
+            "axes.edgecolor": GRID,
+            "axes.labelcolor": MUTED,
+            "axes.titlecolor": INK,
+            "axes.titlesize": 12,
+            "axes.titleweight": "bold",
+            "axes.titlelocation": "left",
+            "axes.spines.top": False,
+            "axes.spines.right": False,
+            "axes.grid": True,
+            "grid.color": GRID,
+            "grid.linewidth": 0.8,
+            "xtick.color": MUTED,
+            "ytick.color": MUTED,
+            "text.color": INK,
+            "legend.frameon": False,
+            "font.size": 10,
+        }
+    )
+    return plt
+
+
+def _in_notebook() -> bool:
+    try:
+        from IPython import get_ipython
+
+        return get_ipython() is not None
+    except ImportError:
+        return False
+
+
+def _results(run: Path) -> dict:
+    f = Path(run) / "results.json"
+    if not f.exists():
+        raise FileNotFoundError(f"{f} is missing: run `tinydistill eval` first")
+    return json.loads(f.read_text(encoding="utf-8"))
+
+
+def _name(label: str) -> str:
+    return label.replace("_", " ")
+
+
+def training(run: str | Path):
+    """Training loss per epoch, and the held-out loss and the kept epoch when `--val-share` was used."""
+    plt = _plt()
+    f = Path(run) / "student" / "history.json"
+    if not f.exists():
+        raise FileNotFoundError(f"{f} is missing: train with tinydistill 0.3.3 or later")
+    h = json.loads(f.read_text(encoding="utf-8"))
+    epochs = [e["epoch"] for e in h["epochs"]]
+    fig, ax = plt.subplots(figsize=(7, 3.6))
+    held = [e.get("held_out_loss") for e in h["epochs"]]
+    has_held = any(v is not None for v in held)
+    ax.plot(  # both lines are the student; the held-out loss decides, so it gets the student's colour
+        epochs,
+        [e["loss"] for e in h["epochs"]],
+        color=REFERENCE if has_held else STUDENT,
+        lw=2,
+        marker="o",
+        ms=5,
+        label="training texts",
+    )
+    if has_held:
+        ax.plot(
+            epochs, held, color=STUDENT, lw=2, marker="o", ms=5, label=f"{h['held_out_texts']} held-out texts"
+        )
+        ax.axvline(h["kept_epoch"], color=REFERENCE, lw=1, ls="--")
+        ax.annotate(
+            f" kept: epoch {h['kept_epoch']}", (h["kept_epoch"], ax.get_ylim()[1]), va="top", color=MUTED
+        )
+    ax.set(xlabel="epoch", ylabel="loss (lower is better)", title="Training")
+    ax.set_xticks(epochs)
+    ax.legend(loc="upper right")
+    fig.tight_layout()
+    return fig
+
+
+def per_label(run: str | Path, limit: int | None = None):
+    """F1 of every label for the teacher and the student, the teacher's weakest labels at the top. With
+    ``limit``, only that many of the labels where the two differ most."""
+    from .analyse import per_label as scores
+    from .analyse import predictions
+
+    plt = _plt()
+    p = predictions(Path(run))
+    names = p["task"].labels
+    t = [s["f1"] for s in scores(p["teacher"], p["gold"], len(names))]
+    s = [x["f1"] for x in scores(p["student"], p["gold"], len(names))]
+    order = sorted(range(len(names)), key=lambda c: t[c])
+    if limit:
+        order = sorted(sorted(order, key=lambda c: -abs(s[c] - t[c]))[:limit], key=lambda c: t[c])
+    fig, ax = plt.subplots(figsize=(7, 0.24 * len(order) + 1.4))
+    y = range(len(order))
+    ax.hlines(y, [min(t[c], s[c]) for c in order], [max(t[c], s[c]) for c in order], color=GRID, lw=2)
+    ax.scatter([t[c] for c in order], y, color=TEACHER, s=36, zorder=3, label="teacher")
+    ax.scatter([s[c] for c in order], y, color=STUDENT, s=36, zorder=3, label="student")
+    ax.set_yticks(list(y), [_name(names[c]) for c in order], fontsize=8)
+    ax.set_ylim(len(order) - 0.5, -0.5)
+    ax.set_xlim(0, 1.02)
+    ax.grid(axis="y", visible=False)
+    ax.set(xlabel="F1 on the test texts", title="Per label: teacher and student")
+    ax.legend(loc="lower right", bbox_to_anchor=(1, 1), ncols=2, borderaxespad=0.2)
+    fig.tight_layout()
+    return fig
+
+
+def confusions(run: str | Path, who: str = "student", labels: int = 12):
+    """Confusion matrix of the student (or ``who="teacher"``) over the ``labels`` labels involved in the most
+    mistakes: rows are the gold labels, columns the predicted ones."""
+    from .analyse import predictions
+
+    plt = _plt()
+    p = predictions(Path(run))
+    names, gold, pred = p["task"].labels, p["gold"], p[who]
+    errors = Counter()
+    for g, x in zip(gold, pred, strict=True):
+        if g != x:
+            errors[g] += 1
+            errors[x] += 1
+    keep = sorted(c for c, _ in errors.most_common(labels)) or list(range(min(labels, len(names))))
+    index = {c: i for i, c in enumerate(keep)}
+    m = [[0] * len(keep) for _ in keep]
+    for g, x in zip(gold, pred, strict=True):
+        if g in index and x in index:
+            m[index[g]][index[x]] += 1
+    fig, ax = plt.subplots(figsize=(0.5 * len(keep) + 3, 0.45 * len(keep) + 2))
+    ax.imshow(m, cmap="Blues")
+    top = max(max(r) for r in m) or 1
+    for i, r in enumerate(m):
+        for j, v in enumerate(r):
+            if v:
+                ax.text(
+                    j, i, v, ha="center", va="center", fontsize=8, color="#ffffff" if v > top / 2 else INK
+                )
+    ticks = [_name(names[c]) for c in keep]
+    ax.set_xticks(range(len(keep)), ticks, rotation=45, ha="right", fontsize=8)
+    ax.set_yticks(range(len(keep)), ticks, fontsize=8)
+    ax.grid(False)
+    ax.set(xlabel="predicted", ylabel="gold", title=f"Most confused labels ({who})")
+    fig.tight_layout()
+    return fig
+
+
+def speed_accuracy(run: str | Path):
+    """Accuracy against latency per text (batch 1, log scale) for every timed model, with the upper bound and
+    the gold-only baseline as reference lines when the run has them."""
+    plt = _plt()
+    rows = _results(Path(run))["models"]
+    timed = [r for r in rows if r.get("latency_ms_b1") and r.get("accuracy") is not None]
+    if not timed:
+        raise ValueError("no timed model with an accuracy in results.json")
+    fig, ax = plt.subplots(figsize=(7, 4))
+    students = [r for r in timed if not r["role"].startswith("teacher")]
+    deployed = min(students, key=lambda r: r["latency_ms_b1"]) if students else None  # e.g. ONNX int8
+    for r in timed:
+        color = TEACHER if r["role"].startswith("teacher") else STUDENT
+        ax.scatter(r["latency_ms_b1"], r["accuracy"], color=color, s=60, zorder=3, edgecolor=SURFACE, lw=1.5)
+        if color == STUDENT and r is not deployed:
+            continue  # label only the teacher and the fastest student, so labels never collide
+        label = r["name"].replace("teacher, ", "teacher ").replace("student, ", "student ")
+        ax.annotate(
+            f"{label} ({r['device']})",
+            (r["latency_ms_b1"], r["accuracy"]),
+            textcoords="offset points",
+            xytext=(6, 4),
+            fontsize=8,
+            color=MUTED,
+        )
+    for r in rows:
+        if r["role"] in ("upper bound", "baseline") and r.get("accuracy") is not None:
+            ls = "--" if r["role"] == "upper bound" else ":"
+            ax.axhline(r["accuracy"], color=REFERENCE, lw=1, ls=ls)
+            ax.annotate(
+                r["role"],
+                (1, r["accuracy"]),
+                xycoords=("axes fraction", "data"),
+                ha="right",
+                va="bottom",
+                fontsize=8,
+                color=MUTED,
+            )
+    from matplotlib.ticker import FuncFormatter
+
+    ax.set_xscale("log")
+    ax.margins(x=0.3)  # room for the labels on both sides
+    ax.xaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:g}"))
+    ax.xaxis.set_minor_formatter(FuncFormatter(lambda v, _: f"{v:g}" if f"{v:g}"[0] in "25" else ""))
+    accs = [r["accuracy"] for r in rows if r.get("accuracy") is not None]
+    ax.set_ylim(max(0.0, min(accs) - 0.1), min(1.0, max(accs) + 0.05))  # never zoom in on a tiny gap
+    ax.set(
+        xlabel="latency per text, batch 1 (ms, log scale)", ylabel="test accuracy", title="Accuracy and speed"
+    )
+    roles = {r["role"].startswith("teacher") for r in timed}
+    for is_teacher, color, label in ((True, TEACHER, "teacher"), (False, STUDENT, "student")):
+        if is_teacher in roles:
+            ax.scatter([], [], color=color, label=label)
+    ax.legend(loc="lower left")
+    fig.tight_layout()
+    return fig
+
+
+def cascade(run: str | Path):
+    """Accuracy as the student passes its least confident texts to the teacher."""
+    plt = _plt()
+    points = _results(Path(run)).get("cascade")
+    if not points:
+        raise ValueError("results.json has no cascade (single-label runs with a gold-labelled test set)")
+    x = [100 * c["share_to_teacher"] for c in points]
+    y = [c["accuracy"] for c in points]
+    fig, ax = plt.subplots(figsize=(7, 3.6))
+    ax.plot(x, y, color=STUDENT, lw=2, marker="o", ms=6)
+    for i, text, ha in ((0, "student only", "left"), (-1, "teacher only", "right")):
+        ax.annotate(
+            text, (x[i], y[i]), textcoords="offset points", xytext=(0, -16), ha=ha, fontsize=8, color=MUTED
+        )
+    ax.margins(y=0.15)
+    ax.set(
+        xlabel="texts sent to the teacher (%)",
+        ylabel="test accuracy",
+        title="Cascade: unsure texts go to the teacher",
+    )
+    fig.tight_layout()
+    return fig
+
+
+def soft_labels(run: str | Path, n: int = 3, k: int = 4):
+    """For ``n`` test texts with different gold labels: the teacher's probabilities (the soft labels the
+    student learned from) next to the student's, over the ``k`` most likely labels."""
+    import math
+
+    from .export import Predictor
+    from .task import load_task
+    from .web import _samples
+
+    plt = _plt()
+    run = Path(run)
+    task = load_task(run / "task.yaml")
+    if task.multi_label:
+        raise ValueError("soft_labels is for single-label tasks")
+    rows = [
+        json.loads(line) for line in (run / "test.jsonl").read_text(encoding="utf-8").splitlines() if line
+    ]
+    rows = _samples([r for r in rows if r.get("scores")], n)
+    if not rows:
+        raise ValueError("the run has no test texts with teacher scores")
+    student = Predictor(run / "onnx").predict_proba([r["text"] for r in rows])
+    fig, axes = plt.subplots(
+        len(rows), 1, figsize=(7, 1.55 * len(rows) * max(1, k / 4) + 0.6), squeeze=False, sharex=True
+    )
+    for ax, r, q in zip(axes[:, 0], rows, student, strict=True):
+        m = max(r["scores"])
+        e = [math.exp(x - m) for x in r["scores"]]
+        p = [x / sum(e) for x in e]
+        top = sorted(range(len(p)), key=lambda c: -(p[c] + q[c]))[:k]
+        y = range(len(top))
+        ax.barh([i - 0.2 for i in y], [p[c] for c in top], height=0.38, color=TEACHER, label="teacher")
+        ax.barh([i + 0.2 for i in y], [q[c] for c in top], height=0.38, color=STUDENT, label="student")
+        gold = r.get("gold")
+        ax.set_yticks(
+            list(y), [_name(task.labels[c]) + ("  (gold)" if c == gold else "") for c in top], fontsize=8
+        )
+        ax.set_ylim(len(top) - 0.5, -0.5)
+        ax.set_xlim(0, 1)
+        ax.grid(axis="y", visible=False)
+        text = r["text"] if len(r["text"]) <= 90 else r["text"][:87] + "..."
+        ax.set_title(f"\u201c{text}\u201d", fontsize=9, fontweight="normal", color=MUTED)
+    axes[0, 0].legend(loc="lower right", bbox_to_anchor=(1, 1.18), ncols=2, borderaxespad=0)
+    axes[-1, 0].set_xlabel("probability")
+    fig.suptitle("Soft labels: teacher and student", x=0.02, ha="left", fontweight="bold", fontsize=12)
+    fig.tight_layout()
+    return fig
+
+
+CHARTS = {
+    "training": training,
+    "speed_accuracy": speed_accuracy,
+    "per_label": per_label,
+    "confusions": confusions,
+    "cascade": cascade,
+    "soft_labels": soft_labels,
+}
+
+
+def save_all(run: str | Path, out: str | Path | None = None) -> list[Path]:
+    """Write every chart the run has data for to ``out`` (default ``RUN/plots``); returns the files."""
+    plt = _plt()
+    out = Path(out) if out else Path(run) / "plots"
+    out.mkdir(parents=True, exist_ok=True)
+    written = []
+    for name, chart in CHARTS.items():
+        try:
+            fig = chart(run)
+        except (FileNotFoundError, ValueError, SystemExit):
+            continue  # this run has no data for the chart (e.g. multi-label, or no cascade)
+        fig.savefig(out / f"{name}.png", dpi=150)
+        plt.close(fig)
+        written.append(out / f"{name}.png")
+    return written

{tinydistill-0.3.2 → tinydistill-0.3.3}/src/tinydistill/student.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 import math
 import random
 import time
@@ -117,15 +118,23 @@ def train_student(
     if static := encoder_of(model_name):
         model, tok = new_static(static, labels, multi_label)
     else:
+        from transformers.utils import logging as hf_logging
+
         tok = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForSequenceClassification.from_pretrained(
-            model_name,
-            num_labels=len(labels),
-            id2label=dict(enumerate(labels)),
-            label2id={name: i for i, name in enumerate(labels)},
-            ignore_mismatched_sizes=True,
-            problem_type="multi_label_classification" if multi_label else "single_label_classification",
-        )
+        level = hf_logging.get_verbosity()
+        hf_logging.set_verbosity_error()  # its report of the new (missing) classifier head is expected noise
+        try:
+            model = AutoModelForSequenceClassification.from_pretrained(
+                model_name,
+                num_labels=len(labels),
+                id2label=dict(enumerate(labels)),
+                label2id={name: i for i, name in enumerate(labels)},
+                ignore_mismatched_sizes=True,
+                problem_type="multi_label_classification" if multi_label else "single_label_classification",
+            )
+        finally:
+            hf_logging.set_verbosity(level)
+        log(f"{model_name}: a new classification head for {len(labels)} labels on the pretrained encoder")
     model = model.to(device)
     gold = gold or [None] * len(texts)
     if multi_label:  # multi-hot gold, and which texts have it
@@ -169,6 +178,7 @@ def train_student(
     )
 
     best_loss, best_state, best_epoch = math.inf, None, epochs
+    history: list[dict] = []  # per epoch, saved as history.json next to the model
     for epoch in range(epochs):
         model.train()
         rng.shuffle(order)
@@ -185,6 +195,7 @@ def train_student(
             if on_step:
                 on_step()
         line = f"epoch {epoch + 1}/{epochs}: loss {total / len(order):.4f}"
+        history.append({"epoch": epoch + 1, "loss": total / len(order)})
         if val:
             model.eval()
             with torch.no_grad():
@@ -193,6 +204,7 @@ def train_student(
                 )
             v /= len(val)
             line += f", held-out loss {v:.4f}"
+            history[-1]["held_out_loss"] = v
             if v < best_loss:
                 best_loss, best_epoch = v, epoch + 1
                 best_state = {k: t.detach().clone() for k, t in model.state_dict().items()}
@@ -206,7 +218,12 @@ def train_student(
     out_dir.mkdir(parents=True, exist_ok=True)
     model.save_pretrained(out_dir)
     tok.save_pretrained(out_dir)
-    return best_epoch
+    kept = best_epoch if best_state is not None else len(history)
+    (out_dir / "history.json").write_text(
+        json.dumps({"epochs": history, "kept_epoch": kept, "held_out_texts": len(val)}, indent=1),
+        encoding="utf-8",
+    )
+    return kept
 
 
 def load_student(path: Path, device: str = "cpu"):

{tinydistill-0.3.2 → tinydistill-0.3.3}/tests/test_cli.py

@@ -293,11 +293,14 @@ def test_push_writes_a_model_card_and_uploads(tmp_path, monkeypatch):
     onnx = tmp_path / "run" / "onnx"
     onnx.mkdir(parents=True)
     (onnx / "labels.json").write_text('["neg", "pos"]')
-    (tmp_path / "run" / "meta.json").write_text('{"teacher": "T", "student": "S"}')
+    (tmp_path / "run" / "meta.json").write_text(
+        '{"teacher": "T", "student": "org/S", "data": "hf:org/ds:train", "test": "hf:org/ds:test"}'
+    )
     (tmp_path / "run" / "results.md").write_text("| Model |\n|---|\n| x |\n")
     main(["push", str(onnx), "--repo", "me/clf", "--private", "--license", "apache-2.0"])
     card = (onnx / "README.md").read_text(encoding="utf-8")
     assert "license: apache-2.0" in card and "`T`" in card and "## Results" in card and '"hf:me/clf"' in card
+    assert "base_model: org/S" in card and "datasets:\n- org/ds\n---" in card
     assert calls == [("create", "me/clf", True), ("upload", str(onnx), "me/clf")]
 
 
@@ -959,3 +962,17 @@ def test_train_keeps_the_best_epoch_on_held_out_texts(tmp_path):
     assert "holds out 20% of its training texts and keeps the epoch" in (run / "results.md").read_text(
         encoding="utf-8"
     )
+
+    history = json.loads((run / "student" / "history.json").read_text(encoding="utf-8"))
+    assert (
+        history["kept_epoch"] == meta["kept_epochs"]["student"][0] and "held_out_loss" in history["epochs"][0]
+    )
+    main(["plot", str(run)])
+    assert {p.name for p in (run / "plots").glob("*.png")} == {
+        "training.png",
+        "speed_accuracy.png",
+        "per_label.png",
+        "confusions.png",
+        "cascade.png",
+        "soft_labels.png",
+    }

{tinydistill-0.3.2 → tinydistill-0.3.3}/tests/test_export.py

@@ -85,3 +85,19 @@ def test_explain_export_matches_the_model(tmp_path):
     np.testing.assert_allclose(attention[0, :, : real[0]].sum(-1), 1, atol=1e-4)  # a distribution over tokens
     assert np.allclose(attention[0, :, real[0] :], 0, atol=1e-6)  # padding gets none
     assert export_explain(tmp_path / "student", tmp_path / "q", int8=True).name == "explain.int8.onnx"
+
+
+def test_int8_scores_do_not_depend_on_the_batch(tmp_path):
+    torch.manual_seed(0)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        TINY_ENCODER, num_labels=3, id2label={0: "a", 1: "b", 2: "c"}, ignore_mismatched_sizes=True
+    )
+    model.save_pretrained(tmp_path / "student")
+    AutoTokenizer.from_pretrained(TINY_ENCODER).save_pretrained(tmp_path / "student")
+    export_onnx(tmp_path / "student", tmp_path / "onnx", int8=True)
+    p = Predictor(tmp_path / "onnx")
+    assert p.int8
+    text = "the same text every time"
+    alone = p.logits([text])[0]
+    for others in (["short"], ["a much longer text with many more words in it than the first one", "x"]):
+        np.testing.assert_allclose(p.logits([text, *others])[0], alone, rtol=0, atol=1e-6)

{tinydistill-0.3.2 → tinydistill-0.3.3}/tests/test_loss.py

@@ -175,7 +175,8 @@ def test_best_epoch_is_kept_and_training_stops_early(tmp_path):
         patience=1,
     )  # nothing is learned, so the held-out loss never improves after epoch 1
     assert kept == 1
-    assert sum(line.startswith("epoch ") for line in lines) == 2 and "held-out loss" in lines[0]
+    epochs = [line for line in lines if line.startswith("epoch ")]
+    assert len(epochs) == 2 and "held-out loss" in epochs[0]
     assert any("keeping epoch 1" in line for line in lines) and (tmp_path / "s" / "config.json").exists()
     assert (
         train_student(texts, probs, ["a", "b"], tmp_path / "t", TINY_ENCODER, epochs=2, log=lines.append) == 2