Instructions to use macmacmacmac/CLEF-Screen with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- LiteRT
How to use macmacmacmac/CLEF-Screen with LiteRT:
# No code snippets available yet for this library. # To use this model, check the repository files and the library's documentation. # Want to help? PRs adding snippets are welcome at: # https://github.com/huggingface/huggingface.js
- Notebooks
- Google Colab
- Kaggle
CLEF-Screen: JointSchemaHead redo FT (seed1, 3 epochs) from jev-retrain MI355X run 2026-11-04
Browse files- .gitattributes +1 -0
- README.md +32 -0
- clef-flash_rowprefill_embeds_L4096_fp16fc.tflite +3 -0
- contract.json +789 -0
- head/joint_head.safetensors +3 -0
- head/joint_head_config.json +8 -0
- host/clef_litert/__init__.py +62 -0
- host/clef_litert/_erf.py +95 -0
- host/clef_litert/errors.py +44 -0
- host/clef_litert/head.py +231 -0
- host/clef_litert/host.py +134 -0
- host/clef_litert/pixels.py +80 -0
- host/clef_litert/positions.py +96 -0
- host/clef_litert/request.py +175 -0
- host/clef_litert/response.py +56 -0
- host/clef_litert/runtime.py +273 -0
- host/clef_litert/tables.py +125 -0
- host/requirements-host.txt +27 -0
- joint_schema_model.py +576 -0
- provenance_ft.json +40 -0
- redo-final-results.json +104 -0
- redo-provenance.json +26 -0
- redo-training-log.jsonl +0 -0
- tables/embed_tokens.safetensors +3 -0
- tables/lm_head.safetensors +3 -0
- tables/vision_pos_g256.safetensors +3 -0
- tables/vision_pos_g448.safetensors +3 -0
- tokenizer/tokenizer.json +3 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
|
README.md
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
language:
|
| 3 |
+
- en
|
| 4 |
+
license: other
|
| 5 |
+
tags:
|
| 6 |
+
- security
|
| 7 |
+
- prompt-injection
|
| 8 |
+
- litert
|
| 9 |
+
---
|
| 10 |
+
|
| 11 |
+
# CLEF-Screen
|
| 12 |
+
|
| 13 |
+
Redo fine-tune of the CLEF-Flash `JointSchemaHead` (121,762,820 parameters,
|
| 14 |
+
122 tensors) on frozen Cloudflare CLEF-Flash LiteRT backbone hidden states.
|
| 15 |
+
Only the head is trained; the backbone is unmodified. This repository root
|
| 16 |
+
is a complete CLEF LiteRT checkpoint directory with the fine-tuned head.
|
| 17 |
+
|
| 18 |
+
## Training
|
| 19 |
+
|
| 20 |
+
- Compute: DigitalOcean GPU droplet `jev-retrain` (AMD Instinct MI355X, ROCm), encode 2026-11-04 11:31Z, head training finished 18:24Z.
|
| 21 |
+
- Method: `ft_clef_bench.py`, 3 epoch(s), lr 3e-05, batch 16,
|
| 22 |
+
pos_weight 3.0. Training details and hashes: `redo-provenance.json`;
|
| 23 |
+
full log: `redo-training-log.jsonl` when present.
|
| 24 |
+
- Head dtype in this repo: BF16 (deployment contract). Trained head was already BF16; no conversion.
|
| 25 |
+
|
| 26 |
+
## Evaluation status
|
| 27 |
+
|
| 28 |
+
Dev (21,977 rows) epoch0 AUC 0.7516 -> epoch3 AUC 0.9899, TPR@0.5 0.9911 FPR@0.5 0.2183, TPR@0.99 0.9583 FPR@0.99 0.0507. Script-run test banking+slack (21,965 rows) AUC 0.9897, TPR@0.99 0.9574 FPR@0.99 0.0537. Caveats: dev/test are 92% positive (~1,695 negatives); FPR@0.5 remains 0.22; test was evaluated by the training script at finish; no independent frozen replication yet. Full numbers: redo-final-results.json.
|
| 29 |
+
|
| 30 |
+
No frozen-test, adaptive-attack, or independent-replication result exists yet
|
| 31 |
+
for this checkpoint unless stated above. Treat all numbers as development
|
| 32 |
+
measurements, not deployment guarantees.
|
clef-flash_rowprefill_embeds_L4096_fp16fc.tflite
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:43a638f2976f8bd8e885bdd7d41b1a342bb583c41c81f4528b1379ad4dee823b
|
| 3 |
+
size 13936871792
|
contract.json
ADDED
|
@@ -0,0 +1,789 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"what": "clef-flash on LiteRT: the contract between the graphs and a host",
|
| 3 |
+
"model": {
|
| 4 |
+
"repo": "Cloudflare/clef-flash",
|
| 5 |
+
"revision": "f0b7d3cd65296abfe459e6c228b618f5d8b26d2f",
|
| 6 |
+
"joint_schema_model_sha256": "0e304cf7c6500e8bb59bef7e2afd2c6373f82596dfb3b57d1aa93c175e2dc3a3",
|
| 7 |
+
"hidden_size": 4096,
|
| 8 |
+
"vocab_rows": 248320,
|
| 9 |
+
"tie_word_embeddings": "false (embed_tokens and lm_head are two tables)"
|
| 10 |
+
},
|
| 11 |
+
"host": {
|
| 12 |
+
"package": "host/clef_litert",
|
| 13 |
+
"version": "0.2.0",
|
| 14 |
+
"requirements": "host/requirements-host.txt (Python 3.14.6, macOS 27.0): numpy, tokenizers, pillow and ai-edge-litert (imported when a graph is compiled); no torch, no transformers",
|
| 15 |
+
"files_sha256": {
|
| 16 |
+
"__init__.py": "4bc079245c456d5ef34b21c3cfd6614e367d5173a2fbf95d95a778279d9e80a2",
|
| 17 |
+
"_erf.py": "cf6a894a8d65d069f64fb4f0f0175d969bd2947c31be89eb0992ab3435a4e66d",
|
| 18 |
+
"errors.py": "43893807154cfa7822cdaabf2eb8db3650b9ccc615e2b4231f9437f527e94744",
|
| 19 |
+
"head.py": "7fe0940018882ab7a83b12e78b884e194c332d5a7fcde9a93419f4bdec57f861",
|
| 20 |
+
"host.py": "b4e39b9f525f2f7fd83ad29c8cb301acfcda3780c4484f577a16fdcd5a50627a",
|
| 21 |
+
"pixels.py": "d87420187e422ff35299f5028236954fbbe63f6653e0d8a90a2ccfb0a500abf2",
|
| 22 |
+
"positions.py": "44a86a2843c550fab76d6e9754518e07e1f142ba252ebe6e584825a2870915d4",
|
| 23 |
+
"request.py": "991a3300bf9240d80b30a886f3d96cb77551840ed53d2af4221a5eaa0e118ab0",
|
| 24 |
+
"response.py": "5658ca395deef7ee20c68e8244f5b1dd42af7c8021c6a8138acd685d89d91445",
|
| 25 |
+
"runtime.py": "11ee9dc121b9c224a693d8c9eab244bde5b0bfb4baa10dcd9499167423f8ee7a",
|
| 26 |
+
"tables.py": "6db98e0069162a8a0d9bbc56111f3676d382b34e0b2a1c4a86794f287fe2a521"
|
| 27 |
+
},
|
| 28 |
+
"entry": "ClefLiteRT.from_dir(folder, accelerator='gpu' | 'cpu').decide(request, tile=448 | 256)",
|
| 29 |
+
"steps": "host/clef_litert/__init__.py docstring 1-6; ClefHost.decide(request, row_graph, tower)"
|
| 30 |
+
},
|
| 31 |
+
"folder": {
|
| 32 |
+
"layout": {
|
| 33 |
+
"row": "clef-flash_rowprefill_embeds_L{length}_fp16fc.tflite",
|
| 34 |
+
"tower": "clef-flash_vision_tower_g{tile}_fp16fc.tflite",
|
| 35 |
+
"tower_pos": "tables/vision_pos_g{tile}.safetensors",
|
| 36 |
+
"embed_tokens": "tables/embed_tokens.safetensors",
|
| 37 |
+
"lm_head": "tables/lm_head.safetensors",
|
| 38 |
+
"joint_head": "head/joint_head.safetensors",
|
| 39 |
+
"joint_head_config": "head/joint_head_config.json",
|
| 40 |
+
"tokenizer": "tokenizer/tokenizer.json"
|
| 41 |
+
},
|
| 42 |
+
"lengths": [
|
| 43 |
+
256,
|
| 44 |
+
512,
|
| 45 |
+
1024,
|
| 46 |
+
2048,
|
| 47 |
+
4096
|
| 48 |
+
],
|
| 49 |
+
"tiles": [
|
| 50 |
+
256,
|
| 51 |
+
448
|
| 52 |
+
],
|
| 53 |
+
"rule": "a request runs on the smallest present row graph that holds it; longer than the largest present one = RequestTooLong; one row graph compiled at a time, the towers stay compiled",
|
| 54 |
+
"accelerators": {
|
| 55 |
+
"gpu": "HardwareAccelerator.GPU, GpuOptions(enforce_f32=True) (Metal on a Mac)",
|
| 56 |
+
"cpu": "HardwareAccelerator.CPU, CpuOptions(num_threads=8) (XNNPACK), no weight cache file"
|
| 57 |
+
},
|
| 58 |
+
"not_supported": "the GPU's default precision (fp16 storage): the row graphs' norms overflow and the towers' rows are wrong (measured, not shipped)"
|
| 59 |
+
},
|
| 60 |
+
"row_graphs": [
|
| 61 |
+
{
|
| 62 |
+
"length": 256,
|
| 63 |
+
"file": "clef-flash_rowprefill_embeds_L256_fp16fc.tflite",
|
| 64 |
+
"bytes": 13847841008,
|
| 65 |
+
"sha256": "97d64beb2ec0a57e1e9f17a2c89096f6a44ed8b13c85216b9641d11f0195e6c7",
|
| 66 |
+
"source": "out/r2/clef_rowprefill_L256_v2_fp16fc.tflite",
|
| 67 |
+
"inputs": {
|
| 68 |
+
"embeds": [
|
| 69 |
+
1,
|
| 70 |
+
256,
|
| 71 |
+
4096
|
| 72 |
+
],
|
| 73 |
+
"valid": [
|
| 74 |
+
1,
|
| 75 |
+
256
|
| 76 |
+
],
|
| 77 |
+
"cos": [
|
| 78 |
+
1,
|
| 79 |
+
256,
|
| 80 |
+
64
|
| 81 |
+
],
|
| 82 |
+
"sin": [
|
| 83 |
+
1,
|
| 84 |
+
256,
|
| 85 |
+
64
|
| 86 |
+
]
|
| 87 |
+
},
|
| 88 |
+
"outputs": {
|
| 89 |
+
"hidden": [
|
| 90 |
+
1,
|
| 91 |
+
256,
|
| 92 |
+
4096
|
| 93 |
+
]
|
| 94 |
+
}
|
| 95 |
+
},
|
| 96 |
+
{
|
| 97 |
+
"length": 512,
|
| 98 |
+
"file": "clef-flash_rowprefill_embeds_L512_fp16fc.tflite",
|
| 99 |
+
"bytes": 13850105024,
|
| 100 |
+
"sha256": "cbff9eab565f921b762875d7e92ff6aadccb4b87b1c1e2c536b48318f130f09e",
|
| 101 |
+
"source": "out/r2/clef_rowprefill_L512_v2_fp16fc.tflite",
|
| 102 |
+
"inputs": {
|
| 103 |
+
"embeds": [
|
| 104 |
+
1,
|
| 105 |
+
512,
|
| 106 |
+
4096
|
| 107 |
+
],
|
| 108 |
+
"valid": [
|
| 109 |
+
1,
|
| 110 |
+
512
|
| 111 |
+
],
|
| 112 |
+
"cos": [
|
| 113 |
+
1,
|
| 114 |
+
512,
|
| 115 |
+
64
|
| 116 |
+
],
|
| 117 |
+
"sin": [
|
| 118 |
+
1,
|
| 119 |
+
512,
|
| 120 |
+
64
|
| 121 |
+
]
|
| 122 |
+
},
|
| 123 |
+
"outputs": {
|
| 124 |
+
"hidden": [
|
| 125 |
+
1,
|
| 126 |
+
512,
|
| 127 |
+
4096
|
| 128 |
+
]
|
| 129 |
+
}
|
| 130 |
+
},
|
| 131 |
+
{
|
| 132 |
+
"length": 1024,
|
| 133 |
+
"file": "clef-flash_rowprefill_embeds_L1024_fp16fc.tflite",
|
| 134 |
+
"bytes": 13856205536,
|
| 135 |
+
"sha256": "b8eff74efec3b2f3d2b289d62f5a5f66721f0926e88790f791fe8b4170834dc5",
|
| 136 |
+
"source": "out/r2/clef_rowprefill_L1024_v2_fp16fc.tflite",
|
| 137 |
+
"inputs": {
|
| 138 |
+
"embeds": [
|
| 139 |
+
1,
|
| 140 |
+
1024,
|
| 141 |
+
4096
|
| 142 |
+
],
|
| 143 |
+
"valid": [
|
| 144 |
+
1,
|
| 145 |
+
1024
|
| 146 |
+
],
|
| 147 |
+
"cos": [
|
| 148 |
+
1,
|
| 149 |
+
1024,
|
| 150 |
+
64
|
| 151 |
+
],
|
| 152 |
+
"sin": [
|
| 153 |
+
1,
|
| 154 |
+
1024,
|
| 155 |
+
64
|
| 156 |
+
]
|
| 157 |
+
},
|
| 158 |
+
"outputs": {
|
| 159 |
+
"hidden": [
|
| 160 |
+
1,
|
| 161 |
+
1024,
|
| 162 |
+
4096
|
| 163 |
+
]
|
| 164 |
+
}
|
| 165 |
+
},
|
| 166 |
+
{
|
| 167 |
+
"length": 2048,
|
| 168 |
+
"file": "clef-flash_rowprefill_embeds_L2048_fp16fc.tflite",
|
| 169 |
+
"bytes": 13874698016,
|
| 170 |
+
"sha256": "6ecdbb68d48a9406aaf0c3c9d789c631abe23c605484f9f34746944cb87ff7a3",
|
| 171 |
+
"source": "out/r2/clef_rowprefill_L2048_v2_fp16fc.tflite",
|
| 172 |
+
"inputs": {
|
| 173 |
+
"embeds": [
|
| 174 |
+
1,
|
| 175 |
+
2048,
|
| 176 |
+
4096
|
| 177 |
+
],
|
| 178 |
+
"valid": [
|
| 179 |
+
1,
|
| 180 |
+
2048
|
| 181 |
+
],
|
| 182 |
+
"cos": [
|
| 183 |
+
1,
|
| 184 |
+
2048,
|
| 185 |
+
64
|
| 186 |
+
],
|
| 187 |
+
"sin": [
|
| 188 |
+
1,
|
| 189 |
+
2048,
|
| 190 |
+
64
|
| 191 |
+
]
|
| 192 |
+
},
|
| 193 |
+
"outputs": {
|
| 194 |
+
"hidden": [
|
| 195 |
+
1,
|
| 196 |
+
2048,
|
| 197 |
+
4096
|
| 198 |
+
]
|
| 199 |
+
}
|
| 200 |
+
},
|
| 201 |
+
{
|
| 202 |
+
"length": 4096,
|
| 203 |
+
"file": "clef-flash_rowprefill_embeds_L4096_fp16fc.tflite",
|
| 204 |
+
"bytes": 13936871792,
|
| 205 |
+
"sha256": "43a638f2976f8bd8e885bdd7d41b1a342bb583c41c81f4528b1379ad4dee823b",
|
| 206 |
+
"source": "out/r2/clef_rowprefill_L4096_v2_fp16fc.tflite",
|
| 207 |
+
"inputs": {
|
| 208 |
+
"embeds": [
|
| 209 |
+
1,
|
| 210 |
+
4096,
|
| 211 |
+
4096
|
| 212 |
+
],
|
| 213 |
+
"valid": [
|
| 214 |
+
1,
|
| 215 |
+
4096
|
| 216 |
+
],
|
| 217 |
+
"cos": [
|
| 218 |
+
1,
|
| 219 |
+
4096,
|
| 220 |
+
64
|
| 221 |
+
],
|
| 222 |
+
"sin": [
|
| 223 |
+
1,
|
| 224 |
+
4096,
|
| 225 |
+
64
|
| 226 |
+
]
|
| 227 |
+
},
|
| 228 |
+
"outputs": {
|
| 229 |
+
"hidden": [
|
| 230 |
+
1,
|
| 231 |
+
4096,
|
| 232 |
+
4096
|
| 233 |
+
]
|
| 234 |
+
}
|
| 235 |
+
}
|
| 236 |
+
],
|
| 237 |
+
"tower_graphs": [
|
| 238 |
+
{
|
| 239 |
+
"tile": 256,
|
| 240 |
+
"file": "clef-flash_vision_tower_g256_fp16fc.tflite",
|
| 241 |
+
"bytes": 907587008,
|
| 242 |
+
"sha256": "c4272400890fa479b6d51792dd51e5e93cb74291a0fca90f94b38a248e57e994",
|
| 243 |
+
"source": "out/r3/clef_tower_g256_v2_fp16fc.tflite",
|
| 244 |
+
"inputs": {
|
| 245 |
+
"patches": [
|
| 246 |
+
1,
|
| 247 |
+
256,
|
| 248 |
+
1536
|
| 249 |
+
],
|
| 250 |
+
"pos": [
|
| 251 |
+
1,
|
| 252 |
+
256,
|
| 253 |
+
1152
|
| 254 |
+
]
|
| 255 |
+
},
|
| 256 |
+
"outputs": {
|
| 257 |
+
"rows": [
|
| 258 |
+
1,
|
| 259 |
+
64,
|
| 260 |
+
4096
|
| 261 |
+
]
|
| 262 |
+
},
|
| 263 |
+
"position_rows": {
|
| 264 |
+
"file": "tables/vision_pos_g256.safetensors",
|
| 265 |
+
"key": "pos",
|
| 266 |
+
"dtype": "F32",
|
| 267 |
+
"shape": [
|
| 268 |
+
256,
|
| 269 |
+
1152
|
| 270 |
+
],
|
| 271 |
+
"sha256": "ecd2eba829853f233b9f71f05da81c459e0ecb5ba7fbf5f5986bf42857c5efe9",
|
| 272 |
+
"source": "cache/r3/consts_g256.npz pos_embeds"
|
| 273 |
+
}
|
| 274 |
+
},
|
| 275 |
+
{
|
| 276 |
+
"tile": 448,
|
| 277 |
+
"file": "clef-flash_vision_tower_g448_fp16fc.tflite",
|
| 278 |
+
"bytes": 907891136,
|
| 279 |
+
"sha256": "f18f8d2eefc7241bcfacf93fee9a8c725c1ad9696c63e215ee5aa5167b2308a6",
|
| 280 |
+
"source": "out/r3/clef_tower_g448_v2_fp16fc.tflite",
|
| 281 |
+
"inputs": {
|
| 282 |
+
"patches": [
|
| 283 |
+
1,
|
| 284 |
+
784,
|
| 285 |
+
1536
|
| 286 |
+
],
|
| 287 |
+
"pos": [
|
| 288 |
+
1,
|
| 289 |
+
784,
|
| 290 |
+
1152
|
| 291 |
+
]
|
| 292 |
+
},
|
| 293 |
+
"outputs": {
|
| 294 |
+
"rows": [
|
| 295 |
+
1,
|
| 296 |
+
196,
|
| 297 |
+
4096
|
| 298 |
+
]
|
| 299 |
+
},
|
| 300 |
+
"position_rows": {
|
| 301 |
+
"file": "tables/vision_pos_g448.safetensors",
|
| 302 |
+
"key": "pos",
|
| 303 |
+
"dtype": "F32",
|
| 304 |
+
"shape": [
|
| 305 |
+
784,
|
| 306 |
+
1152
|
| 307 |
+
],
|
| 308 |
+
"sha256": "65f76c320150de7d3e53c6fa0ff7ffd048079153a8d6e114b6d48141c6f9ae26",
|
| 309 |
+
"source": "cache/r3/consts_g448.npz pos_embeds"
|
| 310 |
+
}
|
| 311 |
+
}
|
| 312 |
+
],
|
| 313 |
+
"files": {
|
| 314 |
+
"tokenizer/tokenizer.json": {
|
| 315 |
+
"sha256": "06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523",
|
| 316 |
+
"bytes": 19989325,
|
| 317 |
+
"source": "the repository's own file, unchanged"
|
| 318 |
+
},
|
| 319 |
+
"head/joint_head.safetensors": {
|
| 320 |
+
"sha256": "19cdcec8c81dc9212be320fff47462ab342fbc1278be4368fb3da71241cf5ba0",
|
| 321 |
+
"bytes": 243538016,
|
| 322 |
+
"source": "the repository's own file, unchanged (122 BF16 tensors)"
|
| 323 |
+
},
|
| 324 |
+
"head/joint_head_config.json": {
|
| 325 |
+
"sha256": "77efe959a38b5b17b241543e129e695f3c77465ece55a25985279bd8176279a0",
|
| 326 |
+
"bytes": 119,
|
| 327 |
+
"content": {
|
| 328 |
+
"hidden_size": 4096,
|
| 329 |
+
"width": 1024,
|
| 330 |
+
"routing_layers": 2,
|
| 331 |
+
"layers": 4,
|
| 332 |
+
"heads": 16,
|
| 333 |
+
"feedforward": 4096
|
| 334 |
+
}
|
| 335 |
+
},
|
| 336 |
+
"tables/embed_tokens.safetensors": {
|
| 337 |
+
"sha256": "48586a2af482671b5ca781a761957e296295a7fd0950af241705cb01ad2079fc",
|
| 338 |
+
"bytes": 2034237544,
|
| 339 |
+
"key": "embed_tokens.weight",
|
| 340 |
+
"dtype": "BF16",
|
| 341 |
+
"shape": [
|
| 342 |
+
248320,
|
| 343 |
+
4096
|
| 344 |
+
],
|
| 345 |
+
"metadata": null,
|
| 346 |
+
"source": {
|
| 347 |
+
"file": "model-00001-of-00004.safetensors",
|
| 348 |
+
"key": "model.language_model.embed_tokens.weight",
|
| 349 |
+
"byte_range": [
|
| 350 |
+
2034242120,
|
| 351 |
+
4068479560
|
| 352 |
+
],
|
| 353 |
+
"data_sha256": "d87d444a38822588c6c1ed07dec4404ccd8d526f57d58565cbbbf60acc758afe"
|
| 354 |
+
},
|
| 355 |
+
"made_by": "conversion/r5a_tables.py (clef_litert.copy_bf16_tensor): the checkpoint's bytes, new name",
|
| 356 |
+
"role": "the row graph's embeds: the rows of the ids, widened to float32 by the host"
|
| 357 |
+
},
|
| 358 |
+
"tables/lm_head.safetensors": {
|
| 359 |
+
"sha256": "106ec926c62c1aac2c8ab808833db89012a9132c53f767fa7609ec69e377d4ab",
|
| 360 |
+
"bytes": 2034237536,
|
| 361 |
+
"key": "lm_head.weight",
|
| 362 |
+
"dtype": "BF16",
|
| 363 |
+
"shape": [
|
| 364 |
+
248320,
|
| 365 |
+
4096
|
| 366 |
+
],
|
| 367 |
+
"metadata": null,
|
| 368 |
+
"source": {
|
| 369 |
+
"file": "model-00001-of-00004.safetensors",
|
| 370 |
+
"key": "lm_head.weight",
|
| 371 |
+
"byte_range": [
|
| 372 |
+
4680,
|
| 373 |
+
2034242120
|
| 374 |
+
],
|
| 375 |
+
"data_sha256": "6d6ff4dbf261f833246d32e7c6088f87f4ab40ddb58516d90c931aa6ae61d5e7"
|
| 376 |
+
},
|
| 377 |
+
"made_by": "conversion/r5a_tables.py (clef_litert.copy_bf16_tensor): the checkpoint's bytes, new name",
|
| 378 |
+
"role": "the joint head's lexical vectors: the mean of the rows of each option span's ids"
|
| 379 |
+
}
|
| 380 |
+
},
|
| 381 |
+
"token_ids": {
|
| 382 |
+
"vision_start": 248053,
|
| 383 |
+
"vision_end": 248054,
|
| 384 |
+
"image_pad": 248056,
|
| 385 |
+
"processor_pad": 248044
|
| 386 |
+
},
|
| 387 |
+
"request": {
|
| 388 |
+
"shape": "SystemOne-compatible request shape: {model: str, state: any JSON, questions: {id: {type, instructions, criteria}}, images: [one picture] (optional)}",
|
| 389 |
+
"question_types": {
|
| 390 |
+
"noul": 0,
|
| 391 |
+
"choice": 1,
|
| 392 |
+
"score": 2
|
| 393 |
+
},
|
| 394 |
+
"checks": "the author's systemone() checks, in order (request.validate_request), then the host's",
|
| 395 |
+
"refusals": {
|
| 396 |
+
"MissingModelOrState": "model is not a string or state is absent: 'model and state are required'",
|
| 397 |
+
"NoQuestions": "questions is not a non-empty object: 'at least one question is required'",
|
| 398 |
+
"BadQuestionType": "type not noul / choice / score: '<id>: type must be noul, choice, or score'",
|
| 399 |
+
"EmptyCriteria": "choice or score without criteria: '<id>: criteria must not be empty'",
|
| 400 |
+
"SchemaTooLong": "prefix + picture + schema + suffix > max_length (16384), encode_record's message",
|
| 401 |
+
"EmptyRecord": "no ids or no questions, encode_record's message",
|
| 402 |
+
"UnsupportedMedia": "host only: more than one picture, videos, media_kwargs, an unreadable picture or URL, a tile other than 256 / 448",
|
| 403 |
+
"RequestTooLong": "host only: more than 4096 tokens (refused, never truncated)",
|
| 404 |
+
"base": "every refusal is clef_litert.RequestError, a ValueError"
|
| 405 |
+
}
|
| 406 |
+
},
|
| 407 |
+
"ids": {
|
| 408 |
+
"encode": "tokenizers.Tokenizer.from_file(tokenizer.json).encode(piece, add_special_tokens=False), pieces concatenated (the author's encode_record)",
|
| 409 |
+
"prefix": "<|im_start|>system\nRead the complete state and schema. Decide every field jointly. Each answer must be exactly one of that field's allowed options.<|im_end|>\n<|im_start|>user\nSTATE:\n",
|
| 410 |
+
"prefix_ids": 36,
|
| 411 |
+
"suffix": "\n<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nJOINT SCHEMA DECISIONS:",
|
| 412 |
+
"suffix_ids": 18,
|
| 413 |
+
"schema_header": "\n\nSCHEMA FIELDS:\n",
|
| 414 |
+
"question_piece": "\\nFIELD {i}\\nID: {id}\\nTYPE: {type}\\nINSTRUCTION: [render(instructions or id)]\\nALLOWED OPTIONS:\\n + per option 'OPTION {j}: [render({option_id, description})]\\n' + 'END FIELD\\n' ([..] = the spans)",
|
| 415 |
+
"render": "a string as is, else json.dumps(ensure_ascii=False, separators=(',', ':'), sort_keys=True)",
|
| 416 |
+
"option_order": "noul (true, false); choice: ids sorted as strings; score: levels 0..n-1",
|
| 417 |
+
"noul_default_criteria": {
|
| 418 |
+
"true": "The proposition is true or the answer is yes.",
|
| 419 |
+
"false": "The proposition is false or the answer is no."
|
| 420 |
+
},
|
| 421 |
+
"picture_block": "<|vision_start|> + N x <|image_pad|> + <|vision_end|> + '\\n' right after the prefix (token_offset = the <|vision_start|> index)",
|
| 422 |
+
"state_cut": "to max_length 16384 exactly as encode_record (the row graphs refuse past 4096 first)",
|
| 423 |
+
"usage_input_tokens": "len(ids)"
|
| 424 |
+
},
|
| 425 |
+
"row_graph": {
|
| 426 |
+
"signature": "serving_default",
|
| 427 |
+
"inputs": [
|
| 428 |
+
{
|
| 429 |
+
"name": "embeds",
|
| 430 |
+
"dtype": "float32",
|
| 431 |
+
"shape": [
|
| 432 |
+
1,
|
| 433 |
+
"L",
|
| 434 |
+
4096
|
| 435 |
+
],
|
| 436 |
+
"meaning": "the bf16 embed_tokens rows of the ids widened to float32; the picture rows at the <|image_pad|> positions; 0 after the record"
|
| 437 |
+
},
|
| 438 |
+
{
|
| 439 |
+
"name": "valid",
|
| 440 |
+
"dtype": "float32",
|
| 441 |
+
"shape": [
|
| 442 |
+
1,
|
| 443 |
+
"L"
|
| 444 |
+
],
|
| 445 |
+
"meaning": "1 on the record, 0 after it"
|
| 446 |
+
},
|
| 447 |
+
{
|
| 448 |
+
"name": "cos",
|
| 449 |
+
"dtype": "float32",
|
| 450 |
+
"shape": [
|
| 451 |
+
1,
|
| 452 |
+
"L",
|
| 453 |
+
64
|
| 454 |
+
],
|
| 455 |
+
"meaning": "M-RoPE cos of the record's position planes (positions.py)"
|
| 456 |
+
},
|
| 457 |
+
{
|
| 458 |
+
"name": "sin",
|
| 459 |
+
"dtype": "float32",
|
| 460 |
+
"shape": [
|
| 461 |
+
1,
|
| 462 |
+
"L",
|
| 463 |
+
64
|
| 464 |
+
],
|
| 465 |
+
"meaning": "M-RoPE sin"
|
| 466 |
+
}
|
| 467 |
+
],
|
| 468 |
+
"outputs": [
|
| 469 |
+
{
|
| 470 |
+
"name": "hidden",
|
| 471 |
+
"dtype": "float32",
|
| 472 |
+
"shape": [
|
| 473 |
+
1,
|
| 474 |
+
"L",
|
| 475 |
+
4096
|
| 476 |
+
],
|
| 477 |
+
"meaning": "after the final RMSNorm, every position; the host reads rows 0..T-1"
|
| 478 |
+
}
|
| 479 |
+
],
|
| 480 |
+
"buckets": [
|
| 481 |
+
256,
|
| 482 |
+
512,
|
| 483 |
+
1024,
|
| 484 |
+
2048,
|
| 485 |
+
4096
|
| 486 |
+
],
|
| 487 |
+
"bucket_rule": "the smallest L >= T; T > 4096 refused (RequestTooLong)",
|
| 488 |
+
"padding": "right; embeds 0, valid 0; pad positions continue max + 1, + 2, ... on all three planes",
|
| 489 |
+
"files": [
|
| 490 |
+
"clef-flash_rowprefill_embeds_L256_fp16fc.tflite",
|
| 491 |
+
"clef-flash_rowprefill_embeds_L512_fp16fc.tflite",
|
| 492 |
+
"clef-flash_rowprefill_embeds_L1024_fp16fc.tflite",
|
| 493 |
+
"clef-flash_rowprefill_embeds_L2048_fp16fc.tflite",
|
| 494 |
+
"clef-flash_rowprefill_embeds_L4096_fp16fc.tflite"
|
| 495 |
+
]
|
| 496 |
+
},
|
| 497 |
+
"positions": {
|
| 498 |
+
"planes": "text-only (i, i, i); picture of H x W merged tokens at token_offset P: i <= P -> (i, i, i), k-th picture token -> (P+1, P+1 + k // W, P+1 + k % W), i >= P+N+1 -> i - (N - max(H, W))",
|
| 499 |
+
"rope": {
|
| 500 |
+
"theta": 10000000.0,
|
| 501 |
+
"rotary_dim": 64,
|
| 502 |
+
"mrope_section": [
|
| 503 |
+
11,
|
| 504 |
+
11,
|
| 505 |
+
10
|
| 506 |
+
],
|
| 507 |
+
"interleave": "j % 3 == 1 and j < 33 -> h, j % 3 == 2 and j < 30 -> w, else t",
|
| 508 |
+
"tables": "angle = float32(inv_freq * pos); cos / sin of cat(angle, angle) in float64, rounded to float32"
|
| 509 |
+
}
|
| 510 |
+
},
|
| 511 |
+
"picture": {
|
| 512 |
+
"tiles": {
|
| 513 |
+
"256": {
|
| 514 |
+
"merged_grid": [
|
| 515 |
+
8,
|
| 516 |
+
8
|
| 517 |
+
],
|
| 518 |
+
"picture_tokens": 64,
|
| 519 |
+
"patches": [
|
| 520 |
+
256,
|
| 521 |
+
1536
|
| 522 |
+
]
|
| 523 |
+
},
|
| 524 |
+
"448": {
|
| 525 |
+
"merged_grid": [
|
| 526 |
+
14,
|
| 527 |
+
14
|
| 528 |
+
],
|
| 529 |
+
"picture_tokens": 196,
|
| 530 |
+
"patches": [
|
| 531 |
+
784,
|
| 532 |
+
1536
|
| 533 |
+
]
|
| 534 |
+
}
|
| 535 |
+
},
|
| 536 |
+
"default_tile": 448,
|
| 537 |
+
"preprocess": "PIL open -> convert('RGB') -> resize((tile, tile), BICUBIC) -> /255 -> (x - 0.5) / 0.5 (float64) -> float32 -> merge-block-major patches (C, T=2, 16, 16)",
|
| 538 |
+
"tower_graph": "patches [1, 4 G^2, 1536] and the position rows pos [1, 4 G^2, 1152] float32 in -> picture rows [1, G^2, 4096] float32 out (tower_graphs)",
|
| 539 |
+
"not_read": "the author's native resolution, more than one picture, videos, media_kwargs, URLs"
|
| 540 |
+
},
|
| 541 |
+
"head": {
|
| 542 |
+
"weights": "joint_head.safetensors (BF16 -> float32)",
|
| 543 |
+
"config": {
|
| 544 |
+
"hidden_size": 4096,
|
| 545 |
+
"width": 1024,
|
| 546 |
+
"routing_layers": 2,
|
| 547 |
+
"layers": 4,
|
| 548 |
+
"heads": 16,
|
| 549 |
+
"feedforward": 4096
|
| 550 |
+
},
|
| 551 |
+
"tensors": 122,
|
| 552 |
+
"parameters": 121762820,
|
| 553 |
+
"compute": "NumPy float32 (head.py docstring); GELU exact (erf in float64, FreeBSD s_erf.c)",
|
| 554 |
+
"lexical": "mean of the lm_head.safetensors rows of each option span's ids (not normalized)"
|
| 555 |
+
},
|
| 556 |
+
"response": {
|
| 557 |
+
"probabilities": "per question float32 softmax in NumPy's order (exp(z - max), float32 sum, float32 division), then Python floats",
|
| 558 |
+
"answer": "the author's systemone_answer (round(x, 4); choice = first maximum in criteria order; score = Python's left-to-right sum of level * p)",
|
| 559 |
+
"body": {
|
| 560 |
+
"model": "request.model",
|
| 561 |
+
"answers": "{question id: answer} in request order",
|
| 562 |
+
"usage": {
|
| 563 |
+
"input_tokens": "len(ids)",
|
| 564 |
+
"output_tokens": 0
|
| 565 |
+
}
|
| 566 |
+
}
|
| 567 |
+
},
|
| 568 |
+
"status": "the graphs, the host package and the folder of this repository",
|
| 569 |
+
"checks": {
|
| 570 |
+
"reference": "the provider's own code (joint_schema_model.py systemone() at the revision above) in float32 on the CPU: 243 runs of fixtures/ (text, and pictures at the 256 / 448 px tiles), fixtures/reference_probs.json",
|
| 571 |
+
"line": "on every question that is not a near-tie (top-2 margin <= 0.02) the most likely option is the reference's; every probability within 0.02; the mean over runs of each run's mean |dp| <= 0.002",
|
| 572 |
+
"end_to_end": {
|
| 573 |
+
"gpu": {
|
| 574 |
+
"complete": true,
|
| 575 |
+
"line": {
|
| 576 |
+
"runs": 243,
|
| 577 |
+
"argmax_equal_non_near_tie": "545/545",
|
| 578 |
+
"near_tie": "5/5",
|
| 579 |
+
"max_abs_dp": 0.0016525983810424805,
|
| 580 |
+
"mean_of_run_means": 7.334260193875453e-06,
|
| 581 |
+
"pass": true
|
| 582 |
+
},
|
| 583 |
+
"red_arms": "3/3",
|
| 584 |
+
"pass": true
|
| 585 |
+
},
|
| 586 |
+
"cpu": {
|
| 587 |
+
"complete": true,
|
| 588 |
+
"line": {
|
| 589 |
+
"runs": 69,
|
| 590 |
+
"argmax_equal_non_near_tie": "207/207",
|
| 591 |
+
"near_tie": "3/3",
|
| 592 |
+
"max_abs_dp": 0.0016247034072875977,
|
| 593 |
+
"mean_of_run_means": 1.3047657266541098e-05,
|
| 594 |
+
"pass": true
|
| 595 |
+
},
|
| 596 |
+
"red_arms": "3/3",
|
| 597 |
+
"pass": true
|
| 598 |
+
}
|
| 599 |
+
},
|
| 600 |
+
"example": {
|
| 601 |
+
"gpu": true,
|
| 602 |
+
"cpu": true
|
| 603 |
+
},
|
| 604 |
+
"delegation_as_compiled": {
|
| 605 |
+
"cpu:clef-flash_rowprefill_embeds_L1024_fp16fc.tflite": {
|
| 606 |
+
"compile_seconds": 82.362,
|
| 607 |
+
"fully_accelerated": false,
|
| 608 |
+
"one_partition_all_nodes": false,
|
| 609 |
+
"replacing": [
|
| 610 |
+
{
|
| 611 |
+
"delegated": 18112,
|
| 612 |
+
"total": 18160,
|
| 613 |
+
"delegate": "TfLiteXNNPackDelegate",
|
| 614 |
+
"partitions": 97
|
| 615 |
+
}
|
| 616 |
+
]
|
| 617 |
+
},
|
| 618 |
+
"cpu:clef-flash_rowprefill_embeds_L2048_fp16fc.tflite": {
|
| 619 |
+
"compile_seconds": 94.098,
|
| 620 |
+
"fully_accelerated": false,
|
| 621 |
+
"one_partition_all_nodes": false,
|
| 622 |
+
"replacing": [
|
| 623 |
+
{
|
| 624 |
+
"delegated": 31552,
|
| 625 |
+
"total": 31600,
|
| 626 |
+
"delegate": "TfLiteXNNPackDelegate",
|
| 627 |
+
"partitions": 97
|
| 628 |
+
}
|
| 629 |
+
]
|
| 630 |
+
},
|
| 631 |
+
"cpu:clef-flash_rowprefill_embeds_L256_fp16fc.tflite": {
|
| 632 |
+
"compile_seconds": 3.553,
|
| 633 |
+
"fully_accelerated": false,
|
| 634 |
+
"one_partition_all_nodes": false,
|
| 635 |
+
"replacing": [
|
| 636 |
+
{
|
| 637 |
+
"delegated": 8056,
|
| 638 |
+
"total": 8080,
|
| 639 |
+
"delegate": "TfLiteXNNPackDelegate",
|
| 640 |
+
"partitions": 49
|
| 641 |
+
}
|
| 642 |
+
]
|
| 643 |
+
},
|
| 644 |
+
"cpu:clef-flash_rowprefill_embeds_L4096_fp16fc.tflite": {
|
| 645 |
+
"compile_seconds": 107.976,
|
| 646 |
+
"fully_accelerated": false,
|
| 647 |
+
"one_partition_all_nodes": false,
|
| 648 |
+
"replacing": [
|
| 649 |
+
{
|
| 650 |
+
"delegated": 58432,
|
| 651 |
+
"total": 58480,
|
| 652 |
+
"delegate": "TfLiteXNNPackDelegate",
|
| 653 |
+
"partitions": 97
|
| 654 |
+
}
|
| 655 |
+
]
|
| 656 |
+
},
|
| 657 |
+
"cpu:clef-flash_rowprefill_embeds_L512_fp16fc.tflite": {
|
| 658 |
+
"compile_seconds": 81.079,
|
| 659 |
+
"fully_accelerated": false,
|
| 660 |
+
"one_partition_all_nodes": false,
|
| 661 |
+
"replacing": [
|
| 662 |
+
{
|
| 663 |
+
"delegated": 11392,
|
| 664 |
+
"total": 11440,
|
| 665 |
+
"delegate": "TfLiteXNNPackDelegate",
|
| 666 |
+
"partitions": 97
|
| 667 |
+
}
|
| 668 |
+
]
|
| 669 |
+
},
|
| 670 |
+
"cpu:clef-flash_vision_tower_g256_fp16fc.tflite": {
|
| 671 |
+
"compile_seconds": 0.247,
|
| 672 |
+
"fully_accelerated": true,
|
| 673 |
+
"one_partition_all_nodes": true,
|
| 674 |
+
"replacing": [
|
| 675 |
+
{
|
| 676 |
+
"delegated": 1940,
|
| 677 |
+
"total": 1940,
|
| 678 |
+
"delegate": "TfLiteXNNPackDelegate",
|
| 679 |
+
"partitions": 1
|
| 680 |
+
}
|
| 681 |
+
]
|
| 682 |
+
},
|
| 683 |
+
"cpu:clef-flash_vision_tower_g448_fp16fc.tflite": {
|
| 684 |
+
"compile_seconds": 0.401,
|
| 685 |
+
"fully_accelerated": true,
|
| 686 |
+
"one_partition_all_nodes": true,
|
| 687 |
+
"replacing": [
|
| 688 |
+
{
|
| 689 |
+
"delegated": 1940,
|
| 690 |
+
"total": 1940,
|
| 691 |
+
"delegate": "TfLiteXNNPackDelegate",
|
| 692 |
+
"partitions": 1
|
| 693 |
+
}
|
| 694 |
+
]
|
| 695 |
+
},
|
| 696 |
+
"gpu:clef-flash_rowprefill_embeds_L1024_fp16fc.tflite": {
|
| 697 |
+
"compile_seconds": 138.695,
|
| 698 |
+
"fully_accelerated": true,
|
| 699 |
+
"one_partition_all_nodes": true,
|
| 700 |
+
"replacing": [
|
| 701 |
+
{
|
| 702 |
+
"delegated": 18160,
|
| 703 |
+
"total": 18160,
|
| 704 |
+
"delegate": "LITERT_METAL",
|
| 705 |
+
"partitions": 1
|
| 706 |
+
}
|
| 707 |
+
]
|
| 708 |
+
},
|
| 709 |
+
"gpu:clef-flash_rowprefill_embeds_L2048_fp16fc.tflite": {
|
| 710 |
+
"compile_seconds": 146.108,
|
| 711 |
+
"fully_accelerated": true,
|
| 712 |
+
"one_partition_all_nodes": true,
|
| 713 |
+
"replacing": [
|
| 714 |
+
{
|
| 715 |
+
"delegated": 31600,
|
| 716 |
+
"total": 31600,
|
| 717 |
+
"delegate": "LITERT_METAL",
|
| 718 |
+
"partitions": 1
|
| 719 |
+
}
|
| 720 |
+
]
|
| 721 |
+
},
|
| 722 |
+
"gpu:clef-flash_rowprefill_embeds_L256_fp16fc.tflite": {
|
| 723 |
+
"compile_seconds": 172.478,
|
| 724 |
+
"fully_accelerated": true,
|
| 725 |
+
"one_partition_all_nodes": true,
|
| 726 |
+
"replacing": [
|
| 727 |
+
{
|
| 728 |
+
"delegated": 8080,
|
| 729 |
+
"total": 8080,
|
| 730 |
+
"delegate": "LITERT_METAL",
|
| 731 |
+
"partitions": 1
|
| 732 |
+
}
|
| 733 |
+
]
|
| 734 |
+
},
|
| 735 |
+
"gpu:clef-flash_rowprefill_embeds_L4096_fp16fc.tflite": {
|
| 736 |
+
"compile_seconds": 187.062,
|
| 737 |
+
"fully_accelerated": true,
|
| 738 |
+
"one_partition_all_nodes": true,
|
| 739 |
+
"replacing": [
|
| 740 |
+
{
|
| 741 |
+
"delegated": 58480,
|
| 742 |
+
"total": 58480,
|
| 743 |
+
"delegate": "LITERT_METAL",
|
| 744 |
+
"partitions": 1
|
| 745 |
+
}
|
| 746 |
+
]
|
| 747 |
+
},
|
| 748 |
+
"gpu:clef-flash_rowprefill_embeds_L512_fp16fc.tflite": {
|
| 749 |
+
"compile_seconds": 135.683,
|
| 750 |
+
"fully_accelerated": true,
|
| 751 |
+
"one_partition_all_nodes": true,
|
| 752 |
+
"replacing": [
|
| 753 |
+
{
|
| 754 |
+
"delegated": 11440,
|
| 755 |
+
"total": 11440,
|
| 756 |
+
"delegate": "LITERT_METAL",
|
| 757 |
+
"partitions": 1
|
| 758 |
+
}
|
| 759 |
+
]
|
| 760 |
+
},
|
| 761 |
+
"gpu:clef-flash_vision_tower_g256_fp16fc.tflite": {
|
| 762 |
+
"compile_seconds": 4.983,
|
| 763 |
+
"fully_accelerated": true,
|
| 764 |
+
"one_partition_all_nodes": true,
|
| 765 |
+
"replacing": [
|
| 766 |
+
{
|
| 767 |
+
"delegated": 1940,
|
| 768 |
+
"total": 1940,
|
| 769 |
+
"delegate": "LITERT_METAL",
|
| 770 |
+
"partitions": 1
|
| 771 |
+
}
|
| 772 |
+
]
|
| 773 |
+
},
|
| 774 |
+
"gpu:clef-flash_vision_tower_g448_fp16fc.tflite": {
|
| 775 |
+
"compile_seconds": 4.709,
|
| 776 |
+
"fully_accelerated": true,
|
| 777 |
+
"one_partition_all_nodes": true,
|
| 778 |
+
"replacing": [
|
| 779 |
+
{
|
| 780 |
+
"delegated": 1940,
|
| 781 |
+
"total": 1940,
|
| 782 |
+
"delegate": "LITERT_METAL",
|
| 783 |
+
"partitions": 1
|
| 784 |
+
}
|
| 785 |
+
]
|
| 786 |
+
}
|
| 787 |
+
}
|
| 788 |
+
}
|
| 789 |
+
}
|
head/joint_head.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:8ca78d0c127755948fb285c90caeefcf6e61a2b5f2f3cc726ffae41ca231ab11
|
| 3 |
+
size 243537984
|
head/joint_head_config.json
ADDED
|
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"hidden_size": 4096,
|
| 3 |
+
"width": 1024,
|
| 4 |
+
"routing_layers": 2,
|
| 5 |
+
"layers": 4,
|
| 6 |
+
"heads": 16,
|
| 7 |
+
"feedforward": 4096
|
| 8 |
+
}
|
host/clef_litert/__init__.py
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""clef-flash on LiteRT: the host in plain Python.
|
| 2 |
+
|
| 3 |
+
Dependencies: NumPy, `tokenizers` and Pillow, and ai-edge-litert for the graphs (requirements-host.txt; the
|
| 4 |
+
safetensors format is read here, the `safetensors` library only checks the table files). No PyTorch, no transformers
|
| 5 |
+
and none of the provider's code at run time. The host answers a request of the SystemOne-compatible request shape the way
|
| 6 |
+
the checkpoint's own `systemone()` does (Cloudflare/clef-flash at f0b7d3cd, joint_schema_model.py sha256 0e304cf7...),
|
| 7 |
+
around two kinds of LiteRT graph that are the caller's: a row graph per length (the Qwen3.5-9B text stack, embeds in,
|
| 8 |
+
final-norm hidden states out) and a tower graph per picture tile. The vocabulary tables and the joint schema head stay
|
| 9 |
+
on the host.
|
| 10 |
+
|
| 11 |
+
1. Request: {"model": str, "state": any JSON value, "questions": {id: {"type": "noul" | "choice" | "score",
|
| 12 |
+
"instructions": ..., "criteria": ...}}, "images": [one picture] (optional)}. Choice criteria are an object
|
| 13 |
+
{option id: description}, score criteria a list of level descriptions (lowest first), noul criteria an optional
|
| 14 |
+
object overriding the "true" / "false" descriptions. Checked as `systemone()` checks it (request.validate_request;
|
| 15 |
+
the named refusals of errors.py carry the author's messages).
|
| 16 |
+
2. Ids (request.encode_record): the author's encode_record on the repository's tokenizer.json, every piece tokenized
|
| 17 |
+
on its own without special tokens; with a picture, its block (<|vision_start|>, N x <|image_pad|>, <|vision_end|>,
|
| 18 |
+
"\\n") right after the 36-id prefix. usage.input_tokens = the number of ids.
|
| 19 |
+
3. Picture (pixels.py): one per request, read at a fixed square tile, 448 (default, 14 x 14 = 196 picture tokens) or
|
| 20 |
+
256 (8 x 8 = 64): RGB, Pillow BICUBIC to the tile, (x/255 - 0.5) / 0.5, merge-block-major patches [4 G^2, 1536]
|
| 21 |
+
-> the tower graph of that tile -> [G^2, 4096] picture rows. More pictures, videos and media_kwargs are refused
|
| 22 |
+
(`UnsupportedMedia`); the author's native resolution is not a tower graph.
|
| 23 |
+
4. Row graph inputs (host.ClefHost.row_inputs, at L = the smallest of 256 / 512 / 1024 / 2048 / 4096 that holds the
|
| 24 |
+
record; longer records are refused, `RequestTooLong`, never cut):
|
| 25 |
+
embeds float32 [1, L, 4096] the bf16 embed_tokens rows of the ids widened to float32, the picture rows at the
|
| 26 |
+
<|image_pad|> positions, zeros after the record
|
| 27 |
+
valid float32 [1, L] 1 on the record, 0 after it
|
| 28 |
+
cos float32 [1, L, 64] the M-RoPE tables of the record's three position planes (positions.py)
|
| 29 |
+
sin float32 [1, L, 64]
|
| 30 |
+
-> hidden float32 [1, L, 4096] (every position, after the final RMSNorm); the host reads the first T rows.
|
| 31 |
+
5. Head (head.py): the author's JointSchemaHead in NumPy float32 on those T rows, the ids, the spans and the bf16
|
| 32 |
+
lm_head rows of the option spans -> one logit per option.
|
| 33 |
+
6. Response (response.py): per question the float32 softmax, then the author's systemone_answer ->
|
| 34 |
+
{"model", "answers": {id: answer}, "usage": {"input_tokens", "output_tokens": 0}}.
|
| 35 |
+
|
| 36 |
+
Files: tokenizer.json (the repository's, sha256 06b95093...), embed_tokens.safetensors and lm_head.safetensors (one BF16
|
| 37 |
+
[248320, 4096] tensor each, the checkpoint's bytes, no metadata), joint_head.safetensors and joint_head_config.json
|
| 38 |
+
(the checkpoint's own). contract.json next to this package records names, shapes, files and checks.
|
| 39 |
+
|
| 40 |
+
The graphs (runtime.py, ai-edge-litert's CompiledModel; imported only when a graph is compiled):
|
| 41 |
+
with ClefLiteRT.from_dir("<folder>", accelerator="gpu") as clef: # "gpu" = float32 precision, "cpu" = XNNPACK
|
| 42 |
+
body = clef.decide(request, tile=448)"""
|
| 43 |
+
from .errors import (BadQuestionType, EmptyCriteria, EmptyRecord, MissingModelOrState, NoQuestions, RequestError,
|
| 44 |
+
RequestTooLong, SchemaTooLong, UnsupportedMedia)
|
| 45 |
+
from .head import JointHead
|
| 46 |
+
from .host import FILES, HIDDEN, ClefHost, Encoded
|
| 47 |
+
from .pixels import TILES, tile_patches
|
| 48 |
+
from .positions import BUCKETS, mrope_tables, padded_positions, pick_length, record_positions
|
| 49 |
+
from .request import EncodedQuestion, EncodedRecord, encode_record, question_options, render, validate_request
|
| 50 |
+
from .response import question_probs, response, systemone_answer
|
| 51 |
+
from .runtime import FOLDER, ClefGraphs, ClefLiteRT, LiteRTGraph, load_f32_tensor
|
| 52 |
+
from .tables import Bf16Table, copy_bf16_tensor, load_bf16_tensors
|
| 53 |
+
|
| 54 |
+
__version__ = "0.2.0"
|
| 55 |
+
__all__ = [
|
| 56 |
+
"BUCKETS", "FILES", "FOLDER", "HIDDEN", "TILES", "BadQuestionType", "Bf16Table", "ClefGraphs", "ClefHost",
|
| 57 |
+
"ClefLiteRT", "EmptyCriteria", "EmptyRecord", "Encoded", "EncodedQuestion", "EncodedRecord", "JointHead",
|
| 58 |
+
"LiteRTGraph", "MissingModelOrState", "NoQuestions", "RequestError", "RequestTooLong", "SchemaTooLong",
|
| 59 |
+
"UnsupportedMedia", "copy_bf16_tensor", "encode_record", "load_bf16_tensors", "load_f32_tensor", "mrope_tables",
|
| 60 |
+
"padded_positions", "pick_length", "question_options", "question_probs", "record_positions", "render", "response",
|
| 61 |
+
"systemone_answer", "tile_patches", "validate_request",
|
| 62 |
+
]
|
host/clef_litert/_erf.py
ADDED
|
@@ -0,0 +1,95 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""erf in float64 over NumPy arrays (NumPy has none): FreeBSD msun s_erf.c as ported by musl and Rust `libm`,
|
| 2 |
+
vectorized region by region. Used by the head's GELU (the author's torch.nn.GELU / F.gelu, the exact erf form).
|
| 3 |
+
|
| 4 |
+
* ====================================================
|
| 5 |
+
* Copyright (C) 1993 by Sun Microsystems, Inc. All rights reserved.
|
| 6 |
+
*
|
| 7 |
+
* Developed at SunPro, a Sun Microsystems, Inc. business.
|
| 8 |
+
* Permission to use, copy, modify, and distribute this
|
| 9 |
+
* software is freely granted, provided that this notice
|
| 10 |
+
* is preserved.
|
| 11 |
+
* ===================================================="""
|
| 12 |
+
from __future__ import annotations
|
| 13 |
+
|
| 14 |
+
import math
|
| 15 |
+
|
| 16 |
+
import numpy as np
|
| 17 |
+
|
| 18 |
+
ERX = 8.45062911510467529297e-01
|
| 19 |
+
EFX8 = 1.02703333676410069053e+00
|
| 20 |
+
PP = (1.28379167095512558561e-01, -3.25042107247001499370e-01, -2.84817495755985104766e-02,
|
| 21 |
+
-5.77027029648944159157e-03, -2.37630166566501626084e-05)
|
| 22 |
+
QQ = (3.97917223959155352819e-01, 6.50222499887672944485e-02, 5.08130628187576562776e-03,
|
| 23 |
+
1.32494738004321644526e-04, -3.96022827877536812320e-06)
|
| 24 |
+
PA = (-2.36211856075265944077e-03, 4.14856118683748331666e-01, -3.72207876035701323847e-01,
|
| 25 |
+
3.18346619901161753674e-01, -1.10894694282396677476e-01, 3.54783043256182359371e-02,
|
| 26 |
+
-2.16637559486879084300e-03)
|
| 27 |
+
QA = (1.06420880400844228286e-01, 5.40397917702171048937e-01, 7.18286544141962662868e-02,
|
| 28 |
+
1.26171219808761642112e-01, 1.36370839120290507362e-02, 1.19844998467991074170e-02)
|
| 29 |
+
RA = (-9.86494403484714822705e-03, -6.93858572707181764372e-01, -1.05586262253232909814e+01,
|
| 30 |
+
-6.23753324503260060396e+01, -1.62396669462573470355e+02, -1.84605092906711035994e+02,
|
| 31 |
+
-8.12874355063065934246e+01, -9.81432934416914548592e+00)
|
| 32 |
+
SA = (1.96512716674392571292e+01, 1.37657754143519042600e+02, 4.34565877475229228821e+02,
|
| 33 |
+
6.45387271733267880336e+02, 4.29008140027567833386e+02, 1.08635005541779435134e+02,
|
| 34 |
+
6.57024977031928170135e+00, -6.04244152148580987438e-02)
|
| 35 |
+
RB = (-9.86494292470009928597e-03, -7.99283237680523006574e-01, -1.77579549177547519889e+01,
|
| 36 |
+
-1.60636384855821916062e+02, -6.37566443368389627722e+02, -1.02509513161107724954e+03,
|
| 37 |
+
-4.83519191608651397019e+02)
|
| 38 |
+
SB = (3.03380607434824582924e+01, 3.25792512996573918826e+02, 1.53672958608443695994e+03,
|
| 39 |
+
3.19985821950859553908e+03, 2.55305040643316442583e+03, 4.74528541206955367215e+02,
|
| 40 |
+
-2.24409524465858183362e+01)
|
| 41 |
+
SQRT1_2 = math.sqrt(0.5)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def _poly(z, c):
|
| 45 |
+
"""c[0] + z * (c[1] + z * (... + z * c[-1])), Horner from the last coefficient (the C source's nesting)."""
|
| 46 |
+
r = np.full_like(z, c[-1])
|
| 47 |
+
for a in c[-2::-1]:
|
| 48 |
+
r = a + z * r
|
| 49 |
+
return r
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def _one_plus_z_poly(z, c):
|
| 53 |
+
"""1 + z * (c[0] + z * (c[1] + ...)) (the denominators Q / S)."""
|
| 54 |
+
return 1.0 + z * _poly(z, c)
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
def erf(x) -> np.ndarray:
|
| 58 |
+
"""Elementwise erf of a float64 array (NaN -> NaN, +-inf -> +-1)."""
|
| 59 |
+
x = np.asarray(x, dtype=np.float64)
|
| 60 |
+
out = np.empty_like(x)
|
| 61 |
+
ax = np.abs(x)
|
| 62 |
+
hi = (ax.view(np.uint64) >> np.uint64(32)).astype(np.uint32) # high word of |x|
|
| 63 |
+
sign = np.signbit(x)
|
| 64 |
+
|
| 65 |
+
m = hi < 0x3e300000 # |x| < 2**-28
|
| 66 |
+
out[m] = 0.125 * (8.0 * x[m] + EFX8 * x[m])
|
| 67 |
+
m = (hi >= 0x3e300000) & (hi < 0x3feb0000) # |x| < 0.84375
|
| 68 |
+
xm = x[m]
|
| 69 |
+
z = xm * xm
|
| 70 |
+
out[m] = xm + xm * (_poly(z, PP) / _one_plus_z_poly(z, QQ))
|
| 71 |
+
m = (hi >= 0x3feb0000) & (hi < 0x3ff40000) # 0.84375 <= |x| < 1.25
|
| 72 |
+
s = ax[m] - 1.0
|
| 73 |
+
erfc1 = 1.0 - ERX - _poly(s, PA) / _one_plus_z_poly(s, QA)
|
| 74 |
+
out[m] = np.where(sign[m], -(1.0 - erfc1), 1.0 - erfc1)
|
| 75 |
+
m = (hi >= 0x3ff40000) & (hi < 0x40180000) # 1.25 <= |x| < 6
|
| 76 |
+
if m.any():
|
| 77 |
+
a = ax[m]
|
| 78 |
+
h = hi[m]
|
| 79 |
+
s = 1.0 / (a * a)
|
| 80 |
+
near = h < 0x4006db6d # |x| < 1 / 0.35
|
| 81 |
+
r = np.where(near, _poly(s, RA), _poly(s, RB))
|
| 82 |
+
big_s = np.where(near, _one_plus_z_poly(s, SA), _one_plus_z_poly(s, SB))
|
| 83 |
+
zz = (a.view(np.uint64) & np.uint64(0xFFFFFFFF00000000)).view(np.float64) # low word cleared
|
| 84 |
+
erfc2 = np.exp(-zz * zz - 0.5625) * np.exp((zz - a) * (zz + a) + r / big_s) / a
|
| 85 |
+
out[m] = np.where(sign[m], -(1.0 - erfc2), 1.0 - erfc2)
|
| 86 |
+
m = hi >= 0x40180000 # |x| >= 6, inf, NaN
|
| 87 |
+
y = 1.0 - math.ldexp(1.0, -1022)
|
| 88 |
+
out[m] = np.where(np.isnan(x[m]), np.nan, np.where(sign[m], -y, y))
|
| 89 |
+
return out
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
def gelu(x: np.ndarray) -> np.ndarray:
|
| 93 |
+
"""GELU, the exact form x / 2 * (1 + erf(x / sqrt 2)), evaluated in float64 and rounded to x's float32."""
|
| 94 |
+
xd = np.asarray(x, dtype=np.float64)
|
| 95 |
+
return (0.5 * xd * (1.0 + erf(xd * SQRT1_2))).astype(np.float32)
|
host/clef_litert/errors.py
ADDED
|
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Named refusals of the clef-flash host.
|
| 2 |
+
|
| 3 |
+
Every refusal is a `RequestError`, itself a `ValueError`: code that catches the ValueError the checkpoint's own
|
| 4 |
+
`systemone()` raises catches these too, and the first six carry the author's message word for word. The last two are
|
| 5 |
+
the host's own: it reads one picture at a fixed tile and runs row graphs of at most 4,096 tokens."""
|
| 6 |
+
from __future__ import annotations
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class RequestError(ValueError):
|
| 10 |
+
"""The request cannot be answered as given."""
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
class MissingModelOrState(RequestError):
|
| 14 |
+
"""`model` is not a string or `state` is absent ("model and state are required")."""
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
class NoQuestions(RequestError):
|
| 18 |
+
"""`questions` is not a non-empty object ("at least one question is required")."""
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class BadQuestionType(RequestError):
|
| 22 |
+
"""A question's `type` is not noul, choice or score ("<id>: type must be noul, choice, or score")."""
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
class EmptyCriteria(RequestError):
|
| 26 |
+
"""A choice or score question without criteria ("<id>: criteria must not be empty")."""
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
class SchemaTooLong(RequestError):
|
| 30 |
+
"""The prefix, picture, schema and suffix alone exceed max_length (encode_record: "schema requires N tokens
|
| 31 |
+
before state; maximum is M")."""
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class EmptyRecord(RequestError):
|
| 35 |
+
"""The record gives no ids or no questions (encode_record: "record produced no model input or questions")."""
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
class UnsupportedMedia(RequestError):
|
| 39 |
+
"""Host only: more than one picture, `videos`, `media_kwargs`, a picture the host cannot open, or a tile other
|
| 40 |
+
than 256 / 448 (the two tower graphs)."""
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
class RequestTooLong(RequestError):
|
| 44 |
+
"""Host only: the record is longer than the largest row graph (4,096 tokens). It is refused, never truncated."""
|
host/clef_litert/head.py
ADDED
|
@@ -0,0 +1,231 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The checkpoint's joint schema head (joint_schema_model.py `JointSchemaHead.forward`, one record) in NumPy float32.
|
| 2 |
+
|
| 3 |
+
Weights: joint_head.safetensors (bfloat16, 122 tensors, widened to float32 exactly; `load_state_dict(strict=True)` names)
|
| 4 |
+
and joint_head_config.json {hidden_size 4096, width 1024, routing_layers 2, layers 4, heads 16, feedforward 4096}.
|
| 5 |
+
Inputs: the row graph's hidden states at the record's T real positions (after the final RMSNorm), the record's ids
|
| 6 |
+
(processor form) and spans, and the lm_head table (`lm_head.weight` rows, for the lexical vectors).
|
| 7 |
+
|
| 8 |
+
x = hidden_norm(hidden) LayerNorm(4096); memory = memory_projection(x) [T, 1024]
|
| 9 |
+
global = x[T - 1]; q_i = mean of x over question i's span; ctx_o = mean of x over option o's span
|
| 10 |
+
lexical_o = mean of the lm_head rows of option o's span ids (not normalized)
|
| 11 |
+
queries_o = option_context_projection(ctx_o) + option_lexical_projection(lexical_o)
|
| 12 |
+
+ option_question_projection(q_i) i = the question of option o
|
| 13 |
+
2 x evidence routing: queries += MHA(query_norm(queries), memory_norm(memory)); queries += FFN(feedforward_norm(.))
|
| 14 |
+
(pre-LN, 16 heads of 64, FFN = Linear 1024->4096, GELU, Linear 4096->1024)
|
| 15 |
+
per question: w = softmax over its options of routed . question_projection(q_i) / 32; summary_i = sum w * routed
|
| 16 |
+
fields = question_projection(q) + option_summary_norm(summary) + global_projection(global) + type_embedding(type)
|
| 17 |
+
4 x TransformerDecoderLayer (norm_first, GELU): self-attention over the questions, cross-attention to the memory
|
| 18 |
+
(not normalized), FFN; fields = field_norm(fields)
|
| 19 |
+
prior_o = e^min(prior_logit_scale, ln 100) * cos(lexical_o, q_i + global)
|
| 20 |
+
joint_o = e^min(joint_logit_scale, ln 100) * cos(field_i, option_norm(routed_o)) + residual_scorer([f, o, f*o, |f-o|])
|
| 21 |
+
logit_o = prior_o + sigmoid(residual_gate) * joint_o
|
| 22 |
+
|
| 23 |
+
Every tensor is float32 as in the author's float32 module; LayerNorm eps 1e-5; attention = softmax(q k^T / 8) v per
|
| 24 |
+
head in float32; GELU = the exact erf form (evaluated in float64, rounded to float32); F.normalize eps 1e-12 and
|
| 25 |
+
F.cosine_similarity eps 1e-8 (each side divided by max(norm, eps), then the dot product, as torch computes it).
|
| 26 |
+
Dropout is the identity (eval). The per-question softmax of the logits is the response's (response.py)."""
|
| 27 |
+
from __future__ import annotations
|
| 28 |
+
|
| 29 |
+
import json
|
| 30 |
+
import math
|
| 31 |
+
from pathlib import Path
|
| 32 |
+
|
| 33 |
+
import numpy as np
|
| 34 |
+
|
| 35 |
+
from ._erf import gelu
|
| 36 |
+
from .tables import load_bf16_tensors
|
| 37 |
+
|
| 38 |
+
CONFIG = {"hidden_size": 4096, "width": 1024, "routing_layers": 2, "layers": 4, "heads": 16, "feedforward": 4096}
|
| 39 |
+
LN_EPS = 1e-5
|
| 40 |
+
F32 = np.float32
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def layer_norm(x: np.ndarray, w: np.ndarray, b: np.ndarray, eps: float = LN_EPS) -> np.ndarray:
|
| 44 |
+
mean = x.mean(axis=-1, keepdims=True, dtype=F32)
|
| 45 |
+
xc = x - mean
|
| 46 |
+
var = (xc * xc).mean(axis=-1, keepdims=True, dtype=F32)
|
| 47 |
+
return (xc * (F32(1.0) / np.sqrt(var + F32(eps)))) * w + b
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def normalize(x: np.ndarray, eps: float) -> np.ndarray:
|
| 51 |
+
"""x / max(||x||_2, eps) over the last axis (F.normalize)."""
|
| 52 |
+
n = np.sqrt((x * x).sum(axis=-1, keepdims=True, dtype=F32))
|
| 53 |
+
return x / np.maximum(n, F32(eps))
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
|
| 57 |
+
e = np.exp(x - x.max(axis=axis, keepdims=True))
|
| 58 |
+
return e / e.sum(axis=axis, keepdims=True, dtype=F32)
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
class _Linear:
|
| 62 |
+
def __init__(self, w: np.ndarray, b: np.ndarray | None = None):
|
| 63 |
+
self.wt = np.ascontiguousarray(w.T, dtype=F32) # [in, out]
|
| 64 |
+
self.b = None if b is None else np.ascontiguousarray(b, dtype=F32)
|
| 65 |
+
|
| 66 |
+
def __call__(self, x: np.ndarray) -> np.ndarray:
|
| 67 |
+
y = x @ self.wt
|
| 68 |
+
return y if self.b is None else y + self.b
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
class _LayerNorm:
|
| 72 |
+
def __init__(self, w: np.ndarray, b: np.ndarray):
|
| 73 |
+
self.w, self.b = w.astype(F32), b.astype(F32)
|
| 74 |
+
|
| 75 |
+
def __call__(self, x: np.ndarray) -> np.ndarray:
|
| 76 |
+
return layer_norm(x, self.w, self.b)
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
class _Attention:
|
| 80 |
+
"""nn.MultiheadAttention (batch_first, eval, no masks) for one sequence: q_in [Lq, E], kv_in [Lk, E] -> [Lq, E]."""
|
| 81 |
+
|
| 82 |
+
def __init__(self, t: dict, prefix: str, heads: int):
|
| 83 |
+
w, b = t[prefix + "in_proj_weight"], t[prefix + "in_proj_bias"]
|
| 84 |
+
e = w.shape[1]
|
| 85 |
+
self.q = _Linear(w[:e], b[:e])
|
| 86 |
+
self.k = _Linear(w[e:2 * e], b[e:2 * e])
|
| 87 |
+
self.v = _Linear(w[2 * e:], b[2 * e:])
|
| 88 |
+
self.out = _Linear(t[prefix + "out_proj.weight"], t[prefix + "out_proj.bias"])
|
| 89 |
+
self.heads, self.dim = heads, e // heads
|
| 90 |
+
self.scale = F32(1.0 / math.sqrt(self.dim))
|
| 91 |
+
|
| 92 |
+
def __call__(self, q_in: np.ndarray, kv_in: np.ndarray) -> np.ndarray:
|
| 93 |
+
h, d = self.heads, self.dim
|
| 94 |
+
q = self.q(q_in).reshape(q_in.shape[0], h, d).transpose(1, 0, 2) # [H, Lq, D]
|
| 95 |
+
k = self.k(kv_in).reshape(kv_in.shape[0], h, d).transpose(1, 0, 2) # [H, Lk, D]
|
| 96 |
+
v = self.v(kv_in).reshape(kv_in.shape[0], h, d).transpose(1, 0, 2)
|
| 97 |
+
p = softmax((q @ k.transpose(0, 2, 1)) * self.scale) # [H, Lq, Lk]
|
| 98 |
+
o = (p @ v).transpose(1, 0, 2).reshape(q_in.shape[0], h * d)
|
| 99 |
+
return self.out(o)
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
class _Tracked(dict):
|
| 103 |
+
"""The weights by name, recording which names were read (load_state_dict(strict=True): none missing, none unused)."""
|
| 104 |
+
|
| 105 |
+
def __init__(self, *a):
|
| 106 |
+
super().__init__(*a)
|
| 107 |
+
self.read = set()
|
| 108 |
+
|
| 109 |
+
def __getitem__(self, key):
|
| 110 |
+
self.read.add(key)
|
| 111 |
+
return super().__getitem__(key)
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
class JointHead:
|
| 115 |
+
"""The author's JointSchemaHead for one record. `scores(hidden, ids, questions, lm_head)` -> one float32 logit
|
| 116 |
+
array per question, in the author's option order."""
|
| 117 |
+
|
| 118 |
+
def __init__(self, tensors: dict[str, np.ndarray], config: dict | None = None):
|
| 119 |
+
cfg = dict(CONFIG if config is None else config)
|
| 120 |
+
if cfg != CONFIG:
|
| 121 |
+
raise ValueError(f"head config {cfg} is not the release's {CONFIG}")
|
| 122 |
+
self.config = cfg
|
| 123 |
+
t = _Tracked({k: np.asarray(v, dtype=F32) for k, v in tensors.items()})
|
| 124 |
+
heads = cfg["heads"]
|
| 125 |
+
ln = lambda p: _LayerNorm(t[p + ".weight"], t[p + ".bias"]) # noqa: E731
|
| 126 |
+
lin = lambda p, bias=True: _Linear(t[p + ".weight"], t[p + ".bias"] if bias else None) # noqa: E731
|
| 127 |
+
self.hidden_norm = ln("hidden_norm")
|
| 128 |
+
self.memory_projection = lin("memory_projection", False)
|
| 129 |
+
self.question_projection = lin("question_projection", False)
|
| 130 |
+
self.option_question_projection = lin("option_question_projection", False)
|
| 131 |
+
self.global_projection = lin("global_projection", False)
|
| 132 |
+
self.option_context_projection = lin("option_context_projection", False)
|
| 133 |
+
self.option_lexical_projection = lin("option_lexical_projection", False)
|
| 134 |
+
self.type_embedding = t["type_embedding.weight"]
|
| 135 |
+
self.evidence = [{"query_norm": ln(f"evidence_layers.{i}.query_norm"),
|
| 136 |
+
"memory_norm": ln(f"evidence_layers.{i}.memory_norm"),
|
| 137 |
+
"attention": _Attention(t, f"evidence_layers.{i}.attention.", heads),
|
| 138 |
+
"feedforward_norm": ln(f"evidence_layers.{i}.feedforward_norm"),
|
| 139 |
+
"ff0": lin(f"evidence_layers.{i}.feedforward.0"),
|
| 140 |
+
"ff3": lin(f"evidence_layers.{i}.feedforward.3")}
|
| 141 |
+
for i in range(cfg["routing_layers"])]
|
| 142 |
+
self.option_summary_norm = ln("option_summary_norm")
|
| 143 |
+
self.layers = [{"norm1": ln(f"layers.{i}.norm1"), "norm2": ln(f"layers.{i}.norm2"),
|
| 144 |
+
"norm3": ln(f"layers.{i}.norm3"),
|
| 145 |
+
"self_attn": _Attention(t, f"layers.{i}.self_attn.", heads),
|
| 146 |
+
"multihead_attn": _Attention(t, f"layers.{i}.multihead_attn.", heads),
|
| 147 |
+
"linear1": lin(f"layers.{i}.linear1"), "linear2": lin(f"layers.{i}.linear2")}
|
| 148 |
+
for i in range(cfg["layers"])]
|
| 149 |
+
self.field_norm = ln("field_norm")
|
| 150 |
+
self.option_norm = ln("option_norm")
|
| 151 |
+
self.scorer0 = lin("residual_scorer.0")
|
| 152 |
+
self.scorer3 = lin("residual_scorer.3")
|
| 153 |
+
cap = F32(math.log(100.0))
|
| 154 |
+
self.prior_scale = np.exp(np.minimum(t["prior_logit_scale"].reshape(()), cap)).astype(F32)
|
| 155 |
+
self.joint_scale = np.exp(np.minimum(t["joint_logit_scale"].reshape(()), cap)).astype(F32)
|
| 156 |
+
g = t["residual_gate"].reshape(())
|
| 157 |
+
self.gate = (F32(1.0) / (F32(1.0) + np.exp(-g))).astype(F32)
|
| 158 |
+
unused = set(t) - t.read # a missing name already raised KeyError when it was read
|
| 159 |
+
if unused:
|
| 160 |
+
raise ValueError(f"head weights not used: {sorted(unused)}")
|
| 161 |
+
self.n_tensors = len(t)
|
| 162 |
+
self.n_params = int(sum(v.size for v in dict.values(t)))
|
| 163 |
+
|
| 164 |
+
@classmethod
|
| 165 |
+
def from_files(cls, weights: str | Path, config: str | Path | None = None) -> "JointHead":
|
| 166 |
+
cfg = json.loads(Path(config).read_text()) if config is not None else None
|
| 167 |
+
return cls(load_bf16_tensors(weights), cfg)
|
| 168 |
+
|
| 169 |
+
@staticmethod
|
| 170 |
+
def lexical_rows(lm_head, ids, questions) -> list[np.ndarray]:
|
| 171 |
+
"""Per question, [n_options, 4096]: the mean of the lm_head rows of each option span's ids (float32)."""
|
| 172 |
+
ids = np.asarray(ids, dtype=np.int64)
|
| 173 |
+
out = []
|
| 174 |
+
for q in questions:
|
| 175 |
+
rows = [lm_head.rows(ids[a:b]).mean(axis=0, dtype=F32) for a, b in q.option_spans]
|
| 176 |
+
out.append(np.stack(rows).astype(F32))
|
| 177 |
+
return out
|
| 178 |
+
|
| 179 |
+
def scores(self, hidden: np.ndarray, ids, questions, lm_head=None, lexical: list | None = None) -> list[np.ndarray]:
|
| 180 |
+
"""hidden [T, 4096] (the record's real positions), ids (len T), questions (request.EncodedQuestion or
|
| 181 |
+
anything with question_span / option_spans / type_id) -> [float32 logits per question].
|
| 182 |
+
`lexical` (one [n_options, 4096] array per question) replaces the lm_head gather (controls)."""
|
| 183 |
+
hidden = np.asarray(hidden, dtype=F32)
|
| 184 |
+
if hidden.ndim != 2 or hidden.shape[0] != len(ids):
|
| 185 |
+
raise ValueError(f"hidden {hidden.shape} for {len(ids)} ids")
|
| 186 |
+
if lexical is None:
|
| 187 |
+
lexical = self.lexical_rows(lm_head, ids, questions)
|
| 188 |
+
x = self.hidden_norm(hidden) # [T, 4096]
|
| 189 |
+
memory = self.memory_projection(x) # [T, 1024]
|
| 190 |
+
global_vector = x[-1]
|
| 191 |
+
question_vectors = np.stack([x[a:b].mean(axis=0, dtype=F32) for a, b in (q.question_span for q in questions)])
|
| 192 |
+
type_rows = self.type_embedding[[q.type_id for q in questions]]
|
| 193 |
+
counts = [len(q.option_spans) for q in questions]
|
| 194 |
+
context = [np.stack([x[a:b].mean(axis=0, dtype=F32) for a, b in q.option_spans]) for q in questions]
|
| 195 |
+
|
| 196 |
+
queries = np.concatenate([
|
| 197 |
+
self.option_context_projection(ctx) + self.option_lexical_projection(lex)
|
| 198 |
+
+ self.option_question_projection(question_vectors[i])[None]
|
| 199 |
+
for i, (ctx, lex) in enumerate(zip(context, lexical))]) # [O, 1024]
|
| 200 |
+
for layer in self.evidence:
|
| 201 |
+
routed = layer["attention"](layer["query_norm"](queries), layer["memory_norm"](memory))
|
| 202 |
+
queries = queries + routed
|
| 203 |
+
queries = queries + layer["ff3"](gelu(layer["ff0"](layer["feedforward_norm"](queries))))
|
| 204 |
+
split = np.split(queries, np.cumsum(counts)[:-1])
|
| 205 |
+
|
| 206 |
+
base = self.question_projection(question_vectors) # [Q, 1024]
|
| 207 |
+
summaries = []
|
| 208 |
+
for field, options in zip(base, split):
|
| 209 |
+
w = softmax((options @ field) / F32(math.sqrt(options.shape[-1])), axis=0)
|
| 210 |
+
summaries.append((w[:, None] * options).sum(axis=0, dtype=F32))
|
| 211 |
+
fields = (base + self.option_summary_norm(np.stack(summaries)) + self.global_projection(global_vector)[None]
|
| 212 |
+
+ type_rows)
|
| 213 |
+
for layer in self.layers:
|
| 214 |
+
f1 = layer["norm1"](fields)
|
| 215 |
+
fields = fields + layer["self_attn"](f1, f1)
|
| 216 |
+
fields = fields + layer["multihead_attn"](layer["norm2"](fields), memory)
|
| 217 |
+
fields = fields + layer["linear2"](gelu(layer["linear1"](layer["norm3"](fields))))
|
| 218 |
+
fields = self.field_norm(fields)
|
| 219 |
+
|
| 220 |
+
out = []
|
| 221 |
+
for i, (field, lex, routed) in enumerate(zip(fields, lexical, split)):
|
| 222 |
+
anchor = normalize(question_vectors[i] + global_vector, 1e-12)
|
| 223 |
+
prior = self.prior_scale * (normalize(lex, 1e-12) @ anchor)
|
| 224 |
+
options = self.option_norm(routed)
|
| 225 |
+
f = np.broadcast_to(field, options.shape)
|
| 226 |
+
cosine = (normalize(f, 1e-8) * normalize(options, 1e-8)).sum(axis=-1, dtype=F32)
|
| 227 |
+
features = np.concatenate([f, options, f * options, np.abs(f - options)], axis=-1)
|
| 228 |
+
residual = self.scorer3(gelu(self.scorer0(features)))[:, 0]
|
| 229 |
+
joint = self.joint_scale * cosine + residual
|
| 230 |
+
out.append((prior + self.gate * joint).astype(F32))
|
| 231 |
+
return out
|
host/clef_litert/host.py
ADDED
|
@@ -0,0 +1,134 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ClefHost: one request through the host steps around the graphs (the graphs themselves are the caller's callables).
|
| 2 |
+
|
| 3 |
+
host = ClefHost.from_dir("<folder>") # tokenizer.json, embed_tokens.safetensors, lm_head.safetensors,
|
| 4 |
+
# joint_head.safetensors, joint_head_config.json
|
| 5 |
+
enc = host.encode(request, tile=448) # checks, ids, spans, bucket L; a picture -> enc.patches
|
| 6 |
+
rows = tower(enc.tile, enc.patches) # the tower graph: [G^2, 4096] picture rows (requests with a picture)
|
| 7 |
+
x = host.row_inputs(enc, rows) # {"embeds", "valid", "cos", "sin"} at length enc.length
|
| 8 |
+
hidden = row_graph(enc.length, x) # the row graph of that length: [1, L, 4096] after the final norm
|
| 9 |
+
logits = host.scores(enc, hidden) # the joint head, NumPy float32: one array per question
|
| 10 |
+
body = host.respond(request, enc, logits) # the SystemOne-compatible response
|
| 11 |
+
|
| 12 |
+
or `host.decide(request, row_graph, tower)` for all of it. Refusals: errors.py (the author's ValueError conditions
|
| 13 |
+
by name, plus the host's own: one picture at a 256 / 448 tile, at most 4,096 tokens)."""
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
from dataclasses import dataclass
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
from typing import Callable
|
| 19 |
+
|
| 20 |
+
import numpy as np
|
| 21 |
+
|
| 22 |
+
from .errors import UnsupportedMedia
|
| 23 |
+
from .head import JointHead
|
| 24 |
+
from .pixels import tile_grid, tile_patches
|
| 25 |
+
from .positions import BUCKETS, mrope_tables, padded_positions, pick_length, record_positions
|
| 26 |
+
from .request import MAX_LENGTH, EncodedRecord, encode_record, validate_request
|
| 27 |
+
from .response import question_probs, response
|
| 28 |
+
from .tables import Bf16Table
|
| 29 |
+
|
| 30 |
+
HIDDEN = 4096
|
| 31 |
+
FILES = {"tokenizer": "tokenizer.json", "embed_tokens": "embed_tokens.safetensors", "lm_head": "lm_head.safetensors",
|
| 32 |
+
"joint_head": "joint_head.safetensors", "joint_head_config": "joint_head_config.json"}
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
@dataclass(frozen=True)
|
| 36 |
+
class Encoded:
|
| 37 |
+
record: EncodedRecord
|
| 38 |
+
length: int # the row graph length: the smallest bucket >= record.tokens
|
| 39 |
+
tile: int | None = None # 256 | 448 for a request with a picture
|
| 40 |
+
patches: np.ndarray | None = None # the tower graph's input [4 G^2, 1536] float32
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def media_of(request: dict) -> list:
|
| 44 |
+
"""The request's one picture as a list ([] or [picture]); UnsupportedMedia for what the host does not read."""
|
| 45 |
+
if request.get("videos"):
|
| 46 |
+
raise UnsupportedMedia("videos are not read by this host")
|
| 47 |
+
if request.get("media_kwargs"):
|
| 48 |
+
raise UnsupportedMedia("media_kwargs are not read by this host (one picture at a fixed tile)")
|
| 49 |
+
images = list(request.get("images") or [])
|
| 50 |
+
if len(images) > 1:
|
| 51 |
+
raise UnsupportedMedia(f"{len(images)} pictures: this host reads one picture per request")
|
| 52 |
+
return images
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
class ClefHost:
|
| 56 |
+
def __init__(self, tokenizer, embed_tokens: Bf16Table, lm_head: Bf16Table, head: JointHead,
|
| 57 |
+
buckets=BUCKETS, max_length: int = MAX_LENGTH):
|
| 58 |
+
if embed_tokens.shape[1] != HIDDEN or lm_head.shape[1] != HIDDEN:
|
| 59 |
+
raise ValueError(f"tables {embed_tokens.shape} / {lm_head.shape}: expected [V, {HIDDEN}]")
|
| 60 |
+
self.tokenizer, self.embed_tokens, self.lm_head, self.head = tokenizer, embed_tokens, lm_head, head
|
| 61 |
+
self.buckets, self.max_length = tuple(sorted(buckets)), max_length
|
| 62 |
+
|
| 63 |
+
@classmethod
|
| 64 |
+
def from_files(cls, tokenizer: str | Path, embed_tokens: str | Path, lm_head: str | Path, joint_head: str | Path,
|
| 65 |
+
joint_head_config: str | Path, **kw) -> "ClefHost":
|
| 66 |
+
from tokenizers import Tokenizer
|
| 67 |
+
|
| 68 |
+
return cls(Tokenizer.from_file(str(tokenizer)), Bf16Table(embed_tokens), Bf16Table(lm_head),
|
| 69 |
+
JointHead.from_files(joint_head, joint_head_config), **kw)
|
| 70 |
+
|
| 71 |
+
@classmethod
|
| 72 |
+
def from_dir(cls, folder: str | Path, **kw) -> "ClefHost":
|
| 73 |
+
folder = Path(folder)
|
| 74 |
+
return cls.from_files(**{k: folder / v for k, v in FILES.items()}, **kw)
|
| 75 |
+
|
| 76 |
+
def encode(self, request: dict, tile: int = 448) -> Encoded:
|
| 77 |
+
"""The author's checks, the host's media checks, encode_record, the bucket (RequestTooLong past 4,096)."""
|
| 78 |
+
validate_request(request)
|
| 79 |
+
images = media_of(request)
|
| 80 |
+
grid = None
|
| 81 |
+
if images:
|
| 82 |
+
g = tile_grid(tile)
|
| 83 |
+
grid = (g, g)
|
| 84 |
+
record = encode_record(self.tokenizer, request, max_length=self.max_length, image_grid=grid)
|
| 85 |
+
length = pick_length(record.tokens, self.buckets)
|
| 86 |
+
if not images:
|
| 87 |
+
return Encoded(record, length)
|
| 88 |
+
return Encoded(record, length, tile, tile_patches(images[0], tile))
|
| 89 |
+
|
| 90 |
+
def row_inputs(self, enc: Encoded, image_rows: np.ndarray | None = None) -> dict[str, np.ndarray]:
|
| 91 |
+
"""The row graph's inputs at enc.length: embeds = the float32 rows of the bf16 table at the ids, the picture's
|
| 92 |
+
rows (the tower's output, [H * W, 4096]) at the <|image_pad|> positions; valid = 1 on the record; the pad
|
| 93 |
+
positions after it: embeds 0, valid 0; cos / sin of the record's planes (positions.py)."""
|
| 94 |
+
rec, length = enc.record, enc.length
|
| 95 |
+
t = rec.tokens
|
| 96 |
+
embeds = np.zeros((1, length, HIDDEN), np.float32)
|
| 97 |
+
embeds[0, :t] = self.embed_tokens.rows(rec.ids)
|
| 98 |
+
if rec.image_grid is not None:
|
| 99 |
+
h, w = rec.image_grid
|
| 100 |
+
if image_rows is None:
|
| 101 |
+
raise ValueError("a request with a picture needs the tower's rows")
|
| 102 |
+
image_rows = np.asarray(image_rows, dtype=np.float32)
|
| 103 |
+
if image_rows.shape != (h * w, HIDDEN):
|
| 104 |
+
raise ValueError(f"picture rows {image_rows.shape}, expected {(h * w, HIDDEN)}")
|
| 105 |
+
p = rec.token_offset
|
| 106 |
+
embeds[0, p + 1:p + 1 + h * w] = image_rows
|
| 107 |
+
elif image_rows is not None:
|
| 108 |
+
raise ValueError("picture rows for a request without a picture")
|
| 109 |
+
valid = np.zeros((1, length), np.float32)
|
| 110 |
+
valid[0, :t] = 1.0
|
| 111 |
+
cos, sin = mrope_tables(padded_positions(record_positions(t, rec.token_offset, rec.image_grid), length))
|
| 112 |
+
return {"embeds": embeds, "valid": valid, "cos": cos, "sin": sin}
|
| 113 |
+
|
| 114 |
+
def scores(self, enc: Encoded, hidden: np.ndarray, lexical: list | None = None) -> list[np.ndarray]:
|
| 115 |
+
"""The row graph's hidden ([1, L, 4096] or [T, 4096]) -> the head's logits, one float32 array per question."""
|
| 116 |
+
hidden = np.asarray(hidden, dtype=np.float32)
|
| 117 |
+
if hidden.ndim == 3:
|
| 118 |
+
hidden = hidden[0]
|
| 119 |
+
t = enc.record.tokens
|
| 120 |
+
return self.head.scores(hidden[:t], enc.record.ids, enc.record.questions, self.lm_head, lexical=lexical)
|
| 121 |
+
|
| 122 |
+
def respond(self, request: dict, enc: Encoded, logits: list[np.ndarray]) -> dict:
|
| 123 |
+
return response(request, enc.record.questions, [question_probs(z) for z in logits], enc.record.tokens)
|
| 124 |
+
|
| 125 |
+
def decide(self, request: dict, row_graph: Callable[[int, dict], np.ndarray],
|
| 126 |
+
tower: Callable[[int, np.ndarray], np.ndarray] | None = None, tile: int = 448) -> dict:
|
| 127 |
+
enc = self.encode(request, tile=tile)
|
| 128 |
+
rows = None
|
| 129 |
+
if enc.patches is not None:
|
| 130 |
+
if tower is None:
|
| 131 |
+
raise UnsupportedMedia("a request with a picture needs a tower")
|
| 132 |
+
rows = tower(enc.tile, enc.patches)
|
| 133 |
+
hidden = row_graph(enc.length, self.row_inputs(enc, rows))
|
| 134 |
+
return self.respond(request, enc, self.scores(enc, hidden))
|
host/clef_litert/pixels.py
ADDED
|
@@ -0,0 +1,80 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Picture -> the tower graph's patches, at one fixed square tile per tower graph (256 -> 8 x 8 = 64 picture tokens,
|
| 2 |
+
448 -> 14 x 14 = 196).
|
| 3 |
+
|
| 4 |
+
open -> convert("RGB") -> resize((tile, tile), BICUBIC) Pillow's own integer resampler; the aspect is NOT kept
|
| 5 |
+
-> /255 -> (x - 0.5) / 0.5 in float64 -> float32
|
| 6 |
+
-> merge-block-major patches [4 G^2, 1536], G = tile // 32, each vector (C, T=2, 16, 16)
|
| 7 |
+
|
| 8 |
+
This is what the reference's g256 / g448 arms fed the author's processor (the picture resized with Pillow before the
|
| 9 |
+
processor, whose own resize is the identity at multiples of 32), and the patches are the processor's `pixel_values`
|
| 10 |
+
bit for bit. The author's native arm (the processor's dynamic resolution) is not a tower graph here. EXIF orientation
|
| 11 |
+
is not applied (the reference passed opened PIL images, which the processor takes as they are)."""
|
| 12 |
+
from __future__ import annotations
|
| 13 |
+
|
| 14 |
+
import base64
|
| 15 |
+
import io
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
|
| 18 |
+
import numpy as np
|
| 19 |
+
|
| 20 |
+
from .errors import UnsupportedMedia
|
| 21 |
+
|
| 22 |
+
PATCH, MERGE, TEMPORAL = 16, 2, 2
|
| 23 |
+
IMAGE_MEAN, IMAGE_STD = 0.5, 0.5
|
| 24 |
+
TILES = (256, 448)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def tile_grid(tile: int) -> int:
|
| 28 |
+
"""Merged rows per side of a square tile (256 -> 8, 448 -> 14)."""
|
| 29 |
+
if tile not in TILES:
|
| 30 |
+
raise UnsupportedMedia(f"tile {tile}: the tower graphs read {' or '.join(map(str, TILES))} px tiles")
|
| 31 |
+
return tile // (PATCH * MERGE)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def open_image(image):
|
| 35 |
+
"""A PIL image, a uint8 [H, W, 3] array, the bytes of an image file, a local path, or a base64 data URI -> PIL."""
|
| 36 |
+
from PIL import Image
|
| 37 |
+
|
| 38 |
+
if isinstance(image, Image.Image):
|
| 39 |
+
return image
|
| 40 |
+
if isinstance(image, np.ndarray):
|
| 41 |
+
return Image.fromarray(np.asarray(image, dtype=np.uint8))
|
| 42 |
+
try:
|
| 43 |
+
if isinstance(image, (bytes, bytearray)):
|
| 44 |
+
return Image.open(io.BytesIO(bytes(image)))
|
| 45 |
+
if isinstance(image, str) and image.startswith("data:image/"):
|
| 46 |
+
return Image.open(io.BytesIO(base64.b64decode(image.split(",", 1)[1])))
|
| 47 |
+
if isinstance(image, (str, Path)) and Path(image).is_file():
|
| 48 |
+
return Image.open(image)
|
| 49 |
+
except Exception as e: # Pillow's UnidentifiedImageError, a bad base64 string, ...
|
| 50 |
+
raise UnsupportedMedia(f"the picture could not be opened: {e}") from e
|
| 51 |
+
raise UnsupportedMedia("a picture is a PIL image, a uint8 array, image file bytes, a local path or a data URI "
|
| 52 |
+
"(the host does not fetch URLs)")
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def patchify(x: np.ndarray) -> np.ndarray:
|
| 56 |
+
"""Normalized [H, W, C] -> [gh*gw, C*T*P*P] in Qwen's merge-block-major order: patches iterate
|
| 57 |
+
(block_row, block_col, y-in-block, x-in-block), channel outermost inside the vector, the still frame repeated at
|
| 58 |
+
both temporal slots (the Qwen2VLImageProcessor reshape / permute)."""
|
| 59 |
+
h, w, c = x.shape
|
| 60 |
+
if h % (PATCH * MERGE) or w % (PATCH * MERGE):
|
| 61 |
+
raise ValueError(f"{h}x{w} is not divisible by patch * merge = {PATCH * MERGE}")
|
| 62 |
+
gh, gw = h // PATCH, w // PATCH
|
| 63 |
+
t = x.transpose(2, 0, 1).reshape(c, gh // MERGE, MERGE, PATCH, gw // MERGE, MERGE, PATCH)
|
| 64 |
+
t = t.transpose(1, 4, 2, 5, 0, 3, 6).reshape(gh * gw, 1, c, PATCH, PATCH)
|
| 65 |
+
t = np.broadcast_to(t[:, :, :, None], (gh * gw, 1, c, TEMPORAL, PATCH, PATCH))
|
| 66 |
+
return t.reshape(gh * gw, c * TEMPORAL * PATCH * PATCH)
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def tile_pixels(image, tile: int) -> np.ndarray:
|
| 70 |
+
"""Picture -> uint8 [tile, tile, 3]: RGB, Pillow BICUBIC to the tile."""
|
| 71 |
+
from PIL import Image
|
| 72 |
+
|
| 73 |
+
tile_grid(tile)
|
| 74 |
+
return np.asarray(open_image(image).convert("RGB").resize((tile, tile), Image.Resampling.BICUBIC))
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def tile_patches(image, tile: int) -> np.ndarray:
|
| 78 |
+
"""Picture -> patches [4 G^2, 1536] float32, G = tile // 32 (the tower graph's input)."""
|
| 79 |
+
x = (tile_pixels(image, tile).astype(np.float64) / 255.0 - IMAGE_MEAN) / IMAGE_STD
|
| 80 |
+
return np.ascontiguousarray(patchify(x).astype(np.float32))
|
host/clef_litert/positions.py
ADDED
|
@@ -0,0 +1,96 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Row lengths and positions: the bucket of a record, its three M-RoPE position planes and the cos / sin tables the row
|
| 2 |
+
graph takes as inputs (the graph holds no COS / SIN / position table of its own).
|
| 3 |
+
|
| 4 |
+
Buckets: one row graph per length L in BUCKETS; a record of T tokens runs on the smallest L >= T, right-padded. A
|
| 5 |
+
record longer than the largest graph is refused (`RequestTooLong`), never truncated.
|
| 6 |
+
|
| 7 |
+
Position planes [3, T] (t, h, w) of the author's model (transformers Qwen3_5Model.get_rope_index; closed form recorded
|
| 8 |
+
with the reference, P = token_offset = the <|vision_start|> index, one picture of H x W merged tokens, N = H * W):
|
| 9 |
+
text-only record: (i, i, i)
|
| 10 |
+
i <= P -> (i, i, i); picture token k at P+1+k -> (P+1, P+1 + k // W, P+1 + k % W);
|
| 11 |
+
i >= P+N+1 -> i - (N - max(H, W)) on all three planes
|
| 12 |
+
Pad positions (right of the record) continue the planes' maximum + 1, + 2, ... on all three planes; their outputs are
|
| 13 |
+
never read and the graph's `valid` mask hides them as keys.
|
| 14 |
+
|
| 15 |
+
cos / sin (the author's Qwen3_5TextRotaryEmbedding: default rope, theta 1e7, rotary dim 64 of the 256-wide heads,
|
| 16 |
+
attention scaling 1, interleaved mrope_section [11, 11, 10]):
|
| 17 |
+
inv_freq[j] = 1 / theta^(2j / 64) j = 0..31, every step in float32 as the author's code
|
| 18 |
+
angle[p][i, j] = float32(inv_freq[j] * pos[p][i]) p = t, h, w
|
| 19 |
+
interleave: j % 3 == 1 and j < 33 -> h plane, j % 3 == 2 and j < 30 -> w plane, every other j -> t plane
|
| 20 |
+
cos = cos(cat(angle, angle)), sin = sin(cat(angle, angle)) float32 [1, L, 64]
|
| 21 |
+
The cosine and sine are taken in float64 on the float32 angle and rounded to float32 (torch's float32 cos / sin sit
|
| 22 |
+
within 1 ulp of these)."""
|
| 23 |
+
from __future__ import annotations
|
| 24 |
+
|
| 25 |
+
import numpy as np
|
| 26 |
+
|
| 27 |
+
from .errors import RequestTooLong
|
| 28 |
+
|
| 29 |
+
BUCKETS = (256, 512, 1024, 2048, 4096)
|
| 30 |
+
ROTARY_DIM = 64
|
| 31 |
+
ROPE_THETA = 10000000.0
|
| 32 |
+
MROPE_SECTION = (11, 11, 10)
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def pick_length(tokens: int, buckets=BUCKETS) -> int:
|
| 36 |
+
"""The smallest row graph length that holds `tokens`; RequestTooLong past the largest."""
|
| 37 |
+
for length in sorted(buckets):
|
| 38 |
+
if tokens <= length:
|
| 39 |
+
return length
|
| 40 |
+
raise RequestTooLong(f"the record has {tokens} tokens; the largest row graph holds {max(buckets)}")
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def text_positions(tokens: int) -> np.ndarray:
|
| 44 |
+
"""[3, T] int64 planes of a record without a picture: (i, i, i)."""
|
| 45 |
+
return np.broadcast_to(np.arange(tokens, dtype=np.int64), (3, tokens)).copy()
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def image_positions(tokens: int, token_offset: int, h: int, w: int) -> np.ndarray:
|
| 49 |
+
"""[3, T] int64 planes of a record with one picture of h x w merged tokens right after the prefix."""
|
| 50 |
+
p, n = int(token_offset), int(h) * int(w)
|
| 51 |
+
if not p + 1 + n < tokens:
|
| 52 |
+
raise ValueError(f"a picture of {n} tokens at {p} does not fit a record of {tokens} tokens")
|
| 53 |
+
pos = text_positions(tokens)
|
| 54 |
+
k = np.arange(n, dtype=np.int64)
|
| 55 |
+
pos[0, p + 1:p + 1 + n] = p + 1
|
| 56 |
+
pos[1, p + 1:p + 1 + n] = p + 1 + k // w
|
| 57 |
+
pos[2, p + 1:p + 1 + n] = p + 1 + k % w
|
| 58 |
+
pos[:, p + n + 1:] = np.arange(p + n + 1, tokens, dtype=np.int64) - (n - max(h, w))
|
| 59 |
+
return pos
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def record_positions(tokens: int, token_offset: int | None = None, image_grid=None) -> np.ndarray:
|
| 63 |
+
"""[3, T] planes of a record (text-only when `image_grid` is None)."""
|
| 64 |
+
if image_grid is None:
|
| 65 |
+
return text_positions(tokens)
|
| 66 |
+
return image_positions(tokens, token_offset, *image_grid)
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def padded_positions(pos3: np.ndarray, length: int) -> np.ndarray:
|
| 70 |
+
"""[3, L]: the record's planes, then maximum + 1, + 2, ... on all three planes for the pad positions."""
|
| 71 |
+
n = pos3.shape[1]
|
| 72 |
+
if n > length:
|
| 73 |
+
raise ValueError(f"{n} positions do not fit length {length}")
|
| 74 |
+
out = np.empty((3, length), dtype=np.int64)
|
| 75 |
+
out[:, :n] = pos3
|
| 76 |
+
out[:, n:] = int(pos3.max()) + 1 + np.arange(length - n, dtype=np.int64)[None, :]
|
| 77 |
+
return out
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def inv_freq(theta: float = ROPE_THETA, rotary_dim: int = ROTARY_DIM) -> np.ndarray:
|
| 81 |
+
"""float32 [rotary_dim / 2] = 1 / theta^(arange(0, dim, 2) / dim), every step in float32."""
|
| 82 |
+
expo = np.arange(0, rotary_dim, 2, dtype=np.float32) / np.float32(rotary_dim)
|
| 83 |
+
return (np.float32(1.0) / np.power(np.float32(theta), expo, dtype=np.float32)).astype(np.float32)
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def mrope_tables(pos3: np.ndarray, freq: np.ndarray | None = None,
|
| 87 |
+
section: tuple[int, int, int] = MROPE_SECTION) -> tuple[np.ndarray, np.ndarray]:
|
| 88 |
+
"""[3, L] position planes -> cos, sin float32 [1, L, 2 * len(freq)]."""
|
| 89 |
+
freq = inv_freq() if freq is None else freq
|
| 90 |
+
ang = (pos3.astype(np.float32)[:, :, None] * freq[None, None, :]).astype(np.float32) # [3, L, d/2]
|
| 91 |
+
a = ang[0].copy()
|
| 92 |
+
for plane, offset in ((1, 1), (2, 2)):
|
| 93 |
+
stop = section[plane] * 3
|
| 94 |
+
a[:, offset:stop:3] = ang[plane][:, offset:stop:3]
|
| 95 |
+
emb = np.concatenate([a, a], axis=-1).astype(np.float64)
|
| 96 |
+
return np.cos(emb).astype(np.float32)[None], np.sin(emb).astype(np.float32)[None]
|
host/clef_litert/request.py
ADDED
|
@@ -0,0 +1,175 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Request -> ids and spans: the checks of the checkpoint's own `systemone()` and its `encode_record()`
|
| 2 |
+
(joint_schema_model.py at Cloudflare/clef-flash f0b7d3cd, sha256 0e304cf7...), written out on a `tokenizers.Tokenizer`
|
| 3 |
+
built from the repository's own tokenizer.json (sha256 06b95093...).
|
| 4 |
+
|
| 5 |
+
Each piece is tokenized on its own without special tokens and the pieces are concatenated (the author's `_tokens`):
|
| 6 |
+
|
| 7 |
+
prefix "<|im_start|>system\\n{SYSTEM_PROMPT}<|im_end|>\\n<|im_start|>user\\nSTATE:\\n" 36 ids
|
| 8 |
+
[image] "<|vision_start|>" + N x "<|image_pad|>" + "<|vision_end|>\\n" N + 3 ids
|
| 9 |
+
state render(state): a string as is, anything else compact JSON (sorted keys, ensure_ascii off)
|
| 10 |
+
schema "\\n\\nSCHEMA FIELDS:\\n", then per question
|
| 11 |
+
"\\nFIELD {i}\\nID: {id}\\nTYPE: {type}\\nINSTRUCTION: " + [render(instructions or id)]
|
| 12 |
+
+ "\\nALLOWED OPTIONS:\\n" + per option "OPTION {j}: " + [render({"option_id", "description"})] + "\\n",
|
| 13 |
+
and "END FIELD\\n" [..] = the head's spans, [start, end)
|
| 14 |
+
suffix "\\n<|im_end|>\\n<|im_start|>assistant\\n<think>\\n\\n</think>\\n\\nJOINT SCHEMA DECISIONS:" 18 ids
|
| 15 |
+
|
| 16 |
+
Option order (the head scores options in this order): noul = (true, false), the author's default descriptions
|
| 17 |
+
overridable by `criteria`; choice = option ids sorted as strings; score = levels 0..n-1. The image block (one picture,
|
| 18 |
+
N = H x W merged tokens) sits right after the prefix, where the processor's media ids go; its index is the
|
| 19 |
+
`token_offset`. The state is cut to fit `max_length` (the author's 16384) exactly as encode_record cuts it; the row
|
| 20 |
+
graphs' own limit (4,096) is the host's check (host.py), never a cut.
|
| 21 |
+
|
| 22 |
+
Malformed requests: the author's ValueError conditions raise the named subclasses of errors.py with the author's
|
| 23 |
+
messages; anything the author's code would stop on otherwise (a question that is not an object, choice criteria that
|
| 24 |
+
are not an object, unsortable state keys) stops here the same way (AttributeError, TypeError, ...)."""
|
| 25 |
+
from __future__ import annotations
|
| 26 |
+
|
| 27 |
+
import json
|
| 28 |
+
from dataclasses import dataclass
|
| 29 |
+
from typing import Any
|
| 30 |
+
|
| 31 |
+
from .errors import BadQuestionType, EmptyCriteria, EmptyRecord, MissingModelOrState, NoQuestions, SchemaTooLong
|
| 32 |
+
|
| 33 |
+
SYSTEM_PROMPT = (
|
| 34 |
+
"Read the complete state and schema. Decide every field jointly. Each answer "
|
| 35 |
+
"must be exactly one of that field's allowed options."
|
| 36 |
+
)
|
| 37 |
+
PREFIX = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\nSTATE:\n"
|
| 38 |
+
SUFFIX = "\n<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nJOINT SCHEMA DECISIONS:"
|
| 39 |
+
SCHEMA_HEADER = "\n\nSCHEMA FIELDS:\n"
|
| 40 |
+
QUESTION_TYPES = {"noul": 0, "choice": 1, "score": 2}
|
| 41 |
+
NOUL_CRITERIA = {
|
| 42 |
+
"true": "The proposition is true or the answer is yes.",
|
| 43 |
+
"false": "The proposition is false or the answer is no.",
|
| 44 |
+
}
|
| 45 |
+
MAX_LENGTH = 16384 # encode_record / systemone default
|
| 46 |
+
VISION_START, VISION_END, IMAGE_PAD = 248053, 248054, 248056
|
| 47 |
+
PAD_ID = 248044 # <|endoftext|>, the processor's pad (never written into a row: pads are zeros)
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def render(value: Any) -> str:
|
| 51 |
+
"""A string as is; anything else compact JSON (the author's `render`)."""
|
| 52 |
+
if isinstance(value, str):
|
| 53 |
+
return value
|
| 54 |
+
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
def question_options(question: dict) -> list[tuple[str, Any]]:
|
| 58 |
+
"""(option id, description) in the order the head scores them (the author's `question_options`)."""
|
| 59 |
+
question_type = str(question["type"])
|
| 60 |
+
if question_type == "noul":
|
| 61 |
+
criteria = dict(NOUL_CRITERIA)
|
| 62 |
+
criteria.update(question.get("criteria") or {})
|
| 63 |
+
return [(key, criteria[key]) for key in ("true", "false")]
|
| 64 |
+
if question_type == "choice":
|
| 65 |
+
return sorted((str(key), value) for key, value in question["criteria"].items())
|
| 66 |
+
return [(str(index), value) for index, value in enumerate(question["criteria"])]
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def validate_request(request: dict) -> None:
|
| 70 |
+
"""The checks `systemone()` makes before it encodes, in its order and with its messages."""
|
| 71 |
+
questions = request.get("questions")
|
| 72 |
+
if not isinstance(request.get("model"), str) or "state" not in request:
|
| 73 |
+
raise MissingModelOrState("model and state are required")
|
| 74 |
+
if not isinstance(questions, dict) or not questions:
|
| 75 |
+
raise NoQuestions("at least one question is required")
|
| 76 |
+
for question_id, question in questions.items():
|
| 77 |
+
if question.get("type") not in QUESTION_TYPES:
|
| 78 |
+
raise BadQuestionType(f"{question_id}: type must be noul, choice, or score")
|
| 79 |
+
if question["type"] != "noul" and not question.get("criteria"):
|
| 80 |
+
raise EmptyCriteria(f"{question_id}: criteria must not be empty")
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
def tokens(tok, text: str) -> list[int]:
|
| 84 |
+
"""`text` -> ids without special tokens (`tokenizers.Tokenizer`)."""
|
| 85 |
+
return list(tok.encode(text, add_special_tokens=False).ids)
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
def image_block(tok, n: int) -> list[int]:
|
| 89 |
+
"""The processor's media ids of one picture of n merged tokens: <|vision_start|>, n pads, <|vision_end|>, "\\n"."""
|
| 90 |
+
ids = tokens(tok, "<|vision_start|>" + "<|image_pad|>" * n + "<|vision_end|>\n")
|
| 91 |
+
if not (ids[:1] == [VISION_START] and ids[1:n + 1] == [IMAGE_PAD] * n and ids[n + 1:n + 2] == [VISION_END]
|
| 92 |
+
and len(ids) == n + 3):
|
| 93 |
+
raise ValueError("the tokenizer does not keep the image placeholder tokens whole")
|
| 94 |
+
return ids
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
@dataclass(frozen=True)
|
| 98 |
+
class EncodedQuestion:
|
| 99 |
+
question_id: str
|
| 100 |
+
type: str # noul | choice | score
|
| 101 |
+
type_id: int # 0 | 1 | 2 (the head's type embedding)
|
| 102 |
+
question_span: tuple[int, int] # [start, end) in ids
|
| 103 |
+
option_spans: tuple[tuple[int, int], ...]
|
| 104 |
+
option_ids: tuple[str, ...]
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
@dataclass(frozen=True)
|
| 108 |
+
class EncodedRecord:
|
| 109 |
+
ids: tuple[int, ...] # the processor-form ids (N x <|image_pad|> for a picture)
|
| 110 |
+
questions: tuple[EncodedQuestion, ...]
|
| 111 |
+
token_offset: int | None = None # index of <|vision_start|> (= len(prefix)), None without a picture
|
| 112 |
+
image_grid: tuple[int, int] | None = None # merged (H, W) of the picture, None without one
|
| 113 |
+
|
| 114 |
+
@property
|
| 115 |
+
def tokens(self) -> int:
|
| 116 |
+
return len(self.ids)
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def encode_record(tok, record: dict, max_length: int = MAX_LENGTH, max_state_tokens: int | None = None,
|
| 120 |
+
image_grid: tuple[int, int] | None = None) -> EncodedRecord:
|
| 121 |
+
"""The author's encode_record. `image_grid` = the merged (H, W) of one picture (its block goes right after the
|
| 122 |
+
prefix), None for a record without one."""
|
| 123 |
+
schema_ids = tokens(tok, SCHEMA_HEADER)
|
| 124 |
+
questions = []
|
| 125 |
+
for question_index, (question_id, question) in enumerate(record["questions"].items()):
|
| 126 |
+
schema_ids += tokens(tok, f"\nFIELD {question_index + 1}\nID: {question_id}\nTYPE: {question['type']}"
|
| 127 |
+
f"\nINSTRUCTION: ")
|
| 128 |
+
question_start = len(schema_ids)
|
| 129 |
+
instructions = question.get("instructions")
|
| 130 |
+
if instructions is None or instructions == "":
|
| 131 |
+
instructions = str(question_id)
|
| 132 |
+
schema_ids += tokens(tok, render(instructions))
|
| 133 |
+
question_end = len(schema_ids)
|
| 134 |
+
schema_ids += tokens(tok, "\nALLOWED OPTIONS:\n")
|
| 135 |
+
option_spans, option_ids = [], []
|
| 136 |
+
for option_index, (option_id, description) in enumerate(question_options(question)):
|
| 137 |
+
schema_ids += tokens(tok, f"OPTION {option_index + 1}: ")
|
| 138 |
+
option_start = len(schema_ids)
|
| 139 |
+
semantics = {"option_id": option_id}
|
| 140 |
+
if description is not None:
|
| 141 |
+
semantics["description"] = description
|
| 142 |
+
schema_ids += tokens(tok, render(semantics))
|
| 143 |
+
option_spans.append((option_start, len(schema_ids)))
|
| 144 |
+
option_ids.append(option_id)
|
| 145 |
+
schema_ids += tokens(tok, "\n")
|
| 146 |
+
schema_ids += tokens(tok, "END FIELD\n")
|
| 147 |
+
questions.append((str(question_id), str(question["type"]), QUESTION_TYPES[str(question["type"])],
|
| 148 |
+
(question_start, question_end), tuple(option_spans), tuple(option_ids)))
|
| 149 |
+
|
| 150 |
+
prefix_ids = tokens(tok, PREFIX)
|
| 151 |
+
suffix_ids = tokens(tok, SUFFIX)
|
| 152 |
+
token_offset = None
|
| 153 |
+
if image_grid is not None:
|
| 154 |
+
h, w = (int(v) for v in image_grid)
|
| 155 |
+
token_offset = len(prefix_ids)
|
| 156 |
+
prefix_ids = prefix_ids + image_block(tok, h * w)
|
| 157 |
+
image_grid = (h, w)
|
| 158 |
+
state_ids = tokens(tok, render(record["state"]))
|
| 159 |
+
if max_state_tokens is not None:
|
| 160 |
+
state_ids = state_ids[:max_state_tokens]
|
| 161 |
+
fixed_length = len(prefix_ids) + len(schema_ids) + len(suffix_ids)
|
| 162 |
+
if fixed_length > max_length:
|
| 163 |
+
raise SchemaTooLong(f"schema requires {fixed_length} tokens before state; maximum is {max_length}")
|
| 164 |
+
state_ids = state_ids[: max_length - fixed_length]
|
| 165 |
+
shift = len(prefix_ids) + len(state_ids)
|
| 166 |
+
shifted = tuple(
|
| 167 |
+
EncodedQuestion(question_id=qid, type=qtype, type_id=type_id,
|
| 168 |
+
question_span=(span[0] + shift, span[1] + shift),
|
| 169 |
+
option_spans=tuple((a + shift, b + shift) for a, b in option_spans),
|
| 170 |
+
option_ids=option_ids)
|
| 171 |
+
for qid, qtype, type_id, span, option_spans, option_ids in questions)
|
| 172 |
+
ids = tuple(prefix_ids + state_ids + schema_ids + suffix_ids)
|
| 173 |
+
if not ids or not shifted:
|
| 174 |
+
raise EmptyRecord("record produced no model input or questions")
|
| 175 |
+
return EncodedRecord(ids=ids, questions=shifted, token_offset=token_offset, image_grid=image_grid)
|
host/clef_litert/response.py
ADDED
|
@@ -0,0 +1,56 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Head logits -> the SystemOne-compatible response body, as the checkpoint's own `systemone()` builds it.
|
| 2 |
+
|
| 3 |
+
Probabilities: per question, the float32 softmax of its option logits, in NumPy's order (exp of z - max z, a float32
|
| 4 |
+
sum, a float32 division); the author's `question_logits.float().softmax(-1)` is the same computation in torch. They
|
| 5 |
+
enter Python as floats (`.tolist()`), as in the author's code, and `systemone_answer` is the author's function:
|
| 6 |
+
|
| 7 |
+
noul {"type": "noul", "noul": round(p(true), 4)}
|
| 8 |
+
choice {"type": "choice", "choice": the most probable option in the request's criteria order (first on a tie),
|
| 9 |
+
"confidence": round(p(choice), 4), "probabilities": {option: round(p, 4)} in the criteria order}
|
| 10 |
+
score {"type": "score", "score": round(sum(level * p(level)), 4) (Python's left-to-right sum),
|
| 11 |
+
"confidence": round(max p, 4), "legend": {level: criterion}, "probabilities": {level: round(p, 4)}}
|
| 12 |
+
|
| 13 |
+
Body: {"model": request["model"], "answers": {question id: answer} in the request's order,
|
| 14 |
+
"usage": {"input_tokens": the record's tokens, "output_tokens": 0}}."""
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import numpy as np
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def question_probs(logits) -> np.ndarray:
|
| 21 |
+
"""float32 softmax over one question's option logits."""
|
| 22 |
+
z = np.asarray(logits, dtype=np.float32)
|
| 23 |
+
e = np.exp(z - z.max())
|
| 24 |
+
return e / e.sum(dtype=np.float32)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def systemone_answer(question: dict, probabilities: dict) -> dict:
|
| 28 |
+
"""Per-option probabilities of one question -> its SystemOne answer (the author's function)."""
|
| 29 |
+
if question["type"] == "noul":
|
| 30 |
+
return {"type": "noul", "noul": round(probabilities["true"], 4)}
|
| 31 |
+
if question["type"] == "choice":
|
| 32 |
+
options = [str(option) for option in question["criteria"]]
|
| 33 |
+
choice = max(options, key=probabilities.__getitem__)
|
| 34 |
+
return {
|
| 35 |
+
"type": "choice",
|
| 36 |
+
"choice": choice,
|
| 37 |
+
"confidence": round(probabilities[choice], 4),
|
| 38 |
+
"probabilities": {option: round(probabilities[option], 4) for option in options},
|
| 39 |
+
}
|
| 40 |
+
levels = [str(index) for index in range(len(question["criteria"]))]
|
| 41 |
+
return {
|
| 42 |
+
"type": "score",
|
| 43 |
+
"score": round(sum(index * probabilities[level] for index, level in enumerate(levels)), 4),
|
| 44 |
+
"confidence": round(max(probabilities[level] for level in levels), 4),
|
| 45 |
+
"legend": dict(zip(levels, question["criteria"])),
|
| 46 |
+
"probabilities": {level: round(probabilities[level], 4) for level in levels},
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def response(request: dict, questions, probs, input_tokens: int) -> dict:
|
| 51 |
+
"""The response body from the encoded questions (request.EncodedQuestion) and their float32 probabilities."""
|
| 52 |
+
answers = {q.question_id: systemone_answer(request["questions"][q.question_id],
|
| 53 |
+
dict(zip(q.option_ids, np.asarray(p, np.float32).tolist())))
|
| 54 |
+
for q, p in zip(questions, probs)}
|
| 55 |
+
return {"model": request["model"], "answers": answers,
|
| 56 |
+
"usage": {"input_tokens": int(input_tokens), "output_tokens": 0}}
|
host/clef_litert/runtime.py
ADDED
|
@@ -0,0 +1,273 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The LiteRT graphs of a clef-flash folder on ai-edge-litert's CompiledModel, and the host around them.
|
| 2 |
+
|
| 3 |
+
with ClefLiteRT.from_dir("<folder>", accelerator="gpu") as clef: # "gpu" = float32 precision, "cpu" = XNNPACK
|
| 4 |
+
body = clef.decide(request, tile=448) # the SystemOne-compatible response
|
| 5 |
+
|
| 6 |
+
Folder layout (FOLDER below; every path relative to the folder):
|
| 7 |
+
clef-flash_rowprefill_embeds_L{256,512,1024,2048,4096}_fp16fc.tflite row graphs: embeds, valid, cos, sin -> hidden
|
| 8 |
+
clef-flash_vision_tower_g{256,448}_fp16fc.tflite tower + merger: patches, pos -> rows
|
| 9 |
+
tables/vision_pos_g{256,448}.safetensors the towers' position rows (float32, key "pos")
|
| 10 |
+
tables/embed_tokens.safetensors, tables/lm_head.safetensors the bf16 vocabulary tables
|
| 11 |
+
head/joint_head.safetensors, head/joint_head_config.json the checkpoint's joint schema head
|
| 12 |
+
tokenizer/tokenizer.json the checkpoint's tokenizer
|
| 13 |
+
|
| 14 |
+
Only the row graphs present are used: a request runs on the smallest present length that holds it, and a request
|
| 15 |
+
longer than the largest present one is refused (RequestTooLong), never cut. Graphs are compiled when a request first
|
| 16 |
+
needs them. One row graph stays compiled at a time (its fp16 weights are about 13.9 GB; on a Mac GPU at float32
|
| 17 |
+
precision a compiled row graph held about 37 GB): a request of another length closes it and compiles the next one.
|
| 18 |
+
The towers (about 0.9 GB each) stay compiled once used.
|
| 19 |
+
|
| 20 |
+
Accelerators: "gpu" = HardwareAccelerator.GPU with GpuOptions(enforce_f32=True) (Metal on a Mac; the graphs are
|
| 21 |
+
checked at float32 precision only: at the default precision (fp16 storage) the row graphs' norms overflow and the
|
| 22 |
+
towers' rows are wrong); "cpu" = HardwareAccelerator.CPU with CpuOptions(num_threads=threads) (XNNPACK; no weight
|
| 23 |
+
cache file is written). No other runtime option is set."""
|
| 24 |
+
from __future__ import annotations
|
| 25 |
+
|
| 26 |
+
import json
|
| 27 |
+
import struct
|
| 28 |
+
from pathlib import Path
|
| 29 |
+
|
| 30 |
+
import numpy as np
|
| 31 |
+
|
| 32 |
+
from .head import JointHead
|
| 33 |
+
from .host import HIDDEN, ClefHost
|
| 34 |
+
from .pixels import TILES, tile_grid
|
| 35 |
+
from .positions import BUCKETS
|
| 36 |
+
from .tables import Bf16Table, read_header
|
| 37 |
+
|
| 38 |
+
SIGNATURE = "serving_default"
|
| 39 |
+
ACCELERATORS = ("gpu", "cpu")
|
| 40 |
+
FOLDER = {
|
| 41 |
+
"row": "clef-flash_rowprefill_embeds_L{length}_fp16fc.tflite",
|
| 42 |
+
"tower": "clef-flash_vision_tower_g{tile}_fp16fc.tflite",
|
| 43 |
+
"tower_pos": "tables/vision_pos_g{tile}.safetensors",
|
| 44 |
+
"embed_tokens": "tables/embed_tokens.safetensors",
|
| 45 |
+
"lm_head": "tables/lm_head.safetensors",
|
| 46 |
+
"joint_head": "head/joint_head.safetensors",
|
| 47 |
+
"joint_head_config": "head/joint_head_config.json",
|
| 48 |
+
"tokenizer": "tokenizer/tokenizer.json",
|
| 49 |
+
}
|
| 50 |
+
ROW_INPUTS = ("embeds", "valid", "cos", "sin")
|
| 51 |
+
TOWER_INPUTS = ("patches", "pos")
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def load_f32_tensor(path: str | Path, key: str | None = None) -> np.ndarray:
|
| 55 |
+
"""One float32 tensor of a safetensors file (the towers' position rows), read without a tensor framework."""
|
| 56 |
+
header, data_start = read_header(path)
|
| 57 |
+
names = [k for k in header if k != "__metadata__"]
|
| 58 |
+
if key is None:
|
| 59 |
+
if len(names) != 1:
|
| 60 |
+
raise ValueError(f"{Path(path).name} holds {len(names)} tensors; name the key")
|
| 61 |
+
key = names[0]
|
| 62 |
+
e = header[key]
|
| 63 |
+
if e["dtype"] != "F32":
|
| 64 |
+
raise ValueError(f"{key}: dtype {e['dtype']}, expected F32")
|
| 65 |
+
start, end = e["data_offsets"]
|
| 66 |
+
n = int(np.prod(e["shape"], dtype=np.int64))
|
| 67 |
+
if end - start != 4 * n:
|
| 68 |
+
raise ValueError(f"{key}: {end - start} bytes for {n} F32 values")
|
| 69 |
+
with open(path, "rb") as f:
|
| 70 |
+
f.seek(data_start + start)
|
| 71 |
+
data = f.read(end - start)
|
| 72 |
+
return np.frombuffer(data, dtype="<f4").reshape(e["shape"]).astype(np.float32)
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
def write_f32_tensor(path: str | Path, key: str, array: np.ndarray) -> dict:
|
| 76 |
+
"""A one-key safetensors file of a float32 array, header without __metadata__ (padded with spaces to 8 bytes, as the
|
| 77 |
+
safetensors writer pads). Refuses to overwrite."""
|
| 78 |
+
path = Path(path)
|
| 79 |
+
if path.exists():
|
| 80 |
+
raise FileExistsError(f"refusing to overwrite {path}")
|
| 81 |
+
a = np.ascontiguousarray(array, dtype="<f4")
|
| 82 |
+
hdr = json.dumps({key: {"dtype": "F32", "shape": list(a.shape), "data_offsets": [0, a.nbytes]}},
|
| 83 |
+
separators=(",", ":")).encode()
|
| 84 |
+
hdr += b" " * (-len(hdr) % 8)
|
| 85 |
+
tmp = path.with_name(path.name + ".part")
|
| 86 |
+
with open(tmp, "wb") as f:
|
| 87 |
+
f.write(struct.pack("<Q", len(hdr)) + hdr)
|
| 88 |
+
f.write(a.tobytes())
|
| 89 |
+
tmp.rename(path)
|
| 90 |
+
return {"file": path.name, "key": key, "dtype": "F32", "shape": list(a.shape), "bytes": path.stat().st_size}
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def compile_options(accelerator: str, threads: int = 8):
|
| 94 |
+
"""-> (CompiledModel options, a plain description of them)."""
|
| 95 |
+
from ai_edge_litert.compiled_model import CpuOptions, GpuOptions, HardwareAccelerator, Options
|
| 96 |
+
|
| 97 |
+
if accelerator == "cpu":
|
| 98 |
+
cpu = CpuOptions(num_threads=threads)
|
| 99 |
+
return (Options(hardware_accelerators=HardwareAccelerator.CPU, cpu_options=cpu),
|
| 100 |
+
{"accelerator": "cpu", "backend": "CPU (XNNPACK)", "num_threads": threads})
|
| 101 |
+
if accelerator == "gpu":
|
| 102 |
+
gpu = GpuOptions(enforce_f32=True)
|
| 103 |
+
return (Options(hardware_accelerators=HardwareAccelerator.GPU, gpu_options=gpu),
|
| 104 |
+
{"accelerator": "gpu", "backend": "GPU, float32 precision", "enforce_f32": True})
|
| 105 |
+
raise ValueError(f"accelerator {accelerator!r}: one of {ACCELERATORS}")
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
class LiteRTGraph:
|
| 109 |
+
"""One .tflite file with one signature, compiled on `accelerator`; call it with its inputs by name."""
|
| 110 |
+
|
| 111 |
+
def __init__(self, path: str | Path, accelerator: str = "gpu", threads: int = 8):
|
| 112 |
+
from ai_edge_litert.compiled_model import CompiledModel
|
| 113 |
+
|
| 114 |
+
self.path, self.accelerator = Path(path), accelerator
|
| 115 |
+
opts, self.options = compile_options(accelerator, threads)
|
| 116 |
+
self.model = CompiledModel.from_file(str(self.path), options=opts)
|
| 117 |
+
signatures = list(self.model.get_signature_list())
|
| 118 |
+
if signatures != [SIGNATURE]:
|
| 119 |
+
self.model.close()
|
| 120 |
+
raise ValueError(f"{self.path.name}: signatures {signatures}, expected [{SIGNATURE!r}]")
|
| 121 |
+
self.input_shapes = {n: tuple(int(x) for x in d["shape"])
|
| 122 |
+
for n, d in self.model.get_input_tensor_details(SIGNATURE).items()}
|
| 123 |
+
self.output_shapes = {n: tuple(int(x) for x in d["shape"])
|
| 124 |
+
for n, d in self.model.get_output_tensor_details(SIGNATURE).items()}
|
| 125 |
+
self.inputs = {n: self.model.create_input_buffer_by_name(SIGNATURE, n) for n in self.input_shapes}
|
| 126 |
+
self.outputs = {n: self.model.create_output_buffer_by_name(SIGNATURE, n) for n in self.output_shapes}
|
| 127 |
+
|
| 128 |
+
def fully_accelerated(self) -> bool:
|
| 129 |
+
return bool(self.model.is_fully_accelerated())
|
| 130 |
+
|
| 131 |
+
def __call__(self, **feeds: np.ndarray) -> dict[str, np.ndarray]:
|
| 132 |
+
if sorted(feeds) != sorted(self.input_shapes):
|
| 133 |
+
raise ValueError(f"{self.path.name}: inputs {sorted(feeds)}, expected {sorted(self.input_shapes)}")
|
| 134 |
+
for name, value in feeds.items():
|
| 135 |
+
value = np.ascontiguousarray(value, dtype=np.float32)
|
| 136 |
+
if value.shape != self.input_shapes[name]:
|
| 137 |
+
raise ValueError(f"{self.path.name}: {name} {value.shape}, expected {self.input_shapes[name]}")
|
| 138 |
+
self.inputs[name].write(value)
|
| 139 |
+
self.model.run_by_name(SIGNATURE, self.inputs, self.outputs)
|
| 140 |
+
return {name: np.asarray(self.outputs[name].read(int(np.prod(shape)), np.float32),
|
| 141 |
+
dtype=np.float32).reshape(shape)
|
| 142 |
+
for name, shape in self.output_shapes.items()}
|
| 143 |
+
|
| 144 |
+
def close(self) -> None:
|
| 145 |
+
for buffer in list(self.inputs.values()) + list(self.outputs.values()):
|
| 146 |
+
try:
|
| 147 |
+
buffer.destroy()
|
| 148 |
+
except Exception:
|
| 149 |
+
pass
|
| 150 |
+
self.inputs, self.outputs = {}, {}
|
| 151 |
+
self.model.close()
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
class ClefGraphs:
|
| 155 |
+
"""The row graphs (one file per length) and the tower graphs (one per tile), compiled when first needed.
|
| 156 |
+
row_graph(length, inputs) and tower(tile, patches) are the callables ClefHost.decide takes."""
|
| 157 |
+
|
| 158 |
+
def __init__(self, row_files: dict, tower_files: dict | None = None, tower_pos: dict | None = None,
|
| 159 |
+
accelerator: str = "gpu", threads: int = 8):
|
| 160 |
+
if accelerator not in ACCELERATORS:
|
| 161 |
+
raise ValueError(f"accelerator {accelerator!r}: one of {ACCELERATORS}")
|
| 162 |
+
self.row_files = {int(L): Path(p) for L, p in row_files.items()}
|
| 163 |
+
if not self.row_files:
|
| 164 |
+
raise ValueError("no row graph")
|
| 165 |
+
self.tower_files = {int(t): Path(p) for t, p in (tower_files or {}).items()}
|
| 166 |
+
self.tower_pos = {}
|
| 167 |
+
for t, pos in (tower_pos or {}).items():
|
| 168 |
+
self.tower_pos[int(t)] = load_f32_tensor(pos) if isinstance(pos, (str, Path)) else np.asarray(pos, np.float32)
|
| 169 |
+
missing = sorted(set(self.tower_files) - set(self.tower_pos))
|
| 170 |
+
if missing:
|
| 171 |
+
raise ValueError(f"tower(s) {missing} without position rows")
|
| 172 |
+
self.accelerator, self.threads = accelerator, threads
|
| 173 |
+
self.row = None # (length, LiteRTGraph): the one compiled row graph
|
| 174 |
+
self.towers: dict[int, LiteRTGraph] = {}
|
| 175 |
+
self.compiled: list[dict] = [] # what was compiled, in order (file, accelerator)
|
| 176 |
+
|
| 177 |
+
@classmethod
|
| 178 |
+
def from_dir(cls, folder: str | Path, accelerator: str = "gpu", threads: int = 8, lengths=None) -> "ClefGraphs":
|
| 179 |
+
"""The graphs of a folder laid out as FOLDER (only the files present; `lengths` limits the row graphs)."""
|
| 180 |
+
folder = Path(folder)
|
| 181 |
+
rows = {L: folder / FOLDER["row"].format(length=L) for L in BUCKETS if lengths is None or L in lengths}
|
| 182 |
+
rows = {L: p for L, p in rows.items() if p.is_file()}
|
| 183 |
+
towers = {t: folder / FOLDER["tower"].format(tile=t) for t in TILES}
|
| 184 |
+
towers = {t: p for t, p in towers.items() if p.is_file()}
|
| 185 |
+
pos = {t: folder / FOLDER["tower_pos"].format(tile=t) for t in towers}
|
| 186 |
+
return cls(rows, towers, pos, accelerator, threads)
|
| 187 |
+
|
| 188 |
+
@property
|
| 189 |
+
def lengths(self) -> tuple[int, ...]:
|
| 190 |
+
return tuple(sorted(self.row_files))
|
| 191 |
+
|
| 192 |
+
@property
|
| 193 |
+
def tiles(self) -> tuple[int, ...]:
|
| 194 |
+
return tuple(sorted(self.tower_files))
|
| 195 |
+
|
| 196 |
+
def _row_graph(self, length: int) -> LiteRTGraph:
|
| 197 |
+
if self.row is not None and self.row[0] == length:
|
| 198 |
+
return self.row[1]
|
| 199 |
+
if length not in self.row_files:
|
| 200 |
+
raise ValueError(f"no row graph of length {length} (present: {list(self.lengths)})")
|
| 201 |
+
if self.row is not None: # one row graph at a time
|
| 202 |
+
self.row[1].close()
|
| 203 |
+
self.row = None
|
| 204 |
+
graph = LiteRTGraph(self.row_files[length], self.accelerator, self.threads)
|
| 205 |
+
if graph.input_shapes["embeds"] != (1, length, HIDDEN):
|
| 206 |
+
graph.close()
|
| 207 |
+
raise ValueError(f"{self.row_files[length].name}: embeds {graph.input_shapes['embeds']}")
|
| 208 |
+
self.row = (length, graph)
|
| 209 |
+
self.compiled.append({"file": self.row_files[length].name, "accelerator": self.accelerator})
|
| 210 |
+
return graph
|
| 211 |
+
|
| 212 |
+
def row_graph(self, length: int, inputs: dict) -> np.ndarray:
|
| 213 |
+
"""The row graph of `length` on the host's row_inputs -> hidden [1, L, 4096]."""
|
| 214 |
+
graph = self._row_graph(int(length))
|
| 215 |
+
return graph(**{n: inputs[n] for n in ROW_INPUTS})["hidden"]
|
| 216 |
+
|
| 217 |
+
def _tower(self, tile: int) -> LiteRTGraph:
|
| 218 |
+
if tile not in self.towers:
|
| 219 |
+
if tile not in self.tower_files:
|
| 220 |
+
raise ValueError(f"no tower graph for a {tile} px tile (present: {list(self.tiles)})")
|
| 221 |
+
self.towers[tile] = LiteRTGraph(self.tower_files[tile], self.accelerator, self.threads)
|
| 222 |
+
self.compiled.append({"file": self.tower_files[tile].name, "accelerator": self.accelerator})
|
| 223 |
+
return self.towers[tile]
|
| 224 |
+
|
| 225 |
+
def tower(self, tile: int, patches: np.ndarray) -> np.ndarray:
|
| 226 |
+
"""The tower of `tile` on the host's patches [4 G^2, 1536] -> picture rows [G^2, 4096]."""
|
| 227 |
+
g = tile_grid(tile)
|
| 228 |
+
rows = self._tower(tile)(patches=np.asarray(patches, np.float32)[None], pos=self.tower_pos[tile][None])["rows"]
|
| 229 |
+
if rows.shape != (1, g * g, HIDDEN):
|
| 230 |
+
raise ValueError(f"tower rows {rows.shape}, expected {(1, g * g, HIDDEN)}")
|
| 231 |
+
return rows[0]
|
| 232 |
+
|
| 233 |
+
def close(self) -> None:
|
| 234 |
+
if self.row is not None:
|
| 235 |
+
self.row[1].close()
|
| 236 |
+
self.row = None
|
| 237 |
+
for graph in self.towers.values():
|
| 238 |
+
graph.close()
|
| 239 |
+
self.towers = {}
|
| 240 |
+
|
| 241 |
+
|
| 242 |
+
class ClefLiteRT:
|
| 243 |
+
"""The host (tokenizer, tables, joint head, response) and the graphs of one folder."""
|
| 244 |
+
|
| 245 |
+
def __init__(self, host: ClefHost, graphs: ClefGraphs):
|
| 246 |
+
self.host, self.graphs = host, graphs
|
| 247 |
+
|
| 248 |
+
@classmethod
|
| 249 |
+
def from_dir(cls, folder: str | Path, accelerator: str = "gpu", threads: int = 8, lengths=None) -> "ClefLiteRT":
|
| 250 |
+
from tokenizers import Tokenizer
|
| 251 |
+
|
| 252 |
+
folder = Path(folder)
|
| 253 |
+
graphs = ClefGraphs.from_dir(folder, accelerator, threads, lengths)
|
| 254 |
+
host = ClefHost(Tokenizer.from_file(str(folder / FOLDER["tokenizer"])),
|
| 255 |
+
Bf16Table(folder / FOLDER["embed_tokens"]), Bf16Table(folder / FOLDER["lm_head"]),
|
| 256 |
+
JointHead.from_files(folder / FOLDER["joint_head"], folder / FOLDER["joint_head_config"]),
|
| 257 |
+
buckets=graphs.lengths)
|
| 258 |
+
return cls(host, graphs)
|
| 259 |
+
|
| 260 |
+
def decide(self, request: dict, tile: int = 448) -> dict:
|
| 261 |
+
"""request -> the SystemOne-compatible response (ClefHost.decide on this folder's graphs). A picture at a tile
|
| 262 |
+
whose tower file is not in the folder is refused (UnsupportedMedia)."""
|
| 263 |
+
tower = self.graphs.tower if tile in self.graphs.tiles else None
|
| 264 |
+
return self.host.decide(request, self.graphs.row_graph, tower, tile=tile)
|
| 265 |
+
|
| 266 |
+
def close(self) -> None:
|
| 267 |
+
self.graphs.close()
|
| 268 |
+
|
| 269 |
+
def __enter__(self) -> "ClefLiteRT":
|
| 270 |
+
return self
|
| 271 |
+
|
| 272 |
+
def __exit__(self, *exc) -> None:
|
| 273 |
+
self.close()
|
host/clef_litert/tables.py
ADDED
|
@@ -0,0 +1,125 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""bfloat16 tensors from safetensors files, read without a tensor framework.
|
| 2 |
+
|
| 3 |
+
A safetensors file is an 8-byte little-endian header length, a JSON header ({name: {dtype, shape, data_offsets}},
|
| 4 |
+
optional "__metadata__") and the data. bfloat16 is the top 16 bits of the float32 with the same value, so a row widens
|
| 5 |
+
to float32 exactly by a 16-bit shift (NumPy has no bfloat16: the safetensors 0.8.0 NumPy reader stops on a BF16 tensor
|
| 6 |
+
with "TypeError: data type 'bfloat16' not understood").
|
| 7 |
+
|
| 8 |
+
Bf16Table(path) one [rows, dim] BF16 tensor, memory-mapped; `rows(ids)` -> float32 [n, dim]
|
| 9 |
+
load_bf16_tensors(path) every tensor of a BF16 file -> {name: float32 array} (the joint head, 243 MB)
|
| 10 |
+
copy_bf16_tensor(...) one tensor's bytes copied into a new one-key file without __metadata__ (the tables)"""
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import hashlib
|
| 14 |
+
import json
|
| 15 |
+
import struct
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
|
| 18 |
+
import numpy as np
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def read_header(path: str | Path) -> tuple[dict, int]:
|
| 22 |
+
"""-> (header dict, file offset of the data region)."""
|
| 23 |
+
with open(path, "rb") as f:
|
| 24 |
+
n = struct.unpack("<Q", f.read(8))[0]
|
| 25 |
+
header = json.loads(f.read(n))
|
| 26 |
+
return header, 8 + n
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def widen_bf16(u16: np.ndarray) -> np.ndarray:
|
| 30 |
+
"""uint16 bfloat16 bits -> float32 (exact)."""
|
| 31 |
+
return (np.asarray(u16, dtype=np.uint32) << 16).view(np.float32)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class Bf16Table:
|
| 35 |
+
"""One BF16 [rows, dim] tensor of a safetensors file, memory-mapped (nothing is read until a row is asked for).
|
| 36 |
+
`key` = the tensor's name (default: the file's only tensor)."""
|
| 37 |
+
|
| 38 |
+
def __init__(self, path: str | Path, key: str | None = None):
|
| 39 |
+
self.path = Path(path)
|
| 40 |
+
header, data_start = read_header(self.path)
|
| 41 |
+
names = [k for k in header if k != "__metadata__"]
|
| 42 |
+
if key is None:
|
| 43 |
+
if len(names) != 1:
|
| 44 |
+
raise ValueError(f"{self.path.name} holds {len(names)} tensors; name the key")
|
| 45 |
+
key = names[0]
|
| 46 |
+
e = header[key]
|
| 47 |
+
if e["dtype"] != "BF16" or len(e["shape"]) != 2:
|
| 48 |
+
raise ValueError(f"{key}: expected a 2-D BF16 tensor, got {e['dtype']} {e['shape']}")
|
| 49 |
+
start, end = e["data_offsets"]
|
| 50 |
+
rows, dim = e["shape"]
|
| 51 |
+
if end - start != rows * dim * 2:
|
| 52 |
+
raise ValueError(f"{key}: {end - start} bytes for {rows} x {dim} BF16")
|
| 53 |
+
self.key, self.shape = key, (rows, dim)
|
| 54 |
+
self.u16 = np.memmap(self.path, dtype=np.uint16, mode="r", offset=data_start + start, shape=(rows, dim))
|
| 55 |
+
self.info = {"file": self.path.name, "key": key, "dtype": "BF16", "shape": [rows, dim],
|
| 56 |
+
"data_offset": data_start + start, "bytes": end - start}
|
| 57 |
+
|
| 58 |
+
def rows(self, ids) -> np.ndarray:
|
| 59 |
+
"""float32 [len(ids), dim]: the rows of `ids`, widened exactly."""
|
| 60 |
+
return widen_bf16(self.u16[np.asarray(ids, dtype=np.int64)])
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def load_bf16_tensors(path: str | Path) -> dict[str, np.ndarray]:
|
| 64 |
+
"""Every tensor of a safetensors file whose tensors are all BF16 -> {name: float32 array of its shape}."""
|
| 65 |
+
header, data_start = read_header(path)
|
| 66 |
+
raw = np.memmap(path, dtype=np.uint8, mode="r")
|
| 67 |
+
out = {}
|
| 68 |
+
for name, e in header.items():
|
| 69 |
+
if name == "__metadata__":
|
| 70 |
+
continue
|
| 71 |
+
if e["dtype"] != "BF16":
|
| 72 |
+
raise ValueError(f"{name}: dtype {e['dtype']}, expected BF16")
|
| 73 |
+
start, end = e["data_offsets"]
|
| 74 |
+
n = int(np.prod(e["shape"], dtype=np.int64))
|
| 75 |
+
if end - start != 2 * n:
|
| 76 |
+
raise ValueError(f"{name}: {end - start} bytes for {n} BF16 values")
|
| 77 |
+
u16 = np.frombuffer(raw[data_start + start:data_start + end].tobytes(), dtype="<u2")
|
| 78 |
+
out[name] = widen_bf16(u16).reshape(e["shape"])
|
| 79 |
+
return out
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def sha256_file(path: str | Path, block: int = 1 << 24) -> str:
|
| 83 |
+
h = hashlib.sha256()
|
| 84 |
+
with open(path, "rb") as f:
|
| 85 |
+
for chunk in iter(lambda: f.read(block), b""):
|
| 86 |
+
h.update(chunk)
|
| 87 |
+
return h.hexdigest()
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
def copy_bf16_tensor(src: str | Path, src_key: str, dst: str | Path, dst_key: str, block: int = 1 << 24) -> dict:
|
| 91 |
+
"""Write `dst` = a safetensors file with the one tensor `dst_key` whose bytes are `src_key`'s bytes in `src`,
|
| 92 |
+
copied by seek + block reads (the source is never loaded whole), header without __metadata__ (padded with spaces
|
| 93 |
+
to 8 bytes, as the safetensors writer pads). Refuses to overwrite. -> record of what was copied."""
|
| 94 |
+
dst = Path(dst)
|
| 95 |
+
if dst.exists():
|
| 96 |
+
raise FileExistsError(f"refusing to overwrite {dst}")
|
| 97 |
+
header, data_start = read_header(src)
|
| 98 |
+
e = header[src_key]
|
| 99 |
+
if e["dtype"] != "BF16":
|
| 100 |
+
raise ValueError(f"{src_key}: dtype {e['dtype']}")
|
| 101 |
+
start, end = e["data_offsets"]
|
| 102 |
+
n = end - start
|
| 103 |
+
if n != 2 * int(np.prod(e["shape"], dtype=np.int64)):
|
| 104 |
+
raise ValueError(f"{src_key}: {n} bytes for shape {e['shape']}")
|
| 105 |
+
hdr = json.dumps({dst_key: {"dtype": "BF16", "shape": list(e["shape"]), "data_offsets": [0, n]}},
|
| 106 |
+
separators=(",", ":")).encode()
|
| 107 |
+
hdr += b" " * (-len(hdr) % 8)
|
| 108 |
+
tmp = dst.with_name(dst.name + ".part")
|
| 109 |
+
h_src = hashlib.sha256()
|
| 110 |
+
copied = 0
|
| 111 |
+
with open(src, "rb") as f, open(tmp, "wb") as g:
|
| 112 |
+
g.write(struct.pack("<Q", len(hdr)) + hdr)
|
| 113 |
+
f.seek(data_start + start)
|
| 114 |
+
while copied < n:
|
| 115 |
+
chunk = f.read(min(block, n - copied))
|
| 116 |
+
if not chunk:
|
| 117 |
+
raise IOError(f"{src}: short read at {copied} of {n} bytes")
|
| 118 |
+
h_src.update(chunk)
|
| 119 |
+
g.write(chunk)
|
| 120 |
+
copied += len(chunk)
|
| 121 |
+
tmp.rename(dst)
|
| 122 |
+
return {"source": {"file": Path(src).name, "key": src_key, "dtype": e["dtype"], "shape": list(e["shape"]),
|
| 123 |
+
"byte_range": [data_start + start, data_start + end], "data_sha256": h_src.hexdigest()},
|
| 124 |
+
"file": dst.name, "key": dst_key, "header_bytes": len(hdr), "data_offset": 8 + len(hdr),
|
| 125 |
+
"bytes": dst.stat().st_size, "bytes_copied": copied}
|
host/requirements-host.txt
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# The clef-flash LiteRT host (host/clef_litert) and examples/run_example.py, as tested: Python 3.14.6, macOS 27.0
|
| 2 |
+
# (arm64), `pip install -r host/requirements-host.txt` into a new venv. The package imports numpy, tokenizers and
|
| 3 |
+
# pillow, and ai-edge-litert when it compiles a graph; the other lines are what pip installed with these four
|
| 4 |
+
# (tokenizers declares huggingface_hub, which the host never imports). No torch, no transformers.
|
| 5 |
+
ai-edge-litert==2.2.0
|
| 6 |
+
anyio==4.15.1
|
| 7 |
+
backports.strenum==1.2.8
|
| 8 |
+
certifi==2026.10.19
|
| 9 |
+
click==8.5.0
|
| 10 |
+
filelock==4.2.4
|
| 11 |
+
flatbuffers==25.12.19
|
| 12 |
+
fsspec==2026.9.0
|
| 13 |
+
h11==0.16.0
|
| 14 |
+
hf-xet==1.7.0
|
| 15 |
+
httpcore==1.0.9
|
| 16 |
+
httpx==0.28.1
|
| 17 |
+
huggingface_hub==1.33.0
|
| 18 |
+
idna==3.20
|
| 19 |
+
ml_dtypes==0.6.0
|
| 20 |
+
numpy==2.5.3
|
| 21 |
+
packaging==26.3
|
| 22 |
+
pillow==12.3.0
|
| 23 |
+
protobuf==7.36.2
|
| 24 |
+
PyYAML==6.0.3
|
| 25 |
+
tokenizers==0.23.2
|
| 26 |
+
tqdm==4.70.1
|
| 27 |
+
typing_extensions==4.16.0
|
joint_schema_model.py
ADDED
|
@@ -0,0 +1,576 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Clef: a multimodal Qwen backbone with a joint schema head for typed decisions.
|
| 2 |
+
|
| 3 |
+
A record provides a ``state`` (any JSON value), optional ``images`` and ``videos``,
|
| 4 |
+
and ``questions``. Each question has a ``type`` (``noul``, ``choice``, or ``score``),
|
| 5 |
+
``instructions``, and, for ``choice`` and ``score``, ``criteria`` describing the
|
| 6 |
+
allowed options. The model returns one logit per allowed option for every question.
|
| 7 |
+
``systemone`` answers a Jev/SystemOne ``/v1/systemone`` request body with the same
|
| 8 |
+
response body.
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import json
|
| 14 |
+
import math
|
| 15 |
+
from dataclasses import dataclass
|
| 16 |
+
from dataclasses import field as dataclass_field
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
from typing import Any
|
| 19 |
+
|
| 20 |
+
import torch
|
| 21 |
+
import torch.nn.functional as functional
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
SYSTEM_PROMPT = (
|
| 25 |
+
"Read the complete state and schema. Decide every field jointly. Each answer "
|
| 26 |
+
"must be exactly one of that field's allowed options."
|
| 27 |
+
)
|
| 28 |
+
IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
|
| 29 |
+
VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"
|
| 30 |
+
MEDIA_BATCH_KEYS = ("pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw")
|
| 31 |
+
MEDIA_TOKEN_KEYS = ("mm_token_type_ids",)
|
| 32 |
+
QUESTION_TYPES = {"noul": 0, "choice": 1, "score": 2}
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def render(value: Any) -> str:
|
| 36 |
+
if isinstance(value, str):
|
| 37 |
+
return value
|
| 38 |
+
return json.dumps(
|
| 39 |
+
value,
|
| 40 |
+
ensure_ascii=False,
|
| 41 |
+
separators=(",", ":"),
|
| 42 |
+
sort_keys=True,
|
| 43 |
+
)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def question_options(question: dict[str, Any]) -> list[tuple[str, Any]]:
|
| 47 |
+
question_type = str(question["type"])
|
| 48 |
+
if question_type == "noul":
|
| 49 |
+
criteria = {
|
| 50 |
+
"true": "The proposition is true or the answer is yes.",
|
| 51 |
+
"false": "The proposition is false or the answer is no.",
|
| 52 |
+
}
|
| 53 |
+
criteria.update(question.get("criteria") or {})
|
| 54 |
+
return [(key, criteria[key]) for key in ("true", "false")]
|
| 55 |
+
if question_type == "choice":
|
| 56 |
+
return sorted((str(key), value) for key, value in question["criteria"].items())
|
| 57 |
+
return [(str(index), value) for index, value in enumerate(question["criteria"])]
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
@dataclass(frozen=True)
|
| 61 |
+
class EncodedQuestion:
|
| 62 |
+
question_id: str
|
| 63 |
+
question_type: int
|
| 64 |
+
question_span: tuple[int, int]
|
| 65 |
+
option_spans: tuple[tuple[int, int], ...]
|
| 66 |
+
option_ids: tuple[str, ...]
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
@dataclass(frozen=True)
|
| 70 |
+
class EncodedRecord:
|
| 71 |
+
input_ids: tuple[int, ...]
|
| 72 |
+
questions: tuple[EncodedQuestion, ...]
|
| 73 |
+
record_id: str
|
| 74 |
+
media: dict[str, Any] | None = dataclass_field(default=None, compare=False, repr=False)
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def _tokens(tokenizer: Any, text: str) -> list[int]:
|
| 78 |
+
return tokenizer(text, add_special_tokens=False).input_ids
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
def _encode_media(processor: Any, record: dict[str, Any]) -> tuple[list[int], dict[str, Any] | None]:
|
| 82 |
+
images = list(record.get("images") or [])
|
| 83 |
+
videos = list(record.get("videos") or [])
|
| 84 |
+
if not images and not videos:
|
| 85 |
+
return [], None
|
| 86 |
+
if processor is None:
|
| 87 |
+
raise ValueError("records with images or videos require a processor")
|
| 88 |
+
text = IMAGE_PLACEHOLDER * len(images) + VIDEO_PLACEHOLDER * len(videos) + "\n"
|
| 89 |
+
encoded = processor(
|
| 90 |
+
text=[text],
|
| 91 |
+
images=images or None,
|
| 92 |
+
videos=videos or None,
|
| 93 |
+
return_tensors="pt",
|
| 94 |
+
**(record.get("media_kwargs") or {}),
|
| 95 |
+
)
|
| 96 |
+
media = {key: encoded[key] for key in MEDIA_BATCH_KEYS if key in encoded}
|
| 97 |
+
for key in MEDIA_TOKEN_KEYS:
|
| 98 |
+
if key in encoded:
|
| 99 |
+
media[key] = encoded[key][0].tolist()
|
| 100 |
+
return encoded["input_ids"][0].tolist(), media
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def encode_record(
|
| 104 |
+
tokenizer: Any,
|
| 105 |
+
record: dict[str, Any],
|
| 106 |
+
max_length: int = 16384,
|
| 107 |
+
max_state_tokens: int | None = None,
|
| 108 |
+
processor: Any | None = None,
|
| 109 |
+
) -> EncodedRecord:
|
| 110 |
+
schema_ids = _tokens(tokenizer, "\n\nSCHEMA FIELDS:\n")
|
| 111 |
+
questions: list[EncodedQuestion] = []
|
| 112 |
+
for question_index, (question_id, question) in enumerate(record["questions"].items()):
|
| 113 |
+
schema_ids.extend(
|
| 114 |
+
_tokens(
|
| 115 |
+
tokenizer,
|
| 116 |
+
f"\nFIELD {question_index + 1}\nID: {question_id}\nTYPE: {question['type']}\nINSTRUCTION: ",
|
| 117 |
+
)
|
| 118 |
+
)
|
| 119 |
+
question_start = len(schema_ids)
|
| 120 |
+
instructions = question.get("instructions")
|
| 121 |
+
if instructions is None or instructions == "":
|
| 122 |
+
instructions = str(question_id)
|
| 123 |
+
schema_ids.extend(_tokens(tokenizer, render(instructions)))
|
| 124 |
+
question_end = len(schema_ids)
|
| 125 |
+
schema_ids.extend(_tokens(tokenizer, "\nALLOWED OPTIONS:\n"))
|
| 126 |
+
|
| 127 |
+
option_spans: list[tuple[int, int]] = []
|
| 128 |
+
option_ids: list[str] = []
|
| 129 |
+
for option_index, (option_id, description) in enumerate(question_options(question)):
|
| 130 |
+
schema_ids.extend(_tokens(tokenizer, f"OPTION {option_index + 1}: "))
|
| 131 |
+
option_start = len(schema_ids)
|
| 132 |
+
semantics = {"option_id": option_id}
|
| 133 |
+
if description is not None:
|
| 134 |
+
semantics["description"] = description
|
| 135 |
+
schema_ids.extend(_tokens(tokenizer, render(semantics)))
|
| 136 |
+
option_spans.append((option_start, len(schema_ids)))
|
| 137 |
+
option_ids.append(option_id)
|
| 138 |
+
schema_ids.extend(_tokens(tokenizer, "\n"))
|
| 139 |
+
schema_ids.extend(_tokens(tokenizer, "END FIELD\n"))
|
| 140 |
+
questions.append(
|
| 141 |
+
EncodedQuestion(
|
| 142 |
+
question_id=str(question_id),
|
| 143 |
+
question_type=QUESTION_TYPES[str(question["type"])],
|
| 144 |
+
question_span=(question_start, question_end),
|
| 145 |
+
option_spans=tuple(option_spans),
|
| 146 |
+
option_ids=tuple(option_ids),
|
| 147 |
+
)
|
| 148 |
+
)
|
| 149 |
+
|
| 150 |
+
prefix_ids = _tokens(
|
| 151 |
+
tokenizer,
|
| 152 |
+
f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\nSTATE:\n",
|
| 153 |
+
)
|
| 154 |
+
suffix_ids = _tokens(
|
| 155 |
+
tokenizer,
|
| 156 |
+
"\n<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nJOINT SCHEMA DECISIONS:",
|
| 157 |
+
)
|
| 158 |
+
media_ids, media = _encode_media(processor, record)
|
| 159 |
+
if media is not None:
|
| 160 |
+
media["token_offset"] = len(prefix_ids)
|
| 161 |
+
prefix_ids = prefix_ids + media_ids
|
| 162 |
+
state_ids = _tokens(tokenizer, render(record["state"]))
|
| 163 |
+
if max_state_tokens is not None:
|
| 164 |
+
state_ids = state_ids[:max_state_tokens]
|
| 165 |
+
fixed_length = len(prefix_ids) + len(schema_ids) + len(suffix_ids)
|
| 166 |
+
if fixed_length > max_length:
|
| 167 |
+
raise ValueError(
|
| 168 |
+
f"schema requires {fixed_length} tokens before state; maximum is {max_length}"
|
| 169 |
+
)
|
| 170 |
+
state_ids = state_ids[: max_length - fixed_length]
|
| 171 |
+
schema_offset = len(prefix_ids) + len(state_ids)
|
| 172 |
+
shifted_questions = tuple(
|
| 173 |
+
EncodedQuestion(
|
| 174 |
+
question_id=question.question_id,
|
| 175 |
+
question_type=question.question_type,
|
| 176 |
+
question_span=(
|
| 177 |
+
question.question_span[0] + schema_offset,
|
| 178 |
+
question.question_span[1] + schema_offset,
|
| 179 |
+
),
|
| 180 |
+
option_spans=tuple(
|
| 181 |
+
(start + schema_offset, end + schema_offset)
|
| 182 |
+
for start, end in question.option_spans
|
| 183 |
+
),
|
| 184 |
+
option_ids=question.option_ids,
|
| 185 |
+
)
|
| 186 |
+
for question in questions
|
| 187 |
+
)
|
| 188 |
+
input_ids = tuple(prefix_ids + state_ids + schema_ids + suffix_ids)
|
| 189 |
+
if not input_ids or not shifted_questions:
|
| 190 |
+
raise ValueError("record produced no model input or questions")
|
| 191 |
+
return EncodedRecord(
|
| 192 |
+
input_ids=input_ids,
|
| 193 |
+
questions=shifted_questions,
|
| 194 |
+
record_id=str(record.get("id", "unknown")),
|
| 195 |
+
media=media,
|
| 196 |
+
)
|
| 197 |
+
|
| 198 |
+
|
| 199 |
+
def collate_records(
|
| 200 |
+
records: list[EncodedRecord],
|
| 201 |
+
pad_token_id: int,
|
| 202 |
+
device: torch.device,
|
| 203 |
+
) -> dict[str, Any]:
|
| 204 |
+
maximum_length = max(len(record.input_ids) for record in records)
|
| 205 |
+
input_ids = torch.full(
|
| 206 |
+
(len(records), maximum_length),
|
| 207 |
+
pad_token_id,
|
| 208 |
+
dtype=torch.long,
|
| 209 |
+
device=device,
|
| 210 |
+
)
|
| 211 |
+
attention_mask = torch.zeros(
|
| 212 |
+
(len(records), maximum_length),
|
| 213 |
+
dtype=torch.long,
|
| 214 |
+
device=device,
|
| 215 |
+
)
|
| 216 |
+
for index, record in enumerate(records):
|
| 217 |
+
length = len(record.input_ids)
|
| 218 |
+
input_ids[index, :length] = torch.tensor(record.input_ids, device=device)
|
| 219 |
+
attention_mask[index, :length] = 1
|
| 220 |
+
media: dict[str, torch.Tensor] = {}
|
| 221 |
+
for key in MEDIA_BATCH_KEYS:
|
| 222 |
+
values = [record.media[key] for record in records if record.media and key in record.media]
|
| 223 |
+
if values:
|
| 224 |
+
media[key] = torch.cat(values, dim=0).to(device)
|
| 225 |
+
for key in MEDIA_TOKEN_KEYS:
|
| 226 |
+
if any(record.media and key in record.media for record in records):
|
| 227 |
+
token_values = torch.zeros((len(records), maximum_length), dtype=torch.long, device=device)
|
| 228 |
+
for index, record in enumerate(records):
|
| 229 |
+
if record.media and key in record.media:
|
| 230 |
+
offset = record.media["token_offset"]
|
| 231 |
+
values = torch.tensor(record.media[key], dtype=torch.long, device=device)
|
| 232 |
+
token_values[index, offset : offset + len(values)] = values
|
| 233 |
+
media[key] = token_values
|
| 234 |
+
return {
|
| 235 |
+
"input_ids": input_ids,
|
| 236 |
+
"attention_mask": attention_mask,
|
| 237 |
+
"records": records,
|
| 238 |
+
"media": media,
|
| 239 |
+
}
|
| 240 |
+
|
| 241 |
+
|
| 242 |
+
class EvidenceRoutingLayer(torch.nn.Module):
|
| 243 |
+
def __init__(
|
| 244 |
+
self,
|
| 245 |
+
width: int,
|
| 246 |
+
heads: int,
|
| 247 |
+
feedforward: int,
|
| 248 |
+
dropout: float = 0.0,
|
| 249 |
+
) -> None:
|
| 250 |
+
super().__init__()
|
| 251 |
+
self.query_norm = torch.nn.LayerNorm(width)
|
| 252 |
+
self.memory_norm = torch.nn.LayerNorm(width)
|
| 253 |
+
self.attention = torch.nn.MultiheadAttention(
|
| 254 |
+
width,
|
| 255 |
+
heads,
|
| 256 |
+
dropout=dropout,
|
| 257 |
+
batch_first=True,
|
| 258 |
+
)
|
| 259 |
+
self.attention_dropout = torch.nn.Dropout(dropout)
|
| 260 |
+
self.feedforward_norm = torch.nn.LayerNorm(width)
|
| 261 |
+
self.feedforward = torch.nn.Sequential(
|
| 262 |
+
torch.nn.Linear(width, feedforward),
|
| 263 |
+
torch.nn.GELU(),
|
| 264 |
+
torch.nn.Dropout(dropout),
|
| 265 |
+
torch.nn.Linear(feedforward, width),
|
| 266 |
+
torch.nn.Dropout(dropout),
|
| 267 |
+
)
|
| 268 |
+
|
| 269 |
+
def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
|
| 270 |
+
normalized_queries = self.query_norm(queries)
|
| 271 |
+
routed, _ = self.attention(
|
| 272 |
+
normalized_queries,
|
| 273 |
+
self.memory_norm(memory),
|
| 274 |
+
self.memory_norm(memory),
|
| 275 |
+
need_weights=False,
|
| 276 |
+
)
|
| 277 |
+
queries = queries + self.attention_dropout(routed)
|
| 278 |
+
return queries + self.feedforward(self.feedforward_norm(queries))
|
| 279 |
+
|
| 280 |
+
|
| 281 |
+
class JointSchemaHead(torch.nn.Module):
|
| 282 |
+
def __init__(
|
| 283 |
+
self,
|
| 284 |
+
hidden_size: int,
|
| 285 |
+
width: int,
|
| 286 |
+
routing_layers: int,
|
| 287 |
+
layers: int,
|
| 288 |
+
heads: int,
|
| 289 |
+
feedforward: int,
|
| 290 |
+
dropout: float = 0.0,
|
| 291 |
+
) -> None:
|
| 292 |
+
super().__init__()
|
| 293 |
+
self.hidden_norm = torch.nn.LayerNorm(hidden_size)
|
| 294 |
+
self.memory_projection = torch.nn.Linear(hidden_size, width, bias=False)
|
| 295 |
+
self.question_projection = torch.nn.Linear(hidden_size, width, bias=False)
|
| 296 |
+
self.option_question_projection = torch.nn.Linear(hidden_size, width, bias=False)
|
| 297 |
+
self.global_projection = torch.nn.Linear(hidden_size, width, bias=False)
|
| 298 |
+
self.option_context_projection = torch.nn.Linear(hidden_size, width, bias=False)
|
| 299 |
+
self.option_lexical_projection = torch.nn.Linear(hidden_size, width, bias=False)
|
| 300 |
+
self.type_embedding = torch.nn.Embedding(3, width)
|
| 301 |
+
self.evidence_layers = torch.nn.ModuleList(
|
| 302 |
+
[
|
| 303 |
+
EvidenceRoutingLayer(
|
| 304 |
+
width=width,
|
| 305 |
+
heads=heads,
|
| 306 |
+
feedforward=feedforward,
|
| 307 |
+
dropout=dropout,
|
| 308 |
+
)
|
| 309 |
+
for _ in range(routing_layers)
|
| 310 |
+
]
|
| 311 |
+
)
|
| 312 |
+
self.option_summary_norm = torch.nn.LayerNorm(width)
|
| 313 |
+
self.layers = torch.nn.ModuleList(
|
| 314 |
+
[
|
| 315 |
+
torch.nn.TransformerDecoderLayer(
|
| 316 |
+
d_model=width,
|
| 317 |
+
nhead=heads,
|
| 318 |
+
dim_feedforward=feedforward,
|
| 319 |
+
dropout=dropout,
|
| 320 |
+
activation="gelu",
|
| 321 |
+
batch_first=True,
|
| 322 |
+
norm_first=True,
|
| 323 |
+
)
|
| 324 |
+
for _ in range(layers)
|
| 325 |
+
]
|
| 326 |
+
)
|
| 327 |
+
self.field_norm = torch.nn.LayerNorm(width)
|
| 328 |
+
self.option_norm = torch.nn.LayerNorm(width)
|
| 329 |
+
self.residual_scorer = torch.nn.Sequential(
|
| 330 |
+
torch.nn.Linear(width * 4, width),
|
| 331 |
+
torch.nn.GELU(),
|
| 332 |
+
torch.nn.Dropout(dropout),
|
| 333 |
+
torch.nn.Linear(width, 1),
|
| 334 |
+
)
|
| 335 |
+
self.prior_logit_scale = torch.nn.Parameter(torch.zeros(()))
|
| 336 |
+
self.joint_logit_scale = torch.nn.Parameter(torch.zeros(()))
|
| 337 |
+
self.residual_gate = torch.nn.Parameter(torch.zeros(()))
|
| 338 |
+
|
| 339 |
+
@staticmethod
|
| 340 |
+
def _mean_span(values: torch.Tensor, span: tuple[int, int]) -> torch.Tensor:
|
| 341 |
+
start, end = span
|
| 342 |
+
return values[start:end].mean(dim=0)
|
| 343 |
+
|
| 344 |
+
def forward(
|
| 345 |
+
self,
|
| 346 |
+
hidden_states: torch.Tensor,
|
| 347 |
+
input_ids: torch.Tensor,
|
| 348 |
+
attention_mask: torch.Tensor,
|
| 349 |
+
records: list[EncodedRecord],
|
| 350 |
+
output_embedding_weight: torch.Tensor,
|
| 351 |
+
) -> list[list[torch.Tensor]]:
|
| 352 |
+
results: list[list[torch.Tensor]] = []
|
| 353 |
+
normalized_hidden = self.hidden_norm(hidden_states)
|
| 354 |
+
for batch_index, record in enumerate(records):
|
| 355 |
+
sequence_length = int(attention_mask[batch_index].sum().item())
|
| 356 |
+
sequence_hidden = normalized_hidden[batch_index, :sequence_length]
|
| 357 |
+
memory = self.memory_projection(sequence_hidden).unsqueeze(0)
|
| 358 |
+
global_vector = sequence_hidden[-1]
|
| 359 |
+
question_vectors = torch.stack(
|
| 360 |
+
[
|
| 361 |
+
self._mean_span(sequence_hidden, question.question_span)
|
| 362 |
+
for question in record.questions
|
| 363 |
+
]
|
| 364 |
+
)
|
| 365 |
+
type_ids = torch.tensor(
|
| 366 |
+
[question.question_type for question in record.questions],
|
| 367 |
+
device=hidden_states.device,
|
| 368 |
+
)
|
| 369 |
+
option_contexts: list[torch.Tensor] = []
|
| 370 |
+
lexical_options: list[torch.Tensor] = []
|
| 371 |
+
option_counts = []
|
| 372 |
+
for question in record.questions:
|
| 373 |
+
context_vectors = torch.stack(
|
| 374 |
+
[
|
| 375 |
+
self._mean_span(sequence_hidden, span)
|
| 376 |
+
for span in question.option_spans
|
| 377 |
+
]
|
| 378 |
+
)
|
| 379 |
+
lexical_vectors = []
|
| 380 |
+
for start, end in question.option_spans:
|
| 381 |
+
token_ids = input_ids[batch_index, start:end]
|
| 382 |
+
lexical_vectors.append(output_embedding_weight[token_ids].mean(dim=0))
|
| 383 |
+
lexical = torch.stack(lexical_vectors)
|
| 384 |
+
option_contexts.append(context_vectors)
|
| 385 |
+
lexical_options.append(lexical)
|
| 386 |
+
option_counts.append(len(question.option_spans))
|
| 387 |
+
|
| 388 |
+
option_queries = []
|
| 389 |
+
for question_index, (context_vectors, lexical) in enumerate(
|
| 390 |
+
zip(option_contexts, lexical_options)
|
| 391 |
+
):
|
| 392 |
+
option_queries.append(
|
| 393 |
+
self.option_context_projection(context_vectors)
|
| 394 |
+
+ self.option_lexical_projection(lexical)
|
| 395 |
+
+ self.option_question_projection(
|
| 396 |
+
question_vectors[question_index]
|
| 397 |
+
).unsqueeze(0)
|
| 398 |
+
)
|
| 399 |
+
routed_options = torch.cat(option_queries, dim=0).unsqueeze(0)
|
| 400 |
+
for layer in self.evidence_layers:
|
| 401 |
+
routed_options = layer(routed_options, memory)
|
| 402 |
+
routed_options = routed_options[0]
|
| 403 |
+
split_options = list(torch.split(routed_options, option_counts, dim=0))
|
| 404 |
+
|
| 405 |
+
base_fields = self.question_projection(question_vectors)
|
| 406 |
+
option_summaries = []
|
| 407 |
+
for field, options in zip(base_fields, split_options):
|
| 408 |
+
routing_weights = torch.softmax(
|
| 409 |
+
torch.matmul(options, field) / math.sqrt(options.shape[-1]),
|
| 410 |
+
dim=0,
|
| 411 |
+
)
|
| 412 |
+
option_summaries.append(
|
| 413 |
+
torch.sum(routing_weights.unsqueeze(-1) * options, dim=0)
|
| 414 |
+
)
|
| 415 |
+
fields = (
|
| 416 |
+
base_fields
|
| 417 |
+
+ self.option_summary_norm(torch.stack(option_summaries))
|
| 418 |
+
+ self.global_projection(global_vector).unsqueeze(0)
|
| 419 |
+
+ self.type_embedding(type_ids)
|
| 420 |
+
)
|
| 421 |
+
fields = fields.unsqueeze(0)
|
| 422 |
+
for layer in self.layers:
|
| 423 |
+
fields = layer(fields, memory)
|
| 424 |
+
fields = self.field_norm(fields[0])
|
| 425 |
+
|
| 426 |
+
record_logits: list[torch.Tensor] = []
|
| 427 |
+
for field, question, lexical, routed in zip(
|
| 428 |
+
fields,
|
| 429 |
+
record.questions,
|
| 430 |
+
lexical_options,
|
| 431 |
+
split_options,
|
| 432 |
+
):
|
| 433 |
+
anchor = functional.normalize(
|
| 434 |
+
question_vectors[len(record_logits)] + global_vector,
|
| 435 |
+
dim=-1,
|
| 436 |
+
)
|
| 437 |
+
lexical_anchor = functional.normalize(lexical, dim=-1)
|
| 438 |
+
prior_scale = self.prior_logit_scale.clamp(max=math.log(100.0)).exp()
|
| 439 |
+
prior = prior_scale * torch.matmul(lexical_anchor, anchor)
|
| 440 |
+
options = self.option_norm(routed)
|
| 441 |
+
repeated_field = field.unsqueeze(0).expand_as(options)
|
| 442 |
+
cosine = functional.cosine_similarity(repeated_field, options, dim=-1)
|
| 443 |
+
features = torch.cat(
|
| 444 |
+
[
|
| 445 |
+
repeated_field,
|
| 446 |
+
options,
|
| 447 |
+
repeated_field * options,
|
| 448 |
+
torch.abs(repeated_field - options),
|
| 449 |
+
],
|
| 450 |
+
dim=-1,
|
| 451 |
+
)
|
| 452 |
+
residual = self.residual_scorer(features).squeeze(-1)
|
| 453 |
+
joint_scale = self.joint_logit_scale.clamp(max=math.log(100.0)).exp()
|
| 454 |
+
joint = joint_scale * cosine + residual
|
| 455 |
+
record_logits.append(
|
| 456 |
+
prior + torch.sigmoid(self.residual_gate) * joint
|
| 457 |
+
)
|
| 458 |
+
results.append(record_logits)
|
| 459 |
+
return results
|
| 460 |
+
|
| 461 |
+
|
| 462 |
+
class ClefModel(torch.nn.Module):
|
| 463 |
+
def __init__(self, language_model: Any, head: JointSchemaHead) -> None:
|
| 464 |
+
super().__init__()
|
| 465 |
+
self.language_model = language_model
|
| 466 |
+
self.head = head
|
| 467 |
+
|
| 468 |
+
def forward(self, batch: dict[str, Any]) -> list[list[torch.Tensor]]:
|
| 469 |
+
base_model = (
|
| 470 |
+
self.language_model.get_base_model()
|
| 471 |
+
if hasattr(self.language_model, "get_base_model")
|
| 472 |
+
else self.language_model
|
| 473 |
+
)
|
| 474 |
+
media = batch.get("media") or {}
|
| 475 |
+
text_model = base_model.model
|
| 476 |
+
if not media and hasattr(text_model, "language_model"):
|
| 477 |
+
text_model = text_model.language_model
|
| 478 |
+
outputs = text_model(
|
| 479 |
+
input_ids=batch["input_ids"],
|
| 480 |
+
attention_mask=batch["attention_mask"],
|
| 481 |
+
use_cache=False,
|
| 482 |
+
return_dict=True,
|
| 483 |
+
**media,
|
| 484 |
+
)
|
| 485 |
+
return self.head(
|
| 486 |
+
outputs.last_hidden_state,
|
| 487 |
+
batch["input_ids"],
|
| 488 |
+
batch["attention_mask"],
|
| 489 |
+
batch["records"],
|
| 490 |
+
base_model.get_output_embeddings().weight,
|
| 491 |
+
)
|
| 492 |
+
|
| 493 |
+
|
| 494 |
+
def load_release_model(
|
| 495 |
+
model_path: str | Path,
|
| 496 |
+
device: str | torch.device = "cuda",
|
| 497 |
+
dtype: torch.dtype = torch.bfloat16,
|
| 498 |
+
**from_pretrained_kwargs: Any,
|
| 499 |
+
) -> tuple[ClefModel, Any]:
|
| 500 |
+
"""Load a Clef release (merged backbone, joint schema head, and processor)."""
|
| 501 |
+
from huggingface_hub import snapshot_download
|
| 502 |
+
from safetensors.torch import load_file
|
| 503 |
+
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
|
| 504 |
+
|
| 505 |
+
path = Path(model_path)
|
| 506 |
+
if not path.is_dir():
|
| 507 |
+
path = Path(snapshot_download(str(model_path)))
|
| 508 |
+
backbone = Qwen3_5ForConditionalGeneration.from_pretrained(
|
| 509 |
+
path,
|
| 510 |
+
dtype=dtype,
|
| 511 |
+
device_map={"": str(device)},
|
| 512 |
+
**from_pretrained_kwargs,
|
| 513 |
+
)
|
| 514 |
+
backbone.config.use_cache = False
|
| 515 |
+
head_config = json.loads((path / "joint_head_config.json").read_text())
|
| 516 |
+
head = JointSchemaHead(**head_config)
|
| 517 |
+
head.load_state_dict(load_file(path / "joint_head.safetensors"), strict=True)
|
| 518 |
+
head = head.to(device=device, dtype=dtype)
|
| 519 |
+
processor = AutoProcessor.from_pretrained(path)
|
| 520 |
+
return ClefModel(backbone, head).eval(), processor
|
| 521 |
+
|
| 522 |
+
|
| 523 |
+
def systemone_answer(question: dict[str, Any], probabilities: dict[str, float]) -> dict[str, Any]:
|
| 524 |
+
"""Convert per-option probabilities for one question into a SystemOne answer."""
|
| 525 |
+
if question["type"] == "noul":
|
| 526 |
+
return {"type": "noul", "noul": round(probabilities["true"], 4)}
|
| 527 |
+
if question["type"] == "choice":
|
| 528 |
+
options = [str(option) for option in question["criteria"]]
|
| 529 |
+
choice = max(options, key=probabilities.__getitem__)
|
| 530 |
+
return {
|
| 531 |
+
"type": "choice",
|
| 532 |
+
"choice": choice,
|
| 533 |
+
"confidence": round(probabilities[choice], 4),
|
| 534 |
+
"probabilities": {option: round(probabilities[option], 4) for option in options},
|
| 535 |
+
}
|
| 536 |
+
levels = [str(index) for index in range(len(question["criteria"]))]
|
| 537 |
+
return {
|
| 538 |
+
"type": "score",
|
| 539 |
+
"score": round(sum(index * probabilities[level] for index, level in enumerate(levels)), 4),
|
| 540 |
+
"confidence": round(max(probabilities[level] for level in levels), 4),
|
| 541 |
+
"legend": dict(zip(levels, question["criteria"])),
|
| 542 |
+
"probabilities": {level: round(probabilities[level], 4) for level in levels},
|
| 543 |
+
}
|
| 544 |
+
|
| 545 |
+
|
| 546 |
+
@torch.inference_mode()
|
| 547 |
+
def systemone(model: ClefModel, processor: Any, request: dict[str, Any], max_length: int = 16384) -> dict[str, Any]:
|
| 548 |
+
"""Answer a Jev/SystemOne ``/v1/systemone`` request body with a SystemOne response body.
|
| 549 |
+
|
| 550 |
+
The request has ``model``, ``state``, and ``questions``, plus optional ``images`` and ``videos``.
|
| 551 |
+
"""
|
| 552 |
+
questions = request.get("questions")
|
| 553 |
+
if not isinstance(request.get("model"), str) or "state" not in request:
|
| 554 |
+
raise ValueError("model and state are required")
|
| 555 |
+
if not isinstance(questions, dict) or not questions:
|
| 556 |
+
raise ValueError("at least one question is required")
|
| 557 |
+
for question_id, question in questions.items():
|
| 558 |
+
if question.get("type") not in QUESTION_TYPES:
|
| 559 |
+
raise ValueError(f"{question_id}: type must be noul, choice, or score")
|
| 560 |
+
if question["type"] != "noul" and not question.get("criteria"):
|
| 561 |
+
raise ValueError(f"{question_id}: criteria must not be empty")
|
| 562 |
+
encoded = encode_record(processor.tokenizer, request, max_length=max_length, processor=processor)
|
| 563 |
+
device = next(model.parameters()).device
|
| 564 |
+
logits = model(collate_records([encoded], processor.tokenizer.pad_token_id, device))[0]
|
| 565 |
+
answers = {
|
| 566 |
+
question.question_id: systemone_answer(
|
| 567 |
+
questions[question.question_id],
|
| 568 |
+
dict(zip(question.option_ids, question_logits.float().softmax(-1).tolist())),
|
| 569 |
+
)
|
| 570 |
+
for question, question_logits in zip(encoded.questions, logits)
|
| 571 |
+
}
|
| 572 |
+
return {
|
| 573 |
+
"model": request["model"],
|
| 574 |
+
"answers": answers,
|
| 575 |
+
"usage": {"input_tokens": len(encoded.input_ids), "output_tokens": 0},
|
| 576 |
+
}
|
provenance_ft.json
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"epoch": 3,
|
| 3 |
+
"step": 19251,
|
| 4 |
+
"dev_eval": {
|
| 5 |
+
"event": "eval",
|
| 6 |
+
"split": "dev",
|
| 7 |
+
"epoch": 3,
|
| 8 |
+
"step": 19251,
|
| 9 |
+
"loss": 0.0777,
|
| 10 |
+
"tool_output_filter": {
|
| 11 |
+
"n": 21977,
|
| 12 |
+
"pos": 20282,
|
| 13 |
+
"neg": 1695,
|
| 14 |
+
"auc": 0.9899,
|
| 15 |
+
"tpr@0.5": 0.9911,
|
| 16 |
+
"fpr@0.5": 0.2183,
|
| 17 |
+
"tpr@0.9": 0.9819,
|
| 18 |
+
"fpr@0.9": 0.1209,
|
| 19 |
+
"tpr@0.99": 0.9583,
|
| 20 |
+
"fpr@0.99": 0.0507
|
| 21 |
+
},
|
| 22 |
+
"action_authorization_guard": {
|
| 23 |
+
"n": 0,
|
| 24 |
+
"pos": 0,
|
| 25 |
+
"neg": 0,
|
| 26 |
+
"auc": 0.0,
|
| 27 |
+
"tpr@0.35": 0.0,
|
| 28 |
+
"fpr@0.35": 0.0,
|
| 29 |
+
"tpr@0.5": 0.0,
|
| 30 |
+
"fpr@0.5": 0.0,
|
| 31 |
+
"tpr@0.9": 0.0,
|
| 32 |
+
"fpr@0.9": 0.0
|
| 33 |
+
},
|
| 34 |
+
"elapsed_min": 193.4,
|
| 35 |
+
"ts": 1793815560.42
|
| 36 |
+
},
|
| 37 |
+
"lr": 3e-05,
|
| 38 |
+
"pos_weight": 3.0,
|
| 39 |
+
"seed": 20261101
|
| 40 |
+
}
|
redo-final-results.json
ADDED
|
@@ -0,0 +1,104 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"dev_baseline": {
|
| 3 |
+
"event": "eval",
|
| 4 |
+
"split": "dev",
|
| 5 |
+
"epoch": 0,
|
| 6 |
+
"step": 0,
|
| 7 |
+
"loss": 0.6314,
|
| 8 |
+
"tool_output_filter": {
|
| 9 |
+
"n": 21977,
|
| 10 |
+
"pos": 20282,
|
| 11 |
+
"neg": 1695,
|
| 12 |
+
"auc": 0.7516,
|
| 13 |
+
"tpr@0.5": 0.7605,
|
| 14 |
+
"fpr@0.5": 0.3687,
|
| 15 |
+
"tpr@0.9": 0.3445,
|
| 16 |
+
"fpr@0.9": 0.1493,
|
| 17 |
+
"tpr@0.99": 0.0,
|
| 18 |
+
"fpr@0.99": 0.0
|
| 19 |
+
},
|
| 20 |
+
"action_authorization_guard": {
|
| 21 |
+
"n": 0,
|
| 22 |
+
"pos": 0,
|
| 23 |
+
"neg": 0,
|
| 24 |
+
"auc": 0.0,
|
| 25 |
+
"tpr@0.35": 0.0,
|
| 26 |
+
"fpr@0.35": 0.0,
|
| 27 |
+
"tpr@0.5": 0.0,
|
| 28 |
+
"fpr@0.5": 0.0,
|
| 29 |
+
"tpr@0.9": 0.0,
|
| 30 |
+
"fpr@0.9": 0.0
|
| 31 |
+
},
|
| 32 |
+
"elapsed_min": 8.99,
|
| 33 |
+
"ts": 1793804496.1
|
| 34 |
+
},
|
| 35 |
+
"dev_best": {
|
| 36 |
+
"event": "eval",
|
| 37 |
+
"split": "dev_best",
|
| 38 |
+
"epoch": 3,
|
| 39 |
+
"step": 19251,
|
| 40 |
+
"loss": 0.0777,
|
| 41 |
+
"tool_output_filter": {
|
| 42 |
+
"n": 21977,
|
| 43 |
+
"pos": 20282,
|
| 44 |
+
"neg": 1695,
|
| 45 |
+
"auc": 0.9899,
|
| 46 |
+
"tpr@0.5": 0.9911,
|
| 47 |
+
"fpr@0.5": 0.2183,
|
| 48 |
+
"tpr@0.9": 0.9819,
|
| 49 |
+
"fpr@0.9": 0.1209,
|
| 50 |
+
"tpr@0.99": 0.9583,
|
| 51 |
+
"fpr@0.99": 0.0507
|
| 52 |
+
},
|
| 53 |
+
"action_authorization_guard": {
|
| 54 |
+
"n": 0,
|
| 55 |
+
"pos": 0,
|
| 56 |
+
"neg": 0,
|
| 57 |
+
"auc": 0.0,
|
| 58 |
+
"tpr@0.35": 0.0,
|
| 59 |
+
"fpr@0.35": 0.0,
|
| 60 |
+
"tpr@0.5": 0.0,
|
| 61 |
+
"fpr@0.5": 0.0,
|
| 62 |
+
"tpr@0.9": 0.0,
|
| 63 |
+
"fpr@0.9": 0.0
|
| 64 |
+
}
|
| 65 |
+
},
|
| 66 |
+
"test_banking_slack": {
|
| 67 |
+
"event": "eval",
|
| 68 |
+
"split": "test_banking_slack",
|
| 69 |
+
"epoch": 3,
|
| 70 |
+
"step": 19251,
|
| 71 |
+
"loss": 0.0779,
|
| 72 |
+
"tool_output_filter": {
|
| 73 |
+
"n": 21965,
|
| 74 |
+
"pos": 20271,
|
| 75 |
+
"neg": 1694,
|
| 76 |
+
"auc": 0.9897,
|
| 77 |
+
"tpr@0.5": 0.9915,
|
| 78 |
+
"fpr@0.5": 0.2161,
|
| 79 |
+
"tpr@0.9": 0.9819,
|
| 80 |
+
"fpr@0.9": 0.1234,
|
| 81 |
+
"tpr@0.99": 0.9574,
|
| 82 |
+
"fpr@0.99": 0.0537
|
| 83 |
+
},
|
| 84 |
+
"action_authorization_guard": {
|
| 85 |
+
"n": 0,
|
| 86 |
+
"pos": 0,
|
| 87 |
+
"neg": 0,
|
| 88 |
+
"auc": 0.0,
|
| 89 |
+
"tpr@0.35": 0.0,
|
| 90 |
+
"fpr@0.35": 0.0,
|
| 91 |
+
"tpr@0.5": 0.0,
|
| 92 |
+
"fpr@0.5": 0.0,
|
| 93 |
+
"tpr@0.9": 0.0,
|
| 94 |
+
"fpr@0.9": 0.0
|
| 95 |
+
},
|
| 96 |
+
"elapsed_min": 210.49,
|
| 97 |
+
"checkpoint": {
|
| 98 |
+
"path": "/root/work/checkpoints/clef-flash-LiteRT-ft-full",
|
| 99 |
+
"epoch": 3,
|
| 100 |
+
"step": 19251
|
| 101 |
+
},
|
| 102 |
+
"ts": 1793816585.6
|
| 103 |
+
}
|
| 104 |
+
}
|
redo-provenance.json
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"host": "jev-retrain (DigitalOcean gpu-mi355x1-288gb-spot, 168.144.60.186)",
|
| 3 |
+
"wall": "2026-11-04, jev-retrain droplet (AMD MI355X), encode 11:31Z, head training finished 18:24Z",
|
| 4 |
+
"epochs": 3,
|
| 5 |
+
"lr": 3e-05,
|
| 6 |
+
"batch": 16,
|
| 7 |
+
"pos_weight": 3.0,
|
| 8 |
+
"checks": {
|
| 9 |
+
"head_tensors": 122,
|
| 10 |
+
"head_params": 121762820,
|
| 11 |
+
"source_head_dtypes": [
|
| 12 |
+
"BF16"
|
| 13 |
+
],
|
| 14 |
+
"bf16_converted_in_staging": false,
|
| 15 |
+
"source_head_sha256": "8ca78d0c127755948fb285c90caeefcf6e61a2b5f2f3cc726ffae41ca231ab11",
|
| 16 |
+
"staged_head_sha256": "8ca78d0c127755948fb285c90caeefcf6e61a2b5f2f3cc726ffae41ca231ab11",
|
| 17 |
+
"tflite": [
|
| 18 |
+
{
|
| 19 |
+
"name": "clef-flash_rowprefill_embeds_L4096_fp16fc.tflite",
|
| 20 |
+
"bytes": 13936871792
|
| 21 |
+
}
|
| 22 |
+
]
|
| 23 |
+
},
|
| 24 |
+
"repo": "macmacmacmac/CLEF-Screen",
|
| 25 |
+
"publication_note": "Uploaded 2026-11-04 on direct user instruction ahead of the Tuesday-v2 frozen-chain gate (no frozen_chain_manifest.json; independent replication pending). Metrics are script-run dev/test plus fresh-process load-check only."
|
| 26 |
+
}
|
redo-training-log.jsonl
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
tables/embed_tokens.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:48586a2af482671b5ca781a761957e296295a7fd0950af241705cb01ad2079fc
|
| 3 |
+
size 2034237544
|
tables/lm_head.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:106ec926c62c1aac2c8ab808833db89012a9132c53f767fa7609ec69e377d4ab
|
| 3 |
+
size 2034237536
|
tables/vision_pos_g256.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ecd2eba829853f233b9f71f05da81c459e0ecb5ba7fbf5f5986bf42857c5efe9
|
| 3 |
+
size 1179728
|
tables/vision_pos_g448.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:65f76c320150de7d3e53c6fa0ff7ffd048079153a8d6e114b6d48141c6f9ae26
|
| 3 |
+
size 3612752
|
tokenizer/tokenizer.json
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523
|
| 3 |
+
size 19989325
|