tinydistill 0.2.0__py3-none-any.whl

tinydistill/evaluate.py

@@ -0,0 +1,140 @@
+"""Quality metrics, latency benchmarks and the results table."""
+
+from __future__ import annotations
+
+import os
+import platform
+import statistics
+import time
+from collections.abc import Callable, Sequence
+from pathlib import Path
+
+import torch
+
+
+def accuracy(pred: Sequence[int], gold: Sequence[int]) -> float:
+    return sum(p == g for p, g in zip(pred, gold, strict=True)) / len(gold) if gold else float("nan")
+
+
+def macro_f1(pred: Sequence[int], gold: Sequence[int]) -> float:
+    """Unweighted mean F1 over the classes that occur in gold or pred (as scikit-learn's 'macro')."""
+    classes = set(gold) | set(pred)
+    if not classes:
+        return float("nan")
+    f1s = []
+    for c in classes:
+        tp = sum(p == c and g == c for p, g in zip(pred, gold, strict=True))
+        fp = sum(p == c and g != c for p, g in zip(pred, gold, strict=True))
+        fn = sum(p != c and g == c for p, g in zip(pred, gold, strict=True))
+        f1s.append(2 * tp / (2 * tp + fp + fn) if tp else 0.0)
+    return sum(f1s) / len(f1s)
+
+
+def time_calls(fn: Callable[[list[str]], object], texts: list[str], batch: int, reps: int) -> float:
+    """Median wall time (ms) of fn on `reps` different batches of `batch` texts, after one warm-up call."""
+
+    def batch_at(i: int) -> list[str]:
+        return [texts[(i * batch + j) % len(texts)] for j in range(batch)]
+
+    fn(batch_at(reps))  # warm-up on a batch that is not timed
+    times = []
+    for i in range(reps):
+        b = batch_at(i)
+        t0 = time.perf_counter()
+        fn(b)
+        times.append((time.perf_counter() - t0) * 1000)
+    return statistics.median(times)
+
+
+def bench(fn: Callable[[list[str]], object], texts: list[str], reps1: int, reps32: int = 0) -> dict:
+    """Latency at batch 1 and (if reps32 > 0) batch 32. Throughput comes from batch 32 when measured,
+    otherwise from batch 1 (for models that process one text at a time)."""
+    out = {"latency_ms_b1": time_calls(fn, texts, 1, reps1)}
+    if reps32:
+        out["latency_ms_b32"] = time_calls(fn, texts, 32, reps32)
+        out["throughput_per_s"] = 32 / (out["latency_ms_b32"] / 1000)
+    else:
+        out["throughput_per_s"] = 1000 / out["latency_ms_b1"]
+    return out
+
+
+def dir_size_mb(paths: Sequence[Path]) -> float:
+    return sum(p.stat().st_size for p in paths if p.exists()) / 1e6
+
+
+def hf_weights_mb(name: str) -> float | None:
+    """Size of a Hub model's weight files in the local cache (or of a local directory)."""
+    p = Path(name)
+    if not p.exists():
+        try:
+            from huggingface_hub import snapshot_download
+
+            p = Path(snapshot_download(name, local_files_only=True))
+        except Exception:
+            return None
+    files = list(p.glob("*.safetensors")) or list(p.glob("*.bin"))
+    return dir_size_mb(files) if files else None
+
+
+def environment(threads: int) -> dict:
+    return {
+        "cpu": platform.processor() or platform.machine(),
+        "cpu_threads": threads,
+        "torch": torch.__version__,
+        "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
+        "omp_num_threads": os.environ.get("OMP_NUM_THREADS"),
+    }
+
+
+def _fmt(v, spec: str) -> str:
+    if v is None or (isinstance(v, float) and v != v):
+        return "n/a"
+    return format(v, spec)
+
+
+def markdown_table(results: dict) -> str:
+    rows = results["models"]
+    ref = next((r for r in rows if r["role"] == "teacher" and r.get("latency_ms_b1")), None)
+    head = (
+        "| 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 |\n"
+        "|---|---|---|---|---|---|---|---|---|---|\n"
+    )
+    lines = []
+    for r in rows:
+        speed = ref["latency_ms_b1"] / r["latency_ms_b1"] if ref and r.get("latency_ms_b1") else None
+        lines.append(
+            f"| {r['name']} | {r['device']} "
+            f"| {_fmt(r.get('accuracy'), '.3f')} | {_fmt(r.get('macro_f1'), '.3f')} "
+            f"| {_fmt(r.get('agreement'), '.3f')} | {_fmt(r.get('latency_ms_b1'), '.1f')} "
+            f"| {_fmt(r.get('latency_ms_b32'), '.1f')} | {_fmt(r.get('throughput_per_s'), '.1f')} "
+            f"| {_fmt(r.get('size_mb'), '.0f')} | {_fmt(speed, '.1f') + 'x' if speed else 'n/a'} |"
+        )
+    ds = results["data"]
+    tr = results.get("training") or {}
+    k = tr.get("gold_per_class") or 0
+    trained = (
+        f"The student learns from the teacher's soft labels on all {ds['n_train']} training texts, "
+        f"plus cross-entropy on {k} gold labels per class "
+        f"({tr.get('n_gold')} texts, weight {tr.get('ce_weight')})."
+        if k
+        else f"The student learns only from the teacher's soft labels on {ds['n_train']} "
+        "unlabelled training texts."
+    )
+    if tr.get("upper_bound"):
+        trained += (
+            " The upper bound is the same student trained with cross-entropy on all "
+            f"{tr.get('n_gold_upper_bound')} gold training labels and no teacher; "
+            "its speed equals the student's."
+        )
+    notes = [
+        "",
+        f"Task `{results['task']}`, {ds['n_test']} gold-labelled test texts. "
+        f"Teacher: `{results['teacher']}`. Student: `{results['student']}`. {trained}",
+        f"CPU rows use {results['environment']['cpu_threads']} 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. "
+        + (f"Speed-up is relative to the {ref['name']} on {ref['device']}." if ref else ""),
+    ]
+    return head + "\n".join(lines) + "\n" + "\n".join(notes) + "\n"

tinydistill/export.py

@@ -0,0 +1,103 @@
+"""Export the student to ONNX (optionally int8 dynamic quantisation) and run it with onnxruntime."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .student import load_student
+
+
+class _Logits(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, input_ids, attention_mask):
+        return self.model(input_ids=input_ids, attention_mask=attention_mask).logits
+
+
+def export_onnx(student_dir: Path, out_dir: Path, int8: bool = False, opset: int = 17) -> list[Path]:
+    model, tok = load_student(student_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    enc = tok(["an example input", "another one"], padding=True, return_tensors="pt")
+    path = out_dir / "model.onnx"
+    torch.onnx.export(
+        _Logits(model),
+        (enc["input_ids"], enc["attention_mask"]),
+        str(path),
+        input_names=["input_ids", "attention_mask"],
+        output_names=["logits"],
+        dynamic_axes={
+            "input_ids": {0: "batch", 1: "seq"},
+            "attention_mask": {0: "batch", 1: "seq"},
+            "logits": {0: "batch"},
+        },
+        opset_version=opset,
+        dynamo=False,
+    )
+    tok.save_pretrained(out_dir)
+    labels = [model.config.id2label[i] for i in range(model.config.num_labels)]
+    (out_dir / "labels.json").write_text(json.dumps(labels, indent=1), encoding="utf-8")
+    written = [path]
+    if int8:
+        from onnxruntime.quantization import QuantType, quantize_dynamic
+
+        q = out_dir / "model.int8.onnx"
+        quantize_dynamic(str(path), str(q), weight_type=QuantType.QInt8)
+        written.append(q)
+    return written
+
+
+class Predictor:
+    """Run an exported model: ``Predictor("runs/x/onnx").predict(["text", ...])``."""
+
+    def __init__(
+        self, path: str | Path, int8: bool | None = None, threads: int | None = None, max_length: int = 128
+    ):
+        import onnxruntime as ort
+        from transformers import AutoTokenizer
+
+        path = Path(path)
+        if int8 is None:
+            int8 = (path / "model.int8.onnx").exists()
+        file = path / ("model.int8.onnx" if int8 else "model.onnx")
+        opts = ort.SessionOptions()
+        threads = threads or int(os.environ.get("OMP_NUM_THREADS", "0"))
+        if threads:
+            opts.intra_op_num_threads = threads
+            opts.inter_op_num_threads = 1
+        self.session = ort.InferenceSession(str(file), opts, providers=["CPUExecutionProvider"])
+        self.tok = AutoTokenizer.from_pretrained(path)
+        self.labels: list[str] = json.loads((path / "labels.json").read_text(encoding="utf-8"))
+        self.max_length = max_length
+        self.file = file
+
+    def logits(self, texts: list[str], batch_size: int = 32) -> np.ndarray:
+        out = []
+        for b in range(0, len(texts), batch_size):
+            enc = self.tok(
+                texts[b : b + batch_size],
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                return_tensors="np",
+            )
+            feed = {
+                "input_ids": enc["input_ids"].astype(np.int64),
+                "attention_mask": enc["attention_mask"].astype(np.int64),
+            }
+            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:
+        z = self.logits(texts, batch_size)
+        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[str]:
+        return [self.labels[i] for i in self.logits(texts, batch_size).argmax(-1)]

tinydistill/student.py

@@ -0,0 +1,140 @@
+"""Student: a small encoder classifier trained on the teacher's soft targets."""
+
+from __future__ import annotations
+
+import math
+import random
+import time
+from collections.abc import Callable
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+DEFAULT_STUDENT = "sentence-transformers/all-MiniLM-L6-v2"
+
+
+def distillation_loss(
+    student_logits: torch.Tensor,
+    teacher_probs: torch.Tensor,
+    gold: torch.Tensor | None = None,
+    ce_weight: float = 0.0,
+) -> torch.Tensor:
+    """Mean over the batch of KL(teacher || student), plus ce_weight * mean CE on the examples with a gold
+    label (gold = -1 marks examples without one). Standard knowledge distillation (Hinton et al., 2015)."""
+    log_q = F.log_softmax(student_logits.float(), dim=-1)
+    p = teacher_probs.float()
+    kl = (p * (torch.log(p.clamp_min(1e-12)) - log_q)).sum(-1).mean()
+    if ce_weight and gold is not None and (gold >= 0).any():
+        has = gold >= 0
+        kl = kl + ce_weight * F.nll_loss(log_q[has], gold[has])
+    return kl
+
+
+def few_gold(gold: list[int | None], k: int, seed: int = 0) -> list[int | None]:
+    """Keep the gold label of k randomly chosen examples per class (fewer if a class has fewer) and drop the
+    rest, so only a few labels per class are used."""
+    by_class: dict[int, list[int]] = {}
+    for i, g in enumerate(gold):
+        if g is not None:
+            by_class.setdefault(g, []).append(i)
+    rng = random.Random(seed)
+    keep: set[int] = set()
+    for c in sorted(by_class):
+        idx = by_class[c]
+        keep.update(rng.sample(idx, min(k, len(idx))))
+    return [g if i in keep else None for i, g in enumerate(gold)]
+
+
+def train_student(
+    texts: list[str],
+    teacher_probs: torch.Tensor | None,
+    labels: list[str],
+    out_dir: Path,
+    model_name: str = DEFAULT_STUDENT,
+    gold: list[int | None] | None = None,
+    ce_weight: float = 0.0,
+    epochs: int = 5,
+    batch_size: int = 32,
+    lr: float = 1e-4,
+    max_length: int = 128,
+    device: str = "cpu",
+    seed: int = 0,
+    log: Callable[[str], None] = print,
+) -> None:
+    torch.manual_seed(seed)
+    rng = random.Random(seed)
+    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,
+    ).to(device)
+    gold_t = torch.tensor([-1 if g is None else g for g in (gold or [None] * len(texts))])
+
+    steps = epochs * math.ceil(len(texts) / batch_size)
+    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
+    warm = max(1, steps // 10)
+    sched = torch.optim.lr_scheduler.LambdaLR(
+        opt, lambda s: min((s + 1) / warm, max(0.0, (steps - s) / (steps - warm + 1e-9)))
+    )
+
+    order = list(range(len(texts)))
+    for epoch in range(epochs):
+        model.train()
+        rng.shuffle(order)
+        t0, total = time.perf_counter(), 0.0
+        for b in range(0, len(order), batch_size):
+            idx = order[b : b + batch_size]
+            enc = tok(
+                [texts[i] for i in idx],
+                padding=True,
+                truncation=True,
+                max_length=max_length,
+                return_tensors="pt",
+            )
+            enc = {k: v.to(device) for k, v in enc.items()}
+            logits = model(**enc).logits
+            if (
+                teacher_probs is None
+            ):  # supervised only (the upper-bound model): every example has a gold label
+                loss = F.cross_entropy(logits.float(), gold_t[idx].to(device))
+            else:
+                loss = distillation_loss(
+                    logits, teacher_probs[idx].to(device), gold_t[idx].to(device), ce_weight
+                )
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step()
+            sched.step()
+            opt.zero_grad(set_to_none=True)
+            total += loss.item() * len(idx)
+        log(f"epoch {epoch + 1}/{epochs}: loss {total / len(texts):.4f} ({time.perf_counter() - t0:.0f}s)")
+    out_dir.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(out_dir)
+    tok.save_pretrained(out_dir)
+
+
+def load_student(path: Path, device: str = "cpu"):
+    tok = AutoTokenizer.from_pretrained(path)
+    model = AutoModelForSequenceClassification.from_pretrained(path).to(device).eval()
+    return model, tok
+
+
+@torch.no_grad()
+def student_logits(model, tok, texts: list[str], batch_size: int = 64, max_length: int = 128) -> torch.Tensor:
+    dev = next(model.parameters()).device
+    out = []
+    for b in range(0, len(texts), batch_size):
+        enc = tok(
+            texts[b : b + batch_size],
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+            return_tensors="pt",
+        )
+        out.append(model(**{k: v.to(dev) for k, v in enc.items()}).logits.float().cpu())
+    return torch.cat(out) if out else torch.empty(0)

tinydistill/task.py

@@ -0,0 +1,108 @@
+"""Task definition: label ids, their natural-language names (what the teacher scores), optional descriptions
+and the instruction prompt."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import yaml
+
+SURFACE_FORMS = ("lowercase", "capitalised", "uppercase")
+
+
+def label_name(label_id: str) -> str:
+    """Natural-language name for a label id: card_arrival -> 'card arrival', top-up -> 'top up'."""
+    return re.sub(r"[\s_\-.]+", " ", label_id).strip()
+
+
+@dataclass(frozen=True)
+class Task:
+    name: str
+    labels: list[str]  # label ids, in a fixed order (class index = position)
+    verbalisations: list[str]  # natural-language label names; the teacher scores these
+    prompt: str  # must contain {labels} and {text}
+    text_field: str = "text"
+    label_field: str | None = "label"
+    descriptions: list[str | None] = field(default_factory=list)  # optional, shown in the prompt only
+    # Extra ways of writing each name that the teacher's score includes (the name as written always counts).
+    surface_forms: tuple[str, ...] = SURFACE_FORMS
+
+    def render(self, text: str) -> str:
+        """The user message for one input text."""
+        descs = self.descriptions or [None] * len(self.labels)
+        options = "\n".join(
+            f"- {v}: {d}" if d else f"- {v}" for v, d in zip(self.verbalisations, descs, strict=True)
+        )
+        return self.prompt.replace("{labels}", options).replace("{text}", text)
+
+    def to_dict(self) -> dict:
+        labels = {}
+        for i, (lid, v) in enumerate(zip(self.labels, self.verbalisations, strict=True)):
+            d = self.descriptions[i] if self.descriptions else None
+            labels[lid] = {"name": v, "description": d} if d else v
+        return {
+            "name": self.name,
+            "text_field": self.text_field,
+            "label_field": self.label_field,
+            "prompt": self.prompt,
+            "surface_forms": list(self.surface_forms),
+            "labels": labels,
+        }
+
+
+def _parse_label(path, lid: str, value) -> tuple[str, str | None]:
+    if value is None:
+        return label_name(lid), None
+    if isinstance(value, str):
+        return value, None
+    if isinstance(value, dict):
+        unknown = set(value) - {"name", "description"}
+        if unknown:
+            raise ValueError(f"{path}: label {lid!r} has unknown keys {sorted(unknown)}")
+        name = value.get("name") or label_name(lid)
+        desc = value.get("description")
+        return str(name), (str(desc) if desc else None)
+    raise ValueError(f"{path}: label {lid!r} must be a name, a mapping with name/description, or empty")
+
+
+def _parse_forms(path, value) -> tuple[str, ...]:
+    if not isinstance(value, list) or not set(value) <= set(SURFACE_FORMS):
+        raise ValueError(f"{path}: 'surface_forms' must be a list drawn from {list(SURFACE_FORMS)}")
+    return tuple(f for f in SURFACE_FORMS if f in value)
+
+
+def load_task(path: str | Path) -> Task:
+    """Load a task YAML file. ``labels`` is a list of label ids, or a mapping from id to either a name or
+    ``{name: ..., description: ...}`` (both optional). A missing name is derived from the id."""
+    raw = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
+    if not isinstance(raw, dict):
+        raise ValueError(f"{path}: expected a mapping at the top level")
+    labels = raw.get("labels")
+    if isinstance(labels, dict):
+        items = [(str(k), v) for k, v in labels.items()]
+    elif isinstance(labels, list):
+        items = [(str(k), None) for k in labels]
+    else:
+        raise ValueError(f"{path}: 'labels' must be a mapping or a list")
+    if len(items) < 2:
+        raise ValueError(f"{path}: need at least two labels")
+    parsed = [_parse_label(path, lid, v) for lid, v in items]
+    verbs = [p[0] for p in parsed]
+    descs = [p[1] for p in parsed]
+    if len(set(verbs)) != len(verbs):
+        raise ValueError(f"{path}: label names must be unique")
+    prompt = raw.get("prompt")
+    if not isinstance(prompt, str) or prompt.count("{text}") != 1 or "{labels}" not in prompt:
+        raise ValueError(f"{path}: 'prompt' must contain {{labels}} and contain {{text}} exactly once")
+    return Task(
+        name=str(raw.get("name", Path(path).stem)),
+        labels=[lid for lid, _ in items],
+        verbalisations=verbs,
+        prompt=prompt,
+        text_field=str(raw.get("text_field", "text")),
+        label_field=raw.get("label_field", "label"),
+        descriptions=descs if any(descs) else [],
+        surface_forms=_parse_forms(path, raw.get("surface_forms", list(SURFACE_FORMS))),
+    )