tinydistill 0.3.0__tar.gz → 0.3.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. {tinydistill-0.3.0 → tinydistill-0.3.1}/PKG-INFO +24 -13
  2. {tinydistill-0.3.0 → tinydistill-0.3.1}/README.md +23 -12
  3. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/run.sh +10 -1
  4. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/students.md +3 -0
  5. tinydistill-0.3.1/examples/financial-tweets/financial-tweets-7b-20-epochs.json +146 -0
  6. tinydistill-0.3.1/examples/financial-tweets/financial-tweets-7b-20-epochs.md +22 -0
  7. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/run.sh +6 -3
  8. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/go-emotions/run.sh +1 -1
  9. {tinydistill-0.3.0 → tinydistill-0.3.1}/pyproject.toml +1 -1
  10. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/__init__.py +1 -1
  11. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/cli.py +34 -2
  12. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/export.py +79 -0
  13. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/student.py +1 -1
  14. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/teacher.py +28 -0
  15. tinydistill-0.3.1/src/tinydistill/web.py +533 -0
  16. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/test_cli.py +29 -2
  17. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/test_export.py +27 -0
  18. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/test_serve.py +20 -10
  19. tinydistill-0.3.0/src/tinydistill/web.py +0 -171
  20. {tinydistill-0.3.0 → tinydistill-0.3.1}/.gitignore +0 -0
  21. {tinydistill-0.3.0 → tinydistill-0.3.1}/LICENSE +0 -0
  22. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/banking77-1.5b-8shot.json +0 -0
  23. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/banking77-1.5b-8shot.md +0 -0
  24. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/banking77-1.5b.json +0 -0
  25. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/banking77-1.5b.md +0 -0
  26. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/banking77-7b-8shot.json +0 -0
  27. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/banking77-7b-8shot.md +0 -0
  28. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/banking77-7b.json +0 -0
  29. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/banking77-7b.md +0 -0
  30. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/task.yaml +0 -0
  31. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/banking77/teachers.md +0 -0
  32. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-1.5b-8shot.json +0 -0
  33. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-1.5b-8shot.md +0 -0
  34. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-1.5b.json +0 -0
  35. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-1.5b.md +0 -0
  36. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-7b.json +0 -0
  37. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-7b.md +0 -0
  38. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-ensemble.json +0 -0
  39. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-ensemble.md +0 -0
  40. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-synthetic.json +0 -0
  41. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/financial-tweets-synthetic.md +0 -0
  42. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/task.yaml +0 -0
  43. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/financial-tweets/teachers.md +0 -0
  44. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/go-emotions/go-emotions-1.5b.json +0 -0
  45. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/go-emotions/go-emotions-1.5b.md +0 -0
  46. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/go-emotions/task.yaml +0 -0
  47. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/update_readme.py +0 -0
  48. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/vllm-smoke.json +0 -0
  49. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/vllm-smoke.md +0 -0
  50. {tinydistill-0.3.0 → tinydistill-0.3.1}/examples/vllm-smoke.sh +0 -0
  51. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/analyse.py +0 -0
  52. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/data.py +0 -0
  53. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/evaluate.py +0 -0
  54. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/hub.py +0 -0
  55. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/progress.py +0 -0
  56. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/serve.py +0 -0
  57. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/shortlist.py +0 -0
  58. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/static.py +0 -0
  59. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/suggest.py +0 -0
  60. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/synth.py +0 -0
  61. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/task.py +0 -0
  62. {tinydistill-0.3.0 → tinydistill-0.3.1}/src/tinydistill/tune.py +0 -0
  63. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/conftest.py +0 -0
  64. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/test_cache.py +0 -0
  65. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/test_loss.py +0 -0
  66. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/test_progress.py +0 -0
  67. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/test_scoring.py +0 -0
  68. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/test_task.py +0 -0
  69. {tinydistill-0.3.0 → tinydistill-0.3.1}/tests/test_vllm.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: tinydistill
3
- Version: 0.3.0
3
+ Version: 0.3.1
4
4
  Summary: Distil an LLM classification prompt into a small, fast model that runs on a CPU.
5
5
  Project-URL: Homepage, https://github.com/Andreas-Hum/tinydistill
6
6
  Author: Andreas-Hum
@@ -70,7 +70,8 @@ results files in [`examples/`](https://github.com/Andreas-Hum/tinydistill/tree/m
70
70
  quotes them.
71
71
 
72
72
  **Try it:** the [live demo](https://andreas-hum.github.io/tinydistill/) runs the Banking77 student (7B 8-shot
73
- teacher, int8, 23 MB) in your browser.
73
+ teacher, int8, 23 MB) in your browser and shows what happens inside it: the word pieces, what each layer would
74
+ answer, where the classifier looks, which words matter, and the teacher's soft labels next to the student's.
74
75
 
75
76
  ### Banking77: the hard case
76
77
 
@@ -124,15 +125,18 @@ The distilled student is far better than the same student trained on the few gol
124
125
  CPU in about a millisecond. Sending its least confident texts to the teacher closes most of the remaining gap
125
126
  to the teacher.
126
127
 
127
- Other students, the same teacher and soft labels. `static` is the static-embedding student
128
- (`--student static`: token vectors from MiniLM-L6, averaged, no attention), trained for 20 epochs at learning
129
- rate 1e-3 instead of 5 epochs at 1e-4:
128
+ Other students, the same teacher and soft labels. Students train for 5 epochs unless the name says otherwise.
129
+ `static` is the static-embedding student (`--student static`: token vectors from MiniLM-L6, averaged, no
130
+ attention), trained for 20 epochs at learning rate 1e-3:
130
131
 
131
132
  <!-- results: banking77/students.md -->
132
133
  | Run | Teacher | Teacher accuracy | Student | Parameters (M) | Accuracy | Macro-F1 | Agreement with teacher | Deployed as | Latency, batch 1 (ms) | Latency, batch 32 (ms) | Size (MB) |
133
134
  |---|---|---|---|---|---|---|---|---|---|---|---|
134
135
  | banking77-7b-8shot | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L6-v2` | 23 | 0.729 ± 0.011 | 0.697 ± 0.015 | 0.800 ± 0.003 | ONNX int8 | 1.4 | 40.2 | 23 |
135
136
  | all-MiniLM-L12-v2 | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L12-v2` | 33 | 0.743 ± 0.011 | 0.710 ± 0.016 | 0.815 ± 0.004 | ONNX int8 | 3.3 | 79.2 | 34 |
137
+ | all-MiniLM-L12-v2-20-epochs | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L12-v2` | 33 | 0.817 ± 0.001 | 0.814 ± 0.001 | 0.875 ± 0.002 | ONNX int8 | 3.2 | 82.0 | 34 |
138
+ | all-MiniLM-L6-v2-10-epochs | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L6-v2` | 23 | 0.808 ± 0.003 | 0.801 ± 0.002 | 0.860 ± 0.004 | ONNX int8 | 1.1 | 38.1 | 23 |
139
+ | all-MiniLM-L6-v2-20-epochs | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L6-v2` | 23 | 0.813 ± 0.002 | 0.811 ± 0.003 | 0.868 ± 0.003 | ONNX int8 | 1.2 | 36.3 | 23 |
136
140
  | all-mpnet-base-v2 | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-mpnet-base-v2` | 110 | 0.814 ± 0.004 | 0.805 ± 0.007 | 0.863 ± 0.001 | ONNX int8 | 7.3 | 220.4 | 110 |
137
141
  | bge-small-en-v1.5 | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `BAAI/bge-small-en-v1.5` | 33 | 0.739 ± 0.008 | 0.706 ± 0.011 | 0.815 ± 0.004 | ONNX int8 | 2.7 | 75.2 | 34 |
138
142
  | static | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `static` | 8 | 0.803 ± 0.001 | 0.801 ± 0.001 | 0.824 ± 0.001 | ONNX int8 | 0.1 | 5.1 | 8 |
@@ -140,10 +144,14 @@ rate 1e-3 instead of 5 epochs at 1e-4:
140
144
  All latencies on the CPU with 2 threads.
141
145
  <!-- /results -->
142
146
 
143
- Two students are more accurate than their own teacher: the 110M-parameter mpnet and the static student.
144
- Learning from the teacher's soft labels on 10,000 texts, plus the 616 gold labels, smooths out some of its
145
- individual mistakes. The static student does this at about a tenth of a millisecond per text and 8 MB. It
146
- also trained four times longer, so the gap to the small transformers is not only about architecture.
147
+ Training longer matters more than a bigger student: with 20 epochs instead of 5, the 23M-parameter MiniLM-L6
148
+ student is as accurate as the 110M-parameter mpnet, at a sixth of its latency, and more accurate than its own
149
+ 7B teacher. Learning from the teacher's soft labels on 10,000 texts, plus the 616 gold labels, smooths out
150
+ some of the teacher's individual mistakes. The static student is about a point behind the MiniLM students
151
+ trained as long, at a tenth of a millisecond per text and 8 MB. Ten epochs get most of the way (MiniLM-L6:
152
+ 0.808), so 10 is the default from 0.3.1. On the easier financial-tweets task, 20 epochs give the same accuracy
153
+ as 5 ([`financial-tweets-7b-20-epochs.md`](https://github.com/Andreas-Hum/tinydistill/blob/main/examples/financial-tweets/financial-tweets-7b-20-epochs.md)).
154
+ The other tables on this page were trained for 5 epochs, the default until 0.3.1.
147
155
 
148
156
  ### Financial tweets: an easier task
149
157
 
@@ -253,7 +261,8 @@ Banking77 splits are sorted by intent, so a limited run only sees a few intents.
253
261
  Data can be `.csv`, `.jsonl` or `hf:<dataset>[:<config>]:<split>` (needs the `hf` extra). Useful options:
254
262
  `--teacher` (any Hugging Face causal LM), `--load-4bit` (bitsandbytes, CUDA only: `pip install "tinydistill[4bit]"`),
255
263
  `--student` (any Hugging Face encoder), `--epochs`, `--lr`, `--device`, `--threads`, `--seed`, and `--cache-dir`
256
- (runs that share it label each text only once). See `tinydistill <command> --help`.
264
+ (runs that share it label each text only once, and a run whose texts are all cached does not load the
265
+ teacher at all). See `tinydistill <command> --help`.
257
266
 
258
267
  Long steps print their progress, speed and time left every few seconds and write it to `RUN/status.json`;
259
268
  `tinydistill status RUN` reads it from another terminal and tells you whether the process is still alive.
@@ -261,8 +270,8 @@ Long steps print their progress, speed and time left every few seconds and write
261
270
  **Static student.** `--student static` (or `static:<encoder>`) trains a student with no attention at all: one
262
271
  vector per token, averaged, then a linear layer, started from the encoder's own token representations (as in
263
272
  Model2Vec). It is much smaller and faster than a transformer student. It ignores word order, so it can miss
264
- meaning that depends on it (negation, "A rather than B"); on Banking77 it beat the small transformer
265
- students and came close to mpnet (see the results). It exports to ONNX and runs with `predict` and `serve` like any other student, but
273
+ meaning that depends on it (negation, "A rather than B"); on Banking77 it came within about a point of
274
+ MiniLM students trained as long (see the results). It exports to ONNX and runs with `predict` and `serve` like any other student, but
266
275
  not in the browser page. Train it longer and at a higher learning rate, such as `--epochs 20 --lr 1e-3`.
267
276
 
268
277
  ### Choosing the teacher
@@ -307,7 +316,9 @@ using only the standard library. It listens on 127.0.0.1 unless you pass `--host
307
316
 
308
317
  **Browser.** `tinydistill web RUN/onnx --out site` writes a static page plus the model in the layout
309
318
  [transformers.js](https://github.com/huggingface/transformers.js) expects; the int8 model then runs in the
310
- visitor's browser with no server at all. Serve the folder with any static host (GitHub Pages works).
319
+ visitor's browser with no server at all. Below the answer, an "Inside the model" view shows each step of the inference
320
+ (tokens, the answer after each layer, attention, probabilities, which words matter, the teacher's soft labels
321
+ for the examples and the cascade), so the page also explains distillation. Serve the folder with any static host (GitHub Pages works).
311
322
 
312
323
  ### A few gold labels
313
324
 
@@ -40,7 +40,8 @@ results files in [`examples/`](https://github.com/Andreas-Hum/tinydistill/tree/m
40
40
  quotes them.
41
41
 
42
42
  **Try it:** the [live demo](https://andreas-hum.github.io/tinydistill/) runs the Banking77 student (7B 8-shot
43
- teacher, int8, 23 MB) in your browser.
43
+ teacher, int8, 23 MB) in your browser and shows what happens inside it: the word pieces, what each layer would
44
+ answer, where the classifier looks, which words matter, and the teacher's soft labels next to the student's.
44
45
 
45
46
  ### Banking77: the hard case
46
47
 
@@ -94,15 +95,18 @@ The distilled student is far better than the same student trained on the few gol
94
95
  CPU in about a millisecond. Sending its least confident texts to the teacher closes most of the remaining gap
95
96
  to the teacher.
96
97
 
97
- Other students, the same teacher and soft labels. `static` is the static-embedding student
98
- (`--student static`: token vectors from MiniLM-L6, averaged, no attention), trained for 20 epochs at learning
99
- rate 1e-3 instead of 5 epochs at 1e-4:
98
+ Other students, the same teacher and soft labels. Students train for 5 epochs unless the name says otherwise.
99
+ `static` is the static-embedding student (`--student static`: token vectors from MiniLM-L6, averaged, no
100
+ attention), trained for 20 epochs at learning rate 1e-3:
100
101
 
101
102
  <!-- results: banking77/students.md -->
102
103
  | Run | Teacher | Teacher accuracy | Student | Parameters (M) | Accuracy | Macro-F1 | Agreement with teacher | Deployed as | Latency, batch 1 (ms) | Latency, batch 32 (ms) | Size (MB) |
103
104
  |---|---|---|---|---|---|---|---|---|---|---|---|
104
105
  | banking77-7b-8shot | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L6-v2` | 23 | 0.729 ± 0.011 | 0.697 ± 0.015 | 0.800 ± 0.003 | ONNX int8 | 1.4 | 40.2 | 23 |
105
106
  | all-MiniLM-L12-v2 | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L12-v2` | 33 | 0.743 ± 0.011 | 0.710 ± 0.016 | 0.815 ± 0.004 | ONNX int8 | 3.3 | 79.2 | 34 |
107
+ | all-MiniLM-L12-v2-20-epochs | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L12-v2` | 33 | 0.817 ± 0.001 | 0.814 ± 0.001 | 0.875 ± 0.002 | ONNX int8 | 3.2 | 82.0 | 34 |
108
+ | all-MiniLM-L6-v2-10-epochs | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L6-v2` | 23 | 0.808 ± 0.003 | 0.801 ± 0.002 | 0.860 ± 0.004 | ONNX int8 | 1.1 | 38.1 | 23 |
109
+ | all-MiniLM-L6-v2-20-epochs | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L6-v2` | 23 | 0.813 ± 0.002 | 0.811 ± 0.003 | 0.868 ± 0.003 | ONNX int8 | 1.2 | 36.3 | 23 |
106
110
  | all-mpnet-base-v2 | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-mpnet-base-v2` | 110 | 0.814 ± 0.004 | 0.805 ± 0.007 | 0.863 ± 0.001 | ONNX int8 | 7.3 | 220.4 | 110 |
107
111
  | bge-small-en-v1.5 | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `BAAI/bge-small-en-v1.5` | 33 | 0.739 ± 0.008 | 0.706 ± 0.011 | 0.815 ± 0.004 | ONNX int8 | 2.7 | 75.2 | 34 |
108
112
  | static | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `static` | 8 | 0.803 ± 0.001 | 0.801 ± 0.001 | 0.824 ± 0.001 | ONNX int8 | 0.1 | 5.1 | 8 |
@@ -110,10 +114,14 @@ rate 1e-3 instead of 5 epochs at 1e-4:
110
114
  All latencies on the CPU with 2 threads.
111
115
  <!-- /results -->
112
116
 
113
- Two students are more accurate than their own teacher: the 110M-parameter mpnet and the static student.
114
- Learning from the teacher's soft labels on 10,000 texts, plus the 616 gold labels, smooths out some of its
115
- individual mistakes. The static student does this at about a tenth of a millisecond per text and 8 MB. It
116
- also trained four times longer, so the gap to the small transformers is not only about architecture.
117
+ Training longer matters more than a bigger student: with 20 epochs instead of 5, the 23M-parameter MiniLM-L6
118
+ student is as accurate as the 110M-parameter mpnet, at a sixth of its latency, and more accurate than its own
119
+ 7B teacher. Learning from the teacher's soft labels on 10,000 texts, plus the 616 gold labels, smooths out
120
+ some of the teacher's individual mistakes. The static student is about a point behind the MiniLM students
121
+ trained as long, at a tenth of a millisecond per text and 8 MB. Ten epochs get most of the way (MiniLM-L6:
122
+ 0.808), so 10 is the default from 0.3.1. On the easier financial-tweets task, 20 epochs give the same accuracy
123
+ as 5 ([`financial-tweets-7b-20-epochs.md`](https://github.com/Andreas-Hum/tinydistill/blob/main/examples/financial-tweets/financial-tweets-7b-20-epochs.md)).
124
+ The other tables on this page were trained for 5 epochs, the default until 0.3.1.
117
125
 
118
126
  ### Financial tweets: an easier task
119
127
 
@@ -223,7 +231,8 @@ Banking77 splits are sorted by intent, so a limited run only sees a few intents.
223
231
  Data can be `.csv`, `.jsonl` or `hf:<dataset>[:<config>]:<split>` (needs the `hf` extra). Useful options:
224
232
  `--teacher` (any Hugging Face causal LM), `--load-4bit` (bitsandbytes, CUDA only: `pip install "tinydistill[4bit]"`),
225
233
  `--student` (any Hugging Face encoder), `--epochs`, `--lr`, `--device`, `--threads`, `--seed`, and `--cache-dir`
226
- (runs that share it label each text only once). See `tinydistill <command> --help`.
234
+ (runs that share it label each text only once, and a run whose texts are all cached does not load the
235
+ teacher at all). See `tinydistill <command> --help`.
227
236
 
228
237
  Long steps print their progress, speed and time left every few seconds and write it to `RUN/status.json`;
229
238
  `tinydistill status RUN` reads it from another terminal and tells you whether the process is still alive.
@@ -231,8 +240,8 @@ Long steps print their progress, speed and time left every few seconds and write
231
240
  **Static student.** `--student static` (or `static:<encoder>`) trains a student with no attention at all: one
232
241
  vector per token, averaged, then a linear layer, started from the encoder's own token representations (as in
233
242
  Model2Vec). It is much smaller and faster than a transformer student. It ignores word order, so it can miss
234
- meaning that depends on it (negation, "A rather than B"); on Banking77 it beat the small transformer
235
- students and came close to mpnet (see the results). It exports to ONNX and runs with `predict` and `serve` like any other student, but
243
+ meaning that depends on it (negation, "A rather than B"); on Banking77 it came within about a point of
244
+ MiniLM students trained as long (see the results). It exports to ONNX and runs with `predict` and `serve` like any other student, but
236
245
  not in the browser page. Train it longer and at a higher learning rate, such as `--epochs 20 --lr 1e-3`.
237
246
 
238
247
  ### Choosing the teacher
@@ -277,7 +286,9 @@ using only the standard library. It listens on 127.0.0.1 unless you pass `--host
277
286
 
278
287
  **Browser.** `tinydistill web RUN/onnx --out site` writes a static page plus the model in the layout
279
288
  [transformers.js](https://github.com/huggingface/transformers.js) expects; the int8 model then runs in the
280
- visitor's browser with no server at all. Serve the folder with any static host (GitHub Pages works).
289
+ visitor's browser with no server at all. Below the answer, an "Inside the model" view shows each step of the inference
290
+ (tokens, the answer after each layer, attention, probabilities, which words matter, the teacher's soft labels
291
+ for the examples and the cascade), so the page also explains distillation. Serve the folder with any static host (GitHub Pages works).
281
292
 
282
293
  ### A few gold labels
283
294
 
@@ -11,7 +11,8 @@ set -euo pipefail
11
11
  cd "$(dirname "$0")/../.."
12
12
  export OMP_NUM_THREADS=${OMP_NUM_THREADS:-2} MKL_NUM_THREADS=${MKL_NUM_THREADS:-2}
13
13
  base=(examples/banking77/task.yaml --data hf:PolyAI/banking77:train --test hf:PolyAI/banking77:test
14
- --cache-dir runs/cache --device cuda --threads "$OMP_NUM_THREADS" --int8 --seeds 3)
14
+ --cache-dir runs/cache --device cuda --threads "$OMP_NUM_THREADS" --int8 --seeds 3
15
+ --epochs 5) # the default before 0.3.1; the lines that pass --epochs 20 override it
15
16
  few=(--few-shot 8 --gold-per-class 8 --gold-only)
16
17
  big=(--teacher Qwen/Qwen2.5-7B-Instruct --load-4bit)
17
18
 
@@ -28,6 +29,14 @@ done
28
29
  # a static-embedding student (no attention), which wants a higher learning rate and more epochs
29
30
  tinydistill run "${base[@]}" "${big[@]}" --few-shot 8 --gold-per-class 8 --student static --lr 1e-3 --epochs 20 \
30
31
  --out runs/banking77-students/static --skip-teacher-bench
32
+ # MiniLM-L6 at 10 epochs (the default from 0.3.1), and the two smallest transformer students trained as long
33
+ # as the static one, for a like-for-like comparison
34
+ tinydistill run "${base[@]}" "${big[@]}" --few-shot 8 --gold-per-class 8 --epochs 10 \
35
+ --out runs/banking77-students/all-MiniLM-L6-v2-10-epochs --skip-teacher-bench
36
+ for student in sentence-transformers/all-MiniLM-L6-v2 sentence-transformers/all-MiniLM-L12-v2; do
37
+ tinydistill run "${base[@]}" "${big[@]}" --few-shot 8 --gold-per-class 8 --student "$student" --epochs 20 \
38
+ --out "runs/banking77-students/${student##*/}-20-epochs" --skip-teacher-bench
39
+ done
31
40
 
32
41
  for run in banking77-1.5b banking77-1.5b-8shot banking77-7b banking77-7b-8shot; do
33
42
  cp "runs/$run/results.json" "examples/banking77/$run.json"
@@ -2,6 +2,9 @@
2
2
  |---|---|---|---|---|---|---|---|---|---|---|---|
3
3
  | banking77-7b-8shot | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L6-v2` | 23 | 0.729 ± 0.011 | 0.697 ± 0.015 | 0.800 ± 0.003 | ONNX int8 | 1.4 | 40.2 | 23 |
4
4
  | all-MiniLM-L12-v2 | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L12-v2` | 33 | 0.743 ± 0.011 | 0.710 ± 0.016 | 0.815 ± 0.004 | ONNX int8 | 3.3 | 79.2 | 34 |
5
+ | all-MiniLM-L12-v2-20-epochs | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L12-v2` | 33 | 0.817 ± 0.001 | 0.814 ± 0.001 | 0.875 ± 0.002 | ONNX int8 | 3.2 | 82.0 | 34 |
6
+ | all-MiniLM-L6-v2-10-epochs | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L6-v2` | 23 | 0.808 ± 0.003 | 0.801 ± 0.002 | 0.860 ± 0.004 | ONNX int8 | 1.1 | 38.1 | 23 |
7
+ | all-MiniLM-L6-v2-20-epochs | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-MiniLM-L6-v2` | 23 | 0.813 ± 0.002 | 0.811 ± 0.003 | 0.868 ± 0.003 | ONNX int8 | 1.2 | 36.3 | 23 |
5
8
  | all-mpnet-base-v2 | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `sentence-transformers/all-mpnet-base-v2` | 110 | 0.814 ± 0.004 | 0.805 ± 0.007 | 0.863 ± 0.001 | ONNX int8 | 7.3 | 220.4 | 110 |
6
9
  | bge-small-en-v1.5 | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `BAAI/bge-small-en-v1.5` | 33 | 0.739 ± 0.008 | 0.706 ± 0.011 | 0.815 ± 0.004 | ONNX int8 | 2.7 | 75.2 | 34 |
7
10
  | static | 8-shot, 4-bit `Qwen/Qwen2.5-7B-Instruct` | 0.783 | `static` | 8 | 0.803 ± 0.001 | 0.801 ± 0.001 | 0.824 ± 0.001 | ONNX int8 | 0.1 | 5.1 | 8 |
@@ -0,0 +1,146 @@
1
+ {
2
+ "task": "financial-tweets",
3
+ "teacher": "Qwen/Qwen2.5-7B-Instruct",
4
+ "student": "sentence-transformers/all-MiniLM-L6-v2",
5
+ "data": {
6
+ "train": "hf:zeroshot/twitter-financial-news-sentiment:train",
7
+ "test": "hf:zeroshot/twitter-financial-news-sentiment:validation",
8
+ "n_train": 9543,
9
+ "n_test": 2388
10
+ },
11
+ "training": {
12
+ "epochs": 20,
13
+ "lr": 0.0001,
14
+ "batch_size": 32,
15
+ "gold_per_class": 0,
16
+ "gold_file": null,
17
+ "n_gold_file": 0,
18
+ "n_gold": 0,
19
+ "upper_bound": false,
20
+ "gold_only": false,
21
+ "seeds": 3,
22
+ "n_gold_upper_bound": null,
23
+ "label_texts_scored": 0,
24
+ "few_shot": 0,
25
+ "n_examples": 0,
26
+ "ce_weight": 0.0,
27
+ "student_device": "cuda",
28
+ "train_seconds": 820.4,
29
+ "teacher_device": "cuda:0",
30
+ "label_seconds": 6.2,
31
+ "load_4bit": true,
32
+ "shortlist": null,
33
+ "shortlist_model": null
34
+ },
35
+ "environment": {
36
+ "cpu": "x86_64",
37
+ "cpu_threads": 2,
38
+ "torch": "2.15.0+cu132",
39
+ "gpu": "NVIDIA GeForce RTX 4090",
40
+ "omp_num_threads": "2"
41
+ },
42
+ "models": [
43
+ {
44
+ "role": "teacher",
45
+ "name": "teacher, zero-shot, 4-bit",
46
+ "device": "cuda:0",
47
+ "model": "Qwen/Qwen2.5-7B-Instruct",
48
+ "accuracy": 0.8320770519262981,
49
+ "macro_f1": 0.7875756953889134,
50
+ "size_mb": null
51
+ },
52
+ {
53
+ "role": "student",
54
+ "name": "student, PyTorch",
55
+ "device": "cpu",
56
+ "accuracy": 0.824818537130095,
57
+ "accuracy_std": 0.004398646358623353,
58
+ "accuracy_seeds": [
59
+ 0.8249581239530989,
60
+ 0.8203517587939698,
61
+ 0.8291457286432161
62
+ ],
63
+ "macro_f1": 0.7645078986696947,
64
+ "macro_f1_std": 0.0037534820261736976,
65
+ "macro_f1_seeds": [
66
+ 0.7618092859102075,
67
+ 0.7629200730807385,
68
+ 0.7687943370181382
69
+ ],
70
+ "agreement": 0.8710217755443885,
71
+ "agreement_std": 0.002931323283082083,
72
+ "agreement_seeds": [
73
+ 0.8697654941373534,
74
+ 0.86892797319933,
75
+ 0.8743718592964824
76
+ ],
77
+ "params_m": 22.714371,
78
+ "size_mb": 90.869492,
79
+ "latency_ms_b1": 4.909115499685868,
80
+ "latency_ms_b32": 150.38988000014797,
81
+ "throughput_per_s": 212.78027484275216
82
+ },
83
+ {
84
+ "role": "student",
85
+ "name": "student, ONNX fp32",
86
+ "device": "cpu",
87
+ "accuracy": 0.8249581239530989,
88
+ "macro_f1": 0.7618092859102075,
89
+ "agreement": 0.8697654941373534,
90
+ "size_mb": 90.990363,
91
+ "latency_ms_b1": 3.9684855019004317,
92
+ "latency_ms_b32": 148.22442249896994,
93
+ "throughput_per_s": 215.88884922268716
94
+ },
95
+ {
96
+ "role": "student",
97
+ "name": "student, ONNX int8",
98
+ "device": "cpu",
99
+ "accuracy": 0.8203517587939698,
100
+ "macro_f1": 0.7589981023943616,
101
+ "agreement": 0.8710217755443886,
102
+ "size_mb": 23.043866,
103
+ "latency_ms_b1": 1.3503949994628783,
104
+ "latency_ms_b32": 74.80496900097933,
105
+ "throughput_per_s": 427.77906905597496
106
+ }
107
+ ],
108
+ "cascade": [
109
+ {
110
+ "share_to_teacher": 0.0,
111
+ "threshold": 0.3775518238544464,
112
+ "accuracy": 0.8203517587939698,
113
+ "macro_f1": 0.7589981023943616
114
+ },
115
+ {
116
+ "share_to_teacher": 0.10008375209380234,
117
+ "threshold": 0.7880504727363586,
118
+ "accuracy": 0.8345896147403685,
119
+ "macro_f1": 0.7804097140755729
120
+ },
121
+ {
122
+ "share_to_teacher": 0.20016750418760468,
123
+ "threshold": 0.953411340713501,
124
+ "accuracy": 0.8375209380234506,
125
+ "macro_f1": 0.7885267426679218
126
+ },
127
+ {
128
+ "share_to_teacher": 0.2998324958123953,
129
+ "threshold": 0.9881662130355835,
130
+ "accuracy": 0.8358458961474037,
131
+ "macro_f1": 0.7879874137344648
132
+ },
133
+ {
134
+ "share_to_teacher": 0.5,
135
+ "threshold": 0.9987922310829163,
136
+ "accuracy": 0.833752093802345,
137
+ "macro_f1": 0.7876367442676434
138
+ },
139
+ {
140
+ "share_to_teacher": 1.0,
141
+ "threshold": 1.0,
142
+ "accuracy": 0.8320770519262981,
143
+ "macro_f1": 0.7875756953889134
144
+ }
145
+ ]
146
+ }
@@ -0,0 +1,22 @@
1
+ | 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 |
2
+ |---|---|---|---|---|---|---|---|---|---|
3
+ | teacher, zero-shot, 4-bit | cuda:0 | 0.832 | 0.788 | n/a | n/a | n/a | n/a | n/a | n/a |
4
+ | student, PyTorch | cpu | 0.825 ± 0.004 | 0.765 ± 0.004 | 0.871 ± 0.003 | 4.9 | 150.4 | 212.8 | 91 | n/a |
5
+ | student, ONNX fp32 | cpu | 0.825 | 0.762 | 0.870 | 4.0 | 148.2 | 215.9 | 91 | n/a |
6
+ | student, ONNX int8 | cpu | 0.820 | 0.759 | 0.871 | 1.4 | 74.8 | 427.8 | 23 | n/a |
7
+
8
+ Task `financial-tweets`, 2388 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 9543 unlabelled training texts. Values with ± are the mean and standard deviation over 3 training seeds; the ONNX rows are the exported first-seed model.
9
+ 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 packs several texts into one forward pass; it is timed at batch 32 only when that takes under three minutes, otherwise its throughput is 1 / batch-1 latency.
10
+
11
+ ### Cascade: student first, unsure texts to the teacher
12
+
13
+ | Share sent to the teacher | Student confidence threshold | Accuracy | Macro-F1 | Mean latency per text (ms) |
14
+ |---|---|---|---|---|
15
+ | 0% | 0.378 | 0.820 | 0.759 | n/a |
16
+ | 10% | 0.788 | 0.835 | 0.780 | n/a |
17
+ | 20% | 0.953 | 0.838 | 0.789 | n/a |
18
+ | 30% | 0.988 | 0.836 | 0.788 | n/a |
19
+ | 50% | 0.999 | 0.834 | 0.788 | n/a |
20
+ | 100% | 1.000 | 0.832 | 0.788 | n/a |
21
+
22
+ The student answers every text and passes its least confident ones to the teacher. Thresholds are set on the test texts to hit each share; in deployment, set them on a validation split. Latency is the student's batch-1 latency plus the share times the batch-1 latency of the teacher.
@@ -12,7 +12,7 @@ export OMP_NUM_THREADS=${OMP_NUM_THREADS:-2} MKL_NUM_THREADS=${MKL_NUM_THREADS:-
12
12
  data=hf:zeroshot/twitter-financial-news-sentiment
13
13
  task=examples/financial-tweets/task.yaml
14
14
  base=("$task" --test "$data:validation" --cache-dir runs/cache --device cuda --threads "$OMP_NUM_THREADS"
15
- --int8 --seeds 3)
15
+ --int8 --seeds 3 --epochs 5) # 5 epochs: the default before 0.3.1
16
16
  q7=Qwen/Qwen2.5-7B-Instruct
17
17
 
18
18
  real=(--data "$data:train" --upper-bound) # the synthetic run has no gold labels, so no upper bound there
@@ -22,6 +22,9 @@ tinydistill run "${base[@]}" "${real[@]}" --few-shot 8 --gold-per-class 8 --gold
22
22
  --out runs/financial-tweets-1.5b-8shot --teacher-devices cpu,cuda
23
23
  tinydistill run "${base[@]}" "${real[@]}" --teacher "$q7" --load-4bit --out runs/financial-tweets-7b \
24
24
  --teacher-devices cuda
25
+ # the same student trained for 20 epochs: does training longer help on this task too?
26
+ tinydistill run "${base[@]}" --data "$data:train" --teacher "$q7" --load-4bit --epochs 20 \
27
+ --out runs/financial-tweets-7b-20-epochs --skip-teacher-bench
25
28
  tinydistill run "${base[@]}" "${real[@]}" --load-4bit --teacher-devices cuda \
26
29
  --teacher "$q7,mistralai/Mistral-7B-Instruct-v0.3,microsoft/Phi-3.5-mini-instruct" \
27
30
  --out runs/financial-tweets-ensemble
@@ -30,8 +33,8 @@ tinydistill synth "$task" --per-label 500 --teacher "$q7" --load-4bit --device c
30
33
  tinydistill run "${base[@]}" --data runs/financial-tweets-synth.jsonl --teacher "$q7" --load-4bit \
31
34
  --out runs/financial-tweets-synthetic --skip-teacher-bench
32
35
 
33
- for run in financial-tweets-1.5b financial-tweets-1.5b-8shot financial-tweets-7b financial-tweets-ensemble \
34
- financial-tweets-synthetic; do
36
+ for run in financial-tweets-1.5b financial-tweets-1.5b-8shot financial-tweets-7b financial-tweets-7b-20-epochs \
37
+ financial-tweets-ensemble financial-tweets-synthetic; do
35
38
  cp "runs/$run/results.json" "examples/financial-tweets/$run.json"
36
39
  cp "runs/$run/results.md" "examples/financial-tweets/$run.md"
37
40
  done
@@ -10,7 +10,7 @@ data=hf:google-research-datasets/go_emotions:simplified
10
10
 
11
11
  tinydistill run examples/go-emotions/task.yaml --data "$data:train" --limit 10000 --test "$data:test" \
12
12
  --cache-dir runs/cache --device cuda --threads "$OMP_NUM_THREADS" --int8 --seeds 3 --upper-bound \
13
- --gold-per-class 8 --gold-only --out runs/go-emotions-1.5b --teacher-devices cpu,cuda
13
+ --gold-per-class 8 --gold-only --epochs 5 --out runs/go-emotions-1.5b --teacher-devices cpu,cuda
14
14
 
15
15
  cp runs/go-emotions-1.5b/results.json examples/go-emotions/go-emotions-1.5b.json
16
16
  cp runs/go-emotions-1.5b/results.md examples/go-emotions/go-emotions-1.5b.md
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "tinydistill"
7
- version = "0.3.0"
7
+ version = "0.3.1"
8
8
  description = "Distil an LLM classification prompt into a small, fast model that runs on a CPU."
9
9
  readme = "README.md"
10
10
  license = "MIT"
@@ -29,4 +29,4 @@ __all__ = [
29
29
  "read_examples",
30
30
  "train_student",
31
31
  ]
32
- __version__ = "0.3.0"
32
+ __version__ = "0.3.1"
@@ -53,7 +53,15 @@ def cmd_label(a: argparse.Namespace) -> None:
53
53
 
54
54
  from .data import count_rows, read_examples
55
55
  from .progress import Progress
56
- from .teacher import ScoreCache, combine_scores, iter_unique, load_teacher
56
+ from .teacher import (
57
+ ScoreCache,
58
+ combine_scores,
59
+ iter_unique,
60
+ load_teacher,
61
+ pick_device,
62
+ quick_key,
63
+ text_hash,
64
+ )
57
65
 
58
66
  task = load_task(a.task)
59
67
  run = Path(a.out)
@@ -85,7 +93,25 @@ def cmd_label(a: argparse.Namespace) -> None:
85
93
 
86
94
  caches, fingerprints, device = [], [], "cpu"
87
95
  seconds, scored = 0.0, 0
96
+ keys_file = cache_dir / "teacher-keys.json" # quick key (tokenizer only) -> full fingerprint
97
+ keys = json.loads(keys_file.read_text()) if keys_file.exists() else {}
88
98
  for name in _teachers(a.teacher): # one teacher in memory at a time
99
+ quick = None
100
+ if not a.vllm_url:
101
+ quick = quick_key(name, task, a.device, a.load_4bit, examples, a.shortlist, a.shortlist_model)
102
+ known = keys.get(quick)
103
+ cache = ScoreCache(cache_dir / f"teacher-{known}.jsonl") if known else None
104
+ if cache and all(
105
+ text_hash(e.text) in cache.scores
106
+ for _, spec, limit in splits
107
+ for e in read_examples(spec, task, limit, with_gold=False)
108
+ ):
109
+ log(f"{name}: every text is in the cache, so the teacher is not loaded")
110
+ fingerprints.append(known)
111
+ caches.append(cache)
112
+ device = pick_device(a.device)
113
+ device = "cuda:0" if device == "cuda" else device # as a loaded model reports it
114
+ continue
89
115
  log(f"loading teacher {name}")
90
116
  scorer = load_teacher(
91
117
  name,
@@ -100,6 +126,12 @@ def cmd_label(a: argparse.Namespace) -> None:
100
126
  )
101
127
  fingerprints.append(scorer.fingerprint())
102
128
  device = str(scorer.device)
129
+ if quick:
130
+ keys = json.loads(keys_file.read_text()) if keys_file.exists() else {}
131
+ keys[quick] = fingerprints[-1]
132
+ tmp = keys_file.with_suffix(".tmp")
133
+ tmp.write_text(json.dumps(keys, indent=1))
134
+ tmp.replace(keys_file)
103
135
  cache = ScoreCache(cache_dir / f"teacher-{fingerprints[-1]}.jsonl")
104
136
  if cache.scores:
105
137
  log(f"resuming: {len(cache.scores)} texts already scored")
@@ -883,7 +915,7 @@ def build_parser() -> argparse.ArgumentParser:
883
915
  help="Hugging Face encoder, or `static` / `static:<encoder>` for a static-embedding student "
884
916
  "(no attention: much smaller and faster; try --epochs 20 --lr 1e-3)",
885
917
  )
886
- train.add_argument("--epochs", type=int, default=5)
918
+ train.add_argument("--epochs", type=int, default=10, help="passes over the training texts (default 10)")
887
919
  train.add_argument("--batch-size", type=int, default=32)
888
920
  train.add_argument("--lr", type=float, default=1e-4)
889
921
  train.add_argument("--max-length", type=int, default=128)
@@ -54,6 +54,85 @@ def export_onnx(student_dir: Path, out_dir: Path, int8: bool = False, opset: int
54
54
  return written
55
55
 
56
56
 
57
+ def _classifier_head(model):
58
+ """The student's classification head as a function of one layer's hidden states (batch, tokens, dim), or
59
+ None when the model has no such head (the static student)."""
60
+ head, base = getattr(model, "classifier", None), getattr(model, "base_model", None)
61
+ if head is None or base is None or not hasattr(base, "encoder"):
62
+ return None
63
+ if isinstance(head, torch.nn.Linear): # BERT style: the pooler reads the first token
64
+ pooler = getattr(base, "pooler", None)
65
+ return (lambda h: head(pooler(h))) if pooler is not None else (lambda h: head(h[:, 0]))
66
+ return head # RoBERTa/MPNet style heads take the whole sequence and read its first token themselves
67
+
68
+
69
+ class _Explain(torch.nn.Module):
70
+ """Logits, plus what the classifier would answer after each layer, and how much the first token (which
71
+ the classifier reads) attends to every token in each layer, averaged over heads."""
72
+
73
+ def __init__(self, model, head):
74
+ super().__init__()
75
+ self.model, self.head = model, head
76
+
77
+ def forward(self, input_ids, attention_mask):
78
+ out = self.model.base_model(
79
+ input_ids=input_ids,
80
+ attention_mask=attention_mask,
81
+ output_hidden_states=True,
82
+ output_attentions=True,
83
+ )
84
+ hs = torch.stack(out.hidden_states[1:], 1) # (batch, layers, tokens, dim)
85
+ b, n, t, d = hs.shape
86
+ # the head runs once on all layers stacked as a batch (one set of weights in the graph)
87
+ layers = self.head(hs.reshape(b * n, t, d)).reshape(b, n, -1) # (batch, layers, labels)
88
+ attention = torch.stack([a[:, :, 0, :].mean(1) for a in out.attentions], 1) # (batch, layers, tokens)
89
+ return layers[:, -1], layers, attention
90
+
91
+
92
+ def export_explain(student_dir: Path, out_dir: Path, int8: bool = False, opset: int = 17) -> Path | None:
93
+ """Export ``explain.onnx`` (and ``explain.int8.onnx``) with outputs ``logits``, ``layer_logits`` and
94
+ ``cls_attention``, for the teaching view of `tinydistill web`. None if the model has no layers to show."""
95
+ from transformers import AutoModelForSequenceClassification
96
+
97
+ from .static import is_static
98
+
99
+ if is_static(student_dir):
100
+ return None
101
+ model = AutoModelForSequenceClassification.from_pretrained(
102
+ student_dir, attn_implementation="eager"
103
+ ).eval()
104
+ head = _classifier_head(model)
105
+ if head is None:
106
+ return None
107
+ _, tok = load_student(student_dir)
108
+ enc = tok(["an example input", "another one"], padding=True, return_tensors="pt")
109
+ out_dir.mkdir(parents=True, exist_ok=True)
110
+ path = out_dir / "explain.onnx"
111
+ torch.onnx.export(
112
+ _Explain(model, head),
113
+ (enc["input_ids"], enc["attention_mask"]),
114
+ str(path),
115
+ input_names=["input_ids", "attention_mask"],
116
+ output_names=["logits", "layer_logits", "cls_attention"],
117
+ dynamic_axes={
118
+ "input_ids": {0: "batch", 1: "seq"},
119
+ "attention_mask": {0: "batch", 1: "seq"},
120
+ "logits": {0: "batch"},
121
+ "layer_logits": {0: "batch"},
122
+ "cls_attention": {0: "batch", 2: "seq"},
123
+ },
124
+ opset_version=opset,
125
+ dynamo=False,
126
+ )
127
+ if not int8:
128
+ return path
129
+ from onnxruntime.quantization import QuantType, quantize_dynamic
130
+
131
+ q = out_dir / "explain.int8.onnx"
132
+ quantize_dynamic(str(path), str(q), weight_type=QuantType.QInt8)
133
+ return q
134
+
135
+
57
136
  class Predictor:
58
137
  """Run an exported model: ``Predictor("runs/x/onnx").predict(["text", ...])``."""
59
138
 
@@ -96,7 +96,7 @@ def train_student(
96
96
  model_name: str = DEFAULT_STUDENT,
97
97
  gold: list[int | None] | None = None,
98
98
  ce_weight: float = 0.0,
99
- epochs: int = 5,
99
+ epochs: int = 10,
100
100
  batch_size: int = 32,
101
101
  lr: float = 1e-4,
102
102
  max_length: int = 128,
@@ -557,6 +557,34 @@ def load_teacher(
557
557
  return LabelScorer(model, tok, task, max_tokens=max_tokens, examples=examples, shortlist=short)
558
558
 
559
559
 
560
+ def quick_key(
561
+ name: str,
562
+ task: Task,
563
+ device: str,
564
+ load_4bit: bool,
565
+ examples: list[tuple[str, int]] | None = None,
566
+ shortlist: int | None = None,
567
+ shortlist_model: str | None = None,
568
+ ) -> str:
569
+ """What decides a local teacher's cache fingerprint, from the tokenizer alone (no model weights): the
570
+ rendered prompt and continuations, and how the model would be loaded. `label` maps it to the full
571
+ fingerprint, so a run whose texts are all cached never loads the teacher."""
572
+ p = _Prompt()
573
+ p._prepare(AutoTokenizer.from_pretrained(name), task, True, True, examples, None)
574
+ key = {
575
+ "model": name,
576
+ "device": pick_device(device),
577
+ "load_4bit": load_4bit,
578
+ "prefix": p.prefix_ids,
579
+ "suffix": p.suffix,
580
+ "forms": p.forms,
581
+ "form_label": p.form_label,
582
+ "form_answer": p.form_answer,
583
+ "shortlist": [shortlist, shortlist_model] if shortlist else None,
584
+ }
585
+ return hashlib.sha256(json.dumps(key, sort_keys=True).encode()).hexdigest()[:16]
586
+
587
+
560
588
  class ScoreCache:
561
589
  """Append-only JSONL of {"h": text hash, "s": label scores}; reopening it resumes where it stopped."""
562
590