tinydistill 0.2.0__py3-none-any.whl

tinydistill/teacher.py

@@ -0,0 +1,272 @@
+"""Teacher: score each label's verbalisation under a causal LM, cache the scores on disk.
+
+For a text x and a continuation c_1..c_m (a way of writing a label, followed by the end-of-turn token),
+log p(c | prompt(x)) = sum_j log p(c_j | prompt(x), c_<j). A label can be written in several surface forms
+(lowercase, Capitalised, UPPERCASE, as written; each with and without a leading space), and its score is the
+log of the summed probability of its distinct forms. A softmax over the label scores gives the teacher's soft
+targets. There is no free-form generation.
+
+Speed: the prompt is split around the text into a constant prefix and a per-text suffix. The prefix's KV
+cache is computed once. Per text, one forward pass covers the text and all labels at once: the labels' tokens
+form a prefix tree (labels that start alike share nodes), and a 4D attention mask lets every tree node attend
+to the cached prompt, the text and its own ancestors only. There is no padding and no per-label copy of the
+cache; the cache is cropped back to the prompt prefix afterwards.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Iterable, Iterator
+from pathlib import Path
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from .task import Task
+
+_SENTINEL = "\x00TEXT\x00"
+
+
+def pick_device(device: str) -> str:
+    if device == "auto":
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    return device
+
+
+def text_hash(text: str) -> str:
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()[:20]
+
+
+def surface_variants(name: str, kinds=("lowercase", "capitalised", "uppercase")) -> list[str]:
+    """Common ways of writing a label: as written, plus the chosen kinds of lowercase, Capitalised and
+    UPPERCASE; each with and without a leading space. The caller drops forms that tokenise identically."""
+    low = name.lower()
+    variant = {"lowercase": low, "capitalised": low[:1].upper() + low[1:], "uppercase": name.upper()}
+    cased = [name, *(variant[k] for k in ("lowercase", "capitalised", "uppercase") if k in kinds)]
+    out: list[str] = []
+    for form in cased:
+        for v in (form, " " + form):
+            if v not in out:
+                out.append(v)
+    return out
+
+
+class LabelScorer:
+    def __init__(
+        self, model, tokenizer, task: Task, use_chat_template: bool = True, surface_forms: bool = True
+    ):
+        self.model = model.eval()
+        self.tok = tokenizer
+        self.task = task
+        self.device = next(model.parameters()).device
+        self.chat = use_chat_template and bool(getattr(tokenizer, "chat_template", None))
+
+        # Split the rendered prompt around the text.
+        user = task.render(_SENTINEL)
+        if self.chat:
+            full = tokenizer.apply_chat_template(
+                [{"role": "user", "content": user}], tokenize=False, add_generation_prompt=True
+            )
+        else:
+            full = user
+        prefix, suffix = full.split(_SENTINEL)
+        self.suffix = suffix
+        ids = tokenizer(prefix, add_special_tokens=False)["input_ids"]
+        if not self.chat and tokenizer.bos_token_id is not None:
+            ids = [tokenizer.bos_token_id, *ids]
+        self.prefix_ids: list[int] = ids
+
+        # Continuations: each surface form of each label name + the end-of-turn token. Forms that tokenise
+        # identically are kept once. Without surface_forms only the name as written is scored (with a leading
+        # space when there is no chat template, since the answer then follows the prompt text directly).
+        end = [tokenizer.eos_token_id] if tokenizer.eos_token_id is not None else []
+        self.forms: list[list[int]] = []  # token ids of every continuation
+        self.form_label: list[int] = []  # which label each continuation belongs to
+        seen: dict[tuple[int, ...], int] = {}
+        for li, name in enumerate(task.verbalisations):
+            texts = (
+                surface_variants(name, task.surface_forms)
+                if surface_forms
+                else [name if self.chat else " " + name]
+            )
+            for text in texts:
+                ids = tuple(tokenizer(text, add_special_tokens=False)["input_ids"] + end)
+                if seen.get(ids, li) != li:
+                    other = task.verbalisations[seen[ids]]
+                    raise ValueError(f"labels {other!r} and {name!r} share the surface form {text!r}")
+                if ids not in seen:
+                    seen[ids] = li
+                    self.forms.append(list(ids))
+                    self.form_label.append(li)
+        self._build_tree()
+        self._prefix_cache = None
+
+    def _build_tree(self) -> None:
+        """Prefix tree over the labels' input tokens (every label token except the last).
+
+        Node i stands for one distinct label prefix (c_1..c_d); its input token is c_d and its output predicts
+        c_(d+1). The first label token is predicted by the last text position (row -1 below).
+        """
+        nodes: dict[tuple[int, ...], int] = {}
+        parent: list[int] = []
+        depth: list[int] = []
+        tokens: list[int] = []
+        rows: list[int] = []  # for every label token, which output row predicts it (-1 = last text position)
+        targets: list[int] = []
+        owner: list[int] = []
+        for li, lab in enumerate(self.forms):
+            for j, tok in enumerate(lab):
+                rows.append(nodes[tuple(lab[:j])] if j else -1)
+                targets.append(tok)
+                owner.append(li)
+                if j < len(lab) - 1 and tuple(lab[: j + 1]) not in nodes:
+                    nodes[tuple(lab[: j + 1])] = len(tokens)
+                    parent.append(nodes[tuple(lab[:j])] if j else -1)
+                    depth.append(j)
+                    tokens.append(tok)
+        n = len(tokens)
+        anc = torch.eye(n, dtype=torch.bool)
+        for i in range(n):  # parents always come before children
+            if parent[i] >= 0:
+                anc[i] |= anc[parent[i]]
+        dev = self.device
+        self._node_tokens = torch.tensor(tokens, dtype=torch.long, device=dev)
+        self._node_depth = torch.tensor(depth, dtype=torch.long, device=dev)
+        self._node_anc = anc.to(dev)
+        self._rows = torch.tensor(rows, dtype=torch.long, device=dev)
+        self._targets = torch.tensor(targets, dtype=torch.long, device=dev)
+        self._owner = torch.tensor(owner, dtype=torch.long, device=dev)
+
+    def text_ids(self, text: str) -> list[int]:
+        return self.tok(text + self.suffix, add_special_tokens=False)["input_ids"]
+
+    def fingerprint(self) -> str:
+        """Identifies everything that changes the scores, so caches are never mixed up."""
+        cfg = getattr(self.model, "config", None)
+        key = {
+            "model": getattr(cfg, "_name_or_path", ""),
+            "dtype": str(next(self.model.parameters()).dtype),
+            "quant": str(getattr(cfg, "quantization_config", "")),
+            "prefix": self.prefix_ids,
+            "suffix": self.suffix,
+            "forms": self.forms,
+            "form_label": self.form_label,
+        }
+        return hashlib.sha256(json.dumps(key, sort_keys=True).encode()).hexdigest()[:16]
+
+    def score(self, text: str) -> list[float]:
+        """Every label's score: log of the summed probability of its surface forms given the prompt."""
+        f = self.form_scores(text).double()
+        idx = torch.tensor(self.form_label, device=f.device)
+        n = len(self.task.labels)
+        top = torch.full((n,), -torch.inf, dtype=f.dtype, device=f.device).scatter_reduce(0, idx, f, "amax")
+        total = torch.zeros(n, dtype=f.dtype, device=f.device).index_add_(0, idx, (f - top[idx]).exp())
+        return (top + total.log()).tolist()
+
+    @torch.no_grad()
+    def form_scores(self, text: str) -> torch.Tensor:
+        """Log-probability of every continuation in ``self.forms`` given the prompt for ``text``."""
+        dev = self.device
+        P = len(self.prefix_ids)
+        if self._prefix_cache is None:
+            out = self.model(torch.tensor([self.prefix_ids], device=dev), use_cache=True)
+            self._prefix_cache = out.past_key_values
+            self._prefix_last = torch.log_softmax(out.logits[0, -1].float(), -1)
+        ids = torch.tensor(self.text_ids(text), dtype=torch.long, device=dev)
+        n, N = len(ids), len(self._node_tokens)
+        if n == 0 and N == 0:
+            return self._prefix_last[self._targets]
+
+        # Query rows: n text tokens, then N tree nodes. Keys: P cached prompt tokens, then the same rows.
+        allow = torch.zeros((n + N, P + n + N), dtype=torch.bool, device=dev)
+        allow[:, :P] = True
+        allow[:n, P : P + n] = torch.ones((n, n), dtype=torch.bool, device=dev).tril()
+        allow[n:, P : P + n] = True
+        allow[n:, P + n :] = self._node_anc
+        positions = torch.cat([torch.arange(P, P + n, device=dev), P + n + self._node_depth])
+        keep = torch.arange(max(n - 1, 0), n + N, device=dev)  # last text position + every node
+        try:
+            out = self.model(
+                torch.cat([ids, self._node_tokens])[None],
+                attention_mask=allow[None, None],
+                position_ids=positions[None],
+                past_key_values=self._prefix_cache,
+                use_cache=True,
+                logits_to_keep=keep,
+            )
+        finally:
+            extra = self._prefix_cache.get_seq_length() - P
+            if extra:
+                self._prefix_cache.crop(-extra)  # back to the prompt prefix
+        lp = torch.log_softmax(out.logits[0].float(), -1)  # row 0: last text position, row 1 + i: node i
+        if n == 0:  # the prompt prefix itself predicts the first label token
+            lp = torch.cat([self._prefix_last[None], lp])
+        per_token = lp[self._rows + 1, self._targets]
+        total = torch.zeros(len(self.forms), dtype=per_token.dtype, device=dev)
+        return total.index_add_(0, self._owner, per_token)
+
+
+def load_teacher(name: str, task: Task, device: str = "auto", load_4bit: bool = False) -> LabelScorer:
+    device = pick_device(device)
+    tok = AutoTokenizer.from_pretrained(name)
+    kwargs: dict = {}
+    if load_4bit:
+        if not device.startswith("cuda"):
+            raise SystemExit("--load-4bit needs a CUDA device (bitsandbytes); use --device cuda")
+        from transformers import BitsAndBytesConfig
+
+        kwargs["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16
+        )
+        kwargs["device_map"] = {"": device}
+    else:
+        kwargs["dtype"] = torch.bfloat16 if device == "cuda" else torch.float32
+    model = AutoModelForCausalLM.from_pretrained(name, **kwargs)
+    if not load_4bit:
+        model.to(device)
+    return LabelScorer(model, tok, task)
+
+
+class ScoreCache:
+    """Append-only JSONL of {"h": text hash, "s": label scores}; reopening it resumes where it stopped."""
+
+    def __init__(self, path: Path):
+        self.path = path
+        self.scores: dict[str, list[float]] = {}
+        if path.exists():
+            with path.open(encoding="utf-8") as f:
+                for line in f:
+                    try:
+                        row = json.loads(line)
+                    except json.JSONDecodeError:  # a line cut off by an interrupted run
+                        continue
+                    self.scores[row["h"]] = row["s"]
+
+    def fill(self, scorer: LabelScorer, texts: Iterable[str], progress=None) -> int:
+        """Score every text not in the cache yet. Returns how many were scored now."""
+        n = 0
+        with self.path.open("a", encoding="utf-8") as f:
+            for text in texts:
+                h = text_hash(text)
+                if h in self.scores:
+                    continue
+                s = scorer.score(text)
+                self.scores[h] = s
+                f.write(json.dumps({"h": h, "s": s}) + "\n")
+                f.flush()
+                n += 1
+                if progress:
+                    progress(n)
+        return n
+
+    def get(self, text: str) -> list[float]:
+        return self.scores[text_hash(text)]
+
+
+def iter_unique(texts: Iterable[str]) -> Iterator[str]:
+    seen: set[str] = set()
+    for t in texts:
+        if t not in seen:
+            seen.add(t)
+            yield t

tinydistill-0.2.0.dist-info/METADATA

@@ -0,0 +1,294 @@
+Metadata-Version: 2.5
+Name: tinydistill
+Version: 0.2.0
+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
+License-Expression: MIT
+License-File: LICENSE
+Keywords: inference,knowledge distillation,llm,onnx,text classification
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: onnx>=1.15
+Requires-Dist: onnxruntime>=1.17
+Requires-Dist: pyyaml>=6
+Requires-Dist: torch>=2.2
+Requires-Dist: transformers>=4.56
+Provides-Extra: 4bit
+Requires-Dist: accelerate>=0.30; extra == '4bit'
+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: pytest>=8; extra == 'dev'
+Requires-Dist: ruff>=0.5; extra == 'dev'
+Provides-Extra: hf
+Requires-Dist: datasets>=2.18; extra == 'hf'
+Description-Content-Type: text/markdown
+
+# tinydistill
+
+**Most of a large model's quality on one task, many times faster, on a CPU.**
+
+`tinydistill` turns an LLM classification prompt into a small encoder classifier by knowledge distillation.
+You describe the task in a short YAML file (labels, how to say each label, an instruction prompt) and give it
+unlabelled texts. An instruction-tuned LLM labels them with soft targets, a small encoder learns to imitate
+those targets, and the result is exported to ONNX (optionally int8) for fast CPU inference.
+
+It continues my bachelor's project on inference optimisation by distilling LLMs.
+
+## Results
+
+Two tasks with `Qwen/Qwen2.5-1.5B-Instruct` as the zero-shot teacher, and Banking77 also with
+`Qwen/Qwen2.5-7B-Instruct`; the student is always `sentence-transformers/all-MiniLM-L6-v2`. Quality is measured on gold-labelled test texts that
+neither model trains on. Every table is written by the task's `run.sh` from its saved `results*.json` file in
+[`examples/`](https://github.com/Andreas-Hum/tinydistill/tree/main/examples); no number is typed by hand. The teacher's latency includes scoring every surface form
+of every label, which costs most on Banking77 with its many labels, so there the GPU teacher row is the fairer
+reference for speed.
+
+### Banking77: the hard case
+
+Intent detection with 77 fine-grained classes ([Banking77](https://huggingface.co/datasets/PolyAI/banking77),
+CC BY 4.0), reproduced by [`examples/banking77/run.sh`](https://github.com/Andreas-Hum/tinydistill/blob/main/examples/banking77/run.sh).
+
+Teacher only, no gold labels:
+
+
+<!-- results: banking77/results.md -->
+| Model | Device | Accuracy | Macro-F1 | Agreement with teacher | Latency, batch 1 (ms) | Latency, batch 32 (ms) | Throughput (texts/s) | Size (MB) | Speed-up, batch 1 |
+|---|---|---|---|---|---|---|---|---|---|
+| teacher, zero-shot | cpu | 0.406 | 0.367 | n/a | 19955.9 | n/a | 0.1 | 6175 | 1.0x |
+| teacher, zero-shot | cuda | 0.406 | 0.367 | n/a | 60.3 | n/a | 16.6 | 3087 | 331.1x |
+| student, PyTorch | cpu | 0.386 | 0.315 | 0.772 | 3.9 | 78.3 | 408.9 | 91 | 5124.6x |
+| student, ONNX fp32 | cpu | 0.386 | 0.315 | 0.772 | 2.6 | 82.4 | 388.4 | 91 | 7606.9x |
+| student, ONNX int8 | cpu | 0.383 | 0.311 | 0.761 | 1.0 | 34.4 | 931.1 | 23 | 20502.9x |
+
+Task `banking77`, 3080 gold-labelled test texts. Teacher: `Qwen/Qwen2.5-1.5B-Instruct`. Student: `sentence-transformers/all-MiniLM-L6-v2`. The student learns only from the teacher's soft labels on 10003 unlabelled training texts.
+CPU rows use 2 threads. Latency is the median over repeated runs. Size is the teacher's weights as loaded and the students' model files. The teacher scores one text at a time, so it is timed at batch 1 and its throughput is 1 / latency. Speed-up is relative to the teacher, zero-shot on cpu.
+<!-- /results -->
+
+With 8 gold labels per class, and the student trained on all gold labels as an upper bound:
+
+<!-- results: banking77/results-k8.md -->
+| Model | Device | Accuracy | Macro-F1 | Agreement with teacher | Latency, batch 1 (ms) | Latency, batch 32 (ms) | Throughput (texts/s) | Size (MB) | Speed-up, batch 1 |
+|---|---|---|---|---|---|---|---|---|---|
+| teacher, zero-shot | cpu | 0.406 | 0.367 | n/a | 19699.4 | n/a | 0.1 | 6175 | 1.0x |
+| teacher, zero-shot | cuda | 0.406 | 0.367 | n/a | 58.3 | n/a | 17.2 | 3087 | 338.1x |
+| student, PyTorch | cpu | 0.404 | 0.324 | 0.699 | 3.9 | 76.7 | 417.4 | 91 | 5056.3x |
+| upper bound: student on all gold labels | cpu | 0.911 | 0.909 | 0.412 | n/a | n/a | n/a | n/a | n/a |
+| student, ONNX fp32 | cpu | 0.404 | 0.324 | 0.699 | 2.7 | 72.1 | 443.6 | 91 | 7306.9x |
+| student, ONNX int8 | cpu | 0.398 | 0.318 | 0.697 | 1.0 | 36.7 | 872.1 | 23 | 18792.2x |
+
+Task `banking77`, 3080 gold-labelled test texts. Teacher: `Qwen/Qwen2.5-1.5B-Instruct`. Student: `sentence-transformers/all-MiniLM-L6-v2`. The student learns from the teacher's soft labels on all 10003 training texts, plus cross-entropy on 8 gold labels per class (616 texts, weight 1.0). The upper bound is the same student trained with cross-entropy on all 10003 gold training labels and no teacher; its speed equals the student's.
+CPU rows use 2 threads. Latency is the median over repeated runs. Size is the teacher's weights as loaded and the students' model files. The teacher scores one text at a time, so it is timed at batch 1 and its throughput is 1 / latency. Speed-up is relative to the teacher, zero-shot on cpu.
+<!-- /results -->
+
+With a stronger teacher, `Qwen/Qwen2.5-7B-Instruct` loaded in 4 bit (timed on the GPU only: 4-bit loading needs
+CUDA), teacher only:
+
+<!-- results: banking77/results-7b.md -->
+| Model | Device | Accuracy | Macro-F1 | Agreement with teacher | Latency, batch 1 (ms) | Latency, batch 32 (ms) | Throughput (texts/s) | Size (MB) | Speed-up, batch 1 |
+|---|---|---|---|---|---|---|---|---|---|
+| teacher, zero-shot, 4-bit | cuda | 0.630 | 0.607 | n/a | 272.2 | n/a | 3.7 | 5443 | 1.0x |
+| student, PyTorch | cpu | 0.623 | 0.587 | 0.827 | 5.5 | 93.0 | 344.2 | 91 | 49.4x |
+| student, ONNX fp32 | cpu | 0.623 | 0.587 | 0.827 | 3.4 | 86.8 | 368.7 | 91 | 80.8x |
+| student, ONNX int8 | cpu | 0.622 | 0.586 | 0.822 | 1.6 | 43.8 | 730.0 | 23 | 171.2x |
+
+Task `banking77`, 3080 gold-labelled test texts. Teacher: `Qwen/Qwen2.5-7B-Instruct`. Student: `sentence-transformers/all-MiniLM-L6-v2`. The student learns only from the teacher's soft labels on 10003 unlabelled training texts.
+CPU rows use 2 threads. Latency is the median over repeated runs. Size is the teacher's weights as loaded and the students' model files. The teacher scores one text at a time, so it is timed at batch 1 and its throughput is 1 / latency. Speed-up is relative to the teacher, zero-shot, 4-bit on cuda.
+<!-- /results -->
+
+The same teacher with 8 gold labels per class and the upper bound:
+
+<!-- results: banking77/results-k8-7b.md -->
+| Model | Device | Accuracy | Macro-F1 | Agreement with teacher | Latency, batch 1 (ms) | Latency, batch 32 (ms) | Throughput (texts/s) | Size (MB) | Speed-up, batch 1 |
+|---|---|---|---|---|---|---|---|---|---|
+| teacher, zero-shot, 4-bit | cuda | 0.630 | 0.607 | n/a | 295.4 | n/a | 3.4 | 5443 | 1.0x |
+| student, PyTorch | cpu | 0.627 | 0.583 | 0.765 | 4.5 | 93.6 | 341.9 | 91 | 65.5x |
+| upper bound: student on all gold labels | cpu | 0.911 | 0.909 | 0.631 | n/a | n/a | n/a | n/a | n/a |
+| student, ONNX fp32 | cpu | 0.627 | 0.583 | 0.765 | 3.2 | 80.9 | 395.7 | 91 | 92.5x |
+| student, ONNX int8 | cpu | 0.622 | 0.576 | 0.762 | 1.1 | 40.7 | 786.4 | 23 | 262.6x |
+
+Task `banking77`, 3080 gold-labelled test texts. Teacher: `Qwen/Qwen2.5-7B-Instruct`. Student: `sentence-transformers/all-MiniLM-L6-v2`. The student learns from the teacher's soft labels on all 10003 training texts, plus cross-entropy on 8 gold labels per class (616 texts, weight 1.0). The upper bound is the same student trained with cross-entropy on all 10003 gold training labels and no teacher; its speed equals the student's.
+CPU rows use 2 threads. Latency is the median over repeated runs. Size is the teacher's weights as loaded and the students' model files. The teacher scores one text at a time, so it is timed at batch 1 and its throughput is 1 / latency. Speed-up is relative to the teacher, zero-shot, 4-bit on cuda.
+<!-- /results -->
+
+### Financial tweets: an easier task
+
+Sentiment of financial news tweets, 3 classes
+([Twitter Financial News Sentiment](https://huggingface.co/datasets/zeroshot/twitter-financial-news-sentiment),
+MIT licence per the dataset card), reproduced by
+[`examples/financial-tweets/run.sh`](https://github.com/Andreas-Hum/tinydistill/blob/main/examples/financial-tweets/run.sh). The validation split is the test set.
+
+<!-- results: financial-tweets/results.md -->
+| Model | Device | Accuracy | Macro-F1 | Agreement with teacher | Latency, batch 1 (ms) | Latency, batch 32 (ms) | Throughput (texts/s) | Size (MB) | Speed-up, batch 1 |
+|---|---|---|---|---|---|---|---|---|---|
+| teacher, zero-shot | cpu | 0.535 | 0.557 | n/a | 976.3 | n/a | 1.0 | 6175 | 1.0x |
+| teacher, zero-shot | cuda | 0.535 | 0.557 | n/a | 33.7 | n/a | 29.6 | 3087 | 28.9x |
+| student, PyTorch | cpu | 0.517 | 0.542 | 0.815 | 4.5 | 145.4 | 220.1 | 91 | 219.3x |
+| upper bound: student on all gold labels | cpu | 0.867 | 0.824 | 0.531 | n/a | n/a | n/a | n/a | n/a |
+| student, ONNX fp32 | cpu | 0.517 | 0.542 | 0.815 | 3.2 | 148.5 | 215.4 | 91 | 303.9x |
+| student, ONNX int8 | cpu | 0.513 | 0.538 | 0.810 | 1.3 | 71.5 | 447.7 | 23 | 754.7x |
+
+Task `financial-tweets`, 2388 gold-labelled test texts. Teacher: `Qwen/Qwen2.5-1.5B-Instruct`. Student: `sentence-transformers/all-MiniLM-L6-v2`. The student learns only from the teacher's soft labels on 9543 unlabelled training texts. The upper bound is the same student trained with cross-entropy on all 9543 gold training labels and no teacher; its speed equals the student's.
+CPU rows use 2 threads. Latency is the median over repeated runs. Size is the teacher's weights as loaded and the students' model files. The teacher scores one text at a time, so it is timed at batch 1 and its throughput is 1 / latency. Speed-up is relative to the teacher, zero-shot on cpu.
+<!-- /results -->
+
+## Quickstart
+
+```sh
+pip install "tinydistill[hf]"
+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 \
+  --out runs/banking77 --int8
+tinydistill predict runs/banking77/onnx "I still haven't received my new card"
+```
+
+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
+Banking77 splits are sorted by intent, so a limited run only sees a few intents.
+
+`run` does every step. Each step also has its own command, and they all read and write the run directory:
+
+| Command | What it does |
+|---|---|
+| `tinydistill label TASK.yaml --data D [--test T] --out RUN` | teacher scores every label of every text; cached, resumable |
+| `tinydistill train RUN` | distils the student from the teacher's soft targets |
+| `tinydistill export RUN [--int8]` | writes `RUN/onnx` (ONNX, tokenizer, labels) |
+| `tinydistill eval RUN` | writes `RUN/results.md` and `RUN/results.json` |
+| `tinydistill predict RUN/onnx TEXT... [--data FILE]` | classifies texts, prints JSON lines |
+
+Data can be `.csv`, `.jsonl` or `hf:<dataset>[:<config>]:<split>` (needs the `hf` extra). Useful options:
+`--teacher` (any Hugging Face causal LM), `--load-4bit` (bitsandbytes, CUDA only: `pip install "tinydistill[4bit]"`),
+`--student` (any Hugging Face encoder), `--epochs`, `--lr`, `--device`, `--threads`, and `--cache-dir` (runs that
+share it label each text only once). See `tinydistill <command> --help`.
+
+**Few labels.** With `--gold-per-class K`, the student also sees the gold labels of K randomly chosen training
+examples per class, through a cross-entropy term (weight `--ce-weight`, default 1) added to the distillation
+loss. `--upper-bound` additionally trains the same student on all gold training labels without the teacher,
+and the results table gets an upper-bound row: how far the distilled student is from fully supervised training.
+
+From Python:
+
+```python
+from tinydistill import Predictor
+
+clf = Predictor("runs/banking77/onnx")  # uses model.int8.onnx when present
+clf.predict(["How do I top up my card?"])  # list of label names
+clf.predict_proba(["How do I top up my card?"])  # numpy array, one row per text
+```
+
+### The task file
+
+```yaml
+name: banking77
+text_field: text        # column holding the input text
+label_field: label      # column holding gold labels (ids, or integers indexing `labels`)
+prompt: |-
+  You are classifying messages that customers sent to their bank's support chat.
+  Choose the single intent from this list that best matches the message:
+  {labels}
+
+  Message: {text}
+
+  Reply with the intent exactly as written in the list.
+labels:
+  - activate_my_card    # a bare id: its name is derived, "activate my card"
+  - card_arrival
+  # ...
+```
+
+`{labels}` (required) becomes a bulleted list of the label names, and `{text}` is the input. The name is what
+the teacher scores. Each label can be a bare id as above, `id: name`, or carry a description that is shown in
+the list (`- bullish: positive for the stock`) but not scored:
+
+```yaml
+labels:
+  bearish:
+    description: negative for the stock, company or market
+  lost_card: lost or stolen card
+```
+
+## How the label scoring works
+
+The teacher never generates free text. For an input x, the prompt is rendered with the model's chat template
+(user message = the instruction with the label list, generation prompt for the assistant). A continuation is one
+way of writing a label followed by the end-of-turn token, `c_1 ... c_m`, with log-probability
+
+    log p(c | prompt(x)) = sum_j log p(c_j | prompt(x), c_1 ... c_(j-1))
+
+Multi-token labels are handled exactly, and the end-of-turn token keeps a label that is a prefix of another
+(`card` versus `card arrival`) from winning by stopping early. Models do not always write a label the way it is
+listed: Qwen2.5 tends to answer a sentiment question with "Bullish" even when the list says "bullish". So each
+label is scored over its common surface forms, the name as written, lowercase, Capitalised and UPPERCASE, each
+with and without a leading space (forms that tokenise identically count once):
+
+    score(label) = log sum over forms c of p(c | prompt(x))
+
+This is verbaliser coverage: it only counts more ways of writing the same answer, and there is no calibration or
+prior correction. All forms share the prefix tree below, but with many labels they still add up; a task file can
+list fewer (`surface_forms: [lowercase, capitalised]` drops UPPERCASE, which tokenises into many pieces, and
+`surface_forms: []` scores the name as written only). A softmax over the label scores gives the soft targets. The prompt lists each label in its
+primary form, the name as written. The scores are cached as JSON lines keyed by a hash of the text, in a file
+named after a fingerprint of the teacher, prompt and continuations, so an interrupted run resumes and a changed
+prompt never reuses old scores.
+
+To keep this fast, the part of the prompt before the text is run once and its key/value cache is kept. Each
+text then needs a single forward pass that covers the text and all labels together: the tokens of all continuations
+(every surface form of every label) form a prefix tree, so continuations that start alike, such as
+`card arrival` and `card linking`, share nodes, and a 4D attention mask lets each tree node attend only to the
+prompt, the text and its own ancestors. There is no padding and no per-label copy of the cache. Tests check the
+scores against a plain forward pass per continuation.
+
+The student minimises KL(teacher || student) over the training texts, plus `--ce-weight` times the
+cross-entropy on the K gold labels per class chosen by `--gold-per-class`, if any (standard knowledge
+distillation, Hinton et al., 2015).
+
+## Limitations
+
+- **The teacher caps the student.** The student learns to imitate the teacher, mistakes included. A small LLM
+  choosing zero-shot among many fine-grained labels can be weak: Qwen2.5-0.5B-Instruct, for example, often
+  confuses related Banking77 intents. A bigger teacher (try `--load-4bit` with a 7B model) or a better prompt
+  usually helps more than anything on the student side.
+- **Scores are sums, not averages.** Longer verbalisations get lower scores, so very uneven label lengths bias the
+  teacher. Keep verbalisations short and of similar length.
+- **Only case and spacing variants are counted.** Plurals, synonyms or rephrasings of a label are not; keep label
+  names short and close to how the teacher would answer.
+- **A too-small teacher can collapse.** Qwen2.5-0.5B-Instruct, for example, puts nearly every text into one or two
+  labels on the financial-tweets task, whatever the scoring; use 1.5B or larger.
+- **Prompt and label order matter.** Zero-shot LLMs are sensitive to wording and to the order of the options.
+- **Teacher cost grows with labels.** Every text is scored against every surface form of every label (shared
+  prefixes are computed once), one text at a time.
+- **Single-label classification only.** Inputs are truncated to 128 tokens by default (`--max-length`).
+- **Timings depend on the machine.** The results table states the threads and devices used; rerun `eval` on your
+  own hardware.
+- ONNX export uses PyTorch's TorchScript-based exporter.
+
+## Development
+
+```sh
+python -m venv .venv && . .venv/bin/activate
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+pip install -e ".[dev]"
+ruff check . && ruff format --check . && pytest
+```
+
+The tests run on a CPU in a few seconds with tiny random models from `hf-internal-testing`. They check the label
+scores against a direct computation, the distillation loss and metrics against hand-computed values, cache
+resumption, that the ONNX export reproduces the PyTorch logits, and the whole CLI end to end.
+
+## References
+
+- G. Hinton, O. Vinyals, J. Dean. [Distilling the Knowledge in a Neural Network](https://arxiv.org/abs/1503.02531). arXiv:1503.02531, 2015.
+- T. Schick, H. Schütze. [Exploiting Cloze Questions for Few Shot Text Classification and Natural Language Inference](https://arxiv.org/abs/2001.07676). EACL, 2021. (Verbalisers.)
+- I. Casanueva, T. Temčinas, D. Gerz, M. Henderson, I. Vulić. [Efficient Intent Detection with Dual Sentence Encoders](https://arxiv.org/abs/2003.04807). NLP4ConvAI Workshop, 2020. (Banking77.)
+- W. Wang, F. Wei, L. Dong, H. Bao, N. Yang, M. Zhou. [MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers](https://arxiv.org/abs/2002.10957). NeurIPS, 2020.
+- N. Reimers, I. Gurevych. [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084). EMNLP, 2019.
+- Qwen Team. [Qwen2.5 Technical Report](https://arxiv.org/abs/2412.15115). arXiv:2412.15115, 2024.
+
+## Licence
+
+[MIT](https://github.com/Andreas-Hum/tinydistill/blob/main/LICENSE). The Banking77 dataset is CC BY 4.0 (PolyAI); models are downloaded from the Hugging Face Hub under
+their own licences.

tinydistill-0.2.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+tinydistill/__init__.py,sha256=h0LIgsGwV4oESxjmRmQNXjxh7xHObNPe1FnwismRWts,159
+tinydistill/cli.py,sha256=cIf2_pJAXeiEKlXiqutUV_XWMgahsutETinwr-e1NZg,16694
+tinydistill/data.py,sha256=VpxIeEGId6birv1fwxb3-r9qKjmg1euLJN3VWWTp5wo,3452
+tinydistill/evaluate.py,sha256=-MoXGPFj6Aj44MLtr-iXQM7d00LTJXrFLA-2757OAxA,5753
+tinydistill/export.py,sha256=qPypHTcDZIZj_Pz2qwR54yPlBBfo2BVAzQHW_Tp4Iqo,3833
+tinydistill/student.py,sha256=zmb-IU67szW5TxPh3jIzT5ykY72UcuPe_ugCWubATFM,5145
+tinydistill/task.py,sha256=XsZHlLrLVW22pLrV9wCVrUdoAMnLE_BtvkbAykTMVFY,4635
+tinydistill/teacher.py,sha256=91aYoYPm_x9Eil_s2mYqDPmv_8RNIuSFZV0fh1CnNKk,12338
+tinydistill-0.2.0.dist-info/METADATA,sha256=E7sTZN-jOSOnFI98TeEymZ2GBFk50s1G088Y_giICSw,19398
+tinydistill-0.2.0.dist-info/WHEEL,sha256=2iZrJkDIkeDh0k9SikKMqFuRP-ua6mb4Ir1AsER4BEw,87
+tinydistill-0.2.0.dist-info/entry_points.txt,sha256=khfieEGrxSXJBWXEN2ATy8J_KO9--eFulic6a4yLEt0,53
+tinydistill-0.2.0.dist-info/licenses/LICENSE,sha256=G7EoDl2dH0OXmLiIMUnEsW3HXrlTA1sVEwVjThlLQC4,1068
+tinydistill-0.2.0.dist-info/RECORD,,

tinydistill-0.2.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.33.0
+Root-Is-Purelib: true
+Tag: py3-none-any

tinydistill-0.2.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+tinydistill = tinydistill.cli:main

tinydistill-0.2.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Andreas-Hum
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.