Python API¶
The functions and classes most programs need are importable from tinydistill directly; the rest are in its
modules.
In a few lines¶
tinydistill.distiller.Distiller ¶
Distil a task's prompt into a small ONNX classifier, scikit-learn style.
fit scores every text with the teacher LLM (the soft labels), trains the student on them and exports
it to ONNX; predict, predict_proba and describe then run the exported model through a
:class:~tinydistill.export.Predictor. It is a thin wrapper over
:func:~tinydistill.teacher.load_teacher, :func:~tinydistill.student.train_student and
:func:~tinydistill.export.export_onnx, with the same defaults as the command line::
d = Distiller("task.yaml").fit(texts)
d.predict(["How do I top up my card?"])
d.save("model") # a folder Predictor and Distiller.load read
Labels are the task's label ids (classes_), in the task's order. get_params and set_params
return and set the constructor's arguments, so tools such as scikit-learn's clone work without
tinydistill depending on scikit-learn.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
task
|
str | Path | Task
|
the task: a task file's path or a :class: |
required |
teacher
|
str
|
the teacher LLM (a Hugging Face model name). |
'Qwen/Qwen2.5-1.5B-Instruct'
|
student
|
str
|
the student encoder (a Hugging Face model name). |
DEFAULT_STUDENT
|
epochs
|
int
|
passes over the texts. |
10
|
lr
|
float
|
the student's learning rate. |
0.0001
|
batch_size
|
int
|
texts per training step. |
32
|
device
|
str
|
|
'auto'
|
load_4bit
|
bool
|
load the teacher in 4 bits (bitsandbytes, CUDA only). |
False
|
cache_dir
|
str | Path | None
|
a folder for the teacher's scores. The cache files are the ones |
None
|
int8
|
bool
|
also write an int8 dynamically quantised model, and predict with it. |
True
|
seed
|
int
|
the training seed. |
0
|
val_share
|
float
|
hold out this share of the texts and keep the epoch with the lowest loss on them (see
:func: |
0.0
|
ce_weight
|
float
|
weight of the cross-entropy term on the gold labels given to |
1.0
|
log
|
Callable[[str], None]
|
called with each progress message. |
print
|
set_params ¶
set_params(**params) -> Distiller
Set constructor arguments by name; returns the Distiller. Takes effect at the next fit.
fit ¶
fit(texts: list[str], gold: list | None = None) -> Distiller
Label texts with the teacher, train the student on its soft labels and export it.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
texts
|
list[str]
|
the training texts. |
required |
gold
|
list | None
|
optional gold labels, one per text (a label id, name or index, a list of them for a
multi-label task, or None for a text without one), used in a cross-entropy term weighted by
|
None
|
Returns:
| Type | Description |
|---|---|
Distiller
|
the Distiller, fitted. |
predict ¶
predict(texts: list[str]) -> list
The label id per text, or for a multi-label task the list of label ids with probability >= 0.5.
predict_proba ¶
predict_proba(texts: list[str]) -> np.ndarray
Label probabilities, one row per text and one column per label in classes_.
describe ¶
describe(texts: list[str]) -> list[dict]
Per text, JSON-ready dicts as :meth:~tinydistill.export.Predictor.describe returns them.
score ¶
score(texts: list[str], gold: list) -> float
Accuracy of predict(texts) against gold (label ids, names or indices; for a multi-label
task lists of them, and the exact-match accuracy: every label of a text right).
save ¶
save(path: str | Path) -> Path
Write the exported model to the folder path (created if needed), with the task as
task.yaml. :class:~tinydistill.export.Predictor and :meth:load read it.
Returns:
| Type | Description |
|---|---|
Path
|
the folder. |
load
classmethod
¶
load(path: str | Path, **params) -> Distiller
A fitted Distiller for a saved or exported model (a folder, or hf:user/name); the teacher is
not needed to predict. The task comes from the folder's task.yaml; a folder without one (such as
a run's onnx folder) gets a task with only the model's labels. params set constructor
arguments, for a later fit.
tinydistill.distiller.NotFittedError ¶
Bases: RuntimeError
A Distiller was asked to predict before fit or load.
Using a model¶
tinydistill.export.Predictor ¶
Run an exported model: Predictor("runs/x/onnx").predict(["text", ...]), or one on the Hugging Face
Hub: Predictor("hf:user/name").
__init__ ¶
__init__(path: str | Path, int8: bool | None = None, threads: int | None = None, max_length: int | None = None)
max_length: tokens per input; by default the length the student was trained with (128 for
models exported before it was recorded).
logits ¶
logits(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.
predict_proba ¶
predict_proba(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.
predict ¶
predict(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.
describe ¶
describe(texts: list[str]) -> list[dict]
Per text, as JSON-ready dicts: the label and its probability, or (multi-label) the labels and every label's probability.
tinydistill.export.Cascade ¶
The student answers; texts where its top probability is below threshold go to the teacher.
teacher is a LabelScorer (see load_teacher) for the same task and label order.
Labelling¶
tinydistill.task.load_task ¶
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.
tinydistill.task.Task
dataclass
¶
render ¶
render(text: str, examples: list[tuple[str, int]] | None = None, label: str | None = None, lettered: bool = False) -> str
The user message for one input text. examples are (text, label index) pairs shown in place of
{examples}; without them the {examples} line is removed. label fills {label} (multi-label).
lettered: the labels are options A, B, C, ... (examples answer with the letter) and the message
ends with an instruction to answer with the letter (for hosted API teachers).
tinydistill.data.read_examples ¶
read_examples(spec: str, task: Task, limit: int | None = None, with_gold: bool = True) -> Iterator[Example]
tinydistill.teacher.load_teacher ¶
load_teacher(name: str, task: Task, device: str = 'auto', load_4bit: bool = False, max_tokens: int = 2048, examples: list[tuple[str, int]] | None = None, vllm_url: str | None = None, shortlist: int | None = None, shortlist_model: str | None = None, shortlist_keep: list[str] | None = None, api_url: str | None = None, api_key_env: str = 'OPENAI_API_KEY') -> LabelScorer | VLLMScorer | ApiScorer
A local teacher, or with vllm_url the same model served by vLLM, or with api_url the model
name at a hosted OpenAI-compatible chat API (key in the environment variable api_key_env); then
device is not used. shortlist K: score only the K labels most similar to each text (for tasks with
very many labels).
tinydistill.teacher.LabelScorer ¶
Bases: _Prompt
fingerprint ¶
fingerprint() -> str
Identifies everything that changes the scores, so caches are never mixed up.
tinydistill.teacher.ScoreCache ¶
Append-only JSONL of {"h": text hash, "s": label scores}; reopening it resumes where it stopped.
fill ¶
fill(scorer: Scorer, texts: Iterable[str], progress=None, batch: int = 64, on_read=None) -> int
Score every text not in the cache yet, batch texts at a time (the scorer packs them into
forward passes). Returns how many were scored now.
Training and export¶
tinydistill.student.train_student ¶
train_student(texts: list[str], teacher_probs: Tensor | None, labels: list[str], out_dir: Path, model_name: str = DEFAULT_STUDENT, gold: Sequence[int | list[int] | None] | None = None, ce_weight: float = 0.0, epochs: int = 10, batch_size: int = 32, lr: float = 0.0001, max_length: int = 128, device: str = 'cpu', seed: int = 0, log: Callable[[str], None] = print, multi_label: bool = False, on_step: Callable[[], None] | None = None, val_share: float = 0.0, patience: int = 3) -> int
Train and save the student; returns the epoch whose model is saved. With val_share > 0 that share
of the texts is held out (texts without gold labels when there is a teacher), the loss on them is measured
after every epoch, the epoch with the lowest one is kept, and training stops after patience epochs
without improvement. Otherwise every epoch is used and the last one is kept.
tinydistill.student.distillation_loss ¶
distillation_loss(student_logits: Tensor, teacher_probs: Tensor, gold: 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).
tinydistill.student.few_gold ¶
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.
tinydistill.student.few_shot_examples ¶
few_shot_examples(texts: list[str], gold: list[int | None], k: int, seed: int = 0) -> list[tuple[str, int]]
The same K examples per class as few_gold, as (text, label) pairs in a seeded random order.
tinydistill.export.export_onnx ¶
export_onnx(student_dir: Path, out_dir: Path, int8: bool = False, opset: int = 17) -> list[Path]
tinydistill.export.export_explain ¶
export_explain(student_dir: Path, out_dir: Path, int8: bool = False, opset: int = 17) -> Path | None
Export explain.onnx (and explain.int8.onnx) with outputs logits, layer_logits and
cls_attention, for the teaching view of tinydistill web. None if the model has no layers to show.
Evaluation and charts¶
tinydistill.evaluate.macro_f1 ¶
macro_f1(pred: Sequence[Hashable], gold: Sequence[Hashable]) -> float
Unweighted mean F1 over the classes that occur in gold or pred (as scikit-learn's 'macro').
tinydistill.evaluate.top_k_accuracy ¶
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.
tinydistill.plots ¶
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.
soft_labels ¶
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.
speed_accuracy ¶
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.
training ¶
training(run: str | Path)
Training loss per epoch, and the held-out loss and the kept epoch when --val-share was used.
per_label ¶
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.
confusions ¶
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.
cascade ¶
cascade(run: str | Path)
Accuracy as the student passes its least confident texts to the teacher.
save_all ¶
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.
Deploying¶
tinydistill.web.build_site ¶
build_site(onnx_dir: Path, out: Path, title: str | None = None, samples: int = 6, note: str | None = None) -> Path
Write out/index.html and out/model-<hash>/ (config, tokenizer, ONNX) for transformers.js. When the run's
PyTorch student is next to the export, the page's model also returns each layer's answer and the
attention of the first token, for the "Inside the model" view. note is an extra line for the footer,
such as the dataset's licence.
tinydistill.hub.push ¶
push(onnx_dir: Path, repo: str, private: bool = False, license: str | None = None) -> str
Upload an exported model directory to the Hub (needs huggingface-cli login or HF_TOKEN).