Image-Text-to-Text
Transformers
Safetensors
lfm3_vl
feature-extraction
liquid
lfm3
edge
decision
classification
extraction
calibration
system-one
multimodal
decision-model
custom_code
conversational
Instructions to use LiquidAI/d1-2B with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use LiquidAI/d1-2B with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("image-text-to-text", model="LiquidAI/d1-2B", trust_remote_code=True) messages = [ { "role": "user", "content": [ {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"}, {"type": "text", "text": "What animal is on the candy?"} ] }, ] pipe(text=messages)# pip install -U transformers accelerate # Load model directly from transformers import AutoModel model = AutoModel.from_pretrained("LiquidAI/d1-2B", trust_remote_code=True, device_map="auto") - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- vLLM
How to use LiquidAI/d1-2B with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "LiquidAI/d1-2B" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "LiquidAI/d1-2B", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }'Use Docker
docker model run hf.co/LiquidAI/d1-2B
- SGLang
How to use LiquidAI/d1-2B with SGLang:
Install from pip and serve model
# Install SGLang from pip: pip install sglang # Start the SGLang server: python3 -m sglang.launch_server \ --model-path "LiquidAI/d1-2B" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "LiquidAI/d1-2B", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }'Use Docker images
docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=<secret>" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server \ --model-path "LiquidAI/d1-2B" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "LiquidAI/d1-2B", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Describe this image in one sentence." }, { "type": "image_url", "image_url": { "url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" } } ] } ] }' - Docker Model Runner
How to use LiquidAI/d1-2B with Docker Model Runner:
docker model run hf.co/LiquidAI/d1-2B
Commit ·
97c5704
0
Parent(s):
d1-2B
Browse filesCo-authored-by: mlabonne <mlabonne@users.noreply.huggingface.co>
Co-authored-by: Aurelien-Lac <Aurelien-Lac@users.noreply.huggingface.co>
- .gitattributes +36 -0
- LICENSE +71 -0
- README.md +593 -0
- api.py +214 -0
- assets/icons/api.svg +1 -0
- assets/icons/blog.svg +1 -0
- assets/icons/docs.svg +1 -0
- assets/icons/try.svg +1 -0
- assets/liquid-ai-banner.webp +0 -0
- attention.py +332 -0
- box.py +150 -0
- calibration.json +380 -0
- calibration.py +77 -0
- chat_template.jinja +286 -0
- config.json +126 -0
- configuration_lfm3.py +140 -0
- configuration_lfm3_vl.py +211 -0
- depthwise.py +53 -0
- grid.py +45 -0
- image_processing_lfm3_vl.py +301 -0
- integrations_lfm3.py +204 -0
- json_spans.py +59 -0
- lfm3.py +201 -0
- media.py +252 -0
- model.safetensors +3 -0
- modeling_d1.py +73 -0
- modeling_d1_lfm3_vl.py +10 -0
- modeling_lfm3.py +881 -0
- modeling_lfm3_vl.py +968 -0
- native_lfm3.py +24 -0
- native_lfm3_vl.py +25 -0
- parsing.py +297 -0
- pointer.py +166 -0
- processing_lfm3_vl.py +254 -0
- processor_config.json +62 -0
- prompt.py +761 -0
- runner.py +379 -0
- slots.py +61 -0
- tokenizer.json +3 -0
- tokenizer_config.json +20 -0
- video_processing_lfm3_vl.py +404 -0
.gitattributes
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
+
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
+
*.xz 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.json filter=lfs diff=lfs merge=lfs -text
|
LICENSE
ADDED
|
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
LFM Open License v1.0
|
| 2 |
+
|
| 3 |
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
| 4 |
+
|
| 5 |
+
1. Definitions.
|
| 6 |
+
|
| 7 |
+
"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by this document.
|
| 8 |
+
|
| 9 |
+
"Licensor" shall mean Liquid AI, Inc.
|
| 10 |
+
|
| 11 |
+
"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
|
| 12 |
+
|
| 13 |
+
"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
|
| 14 |
+
|
| 15 |
+
"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
|
| 16 |
+
|
| 17 |
+
"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
|
| 18 |
+
|
| 19 |
+
"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work.
|
| 20 |
+
|
| 21 |
+
"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
|
| 22 |
+
|
| 23 |
+
"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
|
| 24 |
+
|
| 25 |
+
"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
|
| 26 |
+
|
| 27 |
+
"Commercial Use" shall mean any use of the Work for direct or indirect commercial advantage or monetary compensation.
|
| 28 |
+
|
| 29 |
+
"Qualified Non-Profit Organization" shall mean a Legal Entity that is organized and operated exclusively for religious, charitable, scientific, testing for public safety, literary, or educational purposes, and which is exempt from federal income tax under Section 501(c)(3) of the United States Internal Revenue Code of 1986, as amended, or any equivalent non-profit or charitable organization in a foreign jurisdiction.
|
| 30 |
+
|
| 31 |
+
"Non-Commercial or Research Purposes" shall mean purposes that do not involve any use of the Work or a Derivative Work for Commercial Use.
|
| 32 |
+
|
| 33 |
+
"Threshold" shall mean annual revenue of 10 million United States dollars ($10,000,000) or more.
|
| 34 |
+
|
| 35 |
+
2. Grant of Copyright License. Subject to the terms and conditions of this License, including the Commercial Use limitation set forth in Section 5, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
|
| 36 |
+
|
| 37 |
+
3. Grant of Patent License. Subject to the terms and conditions of this License, including the Commercial Use limitation set forth in Section 5, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
|
| 38 |
+
|
| 39 |
+
4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
|
| 40 |
+
|
| 41 |
+
(a) You must give any other recipients of the Work or Derivative Works a copy of this License; and
|
| 42 |
+
|
| 43 |
+
(b) You must cause any modified files to carry prominent notices stating that You changed the files; and
|
| 44 |
+
|
| 45 |
+
(c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and
|
| 46 |
+
|
| 47 |
+
(d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License.
|
| 48 |
+
|
| 49 |
+
You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
|
| 50 |
+
|
| 51 |
+
5. Commercial Use Limitation.
|
| 52 |
+
|
| 53 |
+
(a) The rights granted under this License for Commercial Use are conditioned upon You or Your Legal Entity not exceeding the Threshold.
|
| 54 |
+
|
| 55 |
+
(b) Any Commercial Use of the Work or a Derivative Work by a Legal Entity that exceeds the Threshold is not licensed under this Agreement.
|
| 56 |
+
|
| 57 |
+
(c) The Threshold shall not apply to a Qualified Non-Profit Organization's use of the Work or a Derivative Work for Non-Commercial or Research Purposes.
|
| 58 |
+
|
| 59 |
+
6. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
|
| 60 |
+
|
| 61 |
+
7. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except for the reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
|
| 62 |
+
|
| 63 |
+
8. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
|
| 64 |
+
|
| 65 |
+
9. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
|
| 66 |
+
|
| 67 |
+
10. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
|
| 68 |
+
|
| 69 |
+
11. Termination. This License will terminate automatically and immediately if You fail to comply with any of its terms and conditions. Upon termination, You must cease all use of the Work and any Derivative Works and delete all copies in Your possession.
|
| 70 |
+
|
| 71 |
+
END OF TERMS AND CONDITIONS
|
README.md
ADDED
|
@@ -0,0 +1,593 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
library_name: transformers
|
| 3 |
+
license: other
|
| 4 |
+
license_name: lfm1.0
|
| 5 |
+
license_link: LICENSE
|
| 6 |
+
pipeline_tag: image-text-to-text
|
| 7 |
+
tags:
|
| 8 |
+
- liquid
|
| 9 |
+
- lfm3
|
| 10 |
+
- edge
|
| 11 |
+
- decision
|
| 12 |
+
- classification
|
| 13 |
+
- extraction
|
| 14 |
+
- calibration
|
| 15 |
+
- system-one
|
| 16 |
+
- multimodal
|
| 17 |
+
- decision-model
|
| 18 |
+
- custom_code
|
| 19 |
+
language:
|
| 20 |
+
- en
|
| 21 |
+
- ar
|
| 22 |
+
- zh
|
| 23 |
+
- fr
|
| 24 |
+
- de
|
| 25 |
+
- it
|
| 26 |
+
- ja
|
| 27 |
+
- ko
|
| 28 |
+
- pt
|
| 29 |
+
- es
|
| 30 |
+
- vi
|
| 31 |
+
- th
|
| 32 |
+
- id
|
| 33 |
+
- hi
|
| 34 |
+
- ru
|
| 35 |
+
- pl
|
| 36 |
+
---
|
| 37 |
+
|
| 38 |
+
<div align="center">
|
| 39 |
+
<img
|
| 40 |
+
src="https://huggingface.co/LiquidAI/d1-2B/resolve/main/assets/liquid-ai-banner.webp"
|
| 41 |
+
alt="Liquid AI"
|
| 42 |
+
style="width: 100%; max-width: 100%; height: auto; display: block; margin: 0;"
|
| 43 |
+
/>
|
| 44 |
+
<div align="center" style="display: flex; justify-content: center; align-items: center; gap: 1.25em; margin: 8px 0 1.7rem; padding: 0; line-height: 1.3;">
|
| 45 |
+
<a href="https://console.liquid.ai/dashboard" style="display: inline-flex; align-items: center; gap: 6px;"><img src="https://huggingface.co/LiquidAI/d1-2B/resolve/main/assets/icons/api.svg" width="16" height="16" alt="" aria-hidden="true" style="display: block; flex: 0 0 16px; width: 16px; height: 16px; margin: 0;" /><strong>API</strong></a>
|
| 46 |
+
<a href="https://playground.liquid.ai/" style="display: inline-flex; align-items: center; gap: 6px;"><img src="https://huggingface.co/LiquidAI/d1-2B/resolve/main/assets/icons/try.svg" width="16" height="16" alt="" aria-hidden="true" style="display: block; flex: 0 0 16px; width: 16px; height: 16px; margin: 0;" /><strong>Try</strong></a>
|
| 47 |
+
<a href="https://docs.liquid.ai/lfm/getting-started/welcome" style="display: inline-flex; align-items: center; gap: 6px;"><img src="https://huggingface.co/LiquidAI/d1-2B/resolve/main/assets/icons/docs.svg" width="16" height="16" alt="" aria-hidden="true" style="display: block; flex: 0 0 16px; width: 16px; height: 16px; margin: 0;" /><strong>Docs</strong></a>
|
| 48 |
+
<a href="https://www.liquid.ai/blog/open-d1" style="display: inline-flex; align-items: center; gap: 6px;"><img src="https://huggingface.co/LiquidAI/d1-2B/resolve/main/assets/icons/blog.svg" width="16" height="16" alt="" aria-hidden="true" style="display: block; flex: 0 0 16px; width: 16px; height: 16px; margin: 0;" /><strong>Blog</strong></a>
|
| 49 |
+
</div>
|
| 50 |
+
</div>
|
| 51 |
+
|
| 52 |
+
# d1-2B
|
| 53 |
+
|
| 54 |
+
d1-2B is a **decision model for text, pictures, and video**, built on [LFM3-VL-2B](https://huggingface.co/LiquidAI/LFM3-VL-2B). It returns calibrated answers in one forward pass, without generating tokens.
|
| 55 |
+
|
| 56 |
+
- **Best-in-class:** it leads the open d1 models on average and on our board of over a hundred decision benchmarks, ahead of d1-3B at under two thirds of its size.
|
| 57 |
+
- **Beyond classification:** it also extracts text, numbers, and dates with their span, locates objects in pictures, and estimates quantities with an interval.
|
| 58 |
+
- **Millisecond decisions:** 5.7 ms per decision on an NVIDIA H100 and 3.8 ms on an AMD MI325X, with any number of questions about one state in one call.
|
| 59 |
+
|
| 60 |
+
Find more information about open d1 in [our release blog post](https://www.liquid.ai/blog/open-d1).
|
| 61 |
+
|
| 62 |
+
## Model details
|
| 63 |
+
<hr style="margin: -0.7em 0 1.2em 0; border: 0; border-top: 1px solid rgba(127, 127, 127, 0.3);" />
|
| 64 |
+
|
| 65 |
+
The [LFM3 family](https://huggingface.co/collections/LiquidAI/lfm3) includes base, text, vision, audio, omni, and decision models. d1-2B decides over text, pictures, and video, and [d1-Omni-2B](https://huggingface.co/LiquidAI/d1-Omni-2B) adds audio.
|
| 66 |
+
|
| 67 |
+
<table style="width: 100%; display: table; border: 0; border-radius: 0; border-collapse: collapse;">
|
| 68 |
+
<thead style="border: 0;">
|
| 69 |
+
<tr style="border: 0;">
|
| 70 |
+
<th scope="col" style="text-align: left; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">Property</th>
|
| 71 |
+
<th scope="col" style="text-align: left; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">Specification</th>
|
| 72 |
+
</tr>
|
| 73 |
+
</thead>
|
| 74 |
+
<tbody style="border: 0;">
|
| 75 |
+
<tr style="border: 0;">
|
| 76 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><strong>Checkpoint</strong></td>
|
| 77 |
+
<td style="border: 0; padding: 0.55em 0.75em;">Post-trained decision model, built on <a href="https://huggingface.co/LiquidAI/LFM3-VL-2B">LFM3-VL-2B</a></td>
|
| 78 |
+
</tr>
|
| 79 |
+
<tr style="border: 0;">
|
| 80 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><strong>Parameters</strong></td>
|
| 81 |
+
<td style="border: 0; padding: 0.55em 0.75em;">1.92B</td>
|
| 82 |
+
</tr>
|
| 83 |
+
<tr style="border: 0;">
|
| 84 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><strong>Architecture</strong></td>
|
| 85 |
+
<td style="border: 0; padding: 0.55em 0.75em;">LFM3 backbone, 24 layers: 13 gated short convolutions, 5 sliding-window attention, 6 global attention</td>
|
| 86 |
+
</tr>
|
| 87 |
+
<tr style="border: 0;">
|
| 88 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><strong>Context length</strong></td>
|
| 89 |
+
<td style="border: 0; padding: 0.55em 0.75em;">128K tokens</td>
|
| 90 |
+
</tr>
|
| 91 |
+
<tr style="border: 0;">
|
| 92 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><strong>Vocabulary</strong></td>
|
| 93 |
+
<td style="border: 0; padding: 0.55em 0.75em;">128,000 tokens</td>
|
| 94 |
+
</tr>
|
| 95 |
+
<tr style="border: 0;">
|
| 96 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><strong>Input</strong></td>
|
| 97 |
+
<td style="border: 0; padding: 0.55em 0.75em;">A state (text or JSON), pictures or a video, and any number of named questions</td>
|
| 98 |
+
</tr>
|
| 99 |
+
<tr style="border: 0;">
|
| 100 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><strong>Output</strong></td>
|
| 101 |
+
<td style="border: 0; padding: 0.55em 0.75em;">Typed answers with their probabilities, with no generated tokens</td>
|
| 102 |
+
</tr>
|
| 103 |
+
<tr style="border: 0;">
|
| 104 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><strong>Languages</strong></td>
|
| 105 |
+
<td style="border: 0; padding: 0.55em 0.75em;">English, Arabic, Chinese, French, German, Italian, Japanese, Korean, Portuguese, Spanish, Vietnamese, Thai, Indonesian, Hindi, Russian, Polish</td>
|
| 106 |
+
</tr>
|
| 107 |
+
<tr style="border: 0;">
|
| 108 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><strong>License</strong></td>
|
| 109 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><a href="LICENSE">LFM Open License v1.0</a></td>
|
| 110 |
+
</tr>
|
| 111 |
+
</tbody>
|
| 112 |
+
</table>
|
| 113 |
+
|
| 114 |
+
d1-2B is designed for any pipeline step that needs a decision: routing and triage, moderation and guardrails, intent and topic classification, extraction from documents and forms, judging and reranking, and visual inspection. It is not a chat model.
|
| 115 |
+
|
| 116 |
+
Knowledge-heavy questions favor larger models, and answers get less reliable past about 64K tokens of state. Estimates are reliable on counts and scores, less so on free quantities.
|
| 117 |
+
|
| 118 |
+
## Inference
|
| 119 |
+
<hr style="margin: -0.7em 0 1.2em 0; border: 0; border-top: 1px solid rgba(127, 127, 127, 0.3);" />
|
| 120 |
+
|
| 121 |
+
d1-2B runs in [Transformers](https://docs.liquid.ai/deployment/gpu-inference/transformers) with its own model code. It reads its answers from a single forward pass, so there are no sampling parameters to set.
|
| 122 |
+
|
| 123 |
+
**Transformers quick start:**
|
| 124 |
+
|
| 125 |
+
```bash
|
| 126 |
+
pip install "transformers>=5.19" torch torchvision pillow
|
| 127 |
+
```
|
| 128 |
+
|
| 129 |
+
```python
|
| 130 |
+
import torch
|
| 131 |
+
from transformers import AutoModel
|
| 132 |
+
|
| 133 |
+
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
| 134 |
+
dtype = torch.float32 if device == "cpu" else torch.bfloat16
|
| 135 |
+
model = AutoModel.from_pretrained("LiquidAI/d1-2B", trust_remote_code=True, dtype=dtype).to(device)
|
| 136 |
+
|
| 137 |
+
# Several named questions over one state, answered in one pass
|
| 138 |
+
questions = {
|
| 139 |
+
"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}, # yes/no: P(yes)
|
| 140 |
+
"team": {"type": "choice", "instructions": "Which team should handle this?",
|
| 141 |
+
"criteria": {"billing": "Charges, refunds, invoices", "shipping": "Delivery and parcels",
|
| 142 |
+
"technical": "App or site faults"}},
|
| 143 |
+
"topics": {"type": "choice", "multiple": True, "instructions": "Which topics does the message raise?",
|
| 144 |
+
"criteria": {"billing": "Payments and charges", "delivery": "Shipping and parcels",
|
| 145 |
+
"login": "Signing in to the account"}},
|
| 146 |
+
"urgency": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["low", "medium", "high"]},
|
| 147 |
+
}
|
| 148 |
+
out = model.system_one("I was charged twice and my parcel never arrived. Please refund one charge.", questions)
|
| 149 |
+
a = out["answers"]
|
| 150 |
+
urgency = a["urgency"]["legend"][str(int(a["urgency"]["score"]))]
|
| 151 |
+
print(round(a["refund"]["noul"], 2), a["team"]["choice"], a["topics"]["choices"], urgency)
|
| 152 |
+
# 1.0 billing ['billing', 'delivery'] medium
|
| 153 |
+
```
|
| 154 |
+
|
| 155 |
+
<details>
|
| 156 |
+
<summary><strong>Values from a document</strong></summary>
|
| 157 |
+
|
| 158 |
+
```python
|
| 159 |
+
invoice = "ACME SA\nInvoice 2026-117\nTotal due: 1 234,50 EUR\nPay before 03/04/2026 (day/month/year)"
|
| 160 |
+
out = model.system_one(invoice, {
|
| 161 |
+
"payee": {"type": "text", "instructions": "Who issued the invoice?"},
|
| 162 |
+
"total": {"type": "number", "instructions": "How much is due?"},
|
| 163 |
+
"due": {"type": "date", "instructions": "By when must it be paid?"},
|
| 164 |
+
"vat": {"type": "number", "instructions": "How much VAT is charged?"},
|
| 165 |
+
})
|
| 166 |
+
a = out["answers"]
|
| 167 |
+
print(a["payee"]["text"], a["total"]["number"], a["due"]["date"], a["vat"]["number"])
|
| 168 |
+
# ACME SA 1234.5 2026-04-03 None
|
| 169 |
+
```
|
| 170 |
+
|
| 171 |
+
</details>
|
| 172 |
+
|
| 173 |
+
<details>
|
| 174 |
+
<summary><strong>Pictures and boxes</strong></summary>
|
| 175 |
+
|
| 176 |
+
Count, locate one object, or list every instance (`max` reads at most that many):
|
| 177 |
+
|
| 178 |
+
```python
|
| 179 |
+
from transformers.image_utils import load_image
|
| 180 |
+
|
| 181 |
+
image = load_image("http://images.cocodataset.org/val2017/000000039769.jpg") # two cats on a sofa
|
| 182 |
+
out = model.system_one(None, {
|
| 183 |
+
"cats": {"type": "number", "instructions": "How many cats are there?", "range": [0, 10], "step": 1},
|
| 184 |
+
"remote": {"type": "box", "instructions": "Where is the remote control on the left?"},
|
| 185 |
+
"every_cat": {"type": "boxes", "instructions": "Every cat.", "max": 4},
|
| 186 |
+
}, media=[image])
|
| 187 |
+
a = out["answers"]
|
| 188 |
+
print(a["cats"]["number"], a["every_cat"]["count"], {k: round(v, 2) for k, v in a["remote"]["box"].items()})
|
| 189 |
+
# 2.0 2 {'top': 0.15, 'bottom': 0.24, 'left': 0.06, 'right': 0.28}
|
| 190 |
+
```
|
| 191 |
+
|
| 192 |
+
</details>
|
| 193 |
+
|
| 194 |
+
<details>
|
| 195 |
+
<summary><strong>Many requests in one call</strong></summary>
|
| 196 |
+
|
| 197 |
+
```python
|
| 198 |
+
tickets = ["Where is my parcel? It was due Monday.", "The app crashes when I open settings."]
|
| 199 |
+
for out in model.system_one_batch([(t, {"team": questions["team"]}) for t in tickets]):
|
| 200 |
+
print(out["answers"]["team"]["choice"])
|
| 201 |
+
# shipping
|
| 202 |
+
# technical
|
| 203 |
+
```
|
| 204 |
+
|
| 205 |
+
</details>
|
| 206 |
+
|
| 207 |
+
<details>
|
| 208 |
+
<summary><strong>Calls</strong></summary>
|
| 209 |
+
|
| 210 |
+
<table style="width: 100%; display: table; border: 0; border-radius: 0; border-collapse: collapse;">
|
| 211 |
+
<thead style="border: 0;">
|
| 212 |
+
<tr style="border: 0;">
|
| 213 |
+
<th scope="col" style="text-align: left; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">Call</th>
|
| 214 |
+
<th scope="col" style="text-align: left; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">What it does</th>
|
| 215 |
+
</tr>
|
| 216 |
+
</thead>
|
| 217 |
+
<tbody style="border: 0;">
|
| 218 |
+
<tr style="border: 0;">
|
| 219 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>system_one(state, questions, media=None)</code></td>
|
| 220 |
+
<td style="border: 0; padding: 0.55em 0.75em;">Named questions over one state, in one pass. <code>media</code> lists pictures (PIL images or <code>{"type": "image", "path": ...}</code>) and videos as their frames (<code>{"type": "video", "video": [PIL images], "fps": 24}</code>).</td>
|
| 221 |
+
</tr>
|
| 222 |
+
<tr style="border: 0;">
|
| 223 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>system_one_batch([(state, questions[, media]), ...])</code></td>
|
| 224 |
+
<td style="border: 0; padding: 0.55em 0.75em;">Many requests in one call, batched.</td>
|
| 225 |
+
</tr>
|
| 226 |
+
<tr style="border: 0;">
|
| 227 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>set_size=</code> (either call)</td>
|
| 228 |
+
<td style="border: 0; padding: 0.55em 0.75em;">Questions read together after the state. The default is the model's trained size (16). Larger is faster and costs some accuracy on text.</td>
|
| 229 |
+
</tr>
|
| 230 |
+
</tbody>
|
| 231 |
+
</table>
|
| 232 |
+
|
| 233 |
+
A state is a string, any JSON value, or `None` when the media are the whole state. Each call returns `{"answers": {name: answer}, "usage": {"input_tokens": n, "output_tokens": 0}}`.
|
| 234 |
+
|
| 235 |
+
</details>
|
| 236 |
+
|
| 237 |
+
<details>
|
| 238 |
+
<summary><strong>Question types</strong></summary>
|
| 239 |
+
|
| 240 |
+
<table style="width: 100%; display: table; border: 0; border-radius: 0; border-collapse: collapse;">
|
| 241 |
+
<thead style="border: 0;">
|
| 242 |
+
<tr style="border: 0;">
|
| 243 |
+
<th scope="col" style="text-align: left; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);"><code>type</code></th>
|
| 244 |
+
<th scope="col" style="text-align: left; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">Question</th>
|
| 245 |
+
<th scope="col" style="text-align: left; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">Answer</th>
|
| 246 |
+
</tr>
|
| 247 |
+
</thead>
|
| 248 |
+
<tbody style="border: 0;">
|
| 249 |
+
<tr style="border: 0;">
|
| 250 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>noul</code></td>
|
| 251 |
+
<td style="border: 0; padding: 0.55em 0.75em;">Yes or no</td>
|
| 252 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>noul</code>: P(yes)</td>
|
| 253 |
+
</tr>
|
| 254 |
+
<tr style="border: 0;">
|
| 255 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>choice</code></td>
|
| 256 |
+
<td style="border: 0; padding: 0.55em 0.75em;">One of the options <code>criteria</code> names (<code>{name: description}</code>)</td>
|
| 257 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>choice</code>, <code>confidence</code>, <code>probabilities</code></td>
|
| 258 |
+
</tr>
|
| 259 |
+
<tr style="border: 0;">
|
| 260 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>choice</code> with <code>multiple: true</code></td>
|
| 261 |
+
<td style="border: 0; padding: 0.55em 0.75em;">Any of the options <code>criteria</code> names</td>
|
| 262 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>choices</code>, <code>probabilities</code> (each option's own P(yes))</td>
|
| 263 |
+
</tr>
|
| 264 |
+
<tr style="border: 0;">
|
| 265 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>score</code></td>
|
| 266 |
+
<td style="border: 0; padding: 0.55em 0.75em;">One of 2 to 10 ordered levels, with <code>criteria</code> listing them lowest first</td>
|
| 267 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>score</code> (the median level), <code>confidence</code>, <code>probabilities</code>, <code>legend</code></td>
|
| 268 |
+
</tr>
|
| 269 |
+
<tr style="border: 0;">
|
| 270 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>text</code>, <code>number</code>, <code>date</code></td>
|
| 271 |
+
<td style="border: 0; padding: 0.55em 0.75em;">A value to read from the state</td>
|
| 272 |
+
<td style="border: 0; padding: 0.55em 0.75em;">The field of its type (a number with its <code>unit</code>, a date in ISO 8601), <code>span</code>, <code>found</code></td>
|
| 273 |
+
</tr>
|
| 274 |
+
<tr style="border: 0;">
|
| 275 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>number</code> with <code>range</code></td>
|
| 276 |
+
<td style="border: 0; padding: 0.55em 0.75em;">An estimate, <code>step</code> apart</td>
|
| 277 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>number</code> (the median), <code>interval</code> (p10, p90)</td>
|
| 278 |
+
</tr>
|
| 279 |
+
<tr style="border: 0;">
|
| 280 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>box</code>, <code>boxes</code></td>
|
| 281 |
+
<td style="border: 0; padding: 0.55em 0.75em;">One object, or every instance (16 at most, or <code>max</code>), in the one picture</td>
|
| 282 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><code>box</code> (<code>top</code>, <code>bottom</code>, <code>left</code>, <code>right</code> as fractions of the picture), <code>found</code>, <code>confidence</code>, or <code>boxes</code> and <code>count</code></td>
|
| 283 |
+
</tr>
|
| 284 |
+
</tbody>
|
| 285 |
+
</table>
|
| 286 |
+
|
| 287 |
+
A value or a box is missing when `found` is under 0.5, and it is then `None`. Calibrate any decision threshold on data from your own use.
|
| 288 |
+
|
| 289 |
+
</details>
|
| 290 |
+
|
| 291 |
+
## Performance
|
| 292 |
+
<hr style="margin: -0.7em 0 1.2em 0; border: 0; border-top: 1px solid rgba(127, 127, 127, 0.3);" />
|
| 293 |
+
|
| 294 |
+
### Benchmarks
|
| 295 |
+
|
| 296 |
+
d1-2B leads the open d1 models on our decision board, on average, and on most single benchmarks, ahead of d1-3B with under two thirds of its parameters.
|
| 297 |
+
|
| 298 |
+
<div role="region" tabindex="0" style="width: 100%; max-width: 100%; overflow-x: auto; -webkit-overflow-scrolling: touch;">
|
| 299 |
+
<table style="width: 100%; display: table; border: 0; border-radius: 0; border-collapse: collapse;">
|
| 300 |
+
<thead style="border: 0;">
|
| 301 |
+
<tr style="border: 0;">
|
| 302 |
+
<th scope="col" style="text-align: left; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">Benchmark</th>
|
| 303 |
+
<th scope="col" style="text-align: center; background-color: rgba(108, 79, 224, 0.08); padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);"><strong>d1<br>2B</strong></th>
|
| 304 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">d1-omni<br>600M</th>
|
| 305 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">d1-Omni<br>2B</th>
|
| 306 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">d1<br>3B</th>
|
| 307 |
+
</tr>
|
| 308 |
+
</thead>
|
| 309 |
+
<tbody style="border: 0;">
|
| 310 |
+
<tr style="border: 0;"><th scope="rowgroup" style="text-align: left; background-color: transparent; padding: 0.55em 0.75em; border: 0; font-weight: 600; padding-top: 1em; padding-bottom: 0.5em;">Decisions</th><td aria-hidden="true" style="border: 0; padding: 0.55em 0.75em; background-color: rgba(108, 79, 224, 0.08);"></td><td aria-hidden="true" style="border: 0; padding: 0.55em 0.75em;"></td><td aria-hidden="true" style="border: 0; padding: 0.55em 0.75em;"></td><td aria-hidden="true" style="border: 0; padding: 0.55em 0.75em;"></td></tr>
|
| 311 |
+
<tr style="border: 0;">
|
| 312 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><a href="https://huggingface.co/spaces/multimodalart/jev-decision-index">Decision Index</a> public suite (v0.3)</td>
|
| 313 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>51.1</strong></td>
|
| 314 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">18.3</td>
|
| 315 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>49.0</u></td>
|
| 316 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>49.0</u></td>
|
| 317 |
+
</tr>
|
| 318 |
+
<tr style="border: 0;">
|
| 319 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><a href="https://huggingface.co/datasets/Hanno-Labs/decision-bench">DecisionBench (Hanno Labs)</a> v1</td>
|
| 320 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>77.1</strong></td>
|
| 321 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">58.7</td>
|
| 322 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>76.4</u></td>
|
| 323 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">73.1</td>
|
| 324 |
+
</tr>
|
| 325 |
+
<tr style="border: 0;">
|
| 326 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><a href="https://decisionbench.ai">Decision Bench (decisionbench.ai)</a> v4</td>
|
| 327 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>85.2</strong></td>
|
| 328 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">57.9</td>
|
| 329 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>84.7</u></td>
|
| 330 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">83.6</td>
|
| 331 |
+
</tr>
|
| 332 |
+
<tr style="border: 0;">
|
| 333 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><a href="https://huggingface.co/datasets/fastino/fast-decisions">Fast Decisions</a></td>
|
| 334 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><u>58.8</u></td>
|
| 335 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">–</td>
|
| 336 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">56.8</td>
|
| 337 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><strong>59.3</strong></td>
|
| 338 |
+
</tr>
|
| 339 |
+
<tr style="border: 0;">
|
| 340 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><a href="https://github.com/fstandhartinger/jevbench">JevBench</a> v1.4.1</td>
|
| 341 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>77.9</strong></td>
|
| 342 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">58.0</td>
|
| 343 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">75.3</td>
|
| 344 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>76.2</u></td>
|
| 345 |
+
</tr>
|
| 346 |
+
<tr style="border: 0;">
|
| 347 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><a href="https://huggingface.co/datasets/mohit67890/imajev-bench">ImajevBench</a> v2.0-lite</td>
|
| 348 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>75.5</strong></td>
|
| 349 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">21.7</td>
|
| 350 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>72.3</u></td>
|
| 351 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">64.0</td>
|
| 352 |
+
</tr>
|
| 353 |
+
<tr style="border: 0;">
|
| 354 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><a href="https://evals.typesafe.ai">WorkflowEvals</a></td>
|
| 355 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>81.6</strong></td>
|
| 356 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">48.8</td>
|
| 357 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>77.1</u></td>
|
| 358 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">74.4</td>
|
| 359 |
+
</tr>
|
| 360 |
+
<tr style="border: 0;">
|
| 361 |
+
<td style="border: 0; padding: 0.55em 0.75em;"><a href="https://huggingface.co/datasets/LocalLLaMA/typed-decisions">Typed Decisions</a></td>
|
| 362 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;">61.3</td>
|
| 363 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">57.5</td>
|
| 364 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>63.1</u></td>
|
| 365 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><strong>65.0</strong></td>
|
| 366 |
+
</tr>
|
| 367 |
+
<tr style="border: 0;">
|
| 368 |
+
<td style="border: 0; padding: 0.55em 0.75em;">Our board (over a hundred public and private benchmarks)</td>
|
| 369 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>63.2</strong></td>
|
| 370 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">–</td>
|
| 371 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>61.8</u></td>
|
| 372 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">58.8</td>
|
| 373 |
+
</tr>
|
| 374 |
+
</tbody>
|
| 375 |
+
<tbody style="border: 0;">
|
| 376 |
+
<tr style="border: 0;"><th scope="rowgroup" style="text-align: left; background-color: transparent; padding: 0.55em 0.75em; border: 0; font-weight: 600; padding-top: 1em; padding-bottom: 0.5em; border-top: 1px solid rgba(127, 127, 127, 0.25);">Vision</th><td aria-hidden="true" style="border: 0; padding: 0.55em 0.75em; background-color: rgba(108, 79, 224, 0.08); border-top: 1px solid rgba(127, 127, 127, 0.25);"></td><td aria-hidden="true" style="border: 0; padding: 0.55em 0.75em; border-top: 1px solid rgba(127, 127, 127, 0.25);"></td><td aria-hidden="true" style="border: 0; padding: 0.55em 0.75em; border-top: 1px solid rgba(127, 127, 127, 0.25);"></td><td aria-hidden="true" style="border: 0; padding: 0.55em 0.75em; border-top: 1px solid rgba(127, 127, 127, 0.25);"></td></tr>
|
| 377 |
+
<tr style="border: 0;">
|
| 378 |
+
<td style="border: 0; padding: 0.55em 0.75em;">AI2D</td>
|
| 379 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><u>76.7</u></td>
|
| 380 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">37.1</td>
|
| 381 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">71.5</td>
|
| 382 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><strong>79.9</strong></td>
|
| 383 |
+
</tr>
|
| 384 |
+
<tr style="border: 0;">
|
| 385 |
+
<td style="border: 0; padding: 0.55em 0.75em;">BLINK</td>
|
| 386 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>66.2</strong></td>
|
| 387 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">37.7</td>
|
| 388 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>63.5</u></td>
|
| 389 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">59.2</td>
|
| 390 |
+
</tr>
|
| 391 |
+
<tr style="border: 0;">
|
| 392 |
+
<td style="border: 0; padding: 0.55em 0.75em;">CountQA</td>
|
| 393 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>51.4</strong></td>
|
| 394 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">26.8</td>
|
| 395 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>45.4</u></td>
|
| 396 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">42.6</td>
|
| 397 |
+
</tr>
|
| 398 |
+
<tr style="border: 0;">
|
| 399 |
+
<td style="border: 0; padding: 0.55em 0.75em;">CV-Bench</td>
|
| 400 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>86.1</strong></td>
|
| 401 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">40.8</td>
|
| 402 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>84.0</u></td>
|
| 403 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">82.1</td>
|
| 404 |
+
</tr>
|
| 405 |
+
<tr style="border: 0;">
|
| 406 |
+
<td style="border: 0; padding: 0.55em 0.75em;">ERQA</td>
|
| 407 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>45.8</strong></td>
|
| 408 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">27.3</td>
|
| 409 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>42.5</u></td>
|
| 410 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">39.8</td>
|
| 411 |
+
</tr>
|
| 412 |
+
<tr style="border: 0;">
|
| 413 |
+
<td style="border: 0; padding: 0.55em 0.75em;">HallusionBench</td>
|
| 414 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>72.1</strong></td>
|
| 415 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">52.6</td>
|
| 416 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>67.5</u></td>
|
| 417 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">65.3</td>
|
| 418 |
+
</tr>
|
| 419 |
+
<tr style="border: 0;">
|
| 420 |
+
<td style="border: 0; padding: 0.55em 0.75em;">MMBench</td>
|
| 421 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>87.9</strong></td>
|
| 422 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">54.2</td>
|
| 423 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>85.2</u></td>
|
| 424 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">84.9</td>
|
| 425 |
+
</tr>
|
| 426 |
+
<tr style="border: 0;">
|
| 427 |
+
<td style="border: 0; padding: 0.55em 0.75em;">MME</td>
|
| 428 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>85.4</strong></td>
|
| 429 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">61.0</td>
|
| 430 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>83.9</u></td>
|
| 431 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">82.1</td>
|
| 432 |
+
</tr>
|
| 433 |
+
<tr style="border: 0;">
|
| 434 |
+
<td style="border: 0; padding: 0.55em 0.75em;">MMMU-Pro</td>
|
| 435 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>26.4</strong></td>
|
| 436 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">11.9</td>
|
| 437 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>24.1</u></td>
|
| 438 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">18.5</td>
|
| 439 |
+
</tr>
|
| 440 |
+
<tr style="border: 0;">
|
| 441 |
+
<td style="border: 0; padding: 0.55em 0.75em;">MMRB2</td>
|
| 442 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>54.2</strong></td>
|
| 443 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">48.3</td>
|
| 444 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>52.6</u></td>
|
| 445 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">49.5</td>
|
| 446 |
+
</tr>
|
| 447 |
+
<tr style="border: 0;">
|
| 448 |
+
<td style="border: 0; padding: 0.55em 0.75em;">MMStar</td>
|
| 449 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>67.5</strong></td>
|
| 450 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">31.9</td>
|
| 451 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>61.3</u></td>
|
| 452 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">59.9</td>
|
| 453 |
+
</tr>
|
| 454 |
+
<tr style="border: 0;">
|
| 455 |
+
<td style="border: 0; padding: 0.55em 0.75em;">MMVP</td>
|
| 456 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>82.7</strong></td>
|
| 457 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">54.0</td>
|
| 458 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>81.3</u></td>
|
| 459 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">77.0</td>
|
| 460 |
+
</tr>
|
| 461 |
+
<tr style="border: 0;">
|
| 462 |
+
<td style="border: 0; padding: 0.55em 0.75em;">NaturalBench</td>
|
| 463 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>81.5</strong></td>
|
| 464 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">58.8</td>
|
| 465 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>80.0</u></td>
|
| 466 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">79.8</td>
|
| 467 |
+
</tr>
|
| 468 |
+
<tr style="border: 0;">
|
| 469 |
+
<td style="border: 0; padding: 0.55em 0.75em;">POPE</td>
|
| 470 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><u>88.7</u></td>
|
| 471 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">82.7</td>
|
| 472 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><strong>89.9</strong></td>
|
| 473 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">88.6</td>
|
| 474 |
+
</tr>
|
| 475 |
+
<tr style="border: 0;">
|
| 476 |
+
<td style="border: 0; padding: 0.55em 0.75em;">V*</td>
|
| 477 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>76.4</strong></td>
|
| 478 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">35.6</td>
|
| 479 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><strong>76.4</strong></td>
|
| 480 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>73.8</u></td>
|
| 481 |
+
</tr>
|
| 482 |
+
<tr style="border: 0;">
|
| 483 |
+
<td style="border: 0; padding: 0.55em 0.75em;">VisualPuzzles</td>
|
| 484 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>29.9</strong></td>
|
| 485 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">26.2</td>
|
| 486 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">26.9</td>
|
| 487 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>29.5</u></td>
|
| 488 |
+
</tr>
|
| 489 |
+
<tr style="border: 0;">
|
| 490 |
+
<td style="border: 0; padding: 0.55em 0.75em;">VisualWebBench</td>
|
| 491 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>96.2</strong></td>
|
| 492 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">14.9</td>
|
| 493 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>88.6</u></td>
|
| 494 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">71.4</td>
|
| 495 |
+
</tr>
|
| 496 |
+
<tr style="border: 0;">
|
| 497 |
+
<td style="border: 0; padding: 0.55em 0.75em;">VL-RewardBench</td>
|
| 498 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em;"><strong>69.3</strong></td>
|
| 499 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">48.1</td>
|
| 500 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;"><u>68.8</u></td>
|
| 501 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">65.0</td>
|
| 502 |
+
</tr>
|
| 503 |
+
</tbody>
|
| 504 |
+
<tbody style="border: 0;">
|
| 505 |
+
<tr style="border: 0;"><td style="border: 0; padding: 0.55em 0.75em; border-top: 1px solid rgba(127, 127, 127, 0.25);"><strong>Average</strong></td><td style="text-align: center; font-variant-numeric: tabular-nums; background-color: rgba(108, 79, 224, 0.08); border: 0; padding: 0.55em 0.75em; border-top: 1px solid rgba(127, 127, 127, 0.25);"><strong>70.2</strong></td><td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em; border-top: 1px solid rgba(127, 127, 127, 0.25);">42.8</td><td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em; border-top: 1px solid rgba(127, 127, 127, 0.25);"><u>67.7</u></td><td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em; border-top: 1px solid rgba(127, 127, 127, 0.25);">65.4</td></tr>
|
| 506 |
+
</tbody>
|
| 507 |
+
</table>
|
| 508 |
+
</div>
|
| 509 |
+
|
| 510 |
+
<details>
|
| 511 |
+
<summary><small>Evaluation details</small></summary>
|
| 512 |
+
|
| 513 |
+
<p><small>All four models are scored on the same items with the same scoring, so a few of d1-3B's and d1-omni-600M's numbers differ slightly from their own cards. d1-omni-600M has no Fast Decisions or board score. Average is the unweighted mean of the public benchmarks that all four models have, so it leaves out Fast Decisions and our board. Image benchmarks are read as decisions over each benchmark's options, on at most 1,000 rows each. Three of them are saturated: POPE, MME, and VisualWebBench.</small></p>
|
| 514 |
+
|
| 515 |
+
</details>
|
| 516 |
+
|
| 517 |
+
### Speed
|
| 518 |
+
|
| 519 |
+
Warm calls after `model.compile()` in bf16, median of 20 calls, and best of three runs on the H100 and the M5 Pro. Compiling takes about a minute on the first calls.
|
| 520 |
+
|
| 521 |
+
<div role="region" tabindex="0" style="width: 100%; max-width: 100%; overflow-x: auto; -webkit-overflow-scrolling: touch;">
|
| 522 |
+
<table style="width: 100%; display: table; border: 0; border-radius: 0; border-collapse: collapse; min-width: 40rem; font-size: calc(0.875em - 1pt);">
|
| 523 |
+
<thead style="border: 0;">
|
| 524 |
+
<tr style="border: 0;">
|
| 525 |
+
<th scope="col" style="text-align: left; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">Hardware</th>
|
| 526 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">One question</th>
|
| 527 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">3 questions</th>
|
| 528 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">5K-token state</th>
|
| 529 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">384 px picture</th>
|
| 530 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">Box</th>
|
| 531 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">Boxes</th>
|
| 532 |
+
<th scope="col" style="text-align: center; padding: 0.55em 0.75em; border: 0; border-bottom: 1px solid rgba(127, 127, 127, 0.55);">64 states in one call</th>
|
| 533 |
+
</tr>
|
| 534 |
+
</thead>
|
| 535 |
+
<tbody style="border: 0;">
|
| 536 |
+
<tr style="border: 0;">
|
| 537 |
+
<td style="border: 0; padding: 0.55em 0.75em;">NVIDIA H100 PCIe</td>
|
| 538 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">5.7 ms</td>
|
| 539 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">7.6 ms</td>
|
| 540 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">52 ms</td>
|
| 541 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">12.3 ms</td>
|
| 542 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">16.1 ms</td>
|
| 543 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">24.4 ms</td>
|
| 544 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">1,132 / s</td>
|
| 545 |
+
</tr>
|
| 546 |
+
<tr style="border: 0;">
|
| 547 |
+
<td style="border: 0; padding: 0.55em 0.75em;">AMD MI325X</td>
|
| 548 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">3.8 ms</td>
|
| 549 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">5.6 ms</td>
|
| 550 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">44 ms</td>
|
| 551 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">8.7 ms</td>
|
| 552 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">11.2 ms</td>
|
| 553 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">16.5 ms</td>
|
| 554 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">1,699 / s</td>
|
| 555 |
+
</tr>
|
| 556 |
+
<tr style="border: 0;">
|
| 557 |
+
<td style="border: 0; padding: 0.55em 0.75em;">Apple M5 Pro (<code>mps</code>)</td>
|
| 558 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">18.6 ms</td>
|
| 559 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">26.5 ms</td>
|
| 560 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">753 ms</td>
|
| 561 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">45.2 ms</td>
|
| 562 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">53.5 ms</td>
|
| 563 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">100 ms</td>
|
| 564 |
+
<td style="text-align: center; font-variant-numeric: tabular-nums; border: 0; padding: 0.55em 0.75em;">109 / s</td>
|
| 565 |
+
</tr>
|
| 566 |
+
</tbody>
|
| 567 |
+
</table>
|
| 568 |
+
</div>
|
| 569 |
+
|
| 570 |
+
## Citation
|
| 571 |
+
<hr style="margin: -0.7em 0 1.2em 0; border: 0; border-top: 1px solid rgba(127, 127, 127, 0.3);" />
|
| 572 |
+
|
| 573 |
+
```bibtex
|
| 574 |
+
@article{liquidAI2026opend1,
|
| 575 |
+
author = {Liquid AI},
|
| 576 |
+
title = {Open d1: Edge decision models for text, vision, and audio},
|
| 577 |
+
journal = {Liquid AI Blog},
|
| 578 |
+
year = {2026},
|
| 579 |
+
note = {https://www.liquid.ai/blog/open-d1},
|
| 580 |
+
}
|
| 581 |
+
```
|
| 582 |
+
|
| 583 |
+
```bibtex
|
| 584 |
+
@article{liquidAI2026LFM3,
|
| 585 |
+
author = {Liquid AI},
|
| 586 |
+
title = {Introducing LFM3: Every Modality for Specialized Intelligence},
|
| 587 |
+
journal = {Liquid AI Blog},
|
| 588 |
+
year = {2026},
|
| 589 |
+
note = {https://www.liquid.ai/blog/lfm3},
|
| 590 |
+
}
|
| 591 |
+
```
|
| 592 |
+
|
| 593 |
+
For questions and feedback, join [Discord](https://discord.com/invite/liquid-ai). For custom deployment and fine-tuning support, [contact Liquid AI](https://www.liquid.ai/contact).
|
api.py
ADDED
|
@@ -0,0 +1,214 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The System One API: named questions over a state in, typed answers out.
|
| 2 |
+
|
| 3 |
+
model.system_one(state, {"refund": {"type": "noul", "instructions": "Is this a refund request?"}})
|
| 4 |
+
# {"answers": {"refund": {"type": "noul", "noul": p}}, "usage": {"input_tokens": n, "output_tokens": 0}}
|
| 5 |
+
|
| 6 |
+
A state is a string, any JSON value, or None with media alone; `media` holds its pictures, audio clips and videos
|
| 7 |
+
(`media`). A question is a dict in the Decision Index's schema (`type`, `instructions`, `criteria`; spoken
|
| 8 |
+
instructions a media dict) or one of `prompt`'s classes.
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
from collections.abc import Mapping, Sequence
|
| 14 |
+
from typing import Any
|
| 15 |
+
|
| 16 |
+
from .media import Media, media_list
|
| 17 |
+
from .prompt import (
|
| 18 |
+
DEFAULT_STATE_STYLE,
|
| 19 |
+
Box,
|
| 20 |
+
Boxes,
|
| 21 |
+
Choice,
|
| 22 |
+
Estimate,
|
| 23 |
+
Noul,
|
| 24 |
+
Question,
|
| 25 |
+
Read,
|
| 26 |
+
as_question,
|
| 27 |
+
check_question,
|
| 28 |
+
state_text,
|
| 29 |
+
typed,
|
| 30 |
+
)
|
| 31 |
+
|
| 32 |
+
# How a score question's distribution becomes its answer. `median`, the default: the first level whose cumulative
|
| 33 |
+
# probability reaches one half, optimal for absolute error (ordinal MAE). `mean`: the expected level, which a caller
|
| 34 |
+
# rounds (the Decision Index's rule), optimal only for squared error. `argmax`: the likeliest level. `median` and
|
| 35 |
+
# `argmax` answer a level, which rounding leaves as it is.
|
| 36 |
+
SCORE_READOUTS = ("mean", "median", "argmax")
|
| 37 |
+
SCORE_READOUT = "median"
|
| 38 |
+
# An explicit "I don't know" option: a choice option whose text (its description, else its name) says the answerer
|
| 39 |
+
# does not know. Not an option saying the evidence is insufficient or that a request cannot be served: those are
|
| 40 |
+
# answers. `abstain_below` (off by default) picks it when the best other option's probability, renormalised over the
|
| 41 |
+
# other options, is below the threshold.
|
| 42 |
+
IDK_TEXTS = frozenset({"i don't know", "i do not know", "i don't know the answer", "i do not know the answer"})
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def abstain_option(criteria: Mapping[str, str | None] | None) -> str | None:
|
| 46 |
+
"""The name of a choice's explicit "I don't know" option (the English `IDK_TEXTS`), or None."""
|
| 47 |
+
if not isinstance(criteria, Mapping):
|
| 48 |
+
return None
|
| 49 |
+
for name, text in criteria.items():
|
| 50 |
+
if str(text or name).strip().lower().replace("\u2019", "'").rstrip(".!") in IDK_TEXTS:
|
| 51 |
+
return name
|
| 52 |
+
return None
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def choice_index(probs: Sequence[float], abstain: int | None = None, abstain_below: float | None = None) -> int:
|
| 56 |
+
"""A choice's answer: the likeliest option; with `abstain` (the "I don't know" option's index) and
|
| 57 |
+
`abstain_below`, that option when the best other option holds less than `abstain_below` of the other options'
|
| 58 |
+
mass."""
|
| 59 |
+
best = max(range(len(probs)), key=probs.__getitem__)
|
| 60 |
+
if abstain is None or abstain_below is None:
|
| 61 |
+
return best
|
| 62 |
+
others = [p for i, p in enumerate(probs) if i != abstain]
|
| 63 |
+
total = sum(others)
|
| 64 |
+
return abstain if total <= 0 or max(others) < abstain_below * total else best
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def score_level(probs: Sequence[float], readout: str = SCORE_READOUT) -> int:
|
| 68 |
+
"""A score question's answer as a level index under `readout` (`mean`: rounded, clamped to the levels)."""
|
| 69 |
+
total = sum(probs) or 1.0
|
| 70 |
+
if readout == "mean":
|
| 71 |
+
return max(0, min(len(probs) - 1, int(round(sum(i * p / total for i, p in enumerate(probs))))))
|
| 72 |
+
if readout == "median":
|
| 73 |
+
mass = 0.0
|
| 74 |
+
for i, p in enumerate(probs):
|
| 75 |
+
mass += p / total
|
| 76 |
+
if mass >= 0.5:
|
| 77 |
+
return i
|
| 78 |
+
return len(probs) - 1
|
| 79 |
+
if readout == "argmax":
|
| 80 |
+
return max(range(len(probs)), key=probs.__getitem__)
|
| 81 |
+
raise ValueError(f"score_readout {readout!r}: one of {SCORE_READOUTS}")
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def answer(q: Question, probs: Sequence[float] | Sequence[Sequence[float]] | Mapping, score_readout: str = SCORE_READOUT,
|
| 85 |
+
abstain_below: float | None = None, context: str = "", calibration=None, edges: str = "median") -> dict:
|
| 86 |
+
"""A noul's P(yes); a choice's pick (`choice_index`) and its probabilities; a score's level under
|
| 87 |
+
`score_readout` (`mean`: the expected level, unrounded); a typed answer's probabilities through `calibration`'s
|
| 88 |
+
tables (a read is found, a box drawn, an option chosen, at 0.5 after them; an estimate's quantiles read at the
|
| 89 |
+
levels its table gives). `edges`: how a box's edges are read (`box.EDGE_READS`, the checkpoint's `d1_box_edges`)."""
|
| 90 |
+
through = calibration.through if calibration is not None else lambda name, value: value
|
| 91 |
+
if isinstance(q, Box):
|
| 92 |
+
from .box import answer as box_answer
|
| 93 |
+
|
| 94 |
+
return box_answer(probs, through, edges)
|
| 95 |
+
if isinstance(q, Boxes):
|
| 96 |
+
from .box import answer_many
|
| 97 |
+
|
| 98 |
+
return answer_many(probs, through, edges)
|
| 99 |
+
if isinstance(q, Estimate):
|
| 100 |
+
value = q.grid.read(probs, None if calibration is None else lambda share: calibration.level("estimate", share))
|
| 101 |
+
return {"type": "number", "number": value["median"], "interval": [value["p10"], value["p90"]], "unit": q.unit}
|
| 102 |
+
if isinstance(q, Read):
|
| 103 |
+
from .parsing import parse_date, parse_number, written_unit
|
| 104 |
+
|
| 105 |
+
found = through(f"found:{q.type}", probs["found"])
|
| 106 |
+
present = found >= 0.5
|
| 107 |
+
text = probs["text"] if present else None
|
| 108 |
+
span = probs["span"] if present else None
|
| 109 |
+
value, unit = text, q.unit
|
| 110 |
+
if text is not None and q.type == "number":
|
| 111 |
+
unit = unit or written_unit(text)
|
| 112 |
+
value = parse_number(text, context, unit=q.unit)
|
| 113 |
+
elif text is not None and q.type == "date":
|
| 114 |
+
value, _ = parse_date(text, context)
|
| 115 |
+
result = {"type": q.type, q.type: value, "span": span, "found": found}
|
| 116 |
+
if q.type != "text":
|
| 117 |
+
result["text"] = text
|
| 118 |
+
if present and "path" in probs:
|
| 119 |
+
result["path"] = probs["path"]
|
| 120 |
+
if q.type == "number":
|
| 121 |
+
result["unit"] = unit
|
| 122 |
+
return result
|
| 123 |
+
if isinstance(q, Noul):
|
| 124 |
+
return {"type": "noul", "noul": probs[0]}
|
| 125 |
+
if isinstance(q, Choice) and q.multiple:
|
| 126 |
+
probs = [through("multiple", p) for p in probs]
|
| 127 |
+
return {"type": "choice", "choices": [name for name, p in zip(q.criteria, probs) if p >= 0.5],
|
| 128 |
+
"probabilities": dict(zip(q.criteria, probs))}
|
| 129 |
+
best = max(range(len(probs)), key=probs.__getitem__)
|
| 130 |
+
if isinstance(q, Choice):
|
| 131 |
+
names = list(q.criteria)
|
| 132 |
+
idk = abstain_option(q.criteria) if abstain_below is not None else None
|
| 133 |
+
pick = choice_index(probs, None if idk is None else names.index(idk), abstain_below)
|
| 134 |
+
return {"type": "choice", "choice": names[pick], "confidence": probs[pick],
|
| 135 |
+
"probabilities": dict(zip(names, probs))}
|
| 136 |
+
level = (sum(i * p for i, p in enumerate(probs)) if score_readout == "mean"
|
| 137 |
+
else float(score_level(probs, score_readout)))
|
| 138 |
+
return {"type": "score", "score": level, "confidence": probs[best],
|
| 139 |
+
"probabilities": {str(i): p for i, p in enumerate(probs)},
|
| 140 |
+
"legend": {str(i): text for i, text in enumerate(q.criteria)}}
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
def request_media(request: tuple) -> list[Media]:
|
| 144 |
+
"""A `(state, questions)` or `(state, questions, media)` request's media."""
|
| 145 |
+
return media_list(request[2] if len(request) > 2 else None)
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
class SystemOneApi:
|
| 149 |
+
"""`system_one` and `system_one_batch`, and the plain probabilities under them, over a model's
|
| 150 |
+
`run(requests)`: each `(state, questions, media)` request's probabilities (in its questions' order)
|
| 151 |
+
and the tokens it read."""
|
| 152 |
+
|
| 153 |
+
serves_typed = True # an engine without the slots and state positions of slots.py says False
|
| 154 |
+
|
| 155 |
+
def ask(self, q: Mapping | Question) -> Question:
|
| 156 |
+
"""A caller's question, refused with what to change, or when this engine cannot serve its type."""
|
| 157 |
+
check_question(q, getattr(self, "tokenizer", None))
|
| 158 |
+
q = as_question(q)
|
| 159 |
+
if typed(q) and not self.serves_typed:
|
| 160 |
+
raise ValueError(f"{type(self).__name__} serves noul, choice and score questions only")
|
| 161 |
+
return q
|
| 162 |
+
|
| 163 |
+
def named(self, questions: Mapping[str, Any]) -> dict[str, Question]:
|
| 164 |
+
"""A request's named questions (`ask`), an error naming the question it is about."""
|
| 165 |
+
if not isinstance(questions, Mapping):
|
| 166 |
+
raise ValueError(f'questions are a dict of name to question, {{"refund": {{"type": "noul", "instructions": '
|
| 167 |
+
f'"Is it a refund request?"}}}}, not {type(questions).__name__}')
|
| 168 |
+
out = {}
|
| 169 |
+
for name, q in questions.items():
|
| 170 |
+
try:
|
| 171 |
+
out[name] = self.ask(q)
|
| 172 |
+
except ValueError as e:
|
| 173 |
+
raise ValueError(f"question {name!r}: {e}") from None
|
| 174 |
+
return out
|
| 175 |
+
|
| 176 |
+
def answer(self, state: Any, q: Question, probs, score_readout: str = SCORE_READOUT,
|
| 177 |
+
abstain_below: float | None = None) -> dict:
|
| 178 |
+
"""One question's answer as `system_one` gives it: a read's value parsed in its state, a typed answer's
|
| 179 |
+
probabilities through the served calibration's tables."""
|
| 180 |
+
context = state_text(state, getattr(self, "state_style", DEFAULT_STATE_STYLE)) if isinstance(q, Read) else ""
|
| 181 |
+
return answer(q, probs, score_readout, abstain_below, context, getattr(self, "calibration", None),
|
| 182 |
+
getattr(getattr(self, "tokenizer", None), "d1_box_edges", "median"))
|
| 183 |
+
|
| 184 |
+
def run(self, requests: Sequence[tuple[Any, list[Question], Sequence[Media]]],
|
| 185 |
+
set_size: int | None = None) -> list[tuple[list[list[float]], int]]:
|
| 186 |
+
raise NotImplementedError
|
| 187 |
+
|
| 188 |
+
def probabilities(self, state: Any, questions: Sequence, media: Sequence | None = None) -> list[list[float]]:
|
| 189 |
+
"""Each question's option probabilities; a multiple choice returns independent P(yes) per option."""
|
| 190 |
+
return self.run([(state, [self.ask(q) for q in questions], media_list(media))])[0][0]
|
| 191 |
+
|
| 192 |
+
def system_one(self, state: Any, questions: Mapping[str, Any], media: Sequence | None = None,
|
| 193 |
+
score_readout: str = SCORE_READOUT, abstain_below: float | None = None,
|
| 194 |
+
set_size: int | None = None) -> dict:
|
| 195 |
+
"""Named questions over one state, and its media if any (pictures, audio clips, videos), in one forward
|
| 196 |
+
pass."""
|
| 197 |
+
return self.system_one_batch([(state, questions, media)], score_readout, abstain_below, set_size)[0]
|
| 198 |
+
|
| 199 |
+
def system_one_batch(self, requests: Sequence[tuple], score_readout: str = SCORE_READOUT,
|
| 200 |
+
abstain_below: float | None = None, set_size: int | None = None) -> list[dict]:
|
| 201 |
+
"""Many `(state, questions)` or `(state, questions, media)` requests in one call, batched as the engine
|
| 202 |
+
reads them (`run`); a request with several questions reads its state once a pass.
|
| 203 |
+
`score_readout` (`SCORE_READOUTS`) sets how a score question answers; `abstain_below` (off by default)
|
| 204 |
+
when a choice with an "I don't know" option picks it (`choice_index`); `set_size` how many questions are
|
| 205 |
+
read together after the state, for this call (the model's trained size by default)."""
|
| 206 |
+
if score_readout not in SCORE_READOUTS:
|
| 207 |
+
raise ValueError(f"score_readout {score_readout!r}: one of {SCORE_READOUTS}")
|
| 208 |
+
if abstain_below is not None and not 0 < abstain_below <= 1:
|
| 209 |
+
raise ValueError(f"abstain_below {abstain_below!r}: a share in (0, 1]")
|
| 210 |
+
named = [(r[0], self.named(r[1]), request_media(r)) for r in requests]
|
| 211 |
+
done = self.run([(state, list(qs.values()), media) for state, qs, media in named], set_size=set_size)
|
| 212 |
+
return [{"answers": {n: self.answer(state, q, p, score_readout, abstain_below) for (n, q), p in zip(qs.items(), probs)},
|
| 213 |
+
"usage": {"input_tokens": read, "output_tokens": 0}}
|
| 214 |
+
for (state, qs, _), (probs, read) in zip(named, done)]
|
assets/icons/api.svg
ADDED
|
|
assets/icons/blog.svg
ADDED
|
|
assets/icons/docs.svg
ADDED
|
|
assets/icons/try.svg
ADDED
|
|
assets/liquid-ai-banner.webp
ADDED
|
attention.py
ADDED
|
@@ -0,0 +1,332 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Which tokens see which on a backbone's full-attention layers, in training and serving alike.
|
| 2 |
+
|
| 3 |
+
| mode | a state token sees | a question token sees |
|
| 4 |
+
|---|---|---|
|
| 5 |
+
| `causal` | the tokens before it | the tokens before it |
|
| 6 |
+
| `state` | the whole state | the whole state, and the question before it |
|
| 7 |
+
| `prompt` | the whole prompt: state, question, options | the whole prompt |
|
| 8 |
+
| `set` | the state and every question | the state and its own question |
|
| 9 |
+
|
| 10 |
+
`state` reads the state both ways and never shows it the question, so one state's keys still serve every question
|
| 11 |
+
about it. `prompt` reads everything before the answer both ways, the state with its question in view: each question is
|
| 12 |
+
then a prompt of its own, with no prefill shared between the questions of one state. `set` keeps `prompt`'s reading
|
| 13 |
+
and packs a state's questions after it, `[state][q1]..[qN]`, each ending at its own answer slot: one pass, the state
|
| 14 |
+
read once. On these layers a question never attends to another; it hears of them only through the state, which reads
|
| 15 |
+
them all. A lone question reads exactly as under `prompt`.
|
| 16 |
+
|
| 17 |
+
Only full-attention layers change; recurrent, convolution and sliding-window layers stay causal. Under `set` they
|
| 18 |
+
run over the packed sequence, so a question also hears the ones before it through them; training packs the same way
|
| 19 |
+
(questions of a state in random order), so the served function is the trained one. A
|
| 20 |
+
checkpoint trains and serves with one mode, recorded in its config (`d1_attention`, with `d1_set_size`, the most
|
| 21 |
+
questions `set` packs after one state).
|
| 22 |
+
"""
|
| 23 |
+
|
| 24 |
+
from __future__ import annotations
|
| 25 |
+
|
| 26 |
+
from contextlib import contextmanager
|
| 27 |
+
from dataclasses import dataclass, field
|
| 28 |
+
|
| 29 |
+
import torch
|
| 30 |
+
import torch.nn.functional as F
|
| 31 |
+
from transformers import AttentionInterface, AttentionMaskInterface
|
| 32 |
+
from transformers.integrations.sdpa_attention import sdpa_attention_forward
|
| 33 |
+
from transformers.masking_utils import sdpa_mask
|
| 34 |
+
|
| 35 |
+
Tensor = torch.Tensor
|
| 36 |
+
SET_SIZE = 4 # the most questions `set` packs after one state when a checkpoint records none (`d1_set_size`)
|
| 37 |
+
TRAINED_SET_SIZE = 128 # the largest cap training takes (`--set-size`), so the most a checkpoint records
|
| 38 |
+
CHUNK = 4096 # queries a `Sees` mask is built for at once
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def _fused(q: torch.Tensor) -> bool:
|
| 42 |
+
"""Whether the flash kernels take these queries: a GPU in half precision, ROCm or CUDA from Ampere on."""
|
| 43 |
+
if not q.is_cuda or q.dtype not in (torch.float16, torch.bfloat16):
|
| 44 |
+
return False
|
| 45 |
+
return bool(torch.version.hip) or torch.cuda.get_device_capability(q.device)[0] >= 8
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def _windowed(q: Tensor, k: Tensor, v: Tensor, window: int, scale: float) -> Tensor:
|
| 49 |
+
"""Causal attention within `window`, as local blocks: each block of `window`
|
| 50 |
+
queries sees its own block and the one before it."""
|
| 51 |
+
batch, heads, length, dim = q.shape
|
| 52 |
+
blocks = -(-length // window)
|
| 53 |
+
|
| 54 |
+
def split(t: Tensor) -> Tensor: # (B, H, T, D) -> (B, blocks, H, window, D)
|
| 55 |
+
t = F.pad(t, (0, 0, 0, blocks * window - length))
|
| 56 |
+
return t.reshape(batch, t.shape[1], blocks, window, dim).transpose(1, 2)
|
| 57 |
+
|
| 58 |
+
def with_previous(t: Tensor) -> Tensor:
|
| 59 |
+
return torch.cat((F.pad(t, (0, 0, 0, 0, 0, 0, 1, 0))[:, :-1], t), dim=3)
|
| 60 |
+
|
| 61 |
+
qb, kb, vb = split(q), with_previous(split(k)), with_previous(split(v))
|
| 62 |
+
qpos = torch.arange(window, device=q.device)[:, None] + window
|
| 63 |
+
kpos = torch.arange(2 * window, device=q.device)[None, :]
|
| 64 |
+
block = torch.arange(blocks, device=q.device)[:, None, None]
|
| 65 |
+
allowed = (kpos <= qpos) & (kpos > qpos - window) & ((block > 0) | (kpos >= window)) # block 0 has none before it
|
| 66 |
+
mask = allowed[:, None].repeat(batch, 1, 1, 1)
|
| 67 |
+
out = F.scaled_dot_product_attention(qb.flatten(0, 1), kb.flatten(0, 1), vb.flatten(0, 1),
|
| 68 |
+
attn_mask=mask, scale=scale, enable_gqa=True)
|
| 69 |
+
out = out.reshape(batch, blocks, heads, window, dim).transpose(1, 2)
|
| 70 |
+
return out.reshape(batch, heads, blocks * window, dim)[:, :, :length]
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def causal_attention(q: Tensor, k: Tensor, v: Tensor, window: int | None, scale: float) -> Tensor:
|
| 74 |
+
"""Right-padded chains: causal, within `window` if set: on the flash kernels their own window at any length (no
|
| 75 |
+
branch on it, so torch.compile keeps one graph across the window), else blocks."""
|
| 76 |
+
if window is None or not _fused(q) and q.shape[2] <= window:
|
| 77 |
+
return F.scaled_dot_product_attention(q, k, v, is_causal=True, scale=scale, enable_gqa=True)
|
| 78 |
+
if not _fused(q):
|
| 79 |
+
return _windowed(q, k, v, window, scale)
|
| 80 |
+
out = torch.ops.aten._flash_attention_forward(*(t.transpose(1, 2) for t in (q, k, v)), None, None, q.shape[2],
|
| 81 |
+
k.shape[2], 0.0, True, False, scale=scale,
|
| 82 |
+
window_size_left=window - 1, window_size_right=0)[0]
|
| 83 |
+
return out.transpose(1, 2)
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def both_ways(mode: str, tokenizer, ids: list[int], suffix: str) -> int | None:
|
| 87 |
+
"""How many of a prompt's first tokens read each other both ways: none (`causal`), its state, or all of it
|
| 88 |
+
(`prompt`, and `set` for a lone question)."""
|
| 89 |
+
if mode == "causal":
|
| 90 |
+
return None
|
| 91 |
+
return len(ids) if mode in ("prompt", "set") else state_length(tokenizer, ids, suffix)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def state_length(tokenizer, ids: list[int], suffix: str) -> int:
|
| 95 |
+
"""A prompt's tokens before its question: the prompt is the state's text then `suffix`, which opens on a token
|
| 96 |
+
boundary (the state ends on a blank line), so its own tokens are the prompt's last ones."""
|
| 97 |
+
return len(ids) - len(tokenizer.encode(suffix, add_special_tokens=False))
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
def pack_set(tokenizer, prefix: str, suffixes: list[str],
|
| 101 |
+
header: str = "", slot_offsets: list[list[int]] | None = None,
|
| 102 |
+
first_ids: list[int] | None = None,
|
| 103 |
+
) -> tuple[list[int], int, list[int], list[list[int]]]:
|
| 104 |
+
"""`set`'s sequence for one state, in any prompt layout: `prefix` (the state, ending with the question
|
| 105 |
+
`header`) and the first question tokenized as its lone prompt is, so one question is exactly its `prompt`
|
| 106 |
+
row; each further question its own group, `header` then its suffix. Returns the ids, the state's length, each
|
| 107 |
+
token's question (1..N, -1 on the state) and each question's answer positions. `slot_offsets` names slots
|
| 108 |
+
relative to its suffix's last token (0); without it, each question has that one slot."""
|
| 109 |
+
if not suffixes:
|
| 110 |
+
raise ValueError("a question set cannot be empty")
|
| 111 |
+
offsets = [[0] for _ in suffixes] if slot_offsets is None else slot_offsets
|
| 112 |
+
if len(offsets) != len(suffixes):
|
| 113 |
+
raise ValueError("one slot layout is required per question")
|
| 114 |
+
|
| 115 |
+
def positions(k: int, length: int, end: int) -> list[int]:
|
| 116 |
+
own = offsets[k]
|
| 117 |
+
if (not own or own != sorted(set(own)) or own[-1] != 0
|
| 118 |
+
or any(isinstance(o, bool) or not isinstance(o, int) or not -length < o <= 0 for o in own)):
|
| 119 |
+
raise ValueError("slots must be increasing positions inside their question, ending at its last token")
|
| 120 |
+
return [end + o for o in own]
|
| 121 |
+
|
| 122 |
+
ids = tokenizer.encode(prefix + suffixes[0], add_special_tokens=False) if first_ids is None else list(first_ids)
|
| 123 |
+
state = state_length(tokenizer, ids, suffixes[0])
|
| 124 |
+
question = [-1] * state + [1] * (len(ids) - state)
|
| 125 |
+
ends = [positions(0, len(ids) - state, len(ids) - 1)]
|
| 126 |
+
for k, suffix in enumerate(suffixes[1:], 2):
|
| 127 |
+
own = tokenizer.encode(header + suffix, add_special_tokens=False)
|
| 128 |
+
ids, question = ids + own, question + [k] * len(own)
|
| 129 |
+
ends.append(positions(k - 1, len(own), len(ids) - 1))
|
| 130 |
+
return ids, state, question, ends
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
def text_backbone(model):
|
| 134 |
+
"""The module that returns hidden states, under any wrapper (DDP, a training wrapper)."""
|
| 135 |
+
m = getattr(model, "module", model)
|
| 136 |
+
if hasattr(m, "get_base_model"):
|
| 137 |
+
m = m.get_base_model()
|
| 138 |
+
return getattr(m, "model", m)
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def language_model(model):
|
| 142 |
+
"""The language model under any wrapper and a vision-language model."""
|
| 143 |
+
m = text_backbone(model)
|
| 144 |
+
return getattr(m, "language_model", m)
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def _runs(*parts: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, int]:
|
| 148 |
+
"""Segments, segment i the runs `(start[i], length[i])` of every part back to back: their flat indices, their
|
| 149 |
+
cumulative lengths (int32, as the varlen kernels take them) and the longest."""
|
| 150 |
+
lengths = sum(n for _, n in parts)
|
| 151 |
+
cu = F.pad(lengths.cumsum(0), (1, 0))
|
| 152 |
+
seg = torch.repeat_interleave(torch.arange(len(lengths), device=lengths.device), lengths)
|
| 153 |
+
off = torch.arange(len(seg), device=seg.device) - cu[seg]
|
| 154 |
+
idx = torch.zeros_like(off)
|
| 155 |
+
for start, n in parts:
|
| 156 |
+
idx = torch.where((off >= 0) & (off < n[seg]), start[seg] + off, idx)
|
| 157 |
+
off = off - n[seg]
|
| 158 |
+
return idx, cu.int(), int(lengths.max()) if len(lengths) else 0
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
def _varlen(q: Tensor, k: Tensor, v: Tensor, cu_q: Tensor, cu_k: Tensor, longest: Tensor, causal: bool,
|
| 162 |
+
scale: float) -> tuple[Tensor, Tensor, Tensor, Tensor]:
|
| 163 |
+
"""Flash attention over packed segments: the output [N, H, D], then what its backward reads (the log-sum-exp
|
| 164 |
+
[H, N], the RNG state and its companion). Causal aligns each segment's last query to its last key. `longest`: the
|
| 165 |
+
longest query and key segment, on the CPU. Compiled, the opaque op `d1::varlen`, so the lengths neither break the
|
| 166 |
+
graph nor recompile it; its backward is flash's own (`_varlen_backward`), as eager's is."""
|
| 167 |
+
max_q, max_k = longest.tolist()
|
| 168 |
+
out, lse, rng, unused, _ = torch.ops.aten._flash_attention_forward(q, k, v, cu_q, cu_k, max_q, max_k, 0.0, causal,
|
| 169 |
+
False, scale=scale)
|
| 170 |
+
return out, lse, rng, unused
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
def _varlen_backward(ctx, grad, *_):
|
| 174 |
+
q, k, v, out, lse, rng, unused, cu_q, cu_k = ctx.saved_tensors
|
| 175 |
+
max_q, max_k = ctx.longest
|
| 176 |
+
dq, dk, dv = torch.ops.aten._flash_attention_backward(grad.contiguous(), q, k, v, out, lse, cu_q, cu_k, max_q,
|
| 177 |
+
max_k, 0.0, ctx.causal, rng, unused, scale=ctx.scale)
|
| 178 |
+
return dq, dk, dv, None, None, None, None, None
|
| 179 |
+
|
| 180 |
+
|
| 181 |
+
def _varlen_context(ctx, inputs, output):
|
| 182 |
+
q, k, v, cu_q, cu_k, longest, causal, scale = inputs
|
| 183 |
+
ctx.save_for_backward(q, k, v, *output, cu_q, cu_k)
|
| 184 |
+
ctx.longest, ctx.causal, ctx.scale = longest.tolist(), causal, scale
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
if not hasattr(torch.ops.d1, "varlen"): # once a process: two repos' copies of this module share the op
|
| 188 |
+
_op = torch.library.custom_op("d1::varlen", _varlen, mutates_args=())
|
| 189 |
+
_op.register_fake(lambda q, k, v, cu_q, cu_k, longest, causal, scale: (
|
| 190 |
+
torch.empty_like(q), q.new_empty(q.shape[1], q.shape[0], dtype=torch.float32),
|
| 191 |
+
q.new_empty((), dtype=torch.int64), q.new_empty((), dtype=torch.int64)))
|
| 192 |
+
_op.register_autograd(_varlen_backward, setup_context=_varlen_context)
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
@dataclass
|
| 196 |
+
class Sees:
|
| 197 |
+
"""The full-attention mask of `state_attention`, never built whole (a [B, 1, T, T] tensor takes 33 GB on a
|
| 198 |
+
180k-token row): `attend` takes the attention under it, `rows` a block of queries' part of it."""
|
| 199 |
+
|
| 200 |
+
state: torch.Tensor # [B] each row's state length
|
| 201 |
+
question: torch.Tensor | None # [B, T] `set`: each token's question, 1..N, and -1 elsewhere
|
| 202 |
+
real: torch.Tensor # [B, T] the tokens that are not padding
|
| 203 |
+
_calls: dict = field(default_factory=dict, init=False, repr=False, compare=False) # once a device, every layer's
|
| 204 |
+
|
| 205 |
+
def __post_init__(self):
|
| 206 |
+
"""On a GPU, the plan on the mask's device, made before any layer: a compiled forward reads it with no graph
|
| 207 |
+
break. Off a GPU, `attend` masks and never plans."""
|
| 208 |
+
if self.real.is_cuda:
|
| 209 |
+
self._calls[self.real.device] = self._plan(self.real.device)
|
| 210 |
+
|
| 211 |
+
def rows(self, start: int, width: int, device) -> torch.Tensor:
|
| 212 |
+
"""Which keys queries `start` to `start + width` see, [B, 1, width, T]."""
|
| 213 |
+
q = torch.arange(start, start + width, device=device)[None, :, None]
|
| 214 |
+
kv = torch.arange(self.real.shape[1], device=device)[None, None, :]
|
| 215 |
+
s = self.state.to(device)[:, None, None]
|
| 216 |
+
if self.question is None: # causal, and the state both ways
|
| 217 |
+
m = (kv <= q) | ((q < s) & (kv < s))
|
| 218 |
+
else: # the state sees every token and every token the state; a question token its own question
|
| 219 |
+
g = self.question.to(device)
|
| 220 |
+
m = (q < s) | (kv < s) | (g[:, start:start + width, None] == g[:, None, :])
|
| 221 |
+
return (m & self.real.to(device)[:, None, :])[:, None]
|
| 222 |
+
|
| 223 |
+
def attend(self, q: Tensor, k: Tensor, v: Tensor, scale: float, window: int | None = None) -> Tensor:
|
| 224 |
+
"""Attention under this mask; q [B, H, T, D], k and v [B, Hkv, T, D]. On the flash kernels, no mask: the
|
| 225 |
+
states over their keys in one varlen call, the questions over their state and themselves in another. Off
|
| 226 |
+
them (CPU, fp32), `CHUNK` masked queries at a time. A sliding layer's `window`: causal within it, as every
|
| 227 |
+
mode keeps it (right-padded rows: a pad never reaches a real token)."""
|
| 228 |
+
if window is not None:
|
| 229 |
+
return causal_attention(q, k, v, window, scale)
|
| 230 |
+
b, h, t, d = q.shape
|
| 231 |
+
if not _fused(q):
|
| 232 |
+
starts = range(0, t, CHUNK) if t > CHUNK else (0,) # a range over t would compile a graph per length
|
| 233 |
+
return torch.cat([F.scaled_dot_product_attention(q[:, :, s:s + CHUNK], k, v, scale=scale, enable_gqa=True,
|
| 234 |
+
attn_mask=self.rows(s, min(CHUNK, t - s), q.device))
|
| 235 |
+
for s in starts], dim=2)
|
| 236 |
+
if q.device not in self._calls: # a model over GPUs: once on each
|
| 237 |
+
self._calls[q.device] = self._plan(q.device)
|
| 238 |
+
if self._calls[q.device] is None: # every row reads all of itself, unpadded: plain attention
|
| 239 |
+
return F.scaled_dot_product_attention(q, k, v, scale=scale, enable_gqa=True)
|
| 240 |
+
flat = [x.transpose(1, 2).reshape(b * t, x.shape[1], d) for x in (q, k, v)]
|
| 241 |
+
out = flat[0].new_zeros(b * t, h, d)
|
| 242 |
+
varlen = torch.ops.d1.varlen if torch.compiler.is_compiling() else _varlen
|
| 243 |
+
for qi, cu_q, ki, cu_k, longest, causal in self._calls[q.device]:
|
| 244 |
+
attended = varlen(flat[0][qi], flat[1][ki], flat[2][ki], cu_q, cu_k, longest, causal, scale)[0]
|
| 245 |
+
out = out.index_copy(0, qi, attended)
|
| 246 |
+
return out.view(b, t, h, d).transpose(1, 2)
|
| 247 |
+
|
| 248 |
+
def _plan(self, device) -> list | None:
|
| 249 |
+
"""The varlen calls, as (queries, their offsets, keys, theirs, the longest of each, causal): each state over
|
| 250 |
+
its row (`set`) or itself (`state`, `prompt`); each question over its state then itself, causally under
|
| 251 |
+
`state`, where the question part is one run a row. Under `set` a row of one question is one segment: all of
|
| 252 |
+
it sees all of it. None when every row is one such segment with no padding."""
|
| 253 |
+
real = self.real.to(device)
|
| 254 |
+
b, t = real.shape
|
| 255 |
+
n = real.sum(1)
|
| 256 |
+
s = torch.minimum(self.state.to(device), n)
|
| 257 |
+
base = torch.arange(b, device=device) * t
|
| 258 |
+
tail = (real & (torch.arange(t, device=device) >= s[:, None])).flatten().nonzero().squeeze(1)
|
| 259 |
+
g = (torch.zeros(b, t, dtype=torch.long, device=device) if self.question is None
|
| 260 |
+
else self.question.to(device)).flatten()[tail]
|
| 261 |
+
new = F.pad((g[1:] != g[:-1]) | (tail[1:] != tail[:-1] + 1), (1, 0), value=True)[:len(tail)]
|
| 262 |
+
first = new.nonzero().squeeze(1)
|
| 263 |
+
start, length = tail[first], torch.diff(F.pad(first, (0, 1), value=len(tail)))
|
| 264 |
+
if self.question is not None:
|
| 265 |
+
lone = torch.bincount(start // t, minlength=b) == 1
|
| 266 |
+
s = torch.where(lone, n, s)
|
| 267 |
+
start, length = start[~lone[start // t]], length[~lone[start // t]]
|
| 268 |
+
if not len(start) and bool((s == t).all()):
|
| 269 |
+
return None
|
| 270 |
+
calls = [(_runs((base, s)), _runs((base, n if self.question is not None else s)), False)]
|
| 271 |
+
if len(start):
|
| 272 |
+
row = start // t
|
| 273 |
+
calls.append((_runs((start, length)), _runs((base[row], s[row]), (start, length)), self.question is None))
|
| 274 |
+
return [(qi, cu_q, ki, cu_k, torch.tensor([max_q, max_k]), causal)
|
| 275 |
+
for (qi, cu_q, max_q), (ki, cu_k, max_k), causal in calls]
|
| 276 |
+
|
| 277 |
+
|
| 278 |
+
def d1_attention(module, q, k, v, attention_mask, scaling=None, sliding_window=None, **kwargs):
|
| 279 |
+
"""transformers' attention function for `attn_implementation="d1"` (registered on import): a `Sees` mask attends
|
| 280 |
+
through `Sees.attend` (a sliding layer within its window), anything else is SDPA's."""
|
| 281 |
+
if isinstance(attention_mask, Sees):
|
| 282 |
+
return attention_mask.attend(q, k, v, scaling, sliding_window).transpose(1, 2), None
|
| 283 |
+
return sdpa_attention_forward(module, q, k, v, attention_mask, scaling=scaling, sliding_window=sliding_window,
|
| 284 |
+
**kwargs)
|
| 285 |
+
|
| 286 |
+
|
| 287 |
+
AttentionInterface.register("d1", d1_attention)
|
| 288 |
+
# the masks transformers builds itself under "d1" (no `state_attention`; a tower, an audio encoder) are SDPA's: an
|
| 289 |
+
# implementation without one gets None for every mask, which drops a sliding window and the padding
|
| 290 |
+
AttentionMaskInterface.register("d1", sdpa_mask)
|
| 291 |
+
ATTENTION_KINDS = ("full_attention", "sliding_attention") # the layer kinds a `Sees` masks; the others take padding
|
| 292 |
+
|
| 293 |
+
|
| 294 |
+
def mask_mapping(layer_types, sees: Sees, pad: Tensor) -> dict:
|
| 295 |
+
"""The per-kind masks of a language model under `Sees`, keyed as its `layer_types`: the `Sees` on attention
|
| 296 |
+
layers, the 2D padding (None without any) on convolution and linear-attention layers."""
|
| 297 |
+
padding = None if pad.all() else pad
|
| 298 |
+
return {kind: sees if kind in ATTENTION_KINDS else padding for kind in set(layer_types)}
|
| 299 |
+
|
| 300 |
+
|
| 301 |
+
@contextmanager
|
| 302 |
+
def state_attention(model, state_len: torch.Tensor | None, question: torch.Tensor | None = None):
|
| 303 |
+
"""Within the block, the model's language model reads each row's first `state_len` tokens both ways; with
|
| 304 |
+
`question` (`set`: each token's question, 1..N, and -1 elsewhere), as `Sees` says. The mask is swapped in at the
|
| 305 |
+
language model, after the multimodal wrapper has placed its media: as the mask of a language model that reads it
|
| 306 |
+
(`reads_sees`), else as transformers' per-kind masks (`mask_mapping`), which `d1_attention` reads."""
|
| 307 |
+
if state_len is None:
|
| 308 |
+
yield
|
| 309 |
+
return
|
| 310 |
+
|
| 311 |
+
def hook(module, args, kwargs):
|
| 312 |
+
pad, x = kwargs.get("attention_mask"), kwargs.get("input_ids")
|
| 313 |
+
if isinstance(pad, dict):
|
| 314 |
+
raise ValueError("state attention replaces a 2D padding mask, not a prepared mapping")
|
| 315 |
+
if pad is None:
|
| 316 |
+
x = kwargs["inputs_embeds"] if x is None else x
|
| 317 |
+
pad = torch.ones(x.shape[:2], dtype=torch.long, device=x.device)
|
| 318 |
+
sees = Sees(state_len.to(pad.device), question, pad.bool())
|
| 319 |
+
if getattr(module, "reads_sees", False):
|
| 320 |
+
kwargs["attention_mask"] = sees
|
| 321 |
+
elif module.config._attn_implementation == "d1":
|
| 322 |
+
kinds = getattr(module.config, "layer_types", None) or ("full_attention",) # a config naming none: all full
|
| 323 |
+
kwargs["attention_mask"] = mask_mapping(kinds, sees, pad)
|
| 324 |
+
else:
|
| 325 |
+
raise ValueError(f"{type(module).__name__} reads `state_attention` with attn_implementation='d1'")
|
| 326 |
+
return args, kwargs
|
| 327 |
+
|
| 328 |
+
handle = language_model(model).register_forward_pre_hook(hook, with_kwargs=True)
|
| 329 |
+
try:
|
| 330 |
+
yield
|
| 331 |
+
finally:
|
| 332 |
+
handle.remove()
|
box.py
ADDED
|
@@ -0,0 +1,150 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Boxes: a presence and four edges, each a distribution over the same grid of fractions of the picture (0 the top
|
| 2 |
+
or left side, 1 the bottom or right): the thousandths 0-999 as the tokenizer writes them on a checkpoint that reads
|
| 3 |
+
numbers (`edge_grid`), else 128 codes.
|
| 4 |
+
|
| 5 |
+
- An edge is its distribution's median: an unsure edge's mean drifts to the picture's middle. A checkpoint that
|
| 6 |
+
records `d1_box_edges: interp` reads it inside the median's bin, where the distribution's CDF
|
| 7 |
+
crosses one half (`EDGE_READS`): the bin's value alone is off by a quarter of a bin on average. One without the
|
| 8 |
+
record reads the bin's value.
|
| 9 |
+
- A box's confidence: the share of boxes drawn from the edges that overlap it (IoU >= 0.5).
|
| 10 |
+
- `boxes`: every instance, blocks of a presence and four edges; block k is present when there are at least k (its
|
| 11 |
+
presence is P(count >= k)), the boxes answered in reading order (left edge, then top). A question reads its `max`
|
| 12 |
+
first blocks, else `INSTANCES`: the count is then at least that when the last is present. Prompt 1 trained 16
|
| 13 |
+
blocks under a prompt saying "16 at most"; prompt 2 states no count and trains 1 to `MAX_INSTANCES` blocks drawn
|
| 14 |
+
per row (`prompt.PROMPTS`).
|
| 15 |
+
- Probabilities go through the served calibration's tables (calibration.py)."""
|
| 16 |
+
|
| 17 |
+
from __future__ import annotations
|
| 18 |
+
|
| 19 |
+
import math
|
| 20 |
+
from bisect import bisect_left
|
| 21 |
+
from collections.abc import Callable
|
| 22 |
+
from functools import lru_cache
|
| 23 |
+
from itertools import accumulate
|
| 24 |
+
|
| 25 |
+
from .grid import Grid
|
| 26 |
+
|
| 27 |
+
EDGES = ("top", "bottom", "left", "right")
|
| 28 |
+
GRID = Grid(0, 1) # an edge's values as codes, 128 of them
|
| 29 |
+
THOUSANDTHS = Grid(0, 0.999, 1000) # an edge's values as numbers: 0-999 thousandths ("1000" takes two tokens)
|
| 30 |
+
DRAWS = 512
|
| 31 |
+
INSTANCES = 16 # a `boxes` question's blocks without `max`: 99.97% of COCO val's (picture, class) pairs hold 16 at most
|
| 32 |
+
MAX_INSTANCES = 64 # the most a `max` asks, and a prompt-2 training row draws (dense scenes: 17-60 instances)
|
| 33 |
+
EDGE_READS = ("median", "interp") # how an edge is read off its distribution; a checkpoint records it (`d1_box_edges`)
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def edge_grid(tokenizer) -> tuple[Grid, int]:
|
| 37 |
+
"""An edge's grid and the factor its values are written at: the thousandths where the checkpoint reads numbers
|
| 38 |
+
(`d1_numbers`) and the tokenizer writes each in one token, else the codes."""
|
| 39 |
+
from .prompt import number_ids
|
| 40 |
+
|
| 41 |
+
numbers = getattr(tokenizer, "d1_numbers", False) and number_ids(tokenizer, THOUSANDTHS, 1000) is not None
|
| 42 |
+
return (THOUSANDTHS, 1000) if numbers else (GRID, 1)
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def read_as(distribution) -> Grid:
|
| 46 |
+
"""The grid an edge's distribution is over, by its length: the codes' 128 values or the thousandths."""
|
| 47 |
+
return THOUSANDTHS if len(distribution) == THOUSANDTHS.size else GRID
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def ordered(a: float, b: float, grid: Grid = GRID) -> tuple[float, float]:
|
| 51 |
+
"""Two edges of one axis, the second below or right of the first: crossed ones become one bin about their middle."""
|
| 52 |
+
if a < b:
|
| 53 |
+
return a, b
|
| 54 |
+
half = grid.step / 2
|
| 55 |
+
middle = min(1 - half, max(half, (a + b) / 2))
|
| 56 |
+
return middle - half, middle + half
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def rectangle(edges: list[float], grid: Grid = GRID) -> dict[str, float]:
|
| 60 |
+
top, bottom = ordered(*edges[:2], grid)
|
| 61 |
+
left, right = ordered(*edges[2:], grid)
|
| 62 |
+
return {"top": top, "bottom": bottom, "left": left, "right": right}
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def area(box: dict) -> float:
|
| 66 |
+
return (box["bottom"] - box["top"]) * (box["right"] - box["left"])
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def iou(a: dict, b: dict) -> float:
|
| 70 |
+
overlap = (max(0.0, min(a["bottom"], b["bottom"]) - max(a["top"], b["top"]))
|
| 71 |
+
* max(0.0, min(a["right"], b["right"]) - max(a["left"], b["left"])))
|
| 72 |
+
return overlap / (area(a) + area(b) - overlap)
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
@lru_cache(maxsize=1)
|
| 76 |
+
def samples() -> tuple[tuple[float, ...], ...]:
|
| 77 |
+
"""`DRAWS` points of a Halton sequence in four dimensions: the same draws every call."""
|
| 78 |
+
def radical_inverse(i: int, base: int) -> float:
|
| 79 |
+
value, scale = 0.0, 1.0
|
| 80 |
+
while i:
|
| 81 |
+
i, digit = divmod(i, base)
|
| 82 |
+
scale /= base
|
| 83 |
+
value += digit * scale
|
| 84 |
+
return value
|
| 85 |
+
|
| 86 |
+
return tuple(tuple(radical_inverse(i, base) for base in (2, 3, 5, 7)) for i in range(1, DRAWS + 1))
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def confidence(edges: list[list[float]], box: dict) -> float:
|
| 90 |
+
"""The share of boxes drawn from the edge distributions (through each one's CDF) with IoU >= 0.5 against `box`."""
|
| 91 |
+
cdfs = [list(accumulate(p / total for p in row)) for row, total in zip(edges, map(math.fsum, edges))]
|
| 92 |
+
grid = read_as(edges[0])
|
| 93 |
+
values = grid.values
|
| 94 |
+
hits = sum(iou(rectangle([values[min(grid.size - 1, bisect_left(cdf, u))] for cdf, u in zip(cdfs, point)], grid),
|
| 95 |
+
box) >= 0.5 for point in samples())
|
| 96 |
+
return hits / DRAWS
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
def answer(probabilities: list[list[float]], through: Callable[[str, float], float] = lambda name, value: value,
|
| 100 |
+
edges: str = "median") -> dict:
|
| 101 |
+
"""The API's box: `probabilities` are presence ([yes, no]) then the four edges' distributions; `through(name, raw)`
|
| 102 |
+
calibrates `found` ("found:box") and the confidence ("box"); `edges`, how each edge is read (`EDGE_READS`)."""
|
| 103 |
+
found = through("found:box", probabilities[0][0])
|
| 104 |
+
if found < 0.5:
|
| 105 |
+
return {"type": "box", "box": None, "found": found, "confidence": 0.0}
|
| 106 |
+
box = medians(probabilities[1:], edges)
|
| 107 |
+
return {"type": "box", "box": box, "found": found, "confidence": through("box", confidence(probabilities[1:], box))}
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def crossing(probabilities, grid: Grid = GRID) -> float:
|
| 111 |
+
"""Where the distribution's CDF crosses one half, the mass spread evenly over each bin (value +- step / 2)."""
|
| 112 |
+
total, mass = math.fsum(probabilities), 0.0
|
| 113 |
+
for value, p in zip(grid.values, probabilities):
|
| 114 |
+
share = p / total
|
| 115 |
+
if share > 0 and mass + share >= 0.5:
|
| 116 |
+
return min(grid.high, max(grid.low, value - grid.step / 2 + (0.5 - mass) / share * grid.step))
|
| 117 |
+
mass += share
|
| 118 |
+
return grid.high
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
def medians(edges: list[list[float]], read: str = "median") -> dict[str, float]:
|
| 122 |
+
"""The box at the four edge distributions' medians: each its bin's value, or where its CDF crosses one half
|
| 123 |
+
inside that bin (`read`)."""
|
| 124 |
+
if read not in EDGE_READS:
|
| 125 |
+
raise ValueError(f"box edges read as one of {EDGE_READS}, not {read!r}")
|
| 126 |
+
grid = read_as(edges[0])
|
| 127 |
+
at = (lambda row: crossing(row, grid)) if read == "interp" else (lambda row: grid.read(row)["median"])
|
| 128 |
+
return rectangle([at(row) for row in edges], grid)
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
def answer_many(probabilities: list[list[float]], through: Callable[[str, float], float] = lambda name, value: value,
|
| 132 |
+
edges: str = "median") -> dict:
|
| 133 |
+
"""The API's `boxes`: `probabilities` are blocks (`INSTANCES`, or a question's `max`) of presence ([yes, no]) and
|
| 134 |
+
four edges. `presence` is P(count >= k) for each k, calibrated ("found:boxes") and kept decreasing; the boxes are
|
| 135 |
+
the blocks at >= 0.5, each with that presence as `found` and its confidence ("boxes"), in reading order; `count`
|
| 136 |
+
is how many, `at_least` when every block read is present (the true count may be higher)."""
|
| 137 |
+
presence, below = [], 1.0
|
| 138 |
+
for k in range(0, len(probabilities), 5):
|
| 139 |
+
below = min(below, through("found:boxes", probabilities[k][0]))
|
| 140 |
+
presence.append(below)
|
| 141 |
+
out = []
|
| 142 |
+
for k, found in enumerate(presence):
|
| 143 |
+
if found < 0.5:
|
| 144 |
+
break
|
| 145 |
+
block = probabilities[5 * k + 1:5 * k + 5]
|
| 146 |
+
box = medians(block, edges)
|
| 147 |
+
out.append({"box": box, "found": found, "confidence": through("boxes", confidence(block, box))})
|
| 148 |
+
out.sort(key=lambda b: (b["box"]["left"], b["box"]["top"]))
|
| 149 |
+
return {"type": "boxes", "boxes": out, "count": len(out), "at_least": bool(out) and len(out) == len(presence),
|
| 150 |
+
"presence": presence}
|
calibration.json
ADDED
|
@@ -0,0 +1,380 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"default": 1.337398453966105,
|
| 3 |
+
"temperatures": {},
|
| 4 |
+
"tables": {
|
| 5 |
+
"found:text": [
|
| 6 |
+
[
|
| 7 |
+
0.005682162615729877,
|
| 8 |
+
0.0
|
| 9 |
+
],
|
| 10 |
+
[
|
| 11 |
+
0.020574209407310595,
|
| 12 |
+
0.011363636363636364
|
| 13 |
+
],
|
| 14 |
+
[
|
| 15 |
+
0.05274184431288053,
|
| 16 |
+
0.022727272727272728
|
| 17 |
+
],
|
| 18 |
+
[
|
| 19 |
+
0.12593568082560191,
|
| 20 |
+
0.11363636363636363
|
| 21 |
+
],
|
| 22 |
+
[
|
| 23 |
+
0.2609023091129281,
|
| 24 |
+
0.25
|
| 25 |
+
],
|
| 26 |
+
[
|
| 27 |
+
0.45711776140061294,
|
| 28 |
+
0.42045454545454547
|
| 29 |
+
],
|
| 30 |
+
[
|
| 31 |
+
0.6460263153368776,
|
| 32 |
+
0.5568181818181818
|
| 33 |
+
],
|
| 34 |
+
[
|
| 35 |
+
0.7728411636569283,
|
| 36 |
+
0.7727272727272727
|
| 37 |
+
],
|
| 38 |
+
[
|
| 39 |
+
0.852111120115627,
|
| 40 |
+
0.8522727272727273
|
| 41 |
+
],
|
| 42 |
+
[
|
| 43 |
+
0.9055508368394591,
|
| 44 |
+
0.9204545454545454
|
| 45 |
+
],
|
| 46 |
+
[
|
| 47 |
+
0.9416752389886163,
|
| 48 |
+
0.9318181818181818
|
| 49 |
+
],
|
| 50 |
+
[
|
| 51 |
+
0.9675597308034246,
|
| 52 |
+
0.9772727272727273
|
| 53 |
+
],
|
| 54 |
+
[
|
| 55 |
+
0.9927885918454691,
|
| 56 |
+
0.9924242424242424
|
| 57 |
+
],
|
| 58 |
+
[
|
| 59 |
+
0.9997205003886156,
|
| 60 |
+
1.0
|
| 61 |
+
]
|
| 62 |
+
],
|
| 63 |
+
"found:number": [
|
| 64 |
+
[
|
| 65 |
+
0.01582507598356401,
|
| 66 |
+
0.0
|
| 67 |
+
],
|
| 68 |
+
[
|
| 69 |
+
0.1465578345315797,
|
| 70 |
+
0.2857142857142857
|
| 71 |
+
],
|
| 72 |
+
[
|
| 73 |
+
0.29448616611106054,
|
| 74 |
+
0.5
|
| 75 |
+
],
|
| 76 |
+
[
|
| 77 |
+
0.6300800165959767,
|
| 78 |
+
0.7142857142857143
|
| 79 |
+
],
|
| 80 |
+
[
|
| 81 |
+
0.9313231366021293,
|
| 82 |
+
0.9285714285714286
|
| 83 |
+
],
|
| 84 |
+
[
|
| 85 |
+
0.9951998972501911,
|
| 86 |
+
1.0
|
| 87 |
+
]
|
| 88 |
+
],
|
| 89 |
+
"found:date": [
|
| 90 |
+
[
|
| 91 |
+
0.008417200646363199,
|
| 92 |
+
0.0
|
| 93 |
+
],
|
| 94 |
+
[
|
| 95 |
+
0.5494083234419426,
|
| 96 |
+
0.5
|
| 97 |
+
],
|
| 98 |
+
[
|
| 99 |
+
0.9654906044403712,
|
| 100 |
+
0.9166666666666666
|
| 101 |
+
],
|
| 102 |
+
[
|
| 103 |
+
0.9966836659113566,
|
| 104 |
+
1.0
|
| 105 |
+
]
|
| 106 |
+
],
|
| 107 |
+
"found:box": [
|
| 108 |
+
[
|
| 109 |
+
0.0010319100776659475,
|
| 110 |
+
0.0027472527472527475
|
| 111 |
+
],
|
| 112 |
+
[
|
| 113 |
+
0.0162678323005357,
|
| 114 |
+
0.11235955056179775
|
| 115 |
+
],
|
| 116 |
+
[
|
| 117 |
+
0.1641463061831407,
|
| 118 |
+
0.4157303370786517
|
| 119 |
+
],
|
| 120 |
+
[
|
| 121 |
+
0.6041711138630492,
|
| 122 |
+
0.8541666666666666
|
| 123 |
+
],
|
| 124 |
+
[
|
| 125 |
+
0.8735690557675277,
|
| 126 |
+
0.978021978021978
|
| 127 |
+
],
|
| 128 |
+
[
|
| 129 |
+
0.9579824398960495,
|
| 130 |
+
0.9887640449438202
|
| 131 |
+
],
|
| 132 |
+
[
|
| 133 |
+
0.9952215192477699,
|
| 134 |
+
1.0
|
| 135 |
+
]
|
| 136 |
+
],
|
| 137 |
+
"box": [
|
| 138 |
+
[
|
| 139 |
+
0.08177083333333333,
|
| 140 |
+
0.18333333333333332
|
| 141 |
+
],
|
| 142 |
+
[
|
| 143 |
+
0.35257161458333336,
|
| 144 |
+
0.4
|
| 145 |
+
],
|
| 146 |
+
[
|
| 147 |
+
0.5962035123966942,
|
| 148 |
+
0.6694214876033058
|
| 149 |
+
],
|
| 150 |
+
[
|
| 151 |
+
0.7902631915983607,
|
| 152 |
+
0.7786885245901639
|
| 153 |
+
],
|
| 154 |
+
[
|
| 155 |
+
0.8755763319672131,
|
| 156 |
+
0.8688524590163934
|
| 157 |
+
],
|
| 158 |
+
[
|
| 159 |
+
0.9086263020833333,
|
| 160 |
+
0.9
|
| 161 |
+
],
|
| 162 |
+
[
|
| 163 |
+
0.9435433650362319,
|
| 164 |
+
0.9492753623188406
|
| 165 |
+
],
|
| 166 |
+
[
|
| 167 |
+
0.9652183219178082,
|
| 168 |
+
0.9726027397260274
|
| 169 |
+
],
|
| 170 |
+
[
|
| 171 |
+
0.9831365659916201,
|
| 172 |
+
0.9888268156424581
|
| 173 |
+
],
|
| 174 |
+
[
|
| 175 |
+
0.9948308270676691,
|
| 176 |
+
1.0
|
| 177 |
+
]
|
| 178 |
+
],
|
| 179 |
+
"estimate": [
|
| 180 |
+
[
|
| 181 |
+
0.030000000000000002,
|
| 182 |
+
0.22420634920634921
|
| 183 |
+
],
|
| 184 |
+
[
|
| 185 |
+
0.09,
|
| 186 |
+
0.28869047619047616
|
| 187 |
+
],
|
| 188 |
+
[
|
| 189 |
+
0.15000000000000002,
|
| 190 |
+
0.3551587301587302
|
| 191 |
+
],
|
| 192 |
+
[
|
| 193 |
+
0.21000000000000002,
|
| 194 |
+
0.4007936507936508
|
| 195 |
+
],
|
| 196 |
+
[
|
| 197 |
+
0.27,
|
| 198 |
+
0.4603174603174603
|
| 199 |
+
],
|
| 200 |
+
[
|
| 201 |
+
0.32999999999999996,
|
| 202 |
+
0.5238095238095238
|
| 203 |
+
],
|
| 204 |
+
[
|
| 205 |
+
0.39,
|
| 206 |
+
0.5922619047619048
|
| 207 |
+
],
|
| 208 |
+
[
|
| 209 |
+
0.44999999999999996,
|
| 210 |
+
0.6448412698412699
|
| 211 |
+
],
|
| 212 |
+
[
|
| 213 |
+
0.51,
|
| 214 |
+
0.7033730158730159
|
| 215 |
+
],
|
| 216 |
+
[
|
| 217 |
+
0.57,
|
| 218 |
+
0.7579365079365079
|
| 219 |
+
],
|
| 220 |
+
[
|
| 221 |
+
0.63,
|
| 222 |
+
0.7986111111111112
|
| 223 |
+
],
|
| 224 |
+
[
|
| 225 |
+
0.69,
|
| 226 |
+
0.8452380952380952
|
| 227 |
+
],
|
| 228 |
+
[
|
| 229 |
+
0.75,
|
| 230 |
+
0.8869047619047619
|
| 231 |
+
],
|
| 232 |
+
[
|
| 233 |
+
0.81,
|
| 234 |
+
0.9246031746031746
|
| 235 |
+
],
|
| 236 |
+
[
|
| 237 |
+
0.87,
|
| 238 |
+
0.9523809523809523
|
| 239 |
+
],
|
| 240 |
+
[
|
| 241 |
+
0.93,
|
| 242 |
+
0.9751984126984127
|
| 243 |
+
],
|
| 244 |
+
[
|
| 245 |
+
0.9799999999999999,
|
| 246 |
+
0.9940476190476191
|
| 247 |
+
]
|
| 248 |
+
],
|
| 249 |
+
"multiple": [
|
| 250 |
+
[
|
| 251 |
+
0.0016627926872626571,
|
| 252 |
+
0.0
|
| 253 |
+
],
|
| 254 |
+
[
|
| 255 |
+
0.008286236297705318,
|
| 256 |
+
0.0024813895781637717
|
| 257 |
+
],
|
| 258 |
+
[
|
| 259 |
+
0.013284336191710702,
|
| 260 |
+
0.0049504950495049506
|
| 261 |
+
],
|
| 262 |
+
[
|
| 263 |
+
0.017877624036637375,
|
| 264 |
+
0.015
|
| 265 |
+
],
|
| 266 |
+
[
|
| 267 |
+
0.023793328700436385,
|
| 268 |
+
0.018518518518518517
|
| 269 |
+
],
|
| 270 |
+
[
|
| 271 |
+
0.03199311782995977,
|
| 272 |
+
0.03398058252427184
|
| 273 |
+
],
|
| 274 |
+
[
|
| 275 |
+
0.04889508202323689,
|
| 276 |
+
0.05459057071960298
|
| 277 |
+
],
|
| 278 |
+
[
|
| 279 |
+
0.07631002315865278,
|
| 280 |
+
0.07407407407407407
|
| 281 |
+
],
|
| 282 |
+
[
|
| 283 |
+
0.10610537429711313,
|
| 284 |
+
0.1262135922330097
|
| 285 |
+
],
|
| 286 |
+
[
|
| 287 |
+
0.1516745984883731,
|
| 288 |
+
0.15865384615384615
|
| 289 |
+
],
|
| 290 |
+
[
|
| 291 |
+
0.2304679394076778,
|
| 292 |
+
0.21393034825870647
|
| 293 |
+
],
|
| 294 |
+
[
|
| 295 |
+
0.374861408323991,
|
| 296 |
+
0.3417085427135678
|
| 297 |
+
],
|
| 298 |
+
[
|
| 299 |
+
0.6689872755703948,
|
| 300 |
+
0.7
|
| 301 |
+
],
|
| 302 |
+
[
|
| 303 |
+
0.976994335126831,
|
| 304 |
+
0.9690721649484536
|
| 305 |
+
]
|
| 306 |
+
],
|
| 307 |
+
"found:boxes": [
|
| 308 |
+
[
|
| 309 |
+
0.0022886859880368398,
|
| 310 |
+
0.0
|
| 311 |
+
],
|
| 312 |
+
[
|
| 313 |
+
0.11311487904813239,
|
| 314 |
+
0.0891089108910891
|
| 315 |
+
],
|
| 316 |
+
[
|
| 317 |
+
0.39857997008840157,
|
| 318 |
+
0.3786407766990291
|
| 319 |
+
],
|
| 320 |
+
[
|
| 321 |
+
0.8182158553298949,
|
| 322 |
+
0.8514851485148515
|
| 323 |
+
],
|
| 324 |
+
[
|
| 325 |
+
0.980104905128495,
|
| 326 |
+
0.9702970297029703
|
| 327 |
+
],
|
| 328 |
+
[
|
| 329 |
+
0.9981022046260327,
|
| 330 |
+
1.0
|
| 331 |
+
]
|
| 332 |
+
],
|
| 333 |
+
"boxes": [
|
| 334 |
+
[
|
| 335 |
+
0.00048828125,
|
| 336 |
+
0.0
|
| 337 |
+
],
|
| 338 |
+
[
|
| 339 |
+
0.008246527777777778,
|
| 340 |
+
0.05555555555555555
|
| 341 |
+
],
|
| 342 |
+
[
|
| 343 |
+
0.0777587890625,
|
| 344 |
+
0.3125
|
| 345 |
+
],
|
| 346 |
+
[
|
| 347 |
+
0.20135498046875,
|
| 348 |
+
0.46875
|
| 349 |
+
],
|
| 350 |
+
[
|
| 351 |
+
0.3692626953125,
|
| 352 |
+
0.5625
|
| 353 |
+
],
|
| 354 |
+
[
|
| 355 |
+
0.50830078125,
|
| 356 |
+
0.625
|
| 357 |
+
],
|
| 358 |
+
[
|
| 359 |
+
0.5980224609375,
|
| 360 |
+
0.6875
|
| 361 |
+
],
|
| 362 |
+
[
|
| 363 |
+
0.6732177734375,
|
| 364 |
+
0.8125
|
| 365 |
+
],
|
| 366 |
+
[
|
| 367 |
+
0.8321126302083334,
|
| 368 |
+
0.9375
|
| 369 |
+
],
|
| 370 |
+
[
|
| 371 |
+
0.9644757699275363,
|
| 372 |
+
0.9710144927536232
|
| 373 |
+
],
|
| 374 |
+
[
|
| 375 |
+
0.9945746527777778,
|
| 376 |
+
1.0
|
| 377 |
+
]
|
| 378 |
+
]
|
| 379 |
+
}
|
| 380 |
+
}
|
calibration.py
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Temperature scaling of option scores, and tables for the typed answers' probabilities.
|
| 2 |
+
|
| 3 |
+
- `calibration.json` holds a global temperature and optional per-type ones; scaling keeps an argmax but can move a
|
| 4 |
+
score's median or expected level.
|
| 5 |
+
- A typed probability goes through its table, a monotone map fitted on held-out rows: `found:text`,
|
| 6 |
+
`found:number`, `found:date`, `found:box`, `box` (a box's confidence), `found:boxes` (P(count >= k)), `boxes`
|
| 7 |
+
(each listed box's confidence), `multiple` (each option's P(yes)), `estimate` (a quantile level, read backwards by
|
| 8 |
+
`Calibration.level`).
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import json
|
| 14 |
+
from bisect import bisect_left
|
| 15 |
+
from dataclasses import dataclass, field
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
|
| 18 |
+
from .prompt import Noul, Question, Score, cardinality
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def bucket(kind: str, k: int) -> str:
|
| 22 |
+
if kind == "noul":
|
| 23 |
+
return "noul"
|
| 24 |
+
if kind == "score":
|
| 25 |
+
return f"score:{k}"
|
| 26 |
+
if k <= 3:
|
| 27 |
+
return "choice:small"
|
| 28 |
+
if k <= 8:
|
| 29 |
+
return "choice:mid"
|
| 30 |
+
if k <= 26:
|
| 31 |
+
return "choice:large"
|
| 32 |
+
return "choice:inventory"
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def question_bucket(q: Question) -> str:
|
| 36 |
+
kind = "noul" if isinstance(q, Noul) else "score" if isinstance(q, Score) else "choice"
|
| 37 |
+
return bucket(kind, cardinality(q))
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def through(value: float, knots: list[list[float]] | None) -> float:
|
| 41 |
+
"""`value` through increasing (raw, calibrated) knots, linear between them, flat past the ends."""
|
| 42 |
+
if not knots:
|
| 43 |
+
return value
|
| 44 |
+
index = bisect_left([raw for raw, _ in knots], value)
|
| 45 |
+
if index == 0:
|
| 46 |
+
return knots[0][1]
|
| 47 |
+
if index == len(knots):
|
| 48 |
+
return knots[-1][1]
|
| 49 |
+
(a, pa), (b, pb) = knots[index - 1], knots[index]
|
| 50 |
+
return pa + (value - a) * (pb - pa) / (b - a)
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
@dataclass
|
| 54 |
+
class Calibration:
|
| 55 |
+
temperatures: dict[str, float] = field(default_factory=dict)
|
| 56 |
+
default: float = 1.0
|
| 57 |
+
counts: dict[str, int] = field(default_factory=dict)
|
| 58 |
+
tables: dict[str, list[list[float]]] = field(default_factory=dict)
|
| 59 |
+
|
| 60 |
+
def temperature(self, q: Question) -> float:
|
| 61 |
+
return self.temperatures.get(question_bucket(q), self.default)
|
| 62 |
+
|
| 63 |
+
def apply(self, q: Question, scores: list[float]) -> list[float]:
|
| 64 |
+
t = self.temperature(q)
|
| 65 |
+
return scores if t == 1.0 else [s / t for s in scores]
|
| 66 |
+
|
| 67 |
+
def through(self, name: str, value: float) -> float:
|
| 68 |
+
return through(value, self.tables.get(name))
|
| 69 |
+
|
| 70 |
+
def level(self, name: str, share: float) -> float:
|
| 71 |
+
"""The level whose observed share is `share`: the table read backwards (both columns increase)."""
|
| 72 |
+
return through(share, [[y, x] for x, y in self.tables.get(name) or ()])
|
| 73 |
+
|
| 74 |
+
@classmethod
|
| 75 |
+
def load(cls, path: Path | str) -> "Calibration":
|
| 76 |
+
d = json.loads(Path(path).read_text())
|
| 77 |
+
return cls(d["temperatures"], d.get("default", 1.0), d.get("counts", {}), d.get("tables", {}))
|
chat_template.jinja
ADDED
|
@@ -0,0 +1,286 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{- bos_token -}}
|
| 2 |
+
{%- set preserve_thinking = preserve_thinking | default(false) -%}
|
| 3 |
+
{%- set enable_thinking = enable_thinking | default(true) -%}
|
| 4 |
+
|
| 5 |
+
{%- macro render_value(v) -%}
|
| 6 |
+
{%- if v is string -%}
|
| 7 |
+
{{- v -}}
|
| 8 |
+
{%- elif v is boolean -%}
|
| 9 |
+
{{- "true" if v else "false" -}}
|
| 10 |
+
{%- elif v is none -%}
|
| 11 |
+
{{- "null" -}}
|
| 12 |
+
{%- elif v is mapping or (v is sequence and v is not string) -%}
|
| 13 |
+
{{- v | tojson -}}
|
| 14 |
+
{%- else -%}
|
| 15 |
+
{{- v -}}
|
| 16 |
+
{%- endif -%}
|
| 17 |
+
{%- endmacro -%}
|
| 18 |
+
|
| 19 |
+
{%- macro parse_content(content) -%}
|
| 20 |
+
{%- if content is string -%}
|
| 21 |
+
{{- content -}}
|
| 22 |
+
{%- elif content is mapping and content.get("type") == "image" -%}
|
| 23 |
+
{{- "<|vision_start|><|image|><|vision_end|>" -}}
|
| 24 |
+
{%- elif content is mapping and content.get("type") == "video" -%}
|
| 25 |
+
{{- "<|vision_start|><|video|><|vision_end|>" -}}
|
| 26 |
+
{%- elif content is mapping and content.get("type") == "text" -%}
|
| 27 |
+
{{- (content.get("text") or "") | string -}}
|
| 28 |
+
{%- elif content is mapping -%}
|
| 29 |
+
{{- content | tojson -}}
|
| 30 |
+
{%- elif content is iterable -%}
|
| 31 |
+
{%- set _ns = namespace(result="") -%}
|
| 32 |
+
{%- for item in content -%}
|
| 33 |
+
{%- if item is string -%}
|
| 34 |
+
{%- set _ns.result = _ns.result + item -%}
|
| 35 |
+
{%- elif item is mapping and item.get("type") == "image" -%}
|
| 36 |
+
{%- set _ns.result = _ns.result + "<|vision_start|><|image|><|vision_end|>" -%}
|
| 37 |
+
{%- elif item is mapping and item.get("type") == "video" -%}
|
| 38 |
+
{%- set _ns.result = _ns.result + "<|vision_start|><|video|><|vision_end|>" -%}
|
| 39 |
+
{%- elif item is mapping and item.get("type") == "text" -%}
|
| 40 |
+
{%- set _ns.result = _ns.result + ((item.get("text") or "") | string) -%}
|
| 41 |
+
{%- else -%}
|
| 42 |
+
{%- set _ns.result = _ns.result + (item | tojson) -%}
|
| 43 |
+
{%- endif -%}
|
| 44 |
+
{%- endfor -%}
|
| 45 |
+
{{- _ns.result -}}
|
| 46 |
+
{%- endif -%}
|
| 47 |
+
{%- endmacro -%}
|
| 48 |
+
|
| 49 |
+
{%- macro render_tool_calls(tool_calls) -%}
|
| 50 |
+
{%- for tool_call in tool_calls -%}
|
| 51 |
+
{%- if not loop.first -%}{{- "\n\n" -}}{%- endif -%}
|
| 52 |
+
{%- set func = tool_call["function"] if "function" in tool_call else tool_call -%}
|
| 53 |
+
{%- set func_args = func.get("arguments") -%}
|
| 54 |
+
{%- if func_args is string and (func_args | trim) not in ["", "{}", "null"] -%}
|
| 55 |
+
{{- raise_exception("Tool call arguments must be a mapping, got a JSON-encoded string: parse arguments with json.loads() before applying the chat template") -}}
|
| 56 |
+
{%- endif -%}
|
| 57 |
+
{{- "<tool_call>\n<function=" + func["name"] + ">\n" -}}
|
| 58 |
+
{%- if func_args is mapping -%}
|
| 59 |
+
{%- for arg_name, arg_value in func_args.items() -%}
|
| 60 |
+
{{- "<parameter=" + arg_name + ">\n" + render_value(arg_value) + "\n</parameter>\n" -}}
|
| 61 |
+
{%- endfor -%}
|
| 62 |
+
{%- endif -%}
|
| 63 |
+
{{- "</function>\n</tool_call>" -}}
|
| 64 |
+
{%- endfor -%}
|
| 65 |
+
{%- endmacro -%}
|
| 66 |
+
|
| 67 |
+
{%- macro tool_to_json(tool) -%}
|
| 68 |
+
{%- set fn = tool.function if tool.function is defined else tool -%}
|
| 69 |
+
{%- set ns_t = namespace(first=true) -%}
|
| 70 |
+
{{- "{" -}}
|
| 71 |
+
{%- for k, v in fn.items() -%}
|
| 72 |
+
{%- if k != "defer_loading" -%}
|
| 73 |
+
{%- if not ns_t.first -%}{{- ", " -}}{%- endif -%}
|
| 74 |
+
{%- set ns_t.first = false -%}
|
| 75 |
+
{{- k | tojson -}}{{- ": " -}}{{- v | tojson -}}
|
| 76 |
+
{%- endif -%}
|
| 77 |
+
{%- endfor -%}
|
| 78 |
+
{{- "}" -}}
|
| 79 |
+
{%- endmacro -%}
|
| 80 |
+
|
| 81 |
+
{#- a tool_reference content part (vLLM, SGLang) names a tool from `tools`,
|
| 82 |
+
usually a deferred one; it renders as that tool's definition -#}
|
| 83 |
+
{%- macro is_tool_reference(item) -%}
|
| 84 |
+
{{- "1" if item is mapping and item.get("type") == "tool_reference" else "" -}}
|
| 85 |
+
{%- endmacro -%}
|
| 86 |
+
|
| 87 |
+
{%- macro referenced_tool_to_json(name) -%}
|
| 88 |
+
{%- set _ns = namespace(found=none) -%}
|
| 89 |
+
{%- for tool in tools or [] -%}
|
| 90 |
+
{%- set fn = tool.function if tool.function is defined else tool -%}
|
| 91 |
+
{%- if fn.name == name -%}
|
| 92 |
+
{%- set _ns.found = fn -%}
|
| 93 |
+
{%- endif -%}
|
| 94 |
+
{%- endfor -%}
|
| 95 |
+
{%- if _ns.found is none -%}
|
| 96 |
+
{{- raise_exception("tool_reference names a tool that is not in tools: " + name | string) -}}
|
| 97 |
+
{%- endif -%}
|
| 98 |
+
{{- tool_to_json(_ns.found) -}}
|
| 99 |
+
{%- endmacro -%}
|
| 100 |
+
|
| 101 |
+
{%- macro render_tool_response(content) -%}
|
| 102 |
+
{#- content that already carries the tag (pre-formatted by the caller) is
|
| 103 |
+
passed through as-is instead of being wrapped a second time -#}
|
| 104 |
+
{%- if content is string and "<tool_response>" in content -%}
|
| 105 |
+
{{- content -}}
|
| 106 |
+
{%- else -%}
|
| 107 |
+
{{- "<tool_response>\n" -}}
|
| 108 |
+
{%- if content is string -%}
|
| 109 |
+
{{- content -}}
|
| 110 |
+
{%- elif content is mapping -%}
|
| 111 |
+
{{- parse_content(content) -}}
|
| 112 |
+
{%- elif content is iterable -%}
|
| 113 |
+
{%- for item in content -%}
|
| 114 |
+
{%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
|
| 115 |
+
{%- if is_tool_reference(item) -%}
|
| 116 |
+
{#- consecutive references share one <additional_tools> block -#}
|
| 117 |
+
{%- if loop.first or not is_tool_reference(loop.previtem) -%}
|
| 118 |
+
{{- "<additional_tools>\n" -}}
|
| 119 |
+
{%- endif -%}
|
| 120 |
+
{{- referenced_tool_to_json(item.name) -}}
|
| 121 |
+
{%- if loop.last or not is_tool_reference(loop.nextitem) -%}
|
| 122 |
+
{{- "\n</additional_tools>" -}}
|
| 123 |
+
{%- endif -%}
|
| 124 |
+
{%- elif item is mapping and item.result is defined and item.result is mapping and item.result.tools is defined -%}
|
| 125 |
+
{{- "<additional_tools>\n" -}}
|
| 126 |
+
{%- set _tl = item.result.tools -%}
|
| 127 |
+
{%- if _tl is string -%}
|
| 128 |
+
{{- _tl -}}
|
| 129 |
+
{%- elif _tl is mapping -%}
|
| 130 |
+
{{- tool_to_json(_tl) -}}
|
| 131 |
+
{%- else -%}
|
| 132 |
+
{%- for t in _tl -%}
|
| 133 |
+
{%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
|
| 134 |
+
{{- tool_to_json(t) -}}
|
| 135 |
+
{%- endfor -%}
|
| 136 |
+
{%- endif -%}
|
| 137 |
+
{{- "\n</additional_tools>" -}}
|
| 138 |
+
{%- else -%}
|
| 139 |
+
{{- parse_content(item) -}}
|
| 140 |
+
{%- endif -%}
|
| 141 |
+
{%- endfor -%}
|
| 142 |
+
{%- endif -%}
|
| 143 |
+
{{- "\n</tool_response>" -}}
|
| 144 |
+
{%- endif -%}
|
| 145 |
+
{%- endmacro -%}
|
| 146 |
+
|
| 147 |
+
{%- set ns = namespace(system_prompt="", last_user_index=-1) -%}
|
| 148 |
+
{%- if messages and messages[0]["role"] == "system" -%}
|
| 149 |
+
{%- if messages[0].get("content") -%}
|
| 150 |
+
{%- set ns.system_prompt = parse_content(messages[0]["content"]) -%}
|
| 151 |
+
{%- endif -%}
|
| 152 |
+
{%- set messages = messages[1:] -%}
|
| 153 |
+
{%- endif -%}
|
| 154 |
+
{%- if tools -%}
|
| 155 |
+
{%- set tools_ns = namespace(lines=[]) -%}
|
| 156 |
+
{%- for tool in tools -%}
|
| 157 |
+
{%- set fn = tool.function if tool.function is defined else tool -%}
|
| 158 |
+
{%- set _dyn = (fn.defer_loading is defined and fn.defer_loading) or (tool.defer_loading is defined and tool.defer_loading) -%}
|
| 159 |
+
{%- if not _dyn -%}
|
| 160 |
+
{%- set tools_ns.lines = tools_ns.lines + [tool_to_json(fn)] -%}
|
| 161 |
+
{%- endif -%}
|
| 162 |
+
{%- endfor -%}
|
| 163 |
+
{%- if tools_ns.lines -%}
|
| 164 |
+
{%- set ns.system_prompt = ns.system_prompt + ("\n\n" if ns.system_prompt else "") + "<tools>\n" + (tools_ns.lines | join("\n")) + "\n</tools>" -%}
|
| 165 |
+
{%- endif -%}
|
| 166 |
+
{%- endif -%}
|
| 167 |
+
{%- if ns.system_prompt -%}
|
| 168 |
+
{{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}}
|
| 169 |
+
{%- endif -%}
|
| 170 |
+
{%- for message in messages -%}
|
| 171 |
+
{%- if message["role"] == "user" -%}
|
| 172 |
+
{%- set ns.last_user_index = loop.index0 -%}
|
| 173 |
+
{%- endif -%}
|
| 174 |
+
{%- endfor -%}
|
| 175 |
+
{#- vLLM's and SGLang's Anthropic /v1/messages adapters split one tool_result that carries
|
| 176 |
+
tool_reference blocks into consecutive tool messages with the same tool_call_id (text first,
|
| 177 |
+
references second). Merge such runs into one tool turn so they render exactly like a single
|
| 178 |
+
OpenAI-format tool message whose content holds the same items. -#}
|
| 179 |
+
{%- set carry = namespace(pending=[]) -%}
|
| 180 |
+
{%- for message in messages -%}
|
| 181 |
+
{%- set _next = messages[loop.index0 + 1] if not loop.last else none -%}
|
| 182 |
+
{%- set _merge = message.role == "tool" and not message.get("additional_tools")
|
| 183 |
+
and _next is not none and _next.role == "tool" and not _next.get("additional_tools")
|
| 184 |
+
and message.get("tool_call_id") and _next.get("tool_call_id") == message.get("tool_call_id") -%}
|
| 185 |
+
{%- set _items = [] -%}
|
| 186 |
+
{%- if message.role == "tool" and not message.get("additional_tools") -%}
|
| 187 |
+
{%- set _c = message.get("content") -%}
|
| 188 |
+
{%- if _c is string -%}
|
| 189 |
+
{%- set _items = [{"type": "text", "text": _c}] if _c else [] -%}
|
| 190 |
+
{%- elif _c is mapping -%}
|
| 191 |
+
{%- set _items = [_c] -%}
|
| 192 |
+
{%- elif _c -%}
|
| 193 |
+
{%- set _items = _c | list -%}
|
| 194 |
+
{%- endif -%}
|
| 195 |
+
{%- endif -%}
|
| 196 |
+
{%- if _merge -%}
|
| 197 |
+
{%- set carry.pending = carry.pending + _items -%}
|
| 198 |
+
{%- else -%}
|
| 199 |
+
{{- "<|im_start|>" + message.role + "\n" -}}
|
| 200 |
+
{%- if message.role == "assistant" -%}
|
| 201 |
+
{%- generation -%}
|
| 202 |
+
{%- set thinking = message.thinking if message.thinking is defined else message.reasoning if message.reasoning is defined else message.reasoning_content -%}
|
| 203 |
+
{%- set keep_thinking = preserve_thinking or loop.index0 > ns.last_user_index -%}
|
| 204 |
+
{%- if thinking is defined and keep_thinking -%}
|
| 205 |
+
{{- "<think>" + (thinking or "") + "</think>" -}}
|
| 206 |
+
{%- endif -%}
|
| 207 |
+
{%- set _cfm_tag = "CONTINUE_FINAL_MESSAGE_TAG " -%}
|
| 208 |
+
{%- set _has_cfm = false -%}
|
| 209 |
+
{%- set content = "" -%}
|
| 210 |
+
{%- if message.get("content") -%}
|
| 211 |
+
{%- set content = parse_content(message.content) -%}
|
| 212 |
+
{%- endif -%}
|
| 213 |
+
{%- if not keep_thinking -%}
|
| 214 |
+
{%- if "</think>" in content -%}
|
| 215 |
+
{%- set content = content.split("</think>")[-1] | trim -%}
|
| 216 |
+
{%- endif -%}
|
| 217 |
+
{%- set content = "<think></think>" + content -%}
|
| 218 |
+
{%- endif -%}
|
| 219 |
+
{%- if content.endswith(_cfm_tag) -%}
|
| 220 |
+
{%- set _has_cfm = true -%}
|
| 221 |
+
{%- set _trunc_len = (content | length) - (_cfm_tag | length) -%}
|
| 222 |
+
{%- set content = content[:_trunc_len] -%}
|
| 223 |
+
{%- endif -%}
|
| 224 |
+
{{- content -}}
|
| 225 |
+
{%- if message.tool_calls -%}
|
| 226 |
+
{{- render_tool_calls(message.tool_calls) -}}
|
| 227 |
+
{%- endif -%}
|
| 228 |
+
{%- if _has_cfm -%}
|
| 229 |
+
{{- _cfm_tag -}}
|
| 230 |
+
{%- endif -%}
|
| 231 |
+
{{- "<|im_end|>\n" -}}
|
| 232 |
+
{%- endgeneration -%}
|
| 233 |
+
{%- elif message.role == "tool" -%}
|
| 234 |
+
{%- if message.get("additional_tools") -%}
|
| 235 |
+
{%- set _at = message["additional_tools"] -%}
|
| 236 |
+
{#- definitions may live in additional_tools itself, or in content
|
| 237 |
+
with additional_tools used as a plain truthy flag -#}
|
| 238 |
+
{%- if _at is mapping or (_at is sequence and _at is not string) -%}
|
| 239 |
+
{%- set _src = _at -%}
|
| 240 |
+
{%- else -%}
|
| 241 |
+
{%- set _src = message.get("content") -%}
|
| 242 |
+
{%- endif -%}
|
| 243 |
+
{#- nothing to reveal: skip the wrapper entirely rather than
|
| 244 |
+
emitting an empty <additional_tools> block -#}
|
| 245 |
+
{%- if _src is string and "<tool_response>" in _src -%}
|
| 246 |
+
{{- _src -}}
|
| 247 |
+
{%- elif _src is string and "<additional_tools>" in _src -%}
|
| 248 |
+
{{- "<tool_response>\n" + _src + "\n</tool_response>" -}}
|
| 249 |
+
{%- elif _src -%}
|
| 250 |
+
{{- "<tool_response>\n<additional_tools>\n" -}}
|
| 251 |
+
{%- if _src is string -%}
|
| 252 |
+
{{- _src -}}
|
| 253 |
+
{%- elif _src is mapping -%}
|
| 254 |
+
{{- tool_to_json(_src) -}}
|
| 255 |
+
{%- else -%}
|
| 256 |
+
{%- for t in _src -%}
|
| 257 |
+
{%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
|
| 258 |
+
{{- tool_to_json(t) -}}
|
| 259 |
+
{%- endfor -%}
|
| 260 |
+
{%- endif -%}
|
| 261 |
+
{{- "\n</additional_tools>\n</tool_response>" -}}
|
| 262 |
+
{%- endif -%}
|
| 263 |
+
{%- elif carry.pending -%}
|
| 264 |
+
{{- render_tool_response(carry.pending + _items) -}}
|
| 265 |
+
{%- set carry.pending = [] -%}
|
| 266 |
+
{%- else -%}
|
| 267 |
+
{#- an empty or missing result still gets the wrapper, so every tool turn has the same shape -#}
|
| 268 |
+
{{- render_tool_response(message.get("content") or "") -}}
|
| 269 |
+
{%- endif -%}
|
| 270 |
+
{{- "<|im_end|>\n" -}}
|
| 271 |
+
{%- else %}
|
| 272 |
+
{%- if message.get("content") -%}
|
| 273 |
+
{{- parse_content(message["content"]) -}}
|
| 274 |
+
{%- endif -%}
|
| 275 |
+
{{- "<|im_end|>\n" -}}
|
| 276 |
+
{%- endif %}
|
| 277 |
+
{%- endif -%}
|
| 278 |
+
{%- endfor -%}
|
| 279 |
+
{%- if add_generation_prompt -%}
|
| 280 |
+
{{- "<|im_start|>assistant\n" -}}
|
| 281 |
+
{%- if enable_thinking -%}
|
| 282 |
+
{{- "<think>" -}}
|
| 283 |
+
{%- else -%}
|
| 284 |
+
{{- "<think></think>" -}}
|
| 285 |
+
{%- endif -%}
|
| 286 |
+
{%- endif -%}
|
config.json
ADDED
|
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"D1Model"
|
| 4 |
+
],
|
| 5 |
+
"bos_token_id": 124894,
|
| 6 |
+
"downsample_factor": 2,
|
| 7 |
+
"dtype": "bfloat16",
|
| 8 |
+
"eos_token_id": 124900,
|
| 9 |
+
"image_token_id": 125030,
|
| 10 |
+
"model_type": "lfm3_vl",
|
| 11 |
+
"pad_token_id": 124893,
|
| 12 |
+
"projector_bias": true,
|
| 13 |
+
"projector_hidden_act": "gelu",
|
| 14 |
+
"projector_hidden_size": 4096,
|
| 15 |
+
"projector_layer_norm_eps": 1e-06,
|
| 16 |
+
"text_config": {
|
| 17 |
+
"bos_token_id": 124894,
|
| 18 |
+
"conv_L_cache": 7,
|
| 19 |
+
"conv_bias": false,
|
| 20 |
+
"dtype": "bfloat16",
|
| 21 |
+
"eos_token_id": 124900,
|
| 22 |
+
"head_dim": 128,
|
| 23 |
+
"hidden_size": 2048,
|
| 24 |
+
"initializer_range": 0.02,
|
| 25 |
+
"intermediate_size": 7680,
|
| 26 |
+
"layer_types": [
|
| 27 |
+
"conv",
|
| 28 |
+
"conv",
|
| 29 |
+
"full_attention",
|
| 30 |
+
"conv",
|
| 31 |
+
"sliding_attention",
|
| 32 |
+
"conv",
|
| 33 |
+
"full_attention",
|
| 34 |
+
"conv",
|
| 35 |
+
"sliding_attention",
|
| 36 |
+
"conv",
|
| 37 |
+
"full_attention",
|
| 38 |
+
"conv",
|
| 39 |
+
"sliding_attention",
|
| 40 |
+
"conv",
|
| 41 |
+
"full_attention",
|
| 42 |
+
"conv",
|
| 43 |
+
"sliding_attention",
|
| 44 |
+
"conv",
|
| 45 |
+
"full_attention",
|
| 46 |
+
"conv",
|
| 47 |
+
"sliding_attention",
|
| 48 |
+
"conv",
|
| 49 |
+
"full_attention",
|
| 50 |
+
"conv"
|
| 51 |
+
],
|
| 52 |
+
"max_position_embeddings": 128000,
|
| 53 |
+
"model_type": "lfm3",
|
| 54 |
+
"norm_eps": 1e-05,
|
| 55 |
+
"num_attention_heads": 16,
|
| 56 |
+
"num_hidden_layers": 24,
|
| 57 |
+
"num_key_value_heads": 4,
|
| 58 |
+
"pad_token_id": 124893,
|
| 59 |
+
"rope_parameters": {
|
| 60 |
+
"full_attention": {
|
| 61 |
+
"partial_rotary_factor": 0.5,
|
| 62 |
+
"rope_theta": 1000000.0,
|
| 63 |
+
"rope_type": "default"
|
| 64 |
+
},
|
| 65 |
+
"sliding_attention": {
|
| 66 |
+
"rope_theta": 10000,
|
| 67 |
+
"rope_type": "default"
|
| 68 |
+
}
|
| 69 |
+
},
|
| 70 |
+
"sliding_window": 256,
|
| 71 |
+
"tie_word_embeddings": true,
|
| 72 |
+
"use_cache": true,
|
| 73 |
+
"use_sandwich_norm": true,
|
| 74 |
+
"vocab_size": 128000
|
| 75 |
+
},
|
| 76 |
+
"tie_word_embeddings": true,
|
| 77 |
+
"transformers_version": "5.19.0",
|
| 78 |
+
"video_token_id": 125029,
|
| 79 |
+
"vision_config": {
|
| 80 |
+
"conv_kernel_size": 3,
|
| 81 |
+
"dtype": "bfloat16",
|
| 82 |
+
"hidden_size": 1024,
|
| 83 |
+
"in_channels": 3,
|
| 84 |
+
"initializer_range": 0.02,
|
| 85 |
+
"intermediate_size": 2560,
|
| 86 |
+
"layer_types": [
|
| 87 |
+
"conv",
|
| 88 |
+
"conv",
|
| 89 |
+
"full_attention",
|
| 90 |
+
"conv",
|
| 91 |
+
"full_attention",
|
| 92 |
+
"conv",
|
| 93 |
+
"full_attention",
|
| 94 |
+
"conv",
|
| 95 |
+
"full_attention",
|
| 96 |
+
"conv",
|
| 97 |
+
"full_attention",
|
| 98 |
+
"conv",
|
| 99 |
+
"full_attention",
|
| 100 |
+
"conv"
|
| 101 |
+
],
|
| 102 |
+
"model_type": "lfm3_vl_vision",
|
| 103 |
+
"norm_eps": 1e-05,
|
| 104 |
+
"num_attention_heads": 16,
|
| 105 |
+
"num_hidden_layers": 14,
|
| 106 |
+
"num_key_value_heads": 8,
|
| 107 |
+
"patch_size": 16,
|
| 108 |
+
"rope_parameters": {
|
| 109 |
+
"rope_theta": 1000000.0,
|
| 110 |
+
"rope_type": "axial"
|
| 111 |
+
},
|
| 112 |
+
"spatial_merge_size": 2,
|
| 113 |
+
"temporal_patch_size": 2
|
| 114 |
+
},
|
| 115 |
+
"vision_end_token_id": 125010,
|
| 116 |
+
"vision_start_token_id": 125009,
|
| 117 |
+
"d1_codes": 2,
|
| 118 |
+
"d1_attention": "set",
|
| 119 |
+
"d1_numbers": true,
|
| 120 |
+
"d1_set_size": 16,
|
| 121 |
+
"d1_vision_max_pixels": 1638400,
|
| 122 |
+
"auto_map": {
|
| 123 |
+
"AutoModel": "modeling_d1_lfm3_vl.D1Model",
|
| 124 |
+
"AutoConfig": "native_lfm3_vl.Lfm3VlConfig"
|
| 125 |
+
}
|
| 126 |
+
}
|
configuration_lfm3.py
ADDED
|
@@ -0,0 +1,140 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Vendored from transformers' LFM3 code (not in a release yet) at 69da1b47a11e362fbb15bfe5d5c8da0790c688c5, src/transformers/models/lfm3/configuration_lfm3.py:
|
| 2 |
+
# its relative imports made absolute, nothing else changed.
|
| 3 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 4 |
+
# This file was automatically generated from src/transformers/models/lfm3/modular_lfm3.py.
|
| 5 |
+
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
| 6 |
+
# the file from the modular. If any change should be done, please apply the change to the
|
| 7 |
+
# modular_lfm3.py file directly. One of our CI enforces this.
|
| 8 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 9 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 10 |
+
#
|
| 11 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 12 |
+
# you may not use this file except in compliance with the License.
|
| 13 |
+
# You may obtain a copy of the License at
|
| 14 |
+
#
|
| 15 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 16 |
+
#
|
| 17 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 18 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 19 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 20 |
+
# See the License for the specific language governing permissions and
|
| 21 |
+
# limitations under the License.
|
| 22 |
+
from huggingface_hub.dataclasses import strict
|
| 23 |
+
|
| 24 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 25 |
+
from transformers.utils import auto_docstring
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-1.7B")
|
| 29 |
+
@strict
|
| 30 |
+
class Lfm3Config(PreTrainedConfig):
|
| 31 |
+
r"""
|
| 32 |
+
rope_parameters (`dict[str, RopeParameters]`, *optional*):
|
| 33 |
+
RoPE parameters keyed by attention layer type, for example `{"full_attention": {"rope_type": "default",
|
| 34 |
+
"rope_theta": 1000000.0, "partial_rotary_factor": 0.5}, "sliding_attention": {"rope_type": "default",
|
| 35 |
+
"rope_theta": 10000.0}}`. `conv` layers do not use RoPE and have no entry. Full-attention layers rotate the
|
| 36 |
+
first `partial_rotary_factor * head_dim` dimensions and leave the rest unrotated.
|
| 37 |
+
conv_bias (`bool`, *optional*, defaults to `False`):
|
| 38 |
+
Whether to use bias in the conv layers.
|
| 39 |
+
conv_L_cache (`int`, *optional*, defaults to 7):
|
| 40 |
+
L_cache dim in the conv layers.
|
| 41 |
+
layer_types (`list[str]`, *optional*):
|
| 42 |
+
Per-layer operator type, each one of `"full_attention"`, `"sliding_attention"` or `"conv"`. Defaults to the
|
| 43 |
+
LFM3 pattern of a full-attention layer every fourth layer starting from the third, a sliding-attention layer
|
| 44 |
+
in between, and gated short convolutions elsewhere.
|
| 45 |
+
sliding_window (`int`, *optional*, defaults to 256):
|
| 46 |
+
Sliding-window size of the `sliding_attention` layers.
|
| 47 |
+
use_sandwich_norm (`bool`, *optional*, defaults to `True`):
|
| 48 |
+
Whether to apply an RMSNorm to the operator and feed-forward outputs before they enter the residual stream.
|
| 49 |
+
|
| 50 |
+
```python
|
| 51 |
+
>>> from transformers import Lfm3Model, Lfm3Config
|
| 52 |
+
|
| 53 |
+
>>> # Initializing a LFM3 model
|
| 54 |
+
>>> configuration = Lfm3Config()
|
| 55 |
+
|
| 56 |
+
>>> # Initializing a model from the LFM3-1.7B style configuration
|
| 57 |
+
>>> model = Lfm3Model(configuration)
|
| 58 |
+
|
| 59 |
+
>>> # Accessing the model configuration
|
| 60 |
+
>>> configuration = model.config
|
| 61 |
+
```
|
| 62 |
+
"""
|
| 63 |
+
|
| 64 |
+
model_type = "lfm3"
|
| 65 |
+
keys_to_ignore_at_inference = ["past_key_values"]
|
| 66 |
+
default_theta = 1000000.0
|
| 67 |
+
|
| 68 |
+
vocab_size: int = 128_000
|
| 69 |
+
hidden_size: int = 2048
|
| 70 |
+
intermediate_size: int = 7680
|
| 71 |
+
num_hidden_layers: int = 24
|
| 72 |
+
num_attention_heads: int = 16
|
| 73 |
+
num_key_value_heads: int = 4
|
| 74 |
+
max_position_embeddings: int = 128_000
|
| 75 |
+
initializer_range: float = 0.02
|
| 76 |
+
norm_eps: float = 0.00001
|
| 77 |
+
use_cache: bool = True
|
| 78 |
+
pad_token_id: int | None = 124_893
|
| 79 |
+
bos_token_id: int | None = 124_894
|
| 80 |
+
eos_token_id: int | list[int] | None = 124_900
|
| 81 |
+
tie_word_embeddings: bool = True
|
| 82 |
+
rope_parameters: dict | None = None
|
| 83 |
+
conv_bias: bool = False
|
| 84 |
+
conv_L_cache: int = 7
|
| 85 |
+
layer_types: list[str] | None = None
|
| 86 |
+
_attention_layer_types = ("full_attention", "sliding_attention")
|
| 87 |
+
head_dim: int | None = None
|
| 88 |
+
sliding_window: int | None = 256
|
| 89 |
+
use_sandwich_norm: bool = True
|
| 90 |
+
|
| 91 |
+
def __post_init__(self, **kwargs):
|
| 92 |
+
if self.head_dim is None:
|
| 93 |
+
self.head_dim = self.hidden_size // self.num_attention_heads
|
| 94 |
+
if self.rope_parameters is None:
|
| 95 |
+
self.rope_parameters = {
|
| 96 |
+
"full_attention": {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 0.5},
|
| 97 |
+
"sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0},
|
| 98 |
+
}
|
| 99 |
+
if self.layer_types is None:
|
| 100 |
+
self.layer_types = [
|
| 101 |
+
"full_attention" if i % 4 == 2 else "sliding_attention" if i % 4 == 0 and i > 0 else "conv"
|
| 102 |
+
for i in range(self.num_hidden_layers)
|
| 103 |
+
]
|
| 104 |
+
|
| 105 |
+
self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings)
|
| 106 |
+
super().__post_init__(**kwargs)
|
| 107 |
+
|
| 108 |
+
def validate_architecture(self):
|
| 109 |
+
"""Part of ``@strict``-powered validation."""
|
| 110 |
+
if len(self.layer_types) != self.num_hidden_layers:
|
| 111 |
+
raise ValueError(
|
| 112 |
+
f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
|
| 113 |
+
)
|
| 114 |
+
allowed_layer_types = (*self._attention_layer_types, "conv")
|
| 115 |
+
invalid_layer_types = [layer_type for layer_type in self.layer_types if layer_type not in allowed_layer_types]
|
| 116 |
+
if invalid_layer_types:
|
| 117 |
+
raise ValueError(f"`layer_types` entries must be one of {allowed_layer_types}; got {invalid_layer_types}.")
|
| 118 |
+
if "sliding_attention" in self.layer_types and self.sliding_window is None:
|
| 119 |
+
raise ValueError("`layer_types` has `sliding_attention` layers but `sliding_window` is None.")
|
| 120 |
+
missing_rope = [
|
| 121 |
+
layer_type
|
| 122 |
+
for layer_type in self._attention_layer_types
|
| 123 |
+
if layer_type in self.layer_types and layer_type not in self.rope_parameters
|
| 124 |
+
]
|
| 125 |
+
if missing_rope:
|
| 126 |
+
raise ValueError(f"`rope_parameters` has no entry for the {missing_rope} layers in `layer_types`.")
|
| 127 |
+
for layer_type, rope_params in self.rope_parameters.items():
|
| 128 |
+
if not isinstance(rope_params, dict):
|
| 129 |
+
continue
|
| 130 |
+
if (
|
| 131 |
+
rope_params.get("rope_type", "default") in ("dynamic", "longrope")
|
| 132 |
+
and rope_params.get("partial_rotary_factor", 1.0) < 1
|
| 133 |
+
):
|
| 134 |
+
raise ValueError(
|
| 135 |
+
f"`{rope_params['rope_type']}` RoPE recomputes its frequencies during the forward pass over the "
|
| 136 |
+
f"rotated dimension only and cannot keep LFM3's partial rotation on `{layer_type}` layers."
|
| 137 |
+
)
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
__all__ = ["Lfm3Config"]
|
configuration_lfm3_vl.py
ADDED
|
@@ -0,0 +1,211 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Vendored from transformers' LFM3 code (not in a release yet) at 69da1b47a11e362fbb15bfe5d5c8da0790c688c5, src/transformers/models/lfm3_vl/configuration_lfm3_vl.py:
|
| 2 |
+
# its relative imports made absolute, nothing else changed.
|
| 3 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 4 |
+
# This file was automatically generated from src/transformers/models/lfm3_vl/modular_lfm3_vl.py.
|
| 5 |
+
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
| 6 |
+
# the file from the modular. If any change should be done, please apply the change to the
|
| 7 |
+
# modular_lfm3_vl.py file directly. One of our CI enforces this.
|
| 8 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 9 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 10 |
+
#
|
| 11 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 12 |
+
# you may not use this file except in compliance with the License.
|
| 13 |
+
# You may obtain a copy of the License at
|
| 14 |
+
#
|
| 15 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 16 |
+
#
|
| 17 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 18 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 19 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 20 |
+
# See the License for the specific language governing permissions and
|
| 21 |
+
# limitations under the License.
|
| 22 |
+
from huggingface_hub.dataclasses import strict
|
| 23 |
+
|
| 24 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 25 |
+
from transformers.utils import auto_docstring
|
| 26 |
+
from transformers.models.auto import CONFIG_MAPPING, AutoConfig
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-VL-1.9B")
|
| 30 |
+
@strict
|
| 31 |
+
class Lfm3VlVisionConfig(PreTrainedConfig):
|
| 32 |
+
r"""
|
| 33 |
+
layer_types (`list[str]`, *optional*):
|
| 34 |
+
Per-layer operator type, each one of `"full_attention"` or `"conv"`. Defaults to two gated convolution
|
| 35 |
+
layers followed by alternating attention and convolution layers.
|
| 36 |
+
patch_size (`int`, *optional*, defaults to 16):
|
| 37 |
+
The spatial size of each image patch.
|
| 38 |
+
temporal_patch_size (`int`, *optional*, defaults to 2):
|
| 39 |
+
The number of video frames packed into one patch.
|
| 40 |
+
spatial_merge_size (`int`, *optional*, defaults to 2):
|
| 41 |
+
The side of the square block of patches merged into one language token.
|
| 42 |
+
conv_kernel_size (`int`, *optional*, defaults to 3):
|
| 43 |
+
The kernel size of the depthwise 2D convolution layers.
|
| 44 |
+
norm_eps (`float`, *optional*, defaults to 1e-05):
|
| 45 |
+
The epsilon used by the RMS normalization layers.
|
| 46 |
+
"""
|
| 47 |
+
|
| 48 |
+
model_type = "lfm3_vl_vision"
|
| 49 |
+
base_config_key = "vision_config"
|
| 50 |
+
default_theta = 1_000_000.0
|
| 51 |
+
default_rope_type = "axial"
|
| 52 |
+
_attention_layer_types = ("full_attention",)
|
| 53 |
+
|
| 54 |
+
hidden_size: int = 1024
|
| 55 |
+
intermediate_size: int = 2560
|
| 56 |
+
num_hidden_layers: int = 14
|
| 57 |
+
num_attention_heads: int = 16
|
| 58 |
+
num_key_value_heads: int = 8
|
| 59 |
+
layer_types: list[str] | None = None
|
| 60 |
+
patch_size: int = 16
|
| 61 |
+
temporal_patch_size: int = 2
|
| 62 |
+
in_channels: int = 3
|
| 63 |
+
spatial_merge_size: int = 2
|
| 64 |
+
conv_kernel_size: int = 3
|
| 65 |
+
norm_eps: float = 1e-5
|
| 66 |
+
initializer_range: float = 0.02
|
| 67 |
+
rope_parameters: dict | None = None
|
| 68 |
+
|
| 69 |
+
def __post_init__(self, **kwargs):
|
| 70 |
+
if self.layer_types is None:
|
| 71 |
+
self.layer_types = ["conv" if i < 2 or i % 2 else "full_attention" for i in range(self.num_hidden_layers)]
|
| 72 |
+
super().__post_init__(**kwargs)
|
| 73 |
+
|
| 74 |
+
def validate_architecture(self):
|
| 75 |
+
"""Part of ``@strict``-powered validation."""
|
| 76 |
+
if len(self.layer_types) != self.num_hidden_layers:
|
| 77 |
+
raise ValueError(
|
| 78 |
+
f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
|
| 79 |
+
)
|
| 80 |
+
allowed_layer_types = (*self._attention_layer_types, "conv")
|
| 81 |
+
invalid_layer_types = [layer_type for layer_type in self.layer_types if layer_type not in allowed_layer_types]
|
| 82 |
+
if invalid_layer_types:
|
| 83 |
+
raise ValueError(f"`layer_types` entries must be one of {allowed_layer_types}; got {invalid_layer_types}.")
|
| 84 |
+
if self.hidden_size % self.num_attention_heads:
|
| 85 |
+
raise ValueError("`hidden_size` must be a multiple of `num_attention_heads`.")
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
# The first checkpoints spell every text and `vision_*` field at the top level of the config
|
| 89 |
+
_LEGACY_TEXT_KEYS = (
|
| 90 |
+
"vocab_size",
|
| 91 |
+
"hidden_size",
|
| 92 |
+
"intermediate_size",
|
| 93 |
+
"num_hidden_layers",
|
| 94 |
+
"num_attention_heads",
|
| 95 |
+
"num_key_value_heads",
|
| 96 |
+
"head_dim",
|
| 97 |
+
"layer_types",
|
| 98 |
+
"conv_bias",
|
| 99 |
+
"conv_L_cache",
|
| 100 |
+
"sliding_window",
|
| 101 |
+
"max_position_embeddings",
|
| 102 |
+
"use_sandwich_norm",
|
| 103 |
+
"use_cache",
|
| 104 |
+
)
|
| 105 |
+
_LEGACY_VISION_KEYS = {
|
| 106 |
+
"vision_hidden_size": "hidden_size",
|
| 107 |
+
"vision_intermediate_size": "intermediate_size",
|
| 108 |
+
"vision_num_hidden_layers": "num_hidden_layers",
|
| 109 |
+
"vision_num_attention_heads": "num_attention_heads",
|
| 110 |
+
"vision_num_key_value_heads": "num_key_value_heads",
|
| 111 |
+
"vision_layer_types": "layer_types",
|
| 112 |
+
"vision_patch_size": "patch_size",
|
| 113 |
+
"vision_temporal_patch_size": "temporal_patch_size",
|
| 114 |
+
"vision_num_channels": "in_channels",
|
| 115 |
+
"vision_spatial_merge_size": "spatial_merge_size",
|
| 116 |
+
}
|
| 117 |
+
_LEGACY_DROPPED_KEYS = ("vision_patch_dim", "vision_max_position_embeddings")
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
@auto_docstring(checkpoint="LiquidAI/LFM3-VL-1.9B")
|
| 121 |
+
@strict
|
| 122 |
+
class Lfm3VlConfig(PreTrainedConfig):
|
| 123 |
+
r"""
|
| 124 |
+
video_token_id (`int`, *optional*, defaults to 125029):
|
| 125 |
+
The id of the placeholder token for video patches.
|
| 126 |
+
vision_start_token_id (`int`, *optional*, defaults to 125009):
|
| 127 |
+
The id of the token opening an image or a video frame.
|
| 128 |
+
vision_end_token_id (`int`, *optional*, defaults to 125010):
|
| 129 |
+
The id of the token closing an image or a video frame.
|
| 130 |
+
projector_layer_norm_eps (`float`, *optional*, defaults to 1e-06):
|
| 131 |
+
The epsilon of the layer normalization applied to the vision features before projection.
|
| 132 |
+
downsample_factor (`int`, *optional*, defaults to 2):
|
| 133 |
+
The side of the square block of vision tokens merged by the projector, equal to the vision encoder's
|
| 134 |
+
`spatial_merge_size`.
|
| 135 |
+
"""
|
| 136 |
+
|
| 137 |
+
model_type = "lfm3_vl"
|
| 138 |
+
sub_configs = {"text_config": AutoConfig, "vision_config": Lfm3VlVisionConfig}
|
| 139 |
+
|
| 140 |
+
text_config: dict | PreTrainedConfig | None = None
|
| 141 |
+
vision_config: dict | PreTrainedConfig | None = None
|
| 142 |
+
image_token_id: int = 125_030
|
| 143 |
+
video_token_id: int = 125_029
|
| 144 |
+
vision_start_token_id: int = 125_009
|
| 145 |
+
vision_end_token_id: int = 125_010
|
| 146 |
+
projector_hidden_act: str = "gelu"
|
| 147 |
+
projector_hidden_size: int = 4096
|
| 148 |
+
projector_bias: bool = True
|
| 149 |
+
projector_layer_norm_eps: float = 1e-6
|
| 150 |
+
downsample_factor: int = 2
|
| 151 |
+
tie_word_embeddings: bool = True
|
| 152 |
+
|
| 153 |
+
def __post_init__(self, **kwargs):
|
| 154 |
+
if "vision_hidden_size" in kwargs:
|
| 155 |
+
kwargs = self._convert_legacy_flat_config(kwargs)
|
| 156 |
+
|
| 157 |
+
if isinstance(self.vision_config, dict):
|
| 158 |
+
self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
|
| 159 |
+
elif self.vision_config is None:
|
| 160 |
+
self.vision_config = self.sub_configs["vision_config"]()
|
| 161 |
+
|
| 162 |
+
if isinstance(self.text_config, dict):
|
| 163 |
+
self.text_config = {"model_type": "lfm3", **self.text_config}
|
| 164 |
+
self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
|
| 165 |
+
elif self.text_config is None:
|
| 166 |
+
self.text_config = CONFIG_MAPPING["lfm3"]()
|
| 167 |
+
self.validate_vision_fields()
|
| 168 |
+
|
| 169 |
+
super().__post_init__(**kwargs)
|
| 170 |
+
|
| 171 |
+
def validate_vision_fields(self):
|
| 172 |
+
"""Part of `@strict`-powered validation, so `validate()` and `save_pretrained` refuse later edits too."""
|
| 173 |
+
if self.downsample_factor != self.vision_config.spatial_merge_size:
|
| 174 |
+
raise ValueError("downsample_factor must match vision_config.spatial_merge_size")
|
| 175 |
+
# an `inputs_embeds` forward embeds both placeholder ids to find them
|
| 176 |
+
vocab_size = self.text_config.vocab_size
|
| 177 |
+
for name in ("image_token_id", "video_token_id"):
|
| 178 |
+
if not 0 <= getattr(self, name) < vocab_size:
|
| 179 |
+
raise ValueError(f"{name}={getattr(self, name)} must lie in [0, text_config.vocab_size={vocab_size})")
|
| 180 |
+
|
| 181 |
+
def _convert_legacy_flat_config(self, kwargs: dict) -> dict:
|
| 182 |
+
text_config = {key: kwargs.pop(key) for key in _LEGACY_TEXT_KEYS if key in kwargs}
|
| 183 |
+
text_config["norm_eps"] = kwargs.pop("rms_norm_eps", 1e-5)
|
| 184 |
+
text_config["rope_parameters"] = {
|
| 185 |
+
"full_attention": {
|
| 186 |
+
"rope_type": "default",
|
| 187 |
+
"rope_theta": kwargs.pop("full_attention_rope_theta", 1_000_000.0),
|
| 188 |
+
"partial_rotary_factor": kwargs.pop("partial_rotary_factor", 0.5),
|
| 189 |
+
},
|
| 190 |
+
"sliding_attention": {
|
| 191 |
+
"rope_type": "default",
|
| 192 |
+
"rope_theta": kwargs.pop("sliding_attention_rope_theta", 10_000.0),
|
| 193 |
+
},
|
| 194 |
+
}
|
| 195 |
+
for key in ("pad_token_id", "bos_token_id", "eos_token_id"):
|
| 196 |
+
if kwargs.get(key) is not None:
|
| 197 |
+
text_config[key] = kwargs[key]
|
| 198 |
+
text_config["tie_word_embeddings"] = self.tie_word_embeddings
|
| 199 |
+
|
| 200 |
+
vision_config = {new: kwargs.pop(old) for old, new in _LEGACY_VISION_KEYS.items() if old in kwargs}
|
| 201 |
+
vision_config["norm_eps"] = text_config["norm_eps"]
|
| 202 |
+
vision_config["rope_parameters"] = {"rope_type": "axial", "rope_theta": kwargs.pop("vision_rope_theta", 1e6)}
|
| 203 |
+
for key in _LEGACY_DROPPED_KEYS:
|
| 204 |
+
kwargs.pop(key, None)
|
| 205 |
+
|
| 206 |
+
self.text_config = text_config
|
| 207 |
+
self.vision_config = vision_config
|
| 208 |
+
return kwargs
|
| 209 |
+
|
| 210 |
+
|
| 211 |
+
__all__ = ["Lfm3VlConfig", "Lfm3VlVisionConfig"]
|
depthwise.py
ADDED
|
@@ -0,0 +1,53 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""LFM's depthwise causal convolution, the same numbers in every process: on ROCm as shifted sums, never MIOpen's
|
| 2 |
+
kernel, whose solver search runs again for every new length."""
|
| 3 |
+
|
| 4 |
+
from __future__ import annotations
|
| 5 |
+
|
| 6 |
+
import torch
|
| 7 |
+
import torch.nn.functional as F
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
def on_rocm(t: torch.Tensor) -> bool:
|
| 11 |
+
return t.is_cuda and bool(torch.version.hip)
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def depthwise(padded: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
|
| 15 |
+
"""Depthwise conv, `(B, C, T + k - 1)` to `(B, C, T)`. On ROCm, `shifted_sum`: MIOpen's depthwise kernel
|
| 16 |
+
searches a solver for every new length, logs a warning each time, may settle on another one in another process,
|
| 17 |
+
and takes 1.25 ms at 5k tokens against 0.07 ms fused (MI325X)."""
|
| 18 |
+
if not on_rocm(padded):
|
| 19 |
+
return F.conv1d(padded, weight, groups=padded.shape[1])
|
| 20 |
+
return shifted_sum(padded, weight)
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def shifted_sum(padded: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
|
| 24 |
+
"""The depthwise conv as k shifted products summed in fp32: one fp32 copy of `padded`, the shifts views of it."""
|
| 25 |
+
k, t = weight.shape[-1], padded.shape[-1] - weight.shape[-1] + 1
|
| 26 |
+
p, w = padded.float(), weight[:, 0, :, None].float()
|
| 27 |
+
return sum(p[..., i:i + t] * w[:, i] for i in range(k)).to(padded.dtype)
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def pinned_conv(conv: torch.nn.Conv1d | torch.nn.Conv2d, x: torch.Tensor) -> torch.Tensor:
|
| 31 |
+
"""`conv(x)` without MIOpen on ROCm: a 1x1 conv as the matmul it is, a 2D conv with one input channel a group
|
| 32 |
+
(`groups` = input channels) as `shifted_sum2d`. MIOpen's solver for the omni downsampler's 1x1 `lin_out` changed
|
| 33 |
+
from one call to the next within a process (MI325X). Elsewhere, and for any other conv, `conv`'s own forward."""
|
| 34 |
+
if on_rocm(x) and conv.padding_mode == "zeros" and set(conv.dilation) == {1}:
|
| 35 |
+
if set(conv.kernel_size) == {1} and set(conv.stride) == {1} and conv.padding in ((0,), (0, 0)) and conv.groups == 1:
|
| 36 |
+
return F.linear(x.movedim(1, -1), conv.weight.flatten(1), conv.bias).movedim(-1, 1)
|
| 37 |
+
if x.ndim == 4 and conv.groups == conv.in_channels and not isinstance(conv.padding, str):
|
| 38 |
+
return shifted_sum2d(x, conv.weight, conv.bias, conv.stride, conv.padding)
|
| 39 |
+
return type(conv).forward(conv, x)
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def shifted_sum2d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, stride: tuple[int, int],
|
| 43 |
+
padding: tuple[int, int]) -> torch.Tensor:
|
| 44 |
+
"""A 2D conv with one input channel a group, `weight` `(C * m, 1, kh, kw)`, as kh x kw strided products summed in
|
| 45 |
+
fp32: each of the C input channels feeds its m output channels."""
|
| 46 |
+
(kh, kw), (sh, sw), (ph, pw) = weight.shape[-2:], stride, padding
|
| 47 |
+
c = x.shape[1]
|
| 48 |
+
ho, wo = (x.shape[-2] + 2 * ph - kh) // sh + 1, (x.shape[-1] + 2 * pw - kw) // sw + 1
|
| 49 |
+
p = F.pad(x.float(), (pw, pw, ph, ph))[:, :, None]
|
| 50 |
+
w = weight.float().view(c, -1, kh, kw)[..., None, None]
|
| 51 |
+
out = sum(p[..., i:i + sh * (ho - 1) + 1:sh, j:j + sw * (wo - 1) + 1:sw] * w[:, :, i, j]
|
| 52 |
+
for i in range(kh) for j in range(kw)).flatten(1, 2)
|
| 53 |
+
return (out if bias is None else out + bias.float()[:, None, None]).to(x.dtype)
|
grid.py
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Ordered values read as a distribution over equally spaced bins."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import math
|
| 6 |
+
from collections.abc import Callable
|
| 7 |
+
from dataclasses import dataclass
|
| 8 |
+
from functools import cached_property
|
| 9 |
+
from typing import Sequence
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
@dataclass(frozen=True)
|
| 13 |
+
class Grid:
|
| 14 |
+
low: float
|
| 15 |
+
high: float
|
| 16 |
+
size: int = 128
|
| 17 |
+
|
| 18 |
+
def __post_init__(self):
|
| 19 |
+
if self.size < 2 or not self.low < self.high:
|
| 20 |
+
raise ValueError("a grid needs two values or more, low to high")
|
| 21 |
+
|
| 22 |
+
@property
|
| 23 |
+
def step(self) -> float:
|
| 24 |
+
return (self.high - self.low) / (self.size - 1)
|
| 25 |
+
|
| 26 |
+
@cached_property
|
| 27 |
+
def values(self) -> tuple[float, ...]:
|
| 28 |
+
return tuple(self.low + i * self.step for i in range(self.size - 1)) + (self.high,)
|
| 29 |
+
|
| 30 |
+
def read(self, probabilities: Sequence[float], level: Callable[[float], float] | None = None) -> dict[str, float]:
|
| 31 |
+
"""The mean, median and central 80 percent interval, on this grid's values; `level`, the level a quantile is
|
| 32 |
+
read at (a calibration's, increasing)."""
|
| 33 |
+
total = math.fsum(probabilities)
|
| 34 |
+
values = self.values
|
| 35 |
+
out = {"mean": math.fsum(v * (p / total) for v, p in zip(values, probabilities))}
|
| 36 |
+
at = [(key, level(q) if level else q) for key, q in (("p10", 0.1), ("median", 0.5), ("p90", 0.9))]
|
| 37 |
+
mass = 0.0
|
| 38 |
+
for value, p in zip(values, probabilities):
|
| 39 |
+
mass += p / total
|
| 40 |
+
for key, quantile in at:
|
| 41 |
+
if key not in out and mass >= quantile:
|
| 42 |
+
out[key] = value
|
| 43 |
+
for key, _ in at: # a level near 1 the summed mass falls short of by rounding: the last value
|
| 44 |
+
out.setdefault(key, values[-1])
|
| 45 |
+
return out
|
image_processing_lfm3_vl.py
ADDED
|
@@ -0,0 +1,301 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Vendored from transformers' LFM3 code (not in a release yet) at 69da1b47a11e362fbb15bfe5d5c8da0790c688c5, src/transformers/models/lfm3_vl/image_processing_lfm3_vl.py:
|
| 2 |
+
# its relative imports made absolute, nothing else changed.
|
| 3 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 4 |
+
# This file was automatically generated from src/transformers/models/lfm3_vl/modular_lfm3_vl.py.
|
| 5 |
+
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
| 6 |
+
# the file from the modular. If any change should be done, please apply the change to the
|
| 7 |
+
# modular_lfm3_vl.py file directly. One of our CI enforces this.
|
| 8 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 9 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 10 |
+
#
|
| 11 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 12 |
+
# you may not use this file except in compliance with the License.
|
| 13 |
+
# You may obtain a copy of the License at
|
| 14 |
+
#
|
| 15 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 16 |
+
#
|
| 17 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 18 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 19 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 20 |
+
# See the License for the specific language governing permissions and
|
| 21 |
+
# limitations under the License.
|
| 22 |
+
import math
|
| 23 |
+
import warnings
|
| 24 |
+
from collections.abc import Iterable
|
| 25 |
+
|
| 26 |
+
import torch
|
| 27 |
+
from torchvision.transforms.v2 import functional as tvF
|
| 28 |
+
|
| 29 |
+
from transformers.image_processing_backends import TorchvisionBackend
|
| 30 |
+
from transformers.image_processing_utils import BatchFeature
|
| 31 |
+
from transformers.image_transforms import group_images_by_shape, reorder_images
|
| 32 |
+
from transformers.image_utils import ImageInput, PILImageResampling, SizeDict
|
| 33 |
+
from transformers.processing_utils import ImagesKwargs, Unpack
|
| 34 |
+
from transformers.utils import TensorType, auto_docstring
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
class Lfm3VlImageProcessorKwargs(ImagesKwargs, total=False):
|
| 38 |
+
r"""
|
| 39 |
+
min_pixels (`int`, *optional*, defaults to `56 * 56`):
|
| 40 |
+
The min pixels of the image to resize the image.
|
| 41 |
+
max_pixels (`int`, *optional*, defaults to `28 * 28 * 1280`):
|
| 42 |
+
The max pixels of the image to resize the image.
|
| 43 |
+
patch_size (`int`, *optional*, defaults to 14):
|
| 44 |
+
The spatial patch size of the vision encoder.
|
| 45 |
+
temporal_patch_size (`int`, *optional*, defaults to 2):
|
| 46 |
+
The temporal patch size of the vision encoder.
|
| 47 |
+
merge_size (`int`, *optional*, defaults to 2):
|
| 48 |
+
The merge size of the vision encoder to llm encoder.
|
| 49 |
+
"""
|
| 50 |
+
|
| 51 |
+
min_pixels: int
|
| 52 |
+
max_pixels: int
|
| 53 |
+
patch_size: int
|
| 54 |
+
temporal_patch_size: int
|
| 55 |
+
merge_size: int
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def smart_resize(
|
| 59 |
+
height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280
|
| 60 |
+
):
|
| 61 |
+
"""Rescales the image so that the following conditions are met:
|
| 62 |
+
|
| 63 |
+
1. Both dimensions (height and width) are divisible by 'factor'.
|
| 64 |
+
|
| 65 |
+
2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
|
| 66 |
+
|
| 67 |
+
3. The aspect ratio of the image is maintained as closely as possible.
|
| 68 |
+
|
| 69 |
+
"""
|
| 70 |
+
if max(height, width) / min(height, width) > 200:
|
| 71 |
+
raise ValueError(
|
| 72 |
+
f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
|
| 73 |
+
)
|
| 74 |
+
h_bar = round(height / factor) * factor
|
| 75 |
+
w_bar = round(width / factor) * factor
|
| 76 |
+
if h_bar * w_bar > max_pixels:
|
| 77 |
+
beta = math.sqrt((height * width) / max_pixels)
|
| 78 |
+
h_bar = max(factor, math.floor(height / beta / factor) * factor)
|
| 79 |
+
w_bar = max(factor, math.floor(width / beta / factor) * factor)
|
| 80 |
+
elif h_bar * w_bar < min_pixels:
|
| 81 |
+
beta = math.sqrt(min_pixels / (height * width))
|
| 82 |
+
h_bar = math.ceil(height * beta / factor) * factor
|
| 83 |
+
w_bar = math.ceil(width * beta / factor) * factor
|
| 84 |
+
return h_bar, w_bar
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
@auto_docstring
|
| 88 |
+
class Lfm3VlImageProcessor(TorchvisionBackend):
|
| 89 |
+
do_resize = True
|
| 90 |
+
resample = PILImageResampling.BICUBIC
|
| 91 |
+
size = {"shortest_edge": 65_536, "longest_edge": 16_777_216}
|
| 92 |
+
default_to_square = False
|
| 93 |
+
do_rescale = True
|
| 94 |
+
do_normalize = True
|
| 95 |
+
image_mean = [0.5, 0.5, 0.5]
|
| 96 |
+
image_std = [0.5, 0.5, 0.5]
|
| 97 |
+
do_convert_rgb = True
|
| 98 |
+
patch_size = 16
|
| 99 |
+
temporal_patch_size = 2
|
| 100 |
+
merge_size = 2
|
| 101 |
+
valid_kwargs = Lfm3VlImageProcessorKwargs
|
| 102 |
+
model_input_names = ["pixel_values", "image_grid_thw"]
|
| 103 |
+
|
| 104 |
+
def __init__(self, **kwargs: Unpack[Lfm3VlImageProcessorKwargs]):
|
| 105 |
+
# backward compatibility: override size with min_pixels and max_pixels if they are provided
|
| 106 |
+
size = kwargs.pop("size", None)
|
| 107 |
+
size = dict(self.size) if size is None else size
|
| 108 |
+
if (min_pixels := kwargs.pop("min_pixels", None)) is not None:
|
| 109 |
+
size["shortest_edge"] = min_pixels
|
| 110 |
+
size.pop("min_pixels", None)
|
| 111 |
+
if (max_pixels := kwargs.pop("max_pixels", None)) is not None:
|
| 112 |
+
size["longest_edge"] = max_pixels
|
| 113 |
+
size.pop("max_pixels", None)
|
| 114 |
+
super().__init__(size=size, **kwargs)
|
| 115 |
+
|
| 116 |
+
def _standardize_kwargs(
|
| 117 |
+
self,
|
| 118 |
+
size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
|
| 119 |
+
min_pixels: int | None = None,
|
| 120 |
+
max_pixels: int | None = None,
|
| 121 |
+
**kwargs,
|
| 122 |
+
) -> dict:
|
| 123 |
+
if min_pixels is not None or max_pixels is not None:
|
| 124 |
+
warnings.warn(
|
| 125 |
+
"Passing `min_pixels` and `max_pixels` to a processor call is deprecated and will be removed in v5.23. "
|
| 126 |
+
"Pass in `size={'longest_edge': xxx, 'shortest_edge': xxx} to override the target size.`",
|
| 127 |
+
FutureWarning,
|
| 128 |
+
)
|
| 129 |
+
|
| 130 |
+
size_dict = dict(size) if isinstance(size, (dict, SizeDict)) else {}
|
| 131 |
+
if min_pixels is not None:
|
| 132 |
+
size_dict["shortest_edge"] = min_pixels
|
| 133 |
+
if max_pixels is not None:
|
| 134 |
+
size_dict["longest_edge"] = max_pixels
|
| 135 |
+
size = SizeDict(**size_dict)
|
| 136 |
+
return super()._standardize_kwargs(size=size, **kwargs)
|
| 137 |
+
|
| 138 |
+
@auto_docstring
|
| 139 |
+
def preprocess(
|
| 140 |
+
self,
|
| 141 |
+
images: ImageInput,
|
| 142 |
+
**kwargs: Unpack[Lfm3VlImageProcessorKwargs],
|
| 143 |
+
) -> BatchFeature:
|
| 144 |
+
return super().preprocess(images, **kwargs)
|
| 145 |
+
|
| 146 |
+
def resize(
|
| 147 |
+
self,
|
| 148 |
+
images: "torch.Tensor",
|
| 149 |
+
size: SizeDict,
|
| 150 |
+
resample: "PILImageResampling | tvF.InterpolationMode | int | None",
|
| 151 |
+
factor: int,
|
| 152 |
+
**kwargs,
|
| 153 |
+
) -> "torch.Tensor":
|
| 154 |
+
"""Resize dynamically based on input image aspect ratio."""
|
| 155 |
+
if not size.shortest_edge or not size.longest_edge:
|
| 156 |
+
raise ValueError(f"`size` dict must contain 'shortest_edge' and 'longest_edge' keys but got {size}.")
|
| 157 |
+
|
| 158 |
+
height, width = images.shape[-2:]
|
| 159 |
+
resized_height, resized_width = smart_resize(
|
| 160 |
+
height,
|
| 161 |
+
width,
|
| 162 |
+
factor=factor,
|
| 163 |
+
min_pixels=size.shortest_edge,
|
| 164 |
+
max_pixels=size.longest_edge,
|
| 165 |
+
)
|
| 166 |
+
return super().resize(
|
| 167 |
+
image=images,
|
| 168 |
+
size=SizeDict(height=resized_height, width=resized_width),
|
| 169 |
+
resample=resample,
|
| 170 |
+
)
|
| 171 |
+
|
| 172 |
+
def patchify(
|
| 173 |
+
self,
|
| 174 |
+
images: "torch.Tensor",
|
| 175 |
+
patch_size: int,
|
| 176 |
+
merge_size: int,
|
| 177 |
+
temporal_patch_size: int,
|
| 178 |
+
) -> tuple["torch.Tensor", int, int]:
|
| 179 |
+
"""Patchifies each image into flat layout of shape (`seq_len`, `patch_dim`) so we can concat dynamically shaped pixels."""
|
| 180 |
+
batch_size, channel, resized_height, resized_width = images.shape
|
| 181 |
+
grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
|
| 182 |
+
patches = images.reshape(
|
| 183 |
+
batch_size,
|
| 184 |
+
channel,
|
| 185 |
+
grid_h // merge_size,
|
| 186 |
+
merge_size,
|
| 187 |
+
patch_size,
|
| 188 |
+
grid_w // merge_size,
|
| 189 |
+
merge_size,
|
| 190 |
+
patch_size,
|
| 191 |
+
)
|
| 192 |
+
# Reorder dimensions to group grid and patch information for subsequent flattening.
|
| 193 |
+
# [batch, grid_h/merge, grid_w/merge, merge, merge, channel, patch, patch]
|
| 194 |
+
patches = patches.permute(0, 2, 5, 3, 6, 1, 4, 7)
|
| 195 |
+
|
| 196 |
+
flatten_patches = (
|
| 197 |
+
patches.unsqueeze(6)
|
| 198 |
+
.expand(-1, -1, -1, -1, -1, -1, temporal_patch_size, -1, -1)
|
| 199 |
+
.reshape(
|
| 200 |
+
batch_size,
|
| 201 |
+
grid_h * grid_w,
|
| 202 |
+
channel * temporal_patch_size * patch_size * patch_size,
|
| 203 |
+
)
|
| 204 |
+
)
|
| 205 |
+
return flatten_patches, grid_h, grid_w
|
| 206 |
+
|
| 207 |
+
def _preprocess(
|
| 208 |
+
self,
|
| 209 |
+
images: list["torch.Tensor"],
|
| 210 |
+
do_resize: bool,
|
| 211 |
+
size: SizeDict,
|
| 212 |
+
resample: "PILImageResampling | tvF.InterpolationMode | int | None",
|
| 213 |
+
do_rescale: bool,
|
| 214 |
+
rescale_factor: float,
|
| 215 |
+
do_normalize: bool,
|
| 216 |
+
image_mean: float | list[float] | None,
|
| 217 |
+
image_std: float | list[float] | None,
|
| 218 |
+
patch_size: int,
|
| 219 |
+
temporal_patch_size: int,
|
| 220 |
+
merge_size: int,
|
| 221 |
+
disable_grouping: bool | None,
|
| 222 |
+
return_tensors: str | TensorType | None,
|
| 223 |
+
**kwargs,
|
| 224 |
+
) -> BatchFeature:
|
| 225 |
+
grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
|
| 226 |
+
resized_images_grouped = {}
|
| 227 |
+
for shape, stacked_images in grouped_images.items():
|
| 228 |
+
if do_resize:
|
| 229 |
+
stacked_images = self.resize(
|
| 230 |
+
images=stacked_images,
|
| 231 |
+
size=size,
|
| 232 |
+
resample=resample,
|
| 233 |
+
factor=patch_size * merge_size,
|
| 234 |
+
)
|
| 235 |
+
resized_images_grouped[shape] = stacked_images
|
| 236 |
+
resized_images = reorder_images(resized_images_grouped, grouped_images_index)
|
| 237 |
+
|
| 238 |
+
grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
|
| 239 |
+
processed_images_grouped = {}
|
| 240 |
+
processed_grids = {}
|
| 241 |
+
for shape, stacked_images in grouped_images.items():
|
| 242 |
+
stacked_images = self.rescale_and_normalize(
|
| 243 |
+
stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
|
| 244 |
+
)
|
| 245 |
+
patches, grid_h, grid_w = self.patchify(
|
| 246 |
+
stacked_images,
|
| 247 |
+
patch_size=patch_size,
|
| 248 |
+
merge_size=merge_size,
|
| 249 |
+
temporal_patch_size=temporal_patch_size,
|
| 250 |
+
)
|
| 251 |
+
|
| 252 |
+
processed_images_grouped[shape] = patches
|
| 253 |
+
processed_grids[shape] = [[1, grid_h, grid_w]] * len(stacked_images)
|
| 254 |
+
|
| 255 |
+
processed_images = reorder_images(processed_images_grouped, grouped_images_index)
|
| 256 |
+
processed_grids_ordered = reorder_images(processed_grids, grouped_images_index)
|
| 257 |
+
pixel_values = processed_images[0] if len(processed_images) == 1 else torch.cat(processed_images, dim=0)
|
| 258 |
+
image_grid_thw = torch.tensor(processed_grids_ordered, dtype=torch.long)
|
| 259 |
+
|
| 260 |
+
return BatchFeature(
|
| 261 |
+
data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, tensor_type=return_tensors
|
| 262 |
+
)
|
| 263 |
+
|
| 264 |
+
def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict | None = None) -> int:
|
| 265 |
+
"""
|
| 266 |
+
A utility that returns number of image patches for a given image size.
|
| 267 |
+
|
| 268 |
+
Note: Do not remove this method! It is used by vLLM to infer the number of patches and placeholders
|
| 269 |
+
without an image input.
|
| 270 |
+
|
| 271 |
+
Args:
|
| 272 |
+
height (`int`):
|
| 273 |
+
Height of the input image.
|
| 274 |
+
width (`int`):
|
| 275 |
+
Width of the input image.
|
| 276 |
+
images_kwargs (`dict`, *optional*)
|
| 277 |
+
Any kwargs to override defaults of the image processor.
|
| 278 |
+
Returns:
|
| 279 |
+
`int`: Number of image patches per image.
|
| 280 |
+
"""
|
| 281 |
+
# `size` is the supported override (`min_pixels` / `max_pixels` are deprecated): count with the budget the
|
| 282 |
+
# preprocessing applies
|
| 283 |
+
images_kwargs = dict(images_kwargs or {})
|
| 284 |
+
size = images_kwargs.pop("size", None) or self.size
|
| 285 |
+
images_kwargs.setdefault("min_pixels", size["shortest_edge"])
|
| 286 |
+
images_kwargs.setdefault("max_pixels", size["longest_edge"])
|
| 287 |
+
images_kwargs = images_kwargs or {}
|
| 288 |
+
min_pixels = images_kwargs["min_pixels"] if "min_pixels" in images_kwargs else self.size["shortest_edge"]
|
| 289 |
+
max_pixels = images_kwargs["max_pixels"] if "max_pixels" in images_kwargs else self.size["longest_edge"]
|
| 290 |
+
patch_size = images_kwargs.get("patch_size", self.patch_size)
|
| 291 |
+
merge_size = images_kwargs.get("merge_size", self.merge_size)
|
| 292 |
+
|
| 293 |
+
factor = patch_size * merge_size
|
| 294 |
+
resized_height, resized_width = smart_resize(
|
| 295 |
+
height, width, factor, min_pixels=min_pixels, max_pixels=max_pixels
|
| 296 |
+
)
|
| 297 |
+
grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
|
| 298 |
+
return grid_h * grid_w
|
| 299 |
+
|
| 300 |
+
|
| 301 |
+
__all__ = ["Lfm3VlImageProcessor"]
|
integrations_lfm3.py
ADDED
|
@@ -0,0 +1,204 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Vendored from transformers' LFM3 code (not in a release yet) at 69da1b47a11e362fbb15bfe5d5c8da0790c688c5, src/transformers/integrations/lfm3.py:
|
| 2 |
+
# its relative imports made absolute, nothing else changed.
|
| 3 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 4 |
+
#
|
| 5 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 6 |
+
# you may not use this file except in compliance with the License.
|
| 7 |
+
# You may obtain a copy of the License at
|
| 8 |
+
#
|
| 9 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 10 |
+
#
|
| 11 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 12 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 13 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 14 |
+
# See the License for the specific language governing permissions and
|
| 15 |
+
# limitations under the License.
|
| 16 |
+
"""Inference kernels for LFM3's width-7 depthwise convolution."""
|
| 17 |
+
|
| 18 |
+
import torch
|
| 19 |
+
|
| 20 |
+
from transformers.utils import is_triton_available
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
if is_triton_available():
|
| 24 |
+
import triton
|
| 25 |
+
import triton.language as tl
|
| 26 |
+
|
| 27 |
+
# the length and strides are runtime arguments, or every prompt length would compile its own kernel
|
| 28 |
+
@triton.jit(do_not_specialize=["T", "X0", "X1", "X2", "S0", "S1", "ST"])
|
| 29 |
+
def _prefill(
|
| 30 |
+
X,
|
| 31 |
+
W,
|
| 32 |
+
BIAS,
|
| 33 |
+
SEQ,
|
| 34 |
+
Y,
|
| 35 |
+
T,
|
| 36 |
+
D: tl.constexpr,
|
| 37 |
+
X0,
|
| 38 |
+
X1,
|
| 39 |
+
X2,
|
| 40 |
+
S0,
|
| 41 |
+
S1,
|
| 42 |
+
ST,
|
| 43 |
+
HAS_BIAS: tl.constexpr,
|
| 44 |
+
PACKED: tl.constexpr,
|
| 45 |
+
BLOCK: tl.constexpr,
|
| 46 |
+
):
|
| 47 |
+
# int64 offsets: batch * channels * length can exceed 2**31 elements
|
| 48 |
+
b, d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
|
| 49 |
+
t = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
|
| 50 |
+
acc = tl.full((BLOCK,), 0, tl.float32)
|
| 51 |
+
if PACKED:
|
| 52 |
+
doc = tl.load(SEQ + b * S0 + (t - T + ST) * S1, (t < T) & (t >= T - ST), other=-1)
|
| 53 |
+
for k in tl.static_range(7):
|
| 54 |
+
pos = t + k - 6
|
| 55 |
+
valid = (t < T) & (pos >= 0)
|
| 56 |
+
if PACKED:
|
| 57 |
+
prev = tl.load(SEQ + b * S0 + (pos - T + ST) * S1, valid & (pos >= T - ST), other=-1)
|
| 58 |
+
valid = valid & (doc == prev)
|
| 59 |
+
x = tl.load(X + b * X0 + d * X1 + pos * X2, valid, other=0).to(tl.float32)
|
| 60 |
+
w = tl.load(W + d * 7 + k).to(tl.float32)
|
| 61 |
+
acc += x * w
|
| 62 |
+
if HAS_BIAS:
|
| 63 |
+
acc += tl.load(BIAS + d).to(tl.float32)
|
| 64 |
+
tl.store(Y + (b * D + d) * T + t, acc, t < T)
|
| 65 |
+
|
| 66 |
+
@triton.jit
|
| 67 |
+
def _decode(
|
| 68 |
+
X,
|
| 69 |
+
W,
|
| 70 |
+
BIAS,
|
| 71 |
+
STATE,
|
| 72 |
+
Y,
|
| 73 |
+
D: tl.constexpr,
|
| 74 |
+
K: tl.constexpr,
|
| 75 |
+
X0: tl.constexpr,
|
| 76 |
+
X1: tl.constexpr,
|
| 77 |
+
S0: tl.constexpr,
|
| 78 |
+
S1: tl.constexpr,
|
| 79 |
+
S2: tl.constexpr,
|
| 80 |
+
HAS_BIAS: tl.constexpr,
|
| 81 |
+
BLOCK: tl.constexpr,
|
| 82 |
+
):
|
| 83 |
+
b = tl.program_id(0).to(tl.int64)
|
| 84 |
+
d = tl.program_id(1).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
|
| 85 |
+
k = tl.arange(0, 8)
|
| 86 |
+
# Load the complete state before writing the shifted window.
|
| 87 |
+
old = tl.load(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, (d[:, None] < D) & (k[None, :] < K), other=0)
|
| 88 |
+
x = tl.load(X + b * X0 + d * X1, d < D, other=0)
|
| 89 |
+
shifted = tl.gather(old, tl.broadcast_to(tl.minimum(k + 1, 7)[None, :], (BLOCK, 8)), 1)
|
| 90 |
+
shifted = tl.where(k[None, :] == K - 1, x[:, None], shifted)
|
| 91 |
+
# A cache may store either K=6 previous inputs or K=7 including the oldest discarded input.
|
| 92 |
+
positions = k + K - 6
|
| 93 |
+
values = tl.gather(old, tl.broadcast_to(tl.minimum(positions, 7)[None, :], (BLOCK, 8)), 1)
|
| 94 |
+
values = tl.where(k[None, :] == 6, x[:, None], values)
|
| 95 |
+
weights = tl.load(W + d[:, None] * 7 + k[None, :], (d[:, None] < D) & (k[None, :] < 7), other=0)
|
| 96 |
+
acc = tl.sum(values.to(tl.float32) * weights.to(tl.float32), 1)
|
| 97 |
+
if HAS_BIAS:
|
| 98 |
+
acc += tl.load(BIAS + d, d < D, other=0).to(tl.float32)
|
| 99 |
+
tl.store(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, shifted, (d[:, None] < D) & (k[None, :] < K))
|
| 100 |
+
tl.store(Y + b * D + d, acc, d < D)
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def can_use_lfm3_conv(hidden_states, weight):
|
| 104 |
+
return (
|
| 105 |
+
is_triton_available()
|
| 106 |
+
and hasattr(tl, "gather") # the decode kernel needs Triton >= 3.3
|
| 107 |
+
and hidden_states.is_cuda
|
| 108 |
+
and not torch.is_grad_enabled()
|
| 109 |
+
and hidden_states.dtype == weight.dtype
|
| 110 |
+
and hidden_states.dtype in (torch.float16, torch.bfloat16, torch.float32)
|
| 111 |
+
and weight.shape[-1] == 7
|
| 112 |
+
)
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
def lfm3_causal_conv1d(hidden_states, weight, bias=None, seq_idx=None):
|
| 116 |
+
batch, channels, length = hidden_states.shape
|
| 117 |
+
# the kernels index every tensor by the batch and channel of `hidden_states` and cannot bound the others
|
| 118 |
+
if weight.shape != (channels, 7) or (bias is not None and bias.shape != (channels,)):
|
| 119 |
+
raise ValueError(
|
| 120 |
+
f"`weight` must have shape ({channels}, 7) and `bias` ({channels},), got {tuple(weight.shape)}"
|
| 121 |
+
)
|
| 122 |
+
if seq_idx is not None and (
|
| 123 |
+
seq_idx.ndim != 2 or seq_idx.shape[0] not in (1, batch) or seq_idx.dtype not in (torch.int32, torch.int64)
|
| 124 |
+
):
|
| 125 |
+
raise ValueError(
|
| 126 |
+
f"`seq_idx` must be an int32 or int64 tensor of shape (1, length) or ({batch}, length), got "
|
| 127 |
+
f"{seq_idx.dtype} {tuple(seq_idx.shape)}"
|
| 128 |
+
)
|
| 129 |
+
if any(t is not None and t.device != hidden_states.device for t in (weight, bias, seq_idx)):
|
| 130 |
+
raise ValueError(
|
| 131 |
+
f"`weight`, `bias` and `seq_idx` must be on the device of `hidden_states` ({hidden_states.device})"
|
| 132 |
+
)
|
| 133 |
+
output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
|
| 134 |
+
# Triton launches on the current device; the model may live on another one (`device_map`, `.to("cuda:1")`).
|
| 135 |
+
with torch.cuda.device(hidden_states.device):
|
| 136 |
+
_prefill[(batch, channels, triton.cdiv(length, 128))](
|
| 137 |
+
hidden_states,
|
| 138 |
+
# the kernels index `bias` and `weight` densely
|
| 139 |
+
weight.contiguous(),
|
| 140 |
+
bias.contiguous() if bias is not None else None,
|
| 141 |
+
seq_idx,
|
| 142 |
+
output,
|
| 143 |
+
length,
|
| 144 |
+
channels,
|
| 145 |
+
*hidden_states.stride(),
|
| 146 |
+
*(
|
| 147 |
+
(0 if seq_idx.shape[0] == 1 else seq_idx.stride(0), seq_idx.stride(1))
|
| 148 |
+
if seq_idx is not None
|
| 149 |
+
else (0, 0)
|
| 150 |
+
),
|
| 151 |
+
seq_idx.shape[-1] if seq_idx is not None else 0,
|
| 152 |
+
bias is not None,
|
| 153 |
+
seq_idx is not None,
|
| 154 |
+
128,
|
| 155 |
+
enable_fp_fusion=False,
|
| 156 |
+
)
|
| 157 |
+
return output
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
def lfm3_causal_conv1d_update(hidden_states, conv_state, weight, bias=None):
|
| 161 |
+
batch, channels, length = hidden_states.shape
|
| 162 |
+
# the kernel addresses one position per row and bounds nothing by the other tensors' shapes
|
| 163 |
+
if length != 1:
|
| 164 |
+
raise ValueError(f"`hidden_states` must hold one position per row, got {tuple(hidden_states.shape)}")
|
| 165 |
+
# the 8-wide tile holds a 6- or 7-wide state (see `positions` in `_decode`)
|
| 166 |
+
if (
|
| 167 |
+
conv_state.ndim != 3
|
| 168 |
+
or conv_state.shape[:2] != (batch, channels)
|
| 169 |
+
or conv_state.shape[-1] not in (6, 7)
|
| 170 |
+
or weight.shape != (channels, 7)
|
| 171 |
+
):
|
| 172 |
+
raise ValueError(
|
| 173 |
+
f"`conv_state` must have shape ({batch}, {channels}, 6 or 7) and `weight` ({channels}, 7), got "
|
| 174 |
+
f"{tuple(conv_state.shape)} and {tuple(weight.shape)}"
|
| 175 |
+
)
|
| 176 |
+
if bias is not None and bias.shape != (channels,):
|
| 177 |
+
raise ValueError(f"`bias` must have shape ({channels},), got {tuple(bias.shape)}")
|
| 178 |
+
# the kernel writes the state in place: rows or channels sharing memory (an `expand`ed cache) would race
|
| 179 |
+
extent = 0
|
| 180 |
+
for stride, size in sorted(zip(conv_state.stride(), conv_state.shape)):
|
| 181 |
+
if size > 1 and stride <= extent:
|
| 182 |
+
raise ValueError(f"`conv_state` must not overlap itself, got strides {conv_state.stride()}")
|
| 183 |
+
extent += stride * (size - 1)
|
| 184 |
+
if any(t is not None and t.device != hidden_states.device for t in (conv_state, weight, bias)):
|
| 185 |
+
raise ValueError(
|
| 186 |
+
f"`conv_state`, `weight` and `bias` must be on the device of `hidden_states` ({hidden_states.device})"
|
| 187 |
+
)
|
| 188 |
+
output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
|
| 189 |
+
with torch.cuda.device(hidden_states.device):
|
| 190 |
+
_decode[(batch, triton.cdiv(channels, 32))](
|
| 191 |
+
hidden_states,
|
| 192 |
+
weight.contiguous(),
|
| 193 |
+
bias.contiguous() if bias is not None else None,
|
| 194 |
+
conv_state,
|
| 195 |
+
output,
|
| 196 |
+
channels,
|
| 197 |
+
conv_state.shape[-1],
|
| 198 |
+
*hidden_states.stride()[:2],
|
| 199 |
+
*conv_state.stride(),
|
| 200 |
+
bias is not None,
|
| 201 |
+
32,
|
| 202 |
+
enable_fp_fusion=False,
|
| 203 |
+
)
|
| 204 |
+
return output
|
json_spans.py
ADDED
|
@@ -0,0 +1,59 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Spans of a rendered JSON state: where its strings sit and what an escaped slice of one says, without torch."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import json
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def string_paths(text: str) -> list[tuple[int, int, str]]:
|
| 9 |
+
"""String-value spans in rendered JSON, excluding quotes; paths use RFC6901 escaping."""
|
| 10 |
+
decoder = json.JSONDecoder()
|
| 11 |
+
ranges = []
|
| 12 |
+
|
| 13 |
+
def skip(at):
|
| 14 |
+
while at < len(text) and text[at].isspace():
|
| 15 |
+
at += 1
|
| 16 |
+
return at
|
| 17 |
+
|
| 18 |
+
def value(at, path):
|
| 19 |
+
at = skip(at)
|
| 20 |
+
if text[at] == "{":
|
| 21 |
+
at = skip(at + 1)
|
| 22 |
+
while text[at] != "}":
|
| 23 |
+
key, end = decoder.raw_decode(text, at)
|
| 24 |
+
if not isinstance(key, str) or text[skip(end)] != ":":
|
| 25 |
+
raise ValueError("invalid JSON object")
|
| 26 |
+
name = key.replace("~", "~0").replace("/", "~1")
|
| 27 |
+
at = skip(value(skip(end) + 1, path + "/" + name))
|
| 28 |
+
if text[at] == ",":
|
| 29 |
+
at = skip(at + 1)
|
| 30 |
+
elif text[at] != "}":
|
| 31 |
+
raise ValueError("invalid JSON object separator")
|
| 32 |
+
return at + 1
|
| 33 |
+
if text[at] == "[":
|
| 34 |
+
at, index = skip(at + 1), 0
|
| 35 |
+
while text[at] != "]":
|
| 36 |
+
at = skip(value(at, path + "/" + str(index)))
|
| 37 |
+
index += 1
|
| 38 |
+
if text[at] == ",":
|
| 39 |
+
at = skip(at + 1)
|
| 40 |
+
elif text[at] != "]":
|
| 41 |
+
raise ValueError("invalid JSON array separator")
|
| 42 |
+
return at + 1
|
| 43 |
+
item, end = decoder.raw_decode(text, at)
|
| 44 |
+
if isinstance(item, str):
|
| 45 |
+
ranges.append((at + 1, end - 1, path))
|
| 46 |
+
return end
|
| 47 |
+
|
| 48 |
+
end = value(0, "")
|
| 49 |
+
if skip(end) != len(text):
|
| 50 |
+
raise ValueError("trailing JSON text")
|
| 51 |
+
return ranges
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def json_text(escaped: str) -> str:
|
| 55 |
+
"""A slice of a rendered JSON string as the characters it encodes (unchanged when it cuts an escape)."""
|
| 56 |
+
try:
|
| 57 |
+
return json.loads(f'"{escaped}"')
|
| 58 |
+
except ValueError:
|
| 59 |
+
return escaped
|
lfm3.py
ADDED
|
@@ -0,0 +1,201 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""LFM3: text (`lfm3`), vision-language (`lfm3_vl`) and omni (`lfm3_omni`) checkpoints on transformers' classes
|
| 2 |
+
(`native_lfm3*`), loaded as their Hub class (`modeling_d1_lfm3*.D1Model`) with d1's attention. d1 adds only:
|
| 3 |
+
|
| 4 |
+
- the masks: `attention.d1_attention` under `attn_implementation="d1"`, its mask function registered;
|
| 5 |
+
- the short convolution's torch path as `depthwise.depthwise` (ROCm: no MIOpen search), taken with autograd and
|
| 6 |
+
under torch.compile (`lfm3_conv`); eager inference keeps the classes' fused kernel; the audio encoder's
|
| 7 |
+
convolutions off MIOpen on ROCm (`lfm3_conv`);
|
| 8 |
+
- the media encoders `SystemOne.compile` compiles beside the language model (`compiled_parts`), and its inductor
|
| 9 |
+
options (`compile_options`); the vision tower in one graph for every picture size (`lfm3_tower`).
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
from __future__ import annotations
|
| 13 |
+
|
| 14 |
+
import functools
|
| 15 |
+
import importlib
|
| 16 |
+
import importlib.util
|
| 17 |
+
import sys
|
| 18 |
+
|
| 19 |
+
import torch
|
| 20 |
+
import torch.nn.functional as F
|
| 21 |
+
|
| 22 |
+
from .depthwise import depthwise, pinned_conv
|
| 23 |
+
from .modeling_d1 import SystemOneModel
|
| 24 |
+
from .prompt import LEADS
|
| 25 |
+
|
| 26 |
+
MODULES = {"lfm3": "modeling_d1_lfm3", "lfm3_vl": "modeling_d1_lfm3_vl", "lfm3_omni": "modeling_d1_lfm3_omni"}
|
| 27 |
+
# the templates open a reasoning block; a non-thinking reply follows it closed at once
|
| 28 |
+
LEADS.update(dict.fromkeys(MODULES, "<think></think>"))
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
class Lfm3SystemOne(SystemOneModel):
|
| 32 |
+
"""What every LFM3 Hub class adds to transformers': d1's attention and its convolution."""
|
| 33 |
+
|
| 34 |
+
attention = "d1"
|
| 35 |
+
compiled_parts: tuple[str, ...] = () # the media encoders `SystemOne.compile` compiles beside the language model
|
| 36 |
+
# compiled for inference with the weights as constants (inductor's freezing): over the fused projections
|
| 37 |
+
# (`fuse_projections`) d1-2B on an MI325X reads 3 questions in 5.5 against 6.2 ms, a 5k-token state in 41.8
|
| 38 |
+
# against 43.3, 64 packed states at 1,427 against 1,390 a second, at the same peak memory (3.65 GiB)
|
| 39 |
+
compile_options = {"freezing": True}
|
| 40 |
+
# the patch grids stay on the host: the tower and the image features read their sizes there, without a sync
|
| 41 |
+
host_inputs = ("image_grid_thw", "video_grid_thw")
|
| 42 |
+
|
| 43 |
+
def __init__(self, config):
|
| 44 |
+
super().__init__(config)
|
| 45 |
+
lfm3_conv(self)
|
| 46 |
+
lfm3_tower(self)
|
| 47 |
+
|
| 48 |
+
def fuse_projections(self) -> None:
|
| 49 |
+
"""For inference (`SystemOne`): the linears that read one input as one matmul (`fuse_projections`)."""
|
| 50 |
+
fuse_projections(self)
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def lfm3_conv(model) -> None:
|
| 54 |
+
"""Each LFM3 short convolution's torch path (`Lfm3ShortConv._conv`) as `depthwise`, and that path under
|
| 55 |
+
torch.compile: inductor fuses its shifted sums with the gates around them, where the classes' fused kernel is a
|
| 56 |
+
call of its own (d1-2B on an MI325X, compiled: 49.0 against 43.2 ms on a 5k-token state, 39.2 against 36.0 ms
|
| 57 |
+
for 64 short states). Eager inference keeps the kernel (10.2 against 11.3 ms a question).
|
| 58 |
+
|
| 59 |
+
The omni's audio encoder, off MIOpen on ROCm as well: its LFM2 convolutions (transformers' `Lfm2ShortConv`, which
|
| 60 |
+
calls its modeling module's `causal_conv1d_fn`, `F.conv1d` without causal-conv1d) as `depthwise`, its downsampler's
|
| 61 |
+
as `pinned_conv`. MIOpen's solver for the downsampler's 1x1 `lin_out` changed from one call to the next, and a
|
| 62 |
+
clip's P(yes) by 3.1e-03 (d1-omni-2B on an MI325X)."""
|
| 63 |
+
for m in model.modules():
|
| 64 |
+
module = sys.modules[type(m).__module__]
|
| 65 |
+
if type(m).__name__ == "Lfm3ShortConv":
|
| 66 |
+
m._conv = functools.partial(_causal_depthwise, m.conv)
|
| 67 |
+
if not getattr(module.can_use_lfm3_conv, "eager_only", False):
|
| 68 |
+
module.can_use_lfm3_conv = _eager_only(module.can_use_lfm3_conv)
|
| 69 |
+
elif type(m).__name__ == "Lfm2ShortConv" and importlib.util.find_spec("causal_conv1d") is None:
|
| 70 |
+
module.causal_conv1d_fn = _lfm2_conv
|
| 71 |
+
elif type(m).__name__ == "Lfm3OmniDownsampler":
|
| 72 |
+
for conv in m.children():
|
| 73 |
+
if isinstance(conv, torch.nn.modules.conv._ConvNd):
|
| 74 |
+
conv.forward = functools.partial(pinned_conv, conv)
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def lfm3_tower(model) -> None:
|
| 78 |
+
"""The vision tower under torch.compile, one graph for every picture size and no graph break:
|
| 79 |
+
|
| 80 |
+
- what its forward derives from the patch grids (positions, `cu_seqlens`, the frames' sizes: a sync and a graph
|
| 81 |
+
break each), made before it on the host once a grid, the way transformers' `vision_utils` takes them
|
| 82 |
+
precomputed; the frames' sizes come in as the grid itself, a list (`_frame_shapes`);
|
| 83 |
+
- its `len(set(frame sizes)) == 1` (one batch for frames of one size) compared by `==`: hashing the sizes makes
|
| 84 |
+
torch.compile guard on their values, a recompile for every picture size (about 10 s on an MI325X).
|
| 85 |
+
|
| 86 |
+
Delete both module patches when transformers' tower compares the sizes and takes the frames' sizes precomputed."""
|
| 87 |
+
for m in model.modules():
|
| 88 |
+
if type(m).__name__ == "Lfm3VlVisionModel":
|
| 89 |
+
m.register_forward_pre_hook(_grid_inputs, with_kwargs=True)
|
| 90 |
+
module = sys.modules[type(m).__module__]
|
| 91 |
+
module.set, module.get_vision_frame_shapes = _distinct, _frame_shapes
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def _grid_inputs(tower, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
|
| 95 |
+
"""The grid's positions, `cu_seqlens`, longest frame and frame sizes, made once a grid."""
|
| 96 |
+
from transformers.vision_utils import get_max_seqlen, get_vision_cu_seqlens, get_vision_position_ids
|
| 97 |
+
|
| 98 |
+
args = list(args)
|
| 99 |
+
grid = (args.pop(1) if len(args) > 1 else kwargs["grid_thw"]).cpu() # on the host already when served
|
| 100 |
+
device = (args[0] if args else kwargs["pixel_values"]).device
|
| 101 |
+
made = tower.__dict__.setdefault("d1_grid_inputs", {})
|
| 102 |
+
key = (tuple(grid.flatten().tolist()), device)
|
| 103 |
+
if key not in made:
|
| 104 |
+
if len(made) >= 256:
|
| 105 |
+
made.clear()
|
| 106 |
+
cu_seqlens = get_vision_cu_seqlens(grid)
|
| 107 |
+
made[key] = {"cu_seqlens": cu_seqlens.to(device), "max_seqlen": get_max_seqlen(cu_seqlens, tower.config),
|
| 108 |
+
"position_ids": get_vision_position_ids(grid, tower.spatial_merge_size).to(device),
|
| 109 |
+
"grid_thw": _grid_frames(grid)}
|
| 110 |
+
kwargs.update(made[key])
|
| 111 |
+
return tuple(args), kwargs
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def _frame_shapes(grid) -> list[list[int]]:
|
| 115 |
+
"""`get_vision_frame_shapes`: the grid `_grid_inputs` passes is its frames' sizes already; a tower without the
|
| 116 |
+
hook (transformers' own model) passes the tensor."""
|
| 117 |
+
return grid if isinstance(grid, list) else _grid_frames(grid)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def _grid_frames(grid: torch.Tensor) -> list[list[int]]:
|
| 121 |
+
return grid[:, 1:].repeat_interleave(grid[:, 0], dim=0).tolist()
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
# the linears each LFM3 module calls on one input, in the order its forward calls them
|
| 125 |
+
SHARED_INPUT = {"Lfm3Attention": ("q_proj", "k_proj", "v_proj", "gate_proj"), "Lfm3MLP": ("w1", "w3"),
|
| 126 |
+
"Lfm3VlVisionAttention": ("q_proj", "k_proj", "v_proj"), "Lfm3VlMLP": ("w1", "w3")}
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
def fuse_projections(model) -> None:
|
| 130 |
+
"""Each module's linears on one input (`SHARED_INPUT`) as one matmul, its output split: the first's call computes
|
| 131 |
+
every part, the others read theirs, the last drops them. Exact, one weight copy, the classes untouched, the
|
| 132 |
+
state dict the native one (a fused model saves as it loaded). For inference only: training keeps the separate
|
| 133 |
+
linears, and a group with any layer not a plain linear stays."""
|
| 134 |
+
for m in model.modules():
|
| 135 |
+
names = SHARED_INPUT.get(type(m).__name__, ())
|
| 136 |
+
linears = [getattr(m, n, None) for n in names]
|
| 137 |
+
if names and all(type(lin) is torch.nn.Linear and lin.bias is None for lin in linears):
|
| 138 |
+
m.fused = Fused(torch.cat([lin.weight for lin in linears]), names, [lin.out_features for lin in linears])
|
| 139 |
+
for i, name in enumerate(names):
|
| 140 |
+
setattr(m, name, Part(m.fused, i))
|
| 141 |
+
m.register_state_dict_post_hook(_unfused)
|
| 142 |
+
|
| 143 |
+
|
| 144 |
+
def _unfused(module, state_dict, prefix, local_metadata) -> None:
|
| 145 |
+
fused = module.fused
|
| 146 |
+
weight = state_dict.pop(prefix + "fused.weight")
|
| 147 |
+
for name, part in zip(fused.names, weight.split(fused.sizes)):
|
| 148 |
+
state_dict[f"{prefix}{name}.weight"] = part.clone()
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
class Fused(torch.nn.Module):
|
| 152 |
+
def __init__(self, weight: torch.Tensor, names: tuple[str, ...], sizes: list[int]):
|
| 153 |
+
super().__init__()
|
| 154 |
+
self.weight = torch.nn.Parameter(weight, requires_grad=False)
|
| 155 |
+
self.names, self.sizes, self.parts = names, sizes, None
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
class Part(torch.nn.Module):
|
| 159 |
+
"""One linear of a `Fused` group: the first part computes the group's matmul."""
|
| 160 |
+
|
| 161 |
+
def __init__(self, fused: Fused, index: int):
|
| 162 |
+
super().__init__()
|
| 163 |
+
self.group, self.index, self.last = [fused], index, index == len(fused.sizes) - 1 # a list: not a submodule
|
| 164 |
+
|
| 165 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 166 |
+
fused = self.group[0]
|
| 167 |
+
if self.index == 0:
|
| 168 |
+
fused.parts = F.linear(x, fused.weight).split(fused.sizes, dim=-1)
|
| 169 |
+
out = fused.parts[self.index]
|
| 170 |
+
if self.last:
|
| 171 |
+
fused.parts = None
|
| 172 |
+
return out
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
def _distinct(items) -> list:
|
| 176 |
+
"""`set(items)`'s members by `==`, without hashing them."""
|
| 177 |
+
seen = []
|
| 178 |
+
for item in items:
|
| 179 |
+
if all(item != other for other in seen):
|
| 180 |
+
seen.append(item)
|
| 181 |
+
return seen
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
def _eager_only(kernel):
|
| 185 |
+
def can_use(hidden_states, weight) -> bool:
|
| 186 |
+
return not torch.compiler.is_compiling() and kernel(hidden_states, weight)
|
| 187 |
+
|
| 188 |
+
can_use.eager_only = True
|
| 189 |
+
return can_use
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
def _causal_depthwise(conv: torch.nn.Conv1d, x: torch.Tensor) -> torch.Tensor:
|
| 193 |
+
"""`conv` over `(B, C, T)`, causal: what `Lfm3ShortConv._conv` computes."""
|
| 194 |
+
return _lfm2_conv(x, conv.weight[:, 0], conv.bias)
|
| 195 |
+
|
| 196 |
+
|
| 197 |
+
def _lfm2_conv(x: torch.Tensor, weight: torch.Tensor, bias=None, seq_idx=None) -> torch.Tensor:
|
| 198 |
+
"""transformers' `causal_conv1d_fn` as `Lfm2ShortConv` calls it: `weight` `(C, k)`, no activation, `seq_idx`
|
| 199 |
+
unused, as in transformers' own torch path."""
|
| 200 |
+
out = depthwise(F.pad(x.to(weight.dtype), (weight.shape[-1] - 1, 0)), weight[:, None])
|
| 201 |
+
return (out if bias is None else out + bias[:, None]).to(x.dtype)
|
media.py
ADDED
|
@@ -0,0 +1,252 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Media: pictures, audio clips and videos, each placed in a prompt as its checkpoint's processor expands it.
|
| 2 |
+
|
| 3 |
+
media=[picture, {"type": "audio", "path": "call.wav"}, {"type": "video", "path": "clip.mp4"}]
|
| 4 |
+
{"type": "noul", "instructions": {"type": "audio", "path": "question.wav"}} # a spoken question
|
| 5 |
+
|
| 6 |
+
- An item is a PIL image (a picture), or `{"type": "image" | "audio" | "video", "path": ...}`, the decoded value under
|
| 7 |
+
its type instead of a path (`{"type": "audio", "audio": waveform, "sampling_rate": 16000}`, a PIL image, a video's
|
| 8 |
+
frames with their `fps`). A spoken state is `state=None` with its clip in `media`.
|
| 9 |
+
- A request's media head its state, in order, as the chat template writes them; a question's spoken instructions sit
|
| 10 |
+
where its text would. Each item's markup is the processor's own expansion, so a prompt is plain text and is packed,
|
| 11 |
+
read both ways and pointed into as a text prompt is (`Reader`).
|
| 12 |
+
- Pictures are capped at the checkpoint's pixel budget first (`cap_pixels`); audio is read as 16-bit PCM WAV or a
|
| 13 |
+
NumPy file and resampled to the feature extractor's rate; a decoded stereo waveform may be (channels, samples) or
|
| 14 |
+
(samples, channels), as soundfile returns it, and its channels are averaged.
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
from __future__ import annotations
|
| 18 |
+
|
| 19 |
+
import math
|
| 20 |
+
from collections.abc import Mapping, Sequence
|
| 21 |
+
from dataclasses import dataclass
|
| 22 |
+
from typing import Any
|
| 23 |
+
|
| 24 |
+
KINDS = ("image", "audio", "video")
|
| 25 |
+
PROCESSORS = {"image": "image_processor", "audio": "feature_extractor", "video": "video_processor"}
|
| 26 |
+
INPUTS = {"image": "images", "audio": "audio", "video": "videos"} # the processor's argument for each kind
|
| 27 |
+
SPOKEN = "" # a spoken question's place in its suffix until its markup replaces it (`Reader.place`)
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
@dataclass(frozen=True, eq=False)
|
| 31 |
+
class Media:
|
| 32 |
+
"""One picture, audio clip or video: a path, or its decoded value (`sampling_rate`: a waveform's, `fps`: a video's
|
| 33 |
+
frames')."""
|
| 34 |
+
|
| 35 |
+
type: str
|
| 36 |
+
value: Any
|
| 37 |
+
sampling_rate: int | None = None
|
| 38 |
+
fps: float | None = None
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def as_media(item: Any) -> Media:
|
| 42 |
+
"""A caller's item (a PIL image, or a dict with its `type`), as `Media`; `Media` passes through."""
|
| 43 |
+
if isinstance(item, Media):
|
| 44 |
+
return item
|
| 45 |
+
if isinstance(item, Mapping):
|
| 46 |
+
kind = item.get("type")
|
| 47 |
+
if kind not in KINDS or kind not in item and "path" not in item:
|
| 48 |
+
raise ValueError(f"a media item is a picture or {{'type': one of {KINDS}, 'path' or <type>: its value}}")
|
| 49 |
+
return Media(kind, item[kind] if kind in item else item["path"], item.get("sampling_rate"), item.get("fps"))
|
| 50 |
+
if hasattr(item, "convert") and hasattr(item, "size"): # a PIL image
|
| 51 |
+
return Media("image", item)
|
| 52 |
+
raise TypeError(f"a media item is a picture or a dict with its type, not {type(item).__name__}")
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def media_list(items: Sequence | None) -> list[Media]:
|
| 56 |
+
return [as_media(x) for x in items or ()]
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def spoken(q) -> Media | None:
|
| 60 |
+
"""A question's spoken instructions, or None for written ones."""
|
| 61 |
+
return q.instructions if isinstance(q.instructions, Media) else None
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def cap_pixels(image, max_pixels: int):
|
| 65 |
+
"""A picture downscaled to at most `max_pixels`, bicubic."""
|
| 66 |
+
image = image.convert("RGB") if hasattr(image, "convert") else image
|
| 67 |
+
w, h = image.size
|
| 68 |
+
if w * h <= max_pixels:
|
| 69 |
+
return image
|
| 70 |
+
from PIL import Image
|
| 71 |
+
|
| 72 |
+
scale = math.sqrt(max_pixels / (w * h))
|
| 73 |
+
return image.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.Resampling.BICUBIC)
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
class Reader:
|
| 77 |
+
"""One checkpoint's media: each item's markup and processor outputs (`encode`), and the model's inputs for a
|
| 78 |
+
sequence holding several (`inputs`)."""
|
| 79 |
+
|
| 80 |
+
def __init__(self, processor, max_pixels: int):
|
| 81 |
+
self.processor, self.max_pixels = processor, max_pixels
|
| 82 |
+
self._texts: dict[tuple[int, ...], str] = {} # an expansion's ids, checked to read back, as text
|
| 83 |
+
|
| 84 |
+
def encode(self, items: Sequence[Media]) -> tuple[str, list[dict]]:
|
| 85 |
+
"""The items' markup back to back, each the template's with its placeholder expanded by the processor, and
|
| 86 |
+
each item's processor outputs."""
|
| 87 |
+
encoded = [self._encode(m) for m in items]
|
| 88 |
+
return "".join(text for text, _ in encoded), [out for _, out in encoded]
|
| 89 |
+
|
| 90 |
+
def place(self, suffix: str, q) -> tuple[str, list[dict]]:
|
| 91 |
+
"""A question's suffix with its spoken instructions' markup in place of `SPOKEN`, and their outputs."""
|
| 92 |
+
voice = spoken(q)
|
| 93 |
+
if voice is None:
|
| 94 |
+
return suffix, []
|
| 95 |
+
markup, parts = self.encode([voice])
|
| 96 |
+
return suffix.replace(SPOKEN, markup), parts
|
| 97 |
+
|
| 98 |
+
def inputs(self, parts: Sequence[dict], ids: Sequence[int]) -> dict:
|
| 99 |
+
"""The model's media inputs for a sequence `ids` holding the items `parts` came from, in this order: their
|
| 100 |
+
processor outputs joined (`join`), and the per-token marks a processor returns (`mm_token_type_ids`) over
|
| 101 |
+
`ids`."""
|
| 102 |
+
out = join([{k: v for k, v in p.items() if k not in ("input_ids", "attention_mask", "mm_token_type_ids")}
|
| 103 |
+
for p in parts])
|
| 104 |
+
if any("mm_token_type_ids" in p for p in parts):
|
| 105 |
+
out["mm_token_type_ids"] = self._marks(ids)
|
| 106 |
+
if "pixel_attention_mask" in out: # pictures padded to a fixed count of patches: cut to the real ones
|
| 107 |
+
n = int(out["pixel_attention_mask"].sum(1).max())
|
| 108 |
+
out["pixel_values"], out["pixel_attention_mask"] = out["pixel_values"][:, :n], out["pixel_attention_mask"][:, :n]
|
| 109 |
+
return out
|
| 110 |
+
|
| 111 |
+
def batch(self, inputs: Sequence[dict], ids) -> dict:
|
| 112 |
+
"""Several sequences' inputs (`inputs`, in row order) as one forward over their right-padded rows `ids`
|
| 113 |
+
[B, T] takes them: `join`ed, and the marks over every row when a processor returns them."""
|
| 114 |
+
out = join([{k: v for k, v in x.items() if k != "mm_token_type_ids"} for x in inputs])
|
| 115 |
+
if any("mm_token_type_ids" in x for x in inputs):
|
| 116 |
+
out["mm_token_type_ids"] = self._marks(ids)
|
| 117 |
+
return out
|
| 118 |
+
|
| 119 |
+
def _marks(self, ids):
|
| 120 |
+
"""Which tokens hold a picture (1) or a video (2), as the processors that return such marks write them: over
|
| 121 |
+
one sequence's ids, [1, T], or a batch's rows, [B, T]."""
|
| 122 |
+
import torch
|
| 123 |
+
|
| 124 |
+
x = torch.as_tensor(ids).cpu()
|
| 125 |
+
x = x[None] if x.dim() == 1 else x
|
| 126 |
+
marks = torch.zeros_like(x)
|
| 127 |
+
for value, name in ((1, "image_token_id"), (2, "video_token_id")):
|
| 128 |
+
if (token := getattr(self.processor, name, None)) is not None:
|
| 129 |
+
marks[x == token] = value
|
| 130 |
+
return marks
|
| 131 |
+
|
| 132 |
+
def _encode(self, item: Media) -> tuple[str, dict]:
|
| 133 |
+
"""One item's expanded markup, as text that encodes back to the processor's ids, and its outputs."""
|
| 134 |
+
if getattr(self.processor, PROCESSORS[item.type], None) is None:
|
| 135 |
+
raise ValueError(f"this checkpoint takes no {item.type}")
|
| 136 |
+
template = template_markup(self.processor, item.type)
|
| 137 |
+
rate = {"sampling_rate": self.processor.feature_extractor.sampling_rate} if item.type == "audio" else {}
|
| 138 |
+
if item.type == "video" and not isinstance(item.value, str): # a path is decoded with its own clock
|
| 139 |
+
if item.fps is None:
|
| 140 |
+
raise ValueError("a video's frames need their fps")
|
| 141 |
+
rate = {"video_metadata": [{"fps": item.fps, "total_num_frames": len(item.value)}]}
|
| 142 |
+
out = self.processor(text=[template], **{INPUTS[item.type]: [[self._decoded(item)]]}, **rate,
|
| 143 |
+
return_tensors="pt", add_special_tokens=False)
|
| 144 |
+
ids = tuple(out["input_ids"][0].tolist())
|
| 145 |
+
if ids not in self._texts: # one check an expansion (a picture size, a clip length), not one a call
|
| 146 |
+
text = self.processor.tokenizer.decode(ids)
|
| 147 |
+
if tuple(self.processor.tokenizer.encode(text, add_special_tokens=False)) != ids:
|
| 148 |
+
raise ValueError(f"{type(self.processor).__name__}'s {item.type} markup does not read back as text")
|
| 149 |
+
if len(self._texts) >= 4096:
|
| 150 |
+
self._texts.clear()
|
| 151 |
+
self._texts[ids] = text
|
| 152 |
+
return self._texts[ids], dict(out)
|
| 153 |
+
|
| 154 |
+
def _decoded(self, item: Media):
|
| 155 |
+
"""What the processor takes: a capped RGB picture, a waveform at the feature extractor's rate, a video."""
|
| 156 |
+
if item.type == "image":
|
| 157 |
+
from PIL import Image
|
| 158 |
+
|
| 159 |
+
image = Image.open(item.value) if isinstance(item.value, str) else item.value
|
| 160 |
+
return cap_pixels(image, self.max_pixels)
|
| 161 |
+
if item.type == "audio":
|
| 162 |
+
rate = self.processor.feature_extractor.sampling_rate
|
| 163 |
+
wave, given = read_audio(item.value, rate) if isinstance(item.value, str) else (item.value, item.sampling_rate)
|
| 164 |
+
if given is None:
|
| 165 |
+
raise ValueError("a waveform needs its sampling_rate")
|
| 166 |
+
return resampled(wave, given, rate)
|
| 167 |
+
return item.value
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
def read_audio(path: str, rate: int):
|
| 171 |
+
"""A clip's waveform and its rate: 16-bit PCM WAV, or a NumPy file (int16 or float, at `rate`), with the standard
|
| 172 |
+
library and NumPy; any other format through transformers' `load_audio` (torchcodec or librosa), at `rate`."""
|
| 173 |
+
import numpy as np
|
| 174 |
+
|
| 175 |
+
if path.endswith(".wav"):
|
| 176 |
+
import wave
|
| 177 |
+
|
| 178 |
+
with wave.open(path, "rb") as f:
|
| 179 |
+
if f.getsampwidth() != 2:
|
| 180 |
+
raise ValueError(f"{path}: 16-bit PCM WAV only")
|
| 181 |
+
pcm = np.frombuffer(f.readframes(f.getnframes()), dtype="<i2").reshape(-1, f.getnchannels())
|
| 182 |
+
return pcm.T.astype(np.float32) / 32768, f.getframerate()
|
| 183 |
+
if path.endswith(".npy"):
|
| 184 |
+
data = np.load(path)
|
| 185 |
+
return (data.astype(np.float32) / 32768 if data.dtype == np.int16 else data.astype(np.float32)), rate
|
| 186 |
+
from transformers.audio_utils import load_audio
|
| 187 |
+
|
| 188 |
+
return load_audio(path, sampling_rate=rate), rate
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
def resampled(wave, rate: int, target: int, zeros: int = 6, rolloff: float = 0.99):
|
| 192 |
+
"""A mono float32 waveform at `target` (channels averaged), resampled in torch as torchaudio's
|
| 193 |
+
`functional.resample` does by default: a Hann-windowed sinc over `zeros` zero crossings, cut at `rolloff` of the
|
| 194 |
+
lower Nyquist, one filter per output phase (no torchaudio: a Hub repo needs no package it does not name)."""
|
| 195 |
+
import numpy as np
|
| 196 |
+
|
| 197 |
+
wave = np.asarray(wave, dtype=np.float32)
|
| 198 |
+
if wave.ndim == 2: # (channels, samples) or (samples, channels): a clip never has more channels than samples
|
| 199 |
+
wave = (wave.T if wave.shape[0] > wave.shape[1] else wave).mean(0)
|
| 200 |
+
if rate == target:
|
| 201 |
+
return wave
|
| 202 |
+
import torch
|
| 203 |
+
|
| 204 |
+
gcd = math.gcd(int(rate), int(target))
|
| 205 |
+
old, new = int(rate) // gcd, int(target) // gcd
|
| 206 |
+
cut = min(old, new) * rolloff
|
| 207 |
+
width = math.ceil(zeros * old / cut)
|
| 208 |
+
t = torch.arange(0, -new, -1, dtype=torch.float32)[:, None, None] / new
|
| 209 |
+
t = ((t + torch.arange(-width, width + old, dtype=torch.float32)[None, None] / old) * cut).clamp(-zeros, zeros)
|
| 210 |
+
window = torch.cos(t * math.pi / zeros / 2) ** 2
|
| 211 |
+
t = t * math.pi
|
| 212 |
+
kernel = torch.where(t == 0, torch.tensor(1.0), t.sin() / t) * (window * (cut / old))
|
| 213 |
+
x = torch.nn.functional.pad(torch.from_numpy(wave)[None, None], (width, width + old))
|
| 214 |
+
out = torch.nn.functional.conv1d(x, kernel, stride=old).transpose(1, 2).reshape(-1)
|
| 215 |
+
return out[:math.ceil(new * len(wave) / old)].numpy()
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
def template_markup(processor, kind: str) -> str:
|
| 219 |
+
"""One item's markup as the processor's chat template writes it at the head of a user turn."""
|
| 220 |
+
messages = [{"role": "user", "content": [{"type": kind}, {"type": "text", "text": "\x00"}]}]
|
| 221 |
+
text = processor.apply_chat_template(messages, add_generation_prompt=False, tokenize=False)
|
| 222 |
+
head = "<|im_start|>user\n"
|
| 223 |
+
return text[text.index(head) + len(head):text.index("\x00")]
|
| 224 |
+
|
| 225 |
+
|
| 226 |
+
def join(inputs: Sequence[dict]) -> dict:
|
| 227 |
+
"""Media inputs of several items, or of several sequences in row order, as one forward takes them: each modality's
|
| 228 |
+
tensors joined on their first axis, padded with zeros to the largest along the others (pictures' patches and
|
| 229 |
+
their mask, clips' mel frames and theirs: `feature_attention_mask`), lists concatenated. A forward fills each
|
| 230 |
+
modality's placeholders in this order, so items of every kind share it."""
|
| 231 |
+
import torch
|
| 232 |
+
|
| 233 |
+
out = {}
|
| 234 |
+
for key in dict.fromkeys(k for i in inputs for k in i):
|
| 235 |
+
values = [i[key] for i in inputs if key in i]
|
| 236 |
+
out[key] = _join(values) if torch.is_tensor(values[0]) else [v for value in values for v in value]
|
| 237 |
+
return out
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
def _join(values: list):
|
| 241 |
+
"""Tensors joined on their first axis, each padded with zeros to the largest along the others."""
|
| 242 |
+
import torch
|
| 243 |
+
|
| 244 |
+
if all(v.shape[1:] == values[0].shape[1:] for v in values):
|
| 245 |
+
return torch.cat(values)
|
| 246 |
+
size = [max(v.shape[d] for v in values) for d in range(1, values[0].dim())]
|
| 247 |
+
out = values[0].new_zeros(sum(v.shape[0] for v in values), *size)
|
| 248 |
+
at = 0
|
| 249 |
+
for v in values:
|
| 250 |
+
out[(slice(at, at + v.shape[0]), *(slice(0, n) for n in v.shape[1:]))] = v
|
| 251 |
+
at += v.shape[0]
|
| 252 |
+
return out
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b5e573687315e82ea80f4f7dd8d76b5cec3b6ee97e3509fd0493cfebb203719f
|
| 3 |
+
size 3836417920
|
modeling_d1.py
ADDED
|
@@ -0,0 +1,73 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The System One API on any backbone: each `modeling_d1_<model_type>.py` serves `D1Model(SystemOneModel, <its
|
| 2 |
+
class>)` to `AutoModel.from_pretrained(repo, trust_remote_code=True)`.
|
| 3 |
+
|
| 4 |
+
model.system_one(state, {name: question}, media=None) # one forward pass; media: pictures, audio, video
|
| 5 |
+
model.system_one_batch([(state, {name: question}), ...]) # many states in one call
|
| 6 |
+
|
| 7 |
+
A release's `calibration.json` (its temperature) is applied to every answer.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
from collections.abc import Mapping, Sequence
|
| 13 |
+
from functools import cached_property
|
| 14 |
+
from typing import Any
|
| 15 |
+
|
| 16 |
+
from .api import SCORE_READOUT
|
| 17 |
+
from .runner import SystemOne
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
class SystemOneModel:
|
| 21 |
+
"""The System One API over any backbone the runner reads. `attention`: the attention function a transformers
|
| 22 |
+
backbone loads with (`attention.d1_attention`), when it has one."""
|
| 23 |
+
|
| 24 |
+
attention: str | None = None
|
| 25 |
+
|
| 26 |
+
@classmethod
|
| 27 |
+
def from_pretrained(cls, *args, revision=None, **kwargs):
|
| 28 |
+
if cls.attention:
|
| 29 |
+
kwargs["attn_implementation"] = cls.attention
|
| 30 |
+
loaded = super().from_pretrained(*args, revision=revision, **kwargs) # with its loading info when asked
|
| 31 |
+
# transformers 5.19 resolves the revision to a commit for the load and keeps no `_commit_hash`
|
| 32 |
+
(loaded[0] if isinstance(loaded, tuple) else loaded).d1_revision = revision
|
| 33 |
+
return loaded
|
| 34 |
+
|
| 35 |
+
def d1_file(self, name: str, required: bool = True) -> str | None:
|
| 36 |
+
"""A file of the loaded revision, not main's: its path, or None for a missing one not `required`."""
|
| 37 |
+
from transformers.utils import cached_file
|
| 38 |
+
|
| 39 |
+
return cached_file(self.name_or_path, name, revision=self._revision(),
|
| 40 |
+
_raise_exceptions_for_missing_entries=required)
|
| 41 |
+
|
| 42 |
+
def _revision(self) -> str | None:
|
| 43 |
+
return getattr(self.config, "_commit_hash", None) or getattr(self, "d1_revision", None)
|
| 44 |
+
|
| 45 |
+
@cached_property
|
| 46 |
+
def engine(self) -> SystemOne:
|
| 47 |
+
"""The runner over this model, with the loaded revision's tokenizer and temperature."""
|
| 48 |
+
from transformers import AutoTokenizer
|
| 49 |
+
|
| 50 |
+
from .calibration import Calibration
|
| 51 |
+
|
| 52 |
+
found = self.d1_file("calibration.json", required=False)
|
| 53 |
+
return SystemOne(model=self.eval(), tokenizer=AutoTokenizer.from_pretrained(self.name_or_path,
|
| 54 |
+
revision=self._revision()),
|
| 55 |
+
calibration=Calibration.load(found) if found else None, revision=self._revision())
|
| 56 |
+
|
| 57 |
+
def compile(self, mode: str | None = None) -> None:
|
| 58 |
+
"""torch.compile what `system_one` runs (`SystemOne.compile`): fused kernels, with no CUDA graphs under `set` attention."""
|
| 59 |
+
self.engine.compile(mode)
|
| 60 |
+
|
| 61 |
+
def system_one(self, state: Any, questions: Mapping[str, Any], media: Sequence | None = None,
|
| 62 |
+
score_readout: str = SCORE_READOUT, set_size: int | None = None) -> dict:
|
| 63 |
+
"""Named questions over a state (text, JSON, or None with media alone) and its media (pictures, audio clips,
|
| 64 |
+
videos: `media`): `{"answers": {name: answer}, "usage": {"input_tokens": n, "output_tokens": 0}}`.
|
| 65 |
+
`score_readout`: a score's `median` level (the default), its expected level (`mean`) or its `argmax` level.
|
| 66 |
+
`set_size`: questions read together after the state, the model's trained size by default; larger is faster
|
| 67 |
+
and costs some accuracy on text."""
|
| 68 |
+
return self.engine.system_one(state, questions, media, score_readout, set_size=set_size)
|
| 69 |
+
|
| 70 |
+
def system_one_batch(self, requests: Sequence[tuple], score_readout: str = SCORE_READOUT,
|
| 71 |
+
set_size: int | None = None) -> list[dict]:
|
| 72 |
+
"""`(state, questions)` or `(state, questions, media)` requests, one response each."""
|
| 73 |
+
return self.engine.system_one_batch(requests, score_readout, set_size=set_size)
|
modeling_d1_lfm3_vl.py
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""`AutoModel.from_pretrained(repo, trust_remote_code=True)`: LFM3-VL with the System One API (`modeling_d1`)."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from .lfm3 import Lfm3SystemOne
|
| 6 |
+
from .native_lfm3_vl import Lfm3VlForConditionalGeneration
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class D1Model(Lfm3SystemOne, Lfm3VlForConditionalGeneration):
|
| 10 |
+
compiled_parts = ("vision_tower", "multi_modal_projector")
|
modeling_lfm3.py
ADDED
|
@@ -0,0 +1,881 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Vendored from transformers' LFM3 code (not in a release yet) at 69da1b47a11e362fbb15bfe5d5c8da0790c688c5, src/transformers/models/lfm3/modeling_lfm3.py:
|
| 2 |
+
# its relative imports made absolute, nothing else changed.
|
| 3 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 4 |
+
# This file was automatically generated from src/transformers/models/lfm3/modular_lfm3.py.
|
| 5 |
+
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
| 6 |
+
# the file from the modular. If any change should be done, please apply the change to the
|
| 7 |
+
# modular_lfm3.py file directly. One of our CI enforces this.
|
| 8 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 9 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 10 |
+
#
|
| 11 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 12 |
+
# you may not use this file except in compliance with the License.
|
| 13 |
+
# You may obtain a copy of the License at
|
| 14 |
+
#
|
| 15 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 16 |
+
#
|
| 17 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 18 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 19 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 20 |
+
# See the License for the specific language governing permissions and
|
| 21 |
+
# limitations under the License.
|
| 22 |
+
import copy
|
| 23 |
+
from collections.abc import Callable
|
| 24 |
+
|
| 25 |
+
import torch
|
| 26 |
+
import torch.nn.functional as F
|
| 27 |
+
from torch import nn
|
| 28 |
+
|
| 29 |
+
from transformers import initialization as init
|
| 30 |
+
from transformers.activations import ACT2FN
|
| 31 |
+
from transformers.cache_utils import Cache, DynamicCache
|
| 32 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 33 |
+
from transformers.generation import GenerationMixin
|
| 34 |
+
from transformers.integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func
|
| 35 |
+
from transformers.integrations.accelerate import force_accelerate_hooks
|
| 36 |
+
from .integrations_lfm3 import can_use_lfm3_conv, lfm3_causal_conv1d, lfm3_causal_conv1d_update
|
| 37 |
+
from transformers.masking_utils import create_causal_mask, create_recurrent_attention_mask, create_sliding_window_causal_mask
|
| 38 |
+
from transformers.modeling_layers import GradientCheckpointingLayer
|
| 39 |
+
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
|
| 40 |
+
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
|
| 41 |
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
| 42 |
+
from transformers.processing_utils import Unpack
|
| 43 |
+
from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
|
| 44 |
+
from transformers.utils.generic import merge_with_config_defaults
|
| 45 |
+
from transformers.utils.import_utils import is_tracing
|
| 46 |
+
from transformers.utils.output_capturing import capture_outputs
|
| 47 |
+
from .configuration_lfm3 import Lfm3Config
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
@use_kernel_forward_from_hub("RMSNorm")
|
| 51 |
+
class Lfm3RMSNorm(nn.Module):
|
| 52 |
+
def __init__(self, hidden_size, eps: float = 1e-6) -> None:
|
| 53 |
+
"""
|
| 54 |
+
Lfm3RMSNorm is equivalent to T5LayerNorm
|
| 55 |
+
"""
|
| 56 |
+
super().__init__()
|
| 57 |
+
self.weight = nn.Parameter(torch.ones(hidden_size))
|
| 58 |
+
self.variance_epsilon = eps
|
| 59 |
+
|
| 60 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 61 |
+
input_dtype = hidden_states.dtype
|
| 62 |
+
hidden_states = hidden_states.to(torch.float32)
|
| 63 |
+
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
| 64 |
+
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
| 65 |
+
return self.weight * hidden_states.to(input_dtype)
|
| 66 |
+
|
| 67 |
+
def extra_repr(self):
|
| 68 |
+
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
class Lfm3RotaryEmbedding(nn.Module):
|
| 72 |
+
"""
|
| 73 |
+
Rotary embedding per attention layer type, like Gemma 3; `conv` layers register no buffers. Partial RoPE keeps the
|
| 74 |
+
frequency progression of the full head dimension and zeroes the tail instead of shrinking the rotated dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(self, config: Lfm3Config, device=None):
|
| 78 |
+
super().__init__()
|
| 79 |
+
self.max_seq_len_cached = config.max_position_embeddings
|
| 80 |
+
self.original_max_seq_len = config.max_position_embeddings
|
| 81 |
+
self.config = config
|
| 82 |
+
self.layer_types = sorted(set(config.layer_types) & set(config.rope_parameters))
|
| 83 |
+
self.rope_type = {}
|
| 84 |
+
for layer_type in self.layer_types:
|
| 85 |
+
rope_params = self.config.rope_parameters[layer_type]
|
| 86 |
+
self.rope_type[layer_type] = rope_params["rope_type"]
|
| 87 |
+
curr_inv_freq, curr_attention_scaling = self.compute_rope_parameters(config, device, layer_type=layer_type)
|
| 88 |
+
setattr(self, f"{layer_type}_inv_freq", nn.Buffer(curr_inv_freq, persistent=False))
|
| 89 |
+
setattr(self, f"{layer_type}_original_inv_freq", nn.Buffer(curr_inv_freq.clone(), persistent=False))
|
| 90 |
+
setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling)
|
| 91 |
+
|
| 92 |
+
@staticmethod
|
| 93 |
+
def compute_default_rope_parameters(
|
| 94 |
+
config: Lfm3Config, device=None, layer_type: str | None = None, **kwargs
|
| 95 |
+
) -> tuple[torch.Tensor, float]:
|
| 96 |
+
"""
|
| 97 |
+
Inverse frequencies of the default RoPE with LFM3's zero-tail partial rotation: the first
|
| 98 |
+
`partial_rotary_factor * head_dim // 2` frequencies follow the progression of the full head dimension and the
|
| 99 |
+
rest are zero, leaving those dimensions unrotated.
|
| 100 |
+
|
| 101 |
+
Returns:
|
| 102 |
+
Tuple of (`torch.Tensor`, `float`): the inverse frequencies and the attention scaling factor (always 1).
|
| 103 |
+
"""
|
| 104 |
+
base = config.rope_parameters[layer_type]["rope_theta"]
|
| 105 |
+
partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
|
| 106 |
+
dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
|
| 107 |
+
rope_angles = int(partial_rotary_factor * dim) // 2
|
| 108 |
+
inv_freq = 1.0 / (base ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float) / dim))
|
| 109 |
+
inv_freq = torch.cat((inv_freq, torch.zeros(dim // 2 - rope_angles, dtype=torch.float)))
|
| 110 |
+
return inv_freq.to(device), 1.0
|
| 111 |
+
|
| 112 |
+
@torch.no_grad()
|
| 113 |
+
@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
|
| 114 |
+
def forward(self, x, position_ids, layer_type):
|
| 115 |
+
inv_freq = getattr(self, f"{layer_type}_inv_freq")
|
| 116 |
+
attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
|
| 117 |
+
|
| 118 |
+
freqs = position_ids[..., None].float() * inv_freq.to(device=x.device, dtype=torch.float)
|
| 119 |
+
emb = torch.cat((freqs, freqs), dim=-1)
|
| 120 |
+
cos = emb.cos() * attention_scaling
|
| 121 |
+
sin = emb.sin() * attention_scaling
|
| 122 |
+
|
| 123 |
+
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
| 124 |
+
|
| 125 |
+
@classmethod
|
| 126 |
+
def compute_rope_parameters(
|
| 127 |
+
cls, config: Lfm3Config, device=None, layer_type: str | None = None
|
| 128 |
+
) -> tuple[torch.Tensor, float]:
|
| 129 |
+
"""Inverse frequencies of any RoPE type in the zero-tail layout: scaled types are evaluated over the full head
|
| 130 |
+
dimension and truncated to the rotated prefix."""
|
| 131 |
+
rope_params = config.rope_parameters[layer_type]
|
| 132 |
+
if rope_params["rope_type"] == "default":
|
| 133 |
+
return cls.compute_default_rope_parameters(config, device, layer_type=layer_type)
|
| 134 |
+
full_config = copy.copy(config)
|
| 135 |
+
full_config.rope_parameters = {
|
| 136 |
+
**config.rope_parameters,
|
| 137 |
+
layer_type: {**rope_params, "partial_rotary_factor": 1.0},
|
| 138 |
+
}
|
| 139 |
+
inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_params["rope_type"]](
|
| 140 |
+
full_config, device, layer_type=layer_type
|
| 141 |
+
)
|
| 142 |
+
dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
|
| 143 |
+
rope_angles = int(rope_params.get("partial_rotary_factor", 1.0) * dim) // 2
|
| 144 |
+
inv_freq = torch.cat((inv_freq[:rope_angles], inv_freq.new_zeros(dim // 2 - rope_angles)))
|
| 145 |
+
return inv_freq, attention_factor
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
class Lfm3MLP(nn.Module):
|
| 149 |
+
def __init__(self, config: Lfm3Config):
|
| 150 |
+
super().__init__()
|
| 151 |
+
self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 152 |
+
self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 153 |
+
self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
| 154 |
+
|
| 155 |
+
def forward(self, x):
|
| 156 |
+
return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
def rotate_half(x):
|
| 160 |
+
"""Rotates half the hidden dims of the input."""
|
| 161 |
+
x1 = x[..., : x.shape[-1] // 2]
|
| 162 |
+
x2 = x[..., x.shape[-1] // 2 :]
|
| 163 |
+
return torch.cat((-x2, x1), dim=-1)
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
| 167 |
+
"""
|
| 168 |
+
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
| 169 |
+
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
| 170 |
+
"""
|
| 171 |
+
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
| 172 |
+
if n_rep == 1:
|
| 173 |
+
return hidden_states
|
| 174 |
+
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
| 175 |
+
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
| 176 |
+
|
| 177 |
+
|
| 178 |
+
def eager_attention_forward(
|
| 179 |
+
module: nn.Module,
|
| 180 |
+
query: torch.Tensor,
|
| 181 |
+
key: torch.Tensor,
|
| 182 |
+
value: torch.Tensor,
|
| 183 |
+
attention_mask: torch.Tensor | None,
|
| 184 |
+
scaling: float,
|
| 185 |
+
dropout: float = 0.0,
|
| 186 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 187 |
+
):
|
| 188 |
+
key_states = repeat_kv(key, module.num_key_value_groups)
|
| 189 |
+
value_states = repeat_kv(value, module.num_key_value_groups)
|
| 190 |
+
|
| 191 |
+
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
| 192 |
+
if attention_mask is not None:
|
| 193 |
+
attn_weights = attn_weights + attention_mask
|
| 194 |
+
|
| 195 |
+
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
|
| 196 |
+
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
|
| 197 |
+
attn_output = torch.matmul(attn_weights, value_states)
|
| 198 |
+
attn_output = attn_output.transpose(1, 2).contiguous()
|
| 199 |
+
|
| 200 |
+
return attn_output, attn_weights
|
| 201 |
+
|
| 202 |
+
|
| 203 |
+
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
|
| 204 |
+
"""Applies RoPE to the query and key tensors in float32, like the reference implementation."""
|
| 205 |
+
cos = cos.unsqueeze(unsqueeze_dim).float()
|
| 206 |
+
sin = sin.unsqueeze(unsqueeze_dim).float()
|
| 207 |
+
q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
|
| 208 |
+
k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
|
| 209 |
+
return q_embed.to(q.dtype), k_embed.to(k.dtype)
|
| 210 |
+
|
| 211 |
+
|
| 212 |
+
@use_kernelized_func(apply_rotary_pos_emb)
|
| 213 |
+
class Lfm3Attention(nn.Module):
|
| 214 |
+
"""Llama attention with QK RMSNorm, a sigmoid output gate per head, a sliding window on `sliding_attention`
|
| 215 |
+
layers and RoPE in float32."""
|
| 216 |
+
|
| 217 |
+
def __init__(self, config: Lfm3Config, layer_idx: int):
|
| 218 |
+
super().__init__()
|
| 219 |
+
self.config = config
|
| 220 |
+
self.layer_idx = layer_idx
|
| 221 |
+
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
| 222 |
+
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
| 223 |
+
self.scaling = self.head_dim**-0.5
|
| 224 |
+
self.is_causal = True
|
| 225 |
+
self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
|
| 226 |
+
self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
|
| 227 |
+
self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
|
| 228 |
+
self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
|
| 229 |
+
self.q_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
|
| 230 |
+
self.k_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
|
| 231 |
+
self.gate_proj = nn.Linear(config.hidden_size, config.num_attention_heads, bias=False)
|
| 232 |
+
self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
|
| 233 |
+
|
| 234 |
+
def forward(
|
| 235 |
+
self,
|
| 236 |
+
hidden_states: torch.Tensor,
|
| 237 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
| 238 |
+
attention_mask: torch.Tensor | None,
|
| 239 |
+
past_key_values: Cache | None = None,
|
| 240 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 241 |
+
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
| 242 |
+
input_shape = hidden_states.shape[:-1]
|
| 243 |
+
hidden_shape = (*input_shape, -1, self.head_dim)
|
| 244 |
+
|
| 245 |
+
query_states = self.q_norm(self.q_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
|
| 246 |
+
key_states = self.k_norm(self.k_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
|
| 247 |
+
value_states = self.v_proj(hidden_states).view(*hidden_shape).transpose(1, 2)
|
| 248 |
+
|
| 249 |
+
cos, sin = position_embeddings
|
| 250 |
+
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
| 251 |
+
|
| 252 |
+
if past_key_values is not None:
|
| 253 |
+
key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
|
| 254 |
+
|
| 255 |
+
attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
|
| 256 |
+
self.config._attn_implementation, eager_attention_forward
|
| 257 |
+
)
|
| 258 |
+
|
| 259 |
+
attn_output, attn_weights = attention_interface(
|
| 260 |
+
self,
|
| 261 |
+
query_states,
|
| 262 |
+
key_states,
|
| 263 |
+
value_states,
|
| 264 |
+
attention_mask,
|
| 265 |
+
dropout=0.0,
|
| 266 |
+
scaling=self.scaling,
|
| 267 |
+
sliding_window=self.sliding_window,
|
| 268 |
+
**kwargs,
|
| 269 |
+
)
|
| 270 |
+
gate = torch.sigmoid(self.gate_proj(hidden_states)).unsqueeze(-1)
|
| 271 |
+
attn_output = (attn_output * gate).reshape(*input_shape, -1).contiguous()
|
| 272 |
+
attn_output = self.o_proj(attn_output)
|
| 273 |
+
return attn_output, attn_weights
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def apply_mask_to_padding_states(hidden_states, attention_mask):
|
| 277 |
+
"""
|
| 278 |
+
Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
|
| 279 |
+
"""
|
| 280 |
+
# NOTE: attention mask is a 2D boolean tensor
|
| 281 |
+
if attention_mask is not None:
|
| 282 |
+
dtype = hidden_states.dtype
|
| 283 |
+
hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
|
| 284 |
+
|
| 285 |
+
return hidden_states
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
|
| 289 |
+
def causal_conv1d_update(
|
| 290 |
+
hidden_states: torch.Tensor,
|
| 291 |
+
conv_state: torch.Tensor,
|
| 292 |
+
weight: nn.Parameter,
|
| 293 |
+
bias: nn.Parameter | None = None,
|
| 294 |
+
activation: str | None = None,
|
| 295 |
+
):
|
| 296 |
+
_, hidden_size, seq_len = hidden_states.shape
|
| 297 |
+
state_len = conv_state.shape[-1]
|
| 298 |
+
|
| 299 |
+
hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
|
| 300 |
+
conv_state.copy_(hidden_states_new[:, :, -state_len:])
|
| 301 |
+
out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
|
| 302 |
+
out = out[:, :, -seq_len:]
|
| 303 |
+
if activation is not None:
|
| 304 |
+
out = ACT2FN[activation](out)
|
| 305 |
+
return out.to(hidden_states.dtype)
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
|
| 309 |
+
def causal_conv1d_fn(
|
| 310 |
+
hidden_states: torch.Tensor,
|
| 311 |
+
weight: nn.Parameter,
|
| 312 |
+
bias: nn.Parameter | None = None,
|
| 313 |
+
activation: str | None = None,
|
| 314 |
+
**kwargs,
|
| 315 |
+
):
|
| 316 |
+
_, hidden_size, seq_len = hidden_states.shape
|
| 317 |
+
padding = weight.shape[-1] - 1
|
| 318 |
+
|
| 319 |
+
out = F.conv1d(
|
| 320 |
+
hidden_states.to(weight.dtype),
|
| 321 |
+
weight=weight.unsqueeze(1),
|
| 322 |
+
bias=bias,
|
| 323 |
+
padding=padding,
|
| 324 |
+
groups=hidden_size,
|
| 325 |
+
)[:, :, :seq_len]
|
| 326 |
+
if activation is not None:
|
| 327 |
+
out = ACT2FN[activation](out)
|
| 328 |
+
return out.to(hidden_states.dtype)
|
| 329 |
+
|
| 330 |
+
|
| 331 |
+
@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
|
| 332 |
+
class Lfm3ShortConv(nn.Module):
|
| 333 |
+
"""LFM2 gated short convolution with packed-document boundaries (`seq_idx`), the padding mask re-applied after a
|
| 334 |
+
biased `in_proj`, and width-7 Triton inference kernels with a PyTorch fallback."""
|
| 335 |
+
|
| 336 |
+
def __init__(
|
| 337 |
+
self,
|
| 338 |
+
config: Lfm3Config,
|
| 339 |
+
layer_idx: int,
|
| 340 |
+
):
|
| 341 |
+
super().__init__()
|
| 342 |
+
self.config = config
|
| 343 |
+
self.layer_idx = layer_idx
|
| 344 |
+
self.conv_kernel_size = config.conv_L_cache
|
| 345 |
+
self.bias = config.conv_bias
|
| 346 |
+
|
| 347 |
+
self.conv = nn.Conv1d(
|
| 348 |
+
in_channels=config.hidden_size,
|
| 349 |
+
out_channels=config.hidden_size,
|
| 350 |
+
kernel_size=self.conv_kernel_size,
|
| 351 |
+
groups=config.hidden_size,
|
| 352 |
+
bias=self.bias,
|
| 353 |
+
padding=self.conv_kernel_size - 1,
|
| 354 |
+
)
|
| 355 |
+
self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=self.bias)
|
| 356 |
+
self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=self.bias)
|
| 357 |
+
|
| 358 |
+
self.layer_type = config.layer_types[layer_idx]
|
| 359 |
+
|
| 360 |
+
@force_accelerate_hooks("conv")
|
| 361 |
+
def forward(
|
| 362 |
+
self,
|
| 363 |
+
hidden_states: torch.Tensor,
|
| 364 |
+
past_key_values: Cache | None = None,
|
| 365 |
+
attention_mask: torch.Tensor | None = None,
|
| 366 |
+
seq_idx: torch.IntTensor | None = None,
|
| 367 |
+
):
|
| 368 |
+
seq_len = hidden_states.shape[1]
|
| 369 |
+
|
| 370 |
+
hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
|
| 371 |
+
BCx = self.in_proj(hidden_states).transpose(-1, -2)
|
| 372 |
+
B, C, x = BCx.chunk(3, dim=-2)
|
| 373 |
+
hidden_states = B * x
|
| 374 |
+
if attention_mask is not None:
|
| 375 |
+
# with `conv_bias` the projection makes the padded positions non-zero again
|
| 376 |
+
hidden_states = hidden_states * attention_mask[:, None, :].to(hidden_states.dtype)
|
| 377 |
+
|
| 378 |
+
use_precomputed_states = past_key_values is not None and past_key_values.has_previous_state(self.layer_idx)
|
| 379 |
+
use_kernels = 2 <= self.conv_kernel_size <= 4
|
| 380 |
+
use_wide_kernel = can_use_lfm3_conv(hidden_states, self.conv.weight)
|
| 381 |
+
|
| 382 |
+
if use_precomputed_states and seq_len == 1 and not past_key_values.layers[self.layer_idx].record_past:
|
| 383 |
+
conv_state = past_key_values.layers[self.layer_idx].conv_states[0]
|
| 384 |
+
if use_wide_kernel and conv_state.shape[-1] in (6, 7):
|
| 385 |
+
hidden_states = lfm3_causal_conv1d_update(
|
| 386 |
+
hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
|
| 387 |
+
)
|
| 388 |
+
elif use_kernels:
|
| 389 |
+
hidden_states = causal_conv1d_update(
|
| 390 |
+
hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
|
| 391 |
+
)
|
| 392 |
+
else:
|
| 393 |
+
hidden_states = torch.cat([conv_state, hidden_states], dim=-1)
|
| 394 |
+
conv_state.copy_(hidden_states[:, :, -conv_state.shape[-1] :])
|
| 395 |
+
hidden_states = self._conv(hidden_states)[:, :, -seq_len:]
|
| 396 |
+
else:
|
| 397 |
+
if past_key_values is not None and seq_idx is not None:
|
| 398 |
+
# `seq_idx` is only set on an empty cache, and the decode steps read the cached window without it:
|
| 399 |
+
# only the last document (after any positional left padding) may enter the window
|
| 400 |
+
last_document = (seq_idx == seq_idx[:, -1:])[:, None].to(hidden_states.dtype)
|
| 401 |
+
past_key_values.update_conv_state(
|
| 402 |
+
hidden_states * last_document, self.layer_idx, conv_kernel_size=self.conv_kernel_size
|
| 403 |
+
)
|
| 404 |
+
elif past_key_values is not None:
|
| 405 |
+
hidden_states = past_key_values.update_conv_state(
|
| 406 |
+
hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
|
| 407 |
+
)
|
| 408 |
+
|
| 409 |
+
if use_wide_kernel:
|
| 410 |
+
hidden_states = lfm3_causal_conv1d(hidden_states, self.conv.weight.squeeze(1), self.conv.bias, seq_idx)
|
| 411 |
+
elif seq_idx is not None:
|
| 412 |
+
hidden_states = self._segmented_conv(hidden_states, seq_idx)
|
| 413 |
+
elif use_kernels:
|
| 414 |
+
hidden_states = causal_conv1d_fn(hidden_states, self.conv.weight.squeeze(1), self.conv.bias)
|
| 415 |
+
else:
|
| 416 |
+
hidden_states = self._conv(hidden_states)
|
| 417 |
+
|
| 418 |
+
# Drop the additional previous states
|
| 419 |
+
if past_key_values is not None:
|
| 420 |
+
hidden_states = hidden_states[:, :, -seq_len:]
|
| 421 |
+
|
| 422 |
+
y = C * hidden_states
|
| 423 |
+
y = y.transpose(-1, -2).contiguous()
|
| 424 |
+
y = self.out_proj(y)
|
| 425 |
+
return y
|
| 426 |
+
|
| 427 |
+
def _conv(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 428 |
+
"""Causal depthwise conv in torch: the `causal_conv1d` kernels only accept widths 2 to 4."""
|
| 429 |
+
conv_out = self.conv(hidden_states.to(self.conv.weight.dtype))[..., : hidden_states.shape[-1]]
|
| 430 |
+
return conv_out.to(hidden_states.dtype)
|
| 431 |
+
|
| 432 |
+
def _segmented_conv(self, hidden_states: torch.Tensor, seq_idx: torch.Tensor) -> torch.Tensor:
|
| 433 |
+
"""Causal depthwise conv that only reads positions sharing the current token's document index."""
|
| 434 |
+
kernel_size = self.conv_kernel_size
|
| 435 |
+
seq_idx = F.pad(seq_idx, (hidden_states.shape[-1] - seq_idx.shape[-1] + kernel_size - 1, 0), value=-1)
|
| 436 |
+
windows = F.pad(hidden_states, (kernel_size - 1, 0)).unfold(-1, kernel_size, 1)
|
| 437 |
+
seq_idx_windows = seq_idx.unfold(-1, kernel_size, 1)
|
| 438 |
+
same_document = (seq_idx_windows == seq_idx_windows[..., -1:]).unsqueeze(1)
|
| 439 |
+
conv_out = torch.einsum("bdtk,dk->bdt", windows * same_document.to(windows.dtype), self.conv.weight.squeeze(1))
|
| 440 |
+
if self.bias:
|
| 441 |
+
conv_out = conv_out + self.conv.bias[None, :, None]
|
| 442 |
+
return conv_out.to(hidden_states.dtype)
|
| 443 |
+
|
| 444 |
+
|
| 445 |
+
class Lfm3DecoderLayer(GradientCheckpointingLayer):
|
| 446 |
+
def __init__(self, config: Lfm3Config, layer_idx: int):
|
| 447 |
+
super().__init__()
|
| 448 |
+
self.is_attention_layer = config.layer_types[layer_idx] in config._attention_layer_types
|
| 449 |
+
self.use_sandwich_norm = config.use_sandwich_norm
|
| 450 |
+
|
| 451 |
+
if self.is_attention_layer:
|
| 452 |
+
self.self_attn = Lfm3Attention(config, layer_idx)
|
| 453 |
+
else:
|
| 454 |
+
self.conv = Lfm3ShortConv(config, layer_idx)
|
| 455 |
+
self.feed_forward = Lfm3MLP(config)
|
| 456 |
+
self.operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 457 |
+
self.ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 458 |
+
if self.use_sandwich_norm:
|
| 459 |
+
self.post_operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 460 |
+
self.post_ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 461 |
+
|
| 462 |
+
def forward(
|
| 463 |
+
self,
|
| 464 |
+
hidden_states: torch.Tensor,
|
| 465 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
| 466 |
+
attention_mask: torch.Tensor | None = None,
|
| 467 |
+
position_ids: torch.LongTensor | None = None,
|
| 468 |
+
past_key_values: Cache | None = None,
|
| 469 |
+
seq_idx: torch.IntTensor | None = None,
|
| 470 |
+
**kwargs,
|
| 471 |
+
) -> torch.Tensor:
|
| 472 |
+
residual = hidden_states
|
| 473 |
+
if self.is_attention_layer:
|
| 474 |
+
hidden_states, _ = self.self_attn(
|
| 475 |
+
hidden_states=self.operator_norm(hidden_states),
|
| 476 |
+
position_embeddings=position_embeddings,
|
| 477 |
+
attention_mask=attention_mask,
|
| 478 |
+
position_ids=position_ids,
|
| 479 |
+
past_key_values=past_key_values,
|
| 480 |
+
**kwargs,
|
| 481 |
+
)
|
| 482 |
+
else:
|
| 483 |
+
hidden_states = self.conv(
|
| 484 |
+
hidden_states=self.operator_norm(hidden_states),
|
| 485 |
+
past_key_values=past_key_values,
|
| 486 |
+
attention_mask=attention_mask,
|
| 487 |
+
seq_idx=seq_idx,
|
| 488 |
+
)
|
| 489 |
+
if self.use_sandwich_norm:
|
| 490 |
+
hidden_states = self.post_operator_norm(hidden_states)
|
| 491 |
+
hidden_states = hidden_states + residual
|
| 492 |
+
|
| 493 |
+
ffn_out = self.feed_forward(self.ffn_norm(hidden_states))
|
| 494 |
+
if self.use_sandwich_norm:
|
| 495 |
+
ffn_out = self.post_ffn_norm(ffn_out)
|
| 496 |
+
hidden_states = hidden_states + ffn_out
|
| 497 |
+
|
| 498 |
+
return hidden_states
|
| 499 |
+
|
| 500 |
+
|
| 501 |
+
def create_masks_by_layer_type(
|
| 502 |
+
config: PreTrainedConfig,
|
| 503 |
+
inputs_embeds: torch.Tensor,
|
| 504 |
+
attention_mask: torch.Tensor | None,
|
| 505 |
+
past_key_values: Cache | None,
|
| 506 |
+
position_ids: torch.Tensor | None = None,
|
| 507 |
+
) -> dict:
|
| 508 |
+
"""The cache sizes each mask from the layer it is asked about, and the generic mask utilities ask about the first
|
| 509 |
+
non-sliding layer, which for a leading `conv` layer redirects to the first attention layer of any type: name the
|
| 510 |
+
first layer of each mask's own type instead."""
|
| 511 |
+
config = config.get_text_config()
|
| 512 |
+
mask_kwargs = {
|
| 513 |
+
"config": config,
|
| 514 |
+
"inputs_embeds": inputs_embeds,
|
| 515 |
+
"attention_mask": attention_mask,
|
| 516 |
+
"past_key_values": past_key_values,
|
| 517 |
+
"position_ids": position_ids,
|
| 518 |
+
}
|
| 519 |
+
layer_types, causal_mask_mapping = config.layer_types, {}
|
| 520 |
+
if "full_attention" in layer_types:
|
| 521 |
+
layer_idx = layer_types.index("full_attention")
|
| 522 |
+
causal_mask_mapping["full_attention"] = create_causal_mask(**mask_kwargs, layer_idx=layer_idx)
|
| 523 |
+
if "sliding_attention" in layer_types:
|
| 524 |
+
layer_idx = layer_types.index("sliding_attention")
|
| 525 |
+
causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
|
| 526 |
+
**mask_kwargs, layer_idx=layer_idx
|
| 527 |
+
)
|
| 528 |
+
if "conv" in layer_types:
|
| 529 |
+
causal_mask_mapping["conv"] = create_recurrent_attention_mask(**mask_kwargs)
|
| 530 |
+
return causal_mask_mapping
|
| 531 |
+
|
| 532 |
+
|
| 533 |
+
@auto_docstring
|
| 534 |
+
class Lfm3PreTrainedModel(PreTrainedModel):
|
| 535 |
+
config: Lfm3Config
|
| 536 |
+
base_model_prefix = "model"
|
| 537 |
+
supports_gradient_checkpointing = True
|
| 538 |
+
_no_split_modules = ["Lfm3DecoderLayer"]
|
| 539 |
+
_skip_keys_device_placement = ["past_key_values"]
|
| 540 |
+
_supports_flash_attn = True
|
| 541 |
+
_supports_sdpa = True
|
| 542 |
+
_supports_flex_attn = True
|
| 543 |
+
|
| 544 |
+
_can_compile_fullgraph = True
|
| 545 |
+
_supports_attention_backend = True
|
| 546 |
+
_can_record_outputs = {
|
| 547 |
+
"hidden_states": Lfm3DecoderLayer,
|
| 548 |
+
"attentions": Lfm3Attention,
|
| 549 |
+
}
|
| 550 |
+
|
| 551 |
+
@torch.no_grad()
|
| 552 |
+
def _init_weights(self, module):
|
| 553 |
+
super()._init_weights(module)
|
| 554 |
+
if isinstance(module, Lfm3RotaryEmbedding):
|
| 555 |
+
for layer_type in module.layer_types:
|
| 556 |
+
curr_inv_freq, _ = module.compute_rope_parameters(module.config, layer_type=layer_type)
|
| 557 |
+
init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
|
| 558 |
+
init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
|
| 559 |
+
|
| 560 |
+
@staticmethod
|
| 561 |
+
def create_masks_for_generate(
|
| 562 |
+
config: PreTrainedConfig,
|
| 563 |
+
inputs_embeds: torch.Tensor,
|
| 564 |
+
attention_mask: torch.Tensor | None,
|
| 565 |
+
past_key_values: Cache | None,
|
| 566 |
+
position_ids: torch.Tensor | None = None,
|
| 567 |
+
**kwargs,
|
| 568 |
+
) -> dict:
|
| 569 |
+
return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
|
| 570 |
+
|
| 571 |
+
|
| 572 |
+
def packed_under_padding(position_ids: torch.Tensor, padding_mask: torch.Tensor, filled) -> torch.Tensor:
|
| 573 |
+
"""Padding repeats position 0, so packing is a restart between two kept tokens, or a kept position 0 on a filled
|
| 574 |
+
cache whose row holds a kept token; attention under a padding mask cannot isolate it."""
|
| 575 |
+
restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
|
| 576 |
+
valid = padding_mask[:, -restarts.shape[-1] :].bool()
|
| 577 |
+
packed = restarts[:, 1:] & valid[:, 1:] & valid[:, :-1]
|
| 578 |
+
if filled:
|
| 579 |
+
# a row whose cached tokens are all padding (chunked prefill of a left-padded batch) has no document to pack
|
| 580 |
+
# against; a mask covering only the chunk keeps the strict rule
|
| 581 |
+
cached = padding_mask[:, : padding_mask.shape[-1] - restarts.shape[-1]]
|
| 582 |
+
cached_valid = cached.bool().any(-1, keepdim=True) if cached.shape[-1] else torch.ones_like(valid[:, :1])
|
| 583 |
+
packed = packed | (valid[:, :1] & (position_ids[:, :1] == 0) & cached_valid)
|
| 584 |
+
return packed
|
| 585 |
+
|
| 586 |
+
|
| 587 |
+
@auto_docstring
|
| 588 |
+
class Lfm3Model(Lfm3PreTrainedModel):
|
| 589 |
+
"""LFM2 model with sliding-attention layers as a third layer type: the mask and the rotary embedding are built per
|
| 590 |
+
layer type, and packed documents (`position_ids` restarting) reset the short-conv window through `seq_idx`."""
|
| 591 |
+
|
| 592 |
+
def __init__(self, config: Lfm3Config):
|
| 593 |
+
super().__init__(config)
|
| 594 |
+
self.padding_idx = config.pad_token_id
|
| 595 |
+
self.vocab_size = config.vocab_size
|
| 596 |
+
|
| 597 |
+
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
| 598 |
+
self.layers = nn.ModuleList(
|
| 599 |
+
[Lfm3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
| 600 |
+
)
|
| 601 |
+
self.rotary_emb = Lfm3RotaryEmbedding(config=config)
|
| 602 |
+
self.gradient_checkpointing = False
|
| 603 |
+
self.embedding_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 604 |
+
|
| 605 |
+
# Initialize weights and apply final processing
|
| 606 |
+
self.post_init()
|
| 607 |
+
|
| 608 |
+
@merge_with_config_defaults
|
| 609 |
+
@capture_outputs
|
| 610 |
+
@auto_docstring
|
| 611 |
+
def forward(
|
| 612 |
+
self,
|
| 613 |
+
input_ids: torch.LongTensor | None = None,
|
| 614 |
+
attention_mask: torch.Tensor | None = None,
|
| 615 |
+
position_ids: torch.LongTensor | None = None,
|
| 616 |
+
past_key_values: Cache | None = None,
|
| 617 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 618 |
+
use_cache: bool | None = None,
|
| 619 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 620 |
+
) -> BaseModelOutputWithPast:
|
| 621 |
+
r"""
|
| 622 |
+
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)` or `dict`, *optional*):
|
| 623 |
+
Mask to avoid performing attention on padding token indices (1 keeps a token, 0 masks it), or the mask
|
| 624 |
+
dict of `create_masks_for_generate`. A mask dict is used as given: `generate` refuses packed
|
| 625 |
+
`position_ids` before its prefill.
|
| 626 |
+
"""
|
| 627 |
+
if (input_ids is None) ^ (inputs_embeds is not None):
|
| 628 |
+
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
| 629 |
+
# `seq_idx` is a shared kwarg (`DataCollatorWithFlattening(return_seq_idx=True)`, read by Lfm2); LFM3 derives
|
| 630 |
+
# the documents from the restarting `position_ids` that come with it
|
| 631 |
+
if kwargs.pop("seq_idx", None) is not None and position_ids is None:
|
| 632 |
+
raise ValueError("Pass packed documents as `position_ids` restarting at 0; LFM3 does not read `seq_idx`")
|
| 633 |
+
|
| 634 |
+
if inputs_embeds is None:
|
| 635 |
+
inputs_embeds = self.embed_tokens(input_ids)
|
| 636 |
+
|
| 637 |
+
if use_cache and past_key_values is None:
|
| 638 |
+
past_key_values = DynamicCache(config=self.config)
|
| 639 |
+
filled = past_key_values is not None and (
|
| 640 |
+
past_key_values.get_seq_length() > 0
|
| 641 |
+
if any(t != "conv" for t in self.config.layer_types)
|
| 642 |
+
else past_key_values.has_previous_state()
|
| 643 |
+
)
|
| 644 |
+
packed_cache = getattr(past_key_values, "lfm3_packed", None)
|
| 645 |
+
if isinstance(packed_cache, torch.Tensor) or packed_cache:
|
| 646 |
+
# Reset layers have no history to mix. Ordinary caches carry a Python `False` and skip this, so the decode
|
| 647 |
+
# steps do no tensor work here; only a compiled packed prefill leaves a tensor flag, and `filled` stays
|
| 648 |
+
# tensor-valued for static caches.
|
| 649 |
+
packed_message = (
|
| 650 |
+
"A cache filled by a packed-document forward cannot be continued: its attention spans the documents"
|
| 651 |
+
)
|
| 652 |
+
packed_cache = packed_cache & filled if isinstance(packed_cache, torch.Tensor) else filled
|
| 653 |
+
if isinstance(packed_cache, torch.Tensor):
|
| 654 |
+
torch_compilable_check(~packed_cache, packed_message)
|
| 655 |
+
if not is_tracing(packed_cache):
|
| 656 |
+
past_key_values.lfm3_packed = bool(packed_cache)
|
| 657 |
+
elif packed_cache:
|
| 658 |
+
raise ValueError(packed_message)
|
| 659 |
+
else:
|
| 660 |
+
past_key_values.lfm3_packed = False
|
| 661 |
+
|
| 662 |
+
# The mask utilities detect packing only with `attention_mask=None`, so a multi-token forward with explicit
|
| 663 |
+
# positions drops the all-ones mask tokenizers emit (a data-dependent check, hence not when traced).
|
| 664 |
+
explicit_positions = position_ids is not None and inputs_embeds.shape[1] > 1
|
| 665 |
+
if (
|
| 666 |
+
explicit_positions
|
| 667 |
+
and isinstance(attention_mask, torch.Tensor)
|
| 668 |
+
and attention_mask.ndim == 2
|
| 669 |
+
and not is_tracing(position_ids)
|
| 670 |
+
and bool(attention_mask.all())
|
| 671 |
+
):
|
| 672 |
+
attention_mask = None
|
| 673 |
+
seq_idx = None
|
| 674 |
+
if position_ids is None:
|
| 675 |
+
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
| 676 |
+
position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
|
| 677 |
+
position_ids = position_ids.unsqueeze(0)
|
| 678 |
+
elif explicit_positions and not isinstance(attention_mask, dict):
|
| 679 |
+
# a mask dict comes from `generate`, which refused packing before its prefill
|
| 680 |
+
if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 2:
|
| 681 |
+
# data-dependent: a traced forward skips it (`generate` refuses packing eagerly before its prefill)
|
| 682 |
+
if not is_tracing(position_ids):
|
| 683 |
+
torch_compilable_check(
|
| 684 |
+
~packed_under_padding(position_ids, attention_mask, filled),
|
| 685 |
+
"Packed documents (`position_ids` restarting) cannot be combined with an attention mask: "
|
| 686 |
+
"attention would not isolate the documents. Pass `attention_mask=None` without a static cache.",
|
| 687 |
+
)
|
| 688 |
+
else:
|
| 689 |
+
restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
|
| 690 |
+
document_ids = restarts.cumsum(-1)
|
| 691 |
+
if is_tracing(position_ids) or bool((document_ids[..., -1] > 0).any()):
|
| 692 |
+
seq_idx = document_ids.to(torch.int32)
|
| 693 |
+
# Packed documents are prefilled in one forward on an empty cache (cached tokens would join the first
|
| 694 |
+
# document), and their cache is never continued
|
| 695 |
+
if filled:
|
| 696 |
+
torch_compilable_check(
|
| 697 |
+
(~restarts[:, 1:]).all() & (position_ids[:, 0] != 0).all(),
|
| 698 |
+
"Packed documents must be prefilled in one forward; the cache already holds tokens",
|
| 699 |
+
)
|
| 700 |
+
# Ordinary continuation must retain the convolution window from the cached document.
|
| 701 |
+
seq_idx = None
|
| 702 |
+
if attention_mask is None:
|
| 703 |
+
# without a mask only the cacheless packed-sequence mask isolates repeated zero positions (left
|
| 704 |
+
# padding): a cache-sized mask and every later decode step would attend them, so they count too
|
| 705 |
+
packed = (document_ids[..., -1] > 0).any()
|
| 706 |
+
else:
|
| 707 |
+
# a prepared 4D mask is the caller's and carries the padding (repeated zero positions), but may
|
| 708 |
+
# or may not isolate documents, so the cache records the packing and is never continued
|
| 709 |
+
steps = torch.diff(position_ids, dim=-1)
|
| 710 |
+
packed = ((steps != 1) & ~((steps == 0) & (position_ids[:, :-1] == 0))).any()
|
| 711 |
+
if past_key_values is not None:
|
| 712 |
+
# resolved here, where the prefill already syncs, so the decode steps carry a Python flag: no
|
| 713 |
+
# per-step host sync, data-dependent op or CUDA-graph output on the cache
|
| 714 |
+
past_key_values.lfm3_packed = packed if is_tracing(position_ids) else bool(packed)
|
| 715 |
+
|
| 716 |
+
if not isinstance(causal_mask_mapping := attention_mask, dict):
|
| 717 |
+
mask_kwargs = {
|
| 718 |
+
"config": self.config,
|
| 719 |
+
"inputs_embeds": inputs_embeds,
|
| 720 |
+
"attention_mask": attention_mask,
|
| 721 |
+
"past_key_values": past_key_values,
|
| 722 |
+
"position_ids": position_ids,
|
| 723 |
+
}
|
| 724 |
+
# The mask utilities only isolate packed documents without a cache, and an empty cache adds nothing.
|
| 725 |
+
layer_types = self.config.layer_types
|
| 726 |
+
if seq_idx is not None and past_key_values is not None:
|
| 727 |
+
attention_layer = next((i for i, t in enumerate(layer_types) if t != "conv"), None)
|
| 728 |
+
if attention_layer is not None and past_key_values.get_seq_length(attention_layer) == 0:
|
| 729 |
+
if past_key_values.is_compileable:
|
| 730 |
+
# static layers always return their full length, so the masks are sized by the cache, which
|
| 731 |
+
# isolates neither documents nor positional-only left padding
|
| 732 |
+
torch_compilable_check(
|
| 733 |
+
~packed,
|
| 734 |
+
"Packed documents, or left padding without an attention mask, cannot be prefilled into a "
|
| 735 |
+
"static cache: its masks are sized by the cache. Pass the padding as `attention_mask`, or "
|
| 736 |
+
"use a `DynamicCache` for packed documents",
|
| 737 |
+
)
|
| 738 |
+
else:
|
| 739 |
+
mask_kwargs["past_key_values"] = None
|
| 740 |
+
causal_mask_mapping = create_masks_by_layer_type(**mask_kwargs)
|
| 741 |
+
|
| 742 |
+
hidden_states = inputs_embeds
|
| 743 |
+
position_embeddings = {
|
| 744 |
+
layer_type: self.rotary_emb(hidden_states, position_ids=position_ids, layer_type=layer_type)
|
| 745 |
+
for layer_type in self.rotary_emb.layer_types
|
| 746 |
+
}
|
| 747 |
+
|
| 748 |
+
for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
|
| 749 |
+
layer_type = self.config.layer_types[i]
|
| 750 |
+
hidden_states = decoder_layer(
|
| 751 |
+
hidden_states,
|
| 752 |
+
attention_mask=causal_mask_mapping[layer_type],
|
| 753 |
+
position_embeddings=position_embeddings.get(layer_type),
|
| 754 |
+
position_ids=position_ids,
|
| 755 |
+
past_key_values=past_key_values,
|
| 756 |
+
seq_idx=seq_idx,
|
| 757 |
+
**kwargs,
|
| 758 |
+
)
|
| 759 |
+
|
| 760 |
+
hidden_states = self.embedding_norm(hidden_states)
|
| 761 |
+
|
| 762 |
+
return BaseModelOutputWithPast(
|
| 763 |
+
last_hidden_state=hidden_states,
|
| 764 |
+
past_key_values=past_key_values,
|
| 765 |
+
)
|
| 766 |
+
|
| 767 |
+
|
| 768 |
+
def reject_packed_generation(config: PreTrainedConfig, input_ids, generation_config, model_kwargs, is_first_iteration):
|
| 769 |
+
"""`generate` decodes from the cache its prefill fills, and no mask it builds isolates packed documents
|
| 770 |
+
(`position_ids` restarting): the prompt is refused before the prefill, one forward or every chunk, writes to it."""
|
| 771 |
+
position_ids, attention_mask = model_kwargs.get("position_ids"), model_kwargs.get("attention_mask")
|
| 772 |
+
cache = model_kwargs.get("past_key_values")
|
| 773 |
+
if cache is None or position_ids is None or position_ids.ndim != 2:
|
| 774 |
+
return
|
| 775 |
+
if attention_mask is not None and attention_mask.ndim != 2:
|
| 776 |
+
return
|
| 777 |
+
has_attention = any(t != "conv" for t in config.get_text_config().layer_types)
|
| 778 |
+
filled = bool(cache.get_seq_length() > 0 if has_attention else cache.has_previous_state())
|
| 779 |
+
inputs_embeds = model_kwargs.get("inputs_embeds") if is_first_iteration else None
|
| 780 |
+
length = input_ids.shape[1] if inputs_embeds is None else inputs_embeds.shape[1]
|
| 781 |
+
if generation_config.prefill_chunk_size is not None:
|
| 782 |
+
# the chunks run from the start of the prompt
|
| 783 |
+
position_ids = position_ids[:, :length]
|
| 784 |
+
else:
|
| 785 |
+
if filled and has_attention:
|
| 786 |
+
# the slice `_prefill` feeds: a prompt that spans the cache is cut to its uncached part
|
| 787 |
+
mask_length = getattr(generation_config, "_mask_length", -1)
|
| 788 |
+
if attention_mask is not None:
|
| 789 |
+
mask_length = attention_mask.shape[1]
|
| 790 |
+
if inputs_embeds is not None or mask_length == length:
|
| 791 |
+
length -= int(cache.get_seq_length())
|
| 792 |
+
position_ids = position_ids[:, position_ids.shape[1] - length :]
|
| 793 |
+
if attention_mask is None:
|
| 794 |
+
attention_mask = torch.ones_like(position_ids)
|
| 795 |
+
if bool(packed_under_padding(position_ids, attention_mask, filled).any()):
|
| 796 |
+
raise ValueError(
|
| 797 |
+
"`generate` cannot decode packed documents (`position_ids` restarting, or starting at 0 on a cache that "
|
| 798 |
+
"holds tokens): attention would span them. Generate each document in its own row of a left-padded batch."
|
| 799 |
+
)
|
| 800 |
+
|
| 801 |
+
|
| 802 |
+
@auto_docstring
|
| 803 |
+
class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
|
| 804 |
+
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
|
| 805 |
+
_tp_plan = {"lm_head": "colwise_gather_output"}
|
| 806 |
+
_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
|
| 807 |
+
_fsdp_plan = {"lm_head": "keep_full_weight"}
|
| 808 |
+
|
| 809 |
+
def __init__(self, config):
|
| 810 |
+
super().__init__(config)
|
| 811 |
+
self.model = Lfm3Model(config)
|
| 812 |
+
self.vocab_size = config.vocab_size
|
| 813 |
+
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
| 814 |
+
|
| 815 |
+
# Initialize weights and apply final processing
|
| 816 |
+
self.post_init()
|
| 817 |
+
|
| 818 |
+
@can_return_tuple
|
| 819 |
+
@auto_docstring
|
| 820 |
+
def forward(
|
| 821 |
+
self,
|
| 822 |
+
input_ids: torch.LongTensor | None = None,
|
| 823 |
+
attention_mask: torch.Tensor | None = None,
|
| 824 |
+
position_ids: torch.LongTensor | None = None,
|
| 825 |
+
past_key_values: Cache | None = None,
|
| 826 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 827 |
+
labels: torch.LongTensor | None = None,
|
| 828 |
+
use_cache: bool | None = None,
|
| 829 |
+
logits_to_keep: int | torch.Tensor = 0,
|
| 830 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 831 |
+
) -> CausalLMOutputWithPast:
|
| 832 |
+
r"""
|
| 833 |
+
Example:
|
| 834 |
+
|
| 835 |
+
```python
|
| 836 |
+
>>> from transformers import AutoTokenizer, Lfm3ForCausalLM
|
| 837 |
+
|
| 838 |
+
>>> model = Lfm3ForCausalLM.from_pretrained("LiquidAI/LFM3-1.7B")
|
| 839 |
+
>>> tokenizer = AutoTokenizer.from_pretrained("LiquidAI/LFM3-1.7B")
|
| 840 |
+
|
| 841 |
+
>>> inputs = tokenizer.apply_chat_template(
|
| 842 |
+
... [{"role": "user", "content": "What is C. elegans?"}], add_generation_prompt=True, return_tensors="pt"
|
| 843 |
+
... )
|
| 844 |
+
>>> generate_ids = model.generate(**inputs, max_new_tokens=64)
|
| 845 |
+
>>> tokenizer.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True)
|
| 846 |
+
```"""
|
| 847 |
+
outputs: BaseModelOutputWithPast = self.model(
|
| 848 |
+
input_ids=input_ids,
|
| 849 |
+
attention_mask=attention_mask,
|
| 850 |
+
position_ids=position_ids,
|
| 851 |
+
past_key_values=past_key_values,
|
| 852 |
+
inputs_embeds=inputs_embeds,
|
| 853 |
+
use_cache=use_cache,
|
| 854 |
+
**kwargs,
|
| 855 |
+
)
|
| 856 |
+
|
| 857 |
+
hidden_states = outputs.last_hidden_state
|
| 858 |
+
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
| 859 |
+
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
| 860 |
+
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
| 861 |
+
|
| 862 |
+
loss = None
|
| 863 |
+
if labels is not None:
|
| 864 |
+
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
|
| 865 |
+
|
| 866 |
+
return CausalLMOutputWithPast(
|
| 867 |
+
loss=loss,
|
| 868 |
+
logits=logits,
|
| 869 |
+
past_key_values=outputs.past_key_values,
|
| 870 |
+
hidden_states=outputs.hidden_states,
|
| 871 |
+
attentions=outputs.attentions,
|
| 872 |
+
)
|
| 873 |
+
|
| 874 |
+
def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
|
| 875 |
+
# every prefill passes here once before its first forward: a check per chunk would leave the earlier chunks
|
| 876 |
+
# in the cache, and a dynamic cache would only refuse packing at the first decode step
|
| 877 |
+
reject_packed_generation(self.config, input_ids, generation_config, model_kwargs, is_first_iteration)
|
| 878 |
+
return super()._prefill(input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration)
|
| 879 |
+
|
| 880 |
+
|
| 881 |
+
__all__ = ["Lfm3ForCausalLM", "Lfm3Model", "Lfm3PreTrainedModel"]
|
modeling_lfm3_vl.py
ADDED
|
@@ -0,0 +1,968 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Vendored from transformers' LFM3 code (not in a release yet) at 69da1b47a11e362fbb15bfe5d5c8da0790c688c5, src/transformers/models/lfm3_vl/modeling_lfm3_vl.py:
|
| 2 |
+
# its relative imports made absolute, nothing else changed.
|
| 3 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 4 |
+
# This file was automatically generated from src/transformers/models/lfm3_vl/modular_lfm3_vl.py.
|
| 5 |
+
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
| 6 |
+
# the file from the modular. If any change should be done, please apply the change to the
|
| 7 |
+
# modular_lfm3_vl.py file directly. One of our CI enforces this.
|
| 8 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 9 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 10 |
+
#
|
| 11 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 12 |
+
# you may not use this file except in compliance with the License.
|
| 13 |
+
# You may obtain a copy of the License at
|
| 14 |
+
#
|
| 15 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 16 |
+
#
|
| 17 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 18 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 19 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 20 |
+
# See the License for the specific language governing permissions and
|
| 21 |
+
# limitations under the License.
|
| 22 |
+
from collections.abc import Callable
|
| 23 |
+
from dataclasses import dataclass
|
| 24 |
+
|
| 25 |
+
import torch
|
| 26 |
+
import torch.nn.functional as F
|
| 27 |
+
from torch import nn
|
| 28 |
+
|
| 29 |
+
from transformers.activations import ACT2FN
|
| 30 |
+
from transformers.cache_utils import Cache
|
| 31 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 32 |
+
from transformers.generation import GenerationMixin
|
| 33 |
+
from transformers.integrations import use_kernel_forward_from_hub
|
| 34 |
+
from transformers.masking_utils import create_causal_mask, create_recurrent_attention_mask, create_sliding_window_causal_mask
|
| 35 |
+
from transformers.modeling_layers import GradientCheckpointingLayer
|
| 36 |
+
from transformers.modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast
|
| 37 |
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
| 38 |
+
from transformers.processing_utils import Unpack
|
| 39 |
+
from transformers.utils import (
|
| 40 |
+
TransformersKwargs,
|
| 41 |
+
auto_docstring,
|
| 42 |
+
can_return_tuple,
|
| 43 |
+
torch_compilable_check,
|
| 44 |
+
)
|
| 45 |
+
from transformers.utils.generic import (
|
| 46 |
+
accepts_precomputed_kwargs,
|
| 47 |
+
is_flash_attention_requested,
|
| 48 |
+
maybe_autocast,
|
| 49 |
+
merge_with_config_defaults,
|
| 50 |
+
)
|
| 51 |
+
from transformers.utils.output_capturing import OutputRecorder, capture_outputs
|
| 52 |
+
from transformers.vision_utils import get_max_seqlen, get_vision_attention_seqlens, get_vision_position_ids
|
| 53 |
+
from transformers.models.auto import AutoModel
|
| 54 |
+
from .configuration_lfm3_vl import Lfm3VlConfig, Lfm3VlVisionConfig
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
@use_kernel_forward_from_hub("RMSNorm")
|
| 58 |
+
class Lfm3VlRMSNorm(nn.Module):
|
| 59 |
+
def __init__(self, hidden_size, eps: float = 1e-6) -> None:
|
| 60 |
+
"""
|
| 61 |
+
Lfm3VlRMSNorm is equivalent to T5LayerNorm
|
| 62 |
+
"""
|
| 63 |
+
super().__init__()
|
| 64 |
+
self.weight = nn.Parameter(torch.ones(hidden_size))
|
| 65 |
+
self.variance_epsilon = eps
|
| 66 |
+
|
| 67 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 68 |
+
input_dtype = hidden_states.dtype
|
| 69 |
+
hidden_states = hidden_states.to(torch.float32)
|
| 70 |
+
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
| 71 |
+
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
| 72 |
+
return self.weight * hidden_states.to(input_dtype)
|
| 73 |
+
|
| 74 |
+
def extra_repr(self):
|
| 75 |
+
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
class Lfm3VlMLP(nn.Module):
|
| 79 |
+
def __init__(self, config: Lfm3VlConfig):
|
| 80 |
+
super().__init__()
|
| 81 |
+
self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 82 |
+
self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 83 |
+
self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
| 84 |
+
|
| 85 |
+
def forward(self, x):
|
| 86 |
+
return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
class Lfm3VlVisionRotaryEmbedding(nn.Module):
|
| 90 |
+
"""
|
| 91 |
+
Simple axial 2D rope with same freqs used for H and W grids. The freqs are
|
| 92 |
+
pre-computed using `head-dim//4` which is later used to concat H and W positions.
|
| 93 |
+
The final angles rotate over the whole head dim, no partial rotation involved.
|
| 94 |
+
"""
|
| 95 |
+
|
| 96 |
+
def __init__(self, config: Lfm3VlVisionConfig):
|
| 97 |
+
super().__init__()
|
| 98 |
+
self.config = config
|
| 99 |
+
|
| 100 |
+
self.rope_type = self.config.rope_parameters["rope_type"]
|
| 101 |
+
rope_init_fn: Callable = self.compute_axial_rope_parameters
|
| 102 |
+
if self.rope_type != "axial":
|
| 103 |
+
raise ValueError(f"{self.__class__.__name__} supports only axial rope, but requested {self.rope_type}")
|
| 104 |
+
inv_freq, self.attention_scaling = rope_init_fn(self.config)
|
| 105 |
+
|
| 106 |
+
self.inv_freq = nn.Buffer(inv_freq, persistent=False)
|
| 107 |
+
self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False)
|
| 108 |
+
|
| 109 |
+
@staticmethod
|
| 110 |
+
def compute_axial_rope_parameters(config: Lfm3VlVisionConfig, **kwargs) -> tuple[torch.Tensor, float]:
|
| 111 |
+
"""Inverse frequencies of the axial RoPE over half the head dimension per axis, with an attention scaling of 1."""
|
| 112 |
+
base = config.rope_parameters["rope_theta"]
|
| 113 |
+
spatial_dim = config.hidden_size // config.num_attention_heads // 2
|
| 114 |
+
inv_freq = 1.0 / (base ** (torch.arange(0, spatial_dim, 2, dtype=torch.float) / spatial_dim))
|
| 115 |
+
return inv_freq, 1.0
|
| 116 |
+
|
| 117 |
+
@torch.no_grad()
|
| 118 |
+
def forward(self, x, position_ids):
|
| 119 |
+
# position_ids: (2, N) — row 0 = h coords, row 1 = w coords
|
| 120 |
+
position_ids_expanded = position_ids[..., None].float()
|
| 121 |
+
device_type = x.device.type if isinstance(x.device.type, str) else "cpu"
|
| 122 |
+
with maybe_autocast(device_type=device_type, enabled=False):
|
| 123 |
+
freqs = position_ids_expanded * self.inv_freq.float()
|
| 124 |
+
cos = freqs.cos() * self.attention_scaling
|
| 125 |
+
sin = freqs.sin() * self.attention_scaling
|
| 126 |
+
|
| 127 |
+
cos = self.recomposition_frequencies(cos)
|
| 128 |
+
sin = self.recomposition_frequencies(sin)
|
| 129 |
+
return cos, sin
|
| 130 |
+
|
| 131 |
+
def recomposition_frequencies(self, freq):
|
| 132 |
+
"""
|
| 133 |
+
Recompose the frequencies into the final spatial layout used per each grid.
|
| 134 |
+
"""
|
| 135 |
+
freq_h, freq_w = freq[:, 0], freq[:, 1]
|
| 136 |
+
freq_hw = torch.cat([freq_h, freq_w], dim=-1)
|
| 137 |
+
return torch.cat([freq_hw, freq_hw], dim=-1)
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
def rotate_half(x):
|
| 141 |
+
"""Rotates half the hidden dims of the input."""
|
| 142 |
+
x1 = x[..., : x.shape[-1] // 2]
|
| 143 |
+
x2 = x[..., x.shape[-1] // 2 :]
|
| 144 |
+
return torch.cat((-x2, x1), dim=-1)
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def apply_rotary_pos_emb_vision(
|
| 148 |
+
q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
|
| 149 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 150 |
+
orig_q_dtype = q.dtype
|
| 151 |
+
orig_k_dtype = k.dtype
|
| 152 |
+
q, k = q.float(), k.float()
|
| 153 |
+
cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
|
| 154 |
+
q_embed = (q * cos) + (rotate_half(q) * sin)
|
| 155 |
+
k_embed = (k * cos) + (rotate_half(k) * sin)
|
| 156 |
+
q_embed = q_embed.to(orig_q_dtype)
|
| 157 |
+
k_embed = k_embed.to(orig_k_dtype)
|
| 158 |
+
return q_embed, k_embed
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
| 162 |
+
"""
|
| 163 |
+
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
| 164 |
+
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
| 165 |
+
"""
|
| 166 |
+
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
| 167 |
+
if n_rep == 1:
|
| 168 |
+
return hidden_states
|
| 169 |
+
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
| 170 |
+
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
def eager_attention_forward(
|
| 174 |
+
module: nn.Module,
|
| 175 |
+
query: torch.Tensor,
|
| 176 |
+
key: torch.Tensor,
|
| 177 |
+
value: torch.Tensor,
|
| 178 |
+
attention_mask: torch.Tensor | None,
|
| 179 |
+
scaling: float,
|
| 180 |
+
dropout: float = 0.0,
|
| 181 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 182 |
+
):
|
| 183 |
+
key_states = repeat_kv(key, module.num_key_value_groups)
|
| 184 |
+
value_states = repeat_kv(value, module.num_key_value_groups)
|
| 185 |
+
|
| 186 |
+
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
| 187 |
+
if attention_mask is not None:
|
| 188 |
+
attn_weights = attn_weights + attention_mask
|
| 189 |
+
|
| 190 |
+
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
|
| 191 |
+
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
|
| 192 |
+
attn_output = torch.matmul(attn_weights, value_states)
|
| 193 |
+
attn_output = attn_output.transpose(1, 2).contiguous()
|
| 194 |
+
|
| 195 |
+
return attn_output, attn_weights
|
| 196 |
+
|
| 197 |
+
|
| 198 |
+
class Lfm3VlVisionAttention(nn.Module):
|
| 199 |
+
"""Bidirectional grouped-query attention over the patches of one frame, with RMSNorm on the query and key heads."""
|
| 200 |
+
|
| 201 |
+
def __init__(self, config: Lfm3VlVisionConfig):
|
| 202 |
+
super().__init__()
|
| 203 |
+
self.config = config
|
| 204 |
+
self.num_heads = config.num_attention_heads
|
| 205 |
+
self.num_key_value_heads = config.num_key_value_heads
|
| 206 |
+
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
| 207 |
+
self.head_dim = config.hidden_size // config.num_attention_heads
|
| 208 |
+
self.scaling = self.head_dim**-0.5
|
| 209 |
+
self.attention_dropout = 0.0
|
| 210 |
+
self.is_causal = False
|
| 211 |
+
self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
|
| 212 |
+
self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
| 213 |
+
self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
| 214 |
+
self.out_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
|
| 215 |
+
self.q_layernorm = Lfm3VlRMSNorm(self.head_dim, eps=config.norm_eps)
|
| 216 |
+
self.k_layernorm = Lfm3VlRMSNorm(self.head_dim, eps=config.norm_eps)
|
| 217 |
+
|
| 218 |
+
def forward(
|
| 219 |
+
self,
|
| 220 |
+
hidden_states: torch.Tensor,
|
| 221 |
+
cu_seqlens: torch.Tensor,
|
| 222 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
| 223 |
+
max_seqlen: int | None = None,
|
| 224 |
+
frame_shapes: list[list[int]] | None = None,
|
| 225 |
+
**kwargs,
|
| 226 |
+
) -> torch.Tensor:
|
| 227 |
+
seq_length = hidden_states.shape[0]
|
| 228 |
+
query_states = self.q_layernorm(self.q_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim))
|
| 229 |
+
key_states = self.k_layernorm(
|
| 230 |
+
self.k_proj(hidden_states).view(seq_length, self.num_key_value_heads, self.head_dim)
|
| 231 |
+
)
|
| 232 |
+
value_states = self.v_proj(hidden_states).view(seq_length, self.num_key_value_heads, self.head_dim)
|
| 233 |
+
|
| 234 |
+
cos, sin = position_embeddings
|
| 235 |
+
query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
|
| 236 |
+
|
| 237 |
+
query_states = query_states.transpose(0, 1).unsqueeze(0)
|
| 238 |
+
key_states = key_states.transpose(0, 1).unsqueeze(0)
|
| 239 |
+
value_states = value_states.transpose(0, 1).unsqueeze(0)
|
| 240 |
+
|
| 241 |
+
attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
|
| 242 |
+
self.config._attn_implementation, eager_attention_forward
|
| 243 |
+
)
|
| 244 |
+
|
| 245 |
+
if is_flash_attention_requested(self.config):
|
| 246 |
+
max_seqlen = get_max_seqlen(cu_seqlens, self.config, kwargs={"max_seqlen": max_seqlen})
|
| 247 |
+
attn_output, attn_weights = attention_interface(
|
| 248 |
+
self,
|
| 249 |
+
query_states,
|
| 250 |
+
key_states,
|
| 251 |
+
value_states,
|
| 252 |
+
attention_mask=None,
|
| 253 |
+
scaling=self.scaling,
|
| 254 |
+
dropout=0.0,
|
| 255 |
+
cu_seq_lens_q=cu_seqlens,
|
| 256 |
+
cu_seq_lens_k=cu_seqlens,
|
| 257 |
+
max_length_q=max_seqlen,
|
| 258 |
+
max_length_k=max_seqlen,
|
| 259 |
+
is_causal=False,
|
| 260 |
+
**kwargs,
|
| 261 |
+
)
|
| 262 |
+
else:
|
| 263 |
+
# Other implementations: frames of one shape (every video, equal-size images) attend as one batch, mixed
|
| 264 |
+
# shapes frame by frame. `frame_shapes` gives the lengths without a device sync.
|
| 265 |
+
if frame_shapes is not None:
|
| 266 |
+
lengths = [height * width for height, width in frame_shapes]
|
| 267 |
+
else:
|
| 268 |
+
lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
|
| 269 |
+
if len(set(lengths)) == 1:
|
| 270 |
+
query_states, key_states, value_states = (
|
| 271 |
+
tensor.reshape(-1, len(lengths), lengths[0], self.head_dim).transpose(0, 1)
|
| 272 |
+
for tensor in (query_states, key_states, value_states)
|
| 273 |
+
)
|
| 274 |
+
attn_output, attn_weights = attention_interface(
|
| 275 |
+
self,
|
| 276 |
+
query_states,
|
| 277 |
+
key_states,
|
| 278 |
+
value_states,
|
| 279 |
+
attention_mask=None,
|
| 280 |
+
scaling=self.scaling,
|
| 281 |
+
dropout=0.0,
|
| 282 |
+
is_causal=False,
|
| 283 |
+
**kwargs,
|
| 284 |
+
)
|
| 285 |
+
else:
|
| 286 |
+
splits = [torch.split(tensor, lengths, dim=2) for tensor in (query_states, key_states, value_states)]
|
| 287 |
+
attn_outputs, attn_weights = [], []
|
| 288 |
+
for q, k, v in zip(*splits):
|
| 289 |
+
output, weights = attention_interface(
|
| 290 |
+
self,
|
| 291 |
+
q,
|
| 292 |
+
k,
|
| 293 |
+
v,
|
| 294 |
+
attention_mask=None,
|
| 295 |
+
scaling=self.scaling,
|
| 296 |
+
dropout=0.0,
|
| 297 |
+
is_causal=False,
|
| 298 |
+
**kwargs,
|
| 299 |
+
)
|
| 300 |
+
attn_outputs.append(output)
|
| 301 |
+
attn_weights.append(weights)
|
| 302 |
+
attn_output = torch.cat(attn_outputs, dim=1)
|
| 303 |
+
# frames of different shapes have differently sized maps: one entry per frame, or None under sdpa
|
| 304 |
+
attn_weights = None if attn_weights[0] is None else tuple(attn_weights)
|
| 305 |
+
|
| 306 |
+
attn_output = attn_output.reshape(seq_length, -1).contiguous()
|
| 307 |
+
attn_output = self.out_proj(attn_output)
|
| 308 |
+
return attn_output, attn_weights
|
| 309 |
+
|
| 310 |
+
|
| 311 |
+
class Lfm3VlVisionConv2d(nn.Module):
|
| 312 |
+
"""Gated depthwise 2D convolution over the patch grid of each frame."""
|
| 313 |
+
|
| 314 |
+
def __init__(self, config: Lfm3VlVisionConfig):
|
| 315 |
+
super().__init__()
|
| 316 |
+
self.spatial_merge_size = config.spatial_merge_size
|
| 317 |
+
self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
|
| 318 |
+
self.spatial_conv = nn.Conv2d(
|
| 319 |
+
config.hidden_size,
|
| 320 |
+
config.hidden_size,
|
| 321 |
+
kernel_size=config.conv_kernel_size,
|
| 322 |
+
padding=config.conv_kernel_size // 2,
|
| 323 |
+
groups=config.hidden_size,
|
| 324 |
+
bias=False,
|
| 325 |
+
)
|
| 326 |
+
self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
|
| 327 |
+
|
| 328 |
+
def forward(self, hidden_states: torch.Tensor, frame_shapes: list[list[int]]):
|
| 329 |
+
B, C, x = self.in_proj(hidden_states).chunk(3, dim=-1)
|
| 330 |
+
gated = B * x
|
| 331 |
+
merge = self.spatial_merge_size
|
| 332 |
+
num_frames = len(frame_shapes)
|
| 333 |
+
if len(set(map(tuple, frame_shapes))) == 1:
|
| 334 |
+
# frames of one shape (every video, equal-size images) convolve as one batch: merge-block order -> raster
|
| 335 |
+
# grid, convolve, back
|
| 336 |
+
height, width = frame_shapes[0]
|
| 337 |
+
frames = gated.view(num_frames, height // merge, width // merge, merge, merge, -1)
|
| 338 |
+
frames = frames.permute(0, 5, 1, 3, 2, 4).reshape(num_frames, -1, height, width)
|
| 339 |
+
frames = self.spatial_conv(frames)
|
| 340 |
+
frames = frames.view(num_frames, -1, height // merge, merge, width // merge, merge).permute(
|
| 341 |
+
0, 2, 4, 3, 5, 1
|
| 342 |
+
)
|
| 343 |
+
conv_output = frames.reshape(gated.shape[0], -1)
|
| 344 |
+
else:
|
| 345 |
+
conv_outputs = []
|
| 346 |
+
lengths = [height * width for height, width in frame_shapes]
|
| 347 |
+
for frame, (height, width) in zip(torch.split(gated, lengths, dim=0), frame_shapes):
|
| 348 |
+
frame = frame.view(height // merge, width // merge, merge, merge, -1)
|
| 349 |
+
frame = frame.permute(4, 0, 2, 1, 3).reshape(1, -1, height, width)
|
| 350 |
+
frame = self.spatial_conv(frame)
|
| 351 |
+
frame = frame.view(-1, height // merge, merge, width // merge, merge).permute(1, 3, 2, 4, 0)
|
| 352 |
+
conv_outputs.append(frame.reshape(height * width, -1))
|
| 353 |
+
conv_output = torch.cat(conv_outputs, dim=0)
|
| 354 |
+
return self.out_proj(C * conv_output)
|
| 355 |
+
|
| 356 |
+
|
| 357 |
+
class Lfm3VlVisionLayer(GradientCheckpointingLayer):
|
| 358 |
+
def __init__(self, config: Lfm3VlVisionConfig, layer_idx: int):
|
| 359 |
+
super().__init__()
|
| 360 |
+
self.is_attention_layer = config.layer_types[layer_idx] == "full_attention"
|
| 361 |
+
if self.is_attention_layer:
|
| 362 |
+
self.self_attn = Lfm3VlVisionAttention(config)
|
| 363 |
+
else:
|
| 364 |
+
self.conv = Lfm3VlVisionConv2d(config)
|
| 365 |
+
self.feed_forward = Lfm3VlMLP(config)
|
| 366 |
+
self.operator_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 367 |
+
self.ffn_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 368 |
+
|
| 369 |
+
def forward(
|
| 370 |
+
self,
|
| 371 |
+
hidden_states: torch.Tensor,
|
| 372 |
+
cu_seqlens: torch.Tensor,
|
| 373 |
+
frame_shapes: list[list[int]],
|
| 374 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
| 375 |
+
**kwargs,
|
| 376 |
+
) -> torch.Tensor:
|
| 377 |
+
if self.is_attention_layer:
|
| 378 |
+
operator_output, _ = self.self_attn(
|
| 379 |
+
self.operator_norm(hidden_states),
|
| 380 |
+
cu_seqlens=cu_seqlens,
|
| 381 |
+
position_embeddings=position_embeddings,
|
| 382 |
+
frame_shapes=frame_shapes,
|
| 383 |
+
**kwargs,
|
| 384 |
+
)
|
| 385 |
+
else:
|
| 386 |
+
operator_output = self.conv(self.operator_norm(hidden_states), frame_shapes)
|
| 387 |
+
hidden_states = hidden_states + operator_output
|
| 388 |
+
hidden_states = hidden_states + self.feed_forward(self.ffn_norm(hidden_states))
|
| 389 |
+
return hidden_states
|
| 390 |
+
|
| 391 |
+
|
| 392 |
+
def create_masks_by_layer_type(
|
| 393 |
+
config: PreTrainedConfig,
|
| 394 |
+
inputs_embeds: torch.Tensor,
|
| 395 |
+
attention_mask: torch.Tensor | None,
|
| 396 |
+
past_key_values: Cache | None,
|
| 397 |
+
position_ids: torch.Tensor | None = None,
|
| 398 |
+
) -> dict:
|
| 399 |
+
"""The cache sizes each mask from the layer it is asked about, and the generic mask utilities ask about the first
|
| 400 |
+
non-sliding layer, which for a leading `conv` layer redirects to the first attention layer of any type: name the
|
| 401 |
+
first layer of each mask's own type instead."""
|
| 402 |
+
config = config.get_text_config()
|
| 403 |
+
mask_kwargs = {
|
| 404 |
+
"config": config,
|
| 405 |
+
"inputs_embeds": inputs_embeds,
|
| 406 |
+
"attention_mask": attention_mask,
|
| 407 |
+
"past_key_values": past_key_values,
|
| 408 |
+
"position_ids": position_ids,
|
| 409 |
+
}
|
| 410 |
+
layer_types, causal_mask_mapping = config.layer_types, {}
|
| 411 |
+
if "full_attention" in layer_types:
|
| 412 |
+
layer_idx = layer_types.index("full_attention")
|
| 413 |
+
causal_mask_mapping["full_attention"] = create_causal_mask(**mask_kwargs, layer_idx=layer_idx)
|
| 414 |
+
if "sliding_attention" in layer_types:
|
| 415 |
+
layer_idx = layer_types.index("sliding_attention")
|
| 416 |
+
causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
|
| 417 |
+
**mask_kwargs, layer_idx=layer_idx
|
| 418 |
+
)
|
| 419 |
+
if "conv" in layer_types:
|
| 420 |
+
causal_mask_mapping["conv"] = create_recurrent_attention_mask(**mask_kwargs)
|
| 421 |
+
return causal_mask_mapping
|
| 422 |
+
|
| 423 |
+
|
| 424 |
+
@auto_docstring
|
| 425 |
+
class Lfm3VlPreTrainedModel(PreTrainedModel):
|
| 426 |
+
config: Lfm3VlConfig
|
| 427 |
+
base_model_prefix = "model"
|
| 428 |
+
input_modalities = ("image", "video", "text")
|
| 429 |
+
supports_gradient_checkpointing = True
|
| 430 |
+
_skip_keys_device_placement = ["past_key_values"]
|
| 431 |
+
|
| 432 |
+
_supports_flash_attn = True
|
| 433 |
+
_supports_sdpa = True
|
| 434 |
+
# the cached decode step traces without graph breaks, like the text model
|
| 435 |
+
_can_compile_fullgraph = True
|
| 436 |
+
_supports_flex_attn = True
|
| 437 |
+
_supports_attention_backend = True
|
| 438 |
+
_no_split_modules = ["Lfm3VlVisionLayer"]
|
| 439 |
+
|
| 440 |
+
@staticmethod
|
| 441 |
+
def create_masks_for_generate(
|
| 442 |
+
config: PreTrainedConfig,
|
| 443 |
+
inputs_embeds: torch.Tensor,
|
| 444 |
+
attention_mask: torch.Tensor | None,
|
| 445 |
+
past_key_values: Cache | None,
|
| 446 |
+
position_ids: torch.Tensor | None = None,
|
| 447 |
+
**kwargs,
|
| 448 |
+
) -> dict:
|
| 449 |
+
# the static-cache masks of the text model are sized by layer type, not by the generic first attention layer
|
| 450 |
+
return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
|
| 451 |
+
|
| 452 |
+
|
| 453 |
+
def get_vision_frame_shapes(grid_thw: torch.Tensor) -> list[list[int]]:
|
| 454 |
+
return grid_thw[:, 1:].repeat_interleave(grid_thw[:, 0], dim=0).tolist()
|
| 455 |
+
|
| 456 |
+
|
| 457 |
+
@auto_docstring(
|
| 458 |
+
custom_intro="""
|
| 459 |
+
The LFM3-VL vision encoder: a stack of gated convolution and bidirectional attention layers over the patches of
|
| 460 |
+
each image or video frame, embedded by a linear patchifier.
|
| 461 |
+
"""
|
| 462 |
+
)
|
| 463 |
+
class Lfm3VlVisionModel(Lfm3VlPreTrainedModel):
|
| 464 |
+
config: Lfm3VlVisionConfig
|
| 465 |
+
input_modalities = ("image", "video")
|
| 466 |
+
_input_embed_layer = "patchifier"
|
| 467 |
+
_can_record_outputs = {
|
| 468 |
+
"hidden_states": Lfm3VlVisionLayer,
|
| 469 |
+
"attentions": OutputRecorder(Lfm3VlVisionAttention, index=1),
|
| 470 |
+
}
|
| 471 |
+
|
| 472 |
+
def __init__(self, config: Lfm3VlVisionConfig):
|
| 473 |
+
super().__init__(config)
|
| 474 |
+
self.spatial_merge_size = config.spatial_merge_size
|
| 475 |
+
patch_dim = config.in_channels * config.temporal_patch_size * config.patch_size**2
|
| 476 |
+
self.patchifier = nn.Linear(patch_dim, config.hidden_size, bias=False)
|
| 477 |
+
self.embedding_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
|
| 478 |
+
self.rotary_pos_emb = Lfm3VlVisionRotaryEmbedding(config)
|
| 479 |
+
self.layers = nn.ModuleList(
|
| 480 |
+
[Lfm3VlVisionLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
| 481 |
+
)
|
| 482 |
+
self.gradient_checkpointing = False
|
| 483 |
+
self.post_init()
|
| 484 |
+
|
| 485 |
+
@merge_with_config_defaults
|
| 486 |
+
@capture_outputs
|
| 487 |
+
@auto_docstring
|
| 488 |
+
def forward(
|
| 489 |
+
self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
|
| 490 |
+
) -> BaseModelOutputWithPooling:
|
| 491 |
+
r"""
|
| 492 |
+
pixel_values (`torch.Tensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 493 |
+
The flattened patches of every image or video, grouped by merge block.
|
| 494 |
+
grid_thw (`torch.LongTensor` of shape `(num_images_or_videos, 3)`):
|
| 495 |
+
The temporal, height and width dimensions of the patch grid of each image or video.
|
| 496 |
+
"""
|
| 497 |
+
position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
|
| 498 |
+
cu_seqlens, max_seqlen = get_vision_attention_seqlens(grid_thw, self.config, kwargs=kwargs)
|
| 499 |
+
frame_shapes = get_vision_frame_shapes(grid_thw)
|
| 500 |
+
|
| 501 |
+
hidden_states = self.patchifier(pixel_values.to(self.patchifier.weight.dtype))
|
| 502 |
+
hidden_states = self.embedding_norm(hidden_states)
|
| 503 |
+
position_embeddings = self.rotary_pos_emb(hidden_states, position_ids)
|
| 504 |
+
|
| 505 |
+
for layer in self.layers:
|
| 506 |
+
hidden_states = layer(
|
| 507 |
+
hidden_states,
|
| 508 |
+
cu_seqlens=cu_seqlens,
|
| 509 |
+
frame_shapes=frame_shapes,
|
| 510 |
+
max_seqlen=max_seqlen,
|
| 511 |
+
position_embeddings=position_embeddings,
|
| 512 |
+
**kwargs,
|
| 513 |
+
)
|
| 514 |
+
|
| 515 |
+
return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
|
| 516 |
+
|
| 517 |
+
|
| 518 |
+
class Lfm3VlMultiModalProjector(nn.Module):
|
| 519 |
+
"""Merges each block of `downsample_factor**2` vision tokens into one language token."""
|
| 520 |
+
|
| 521 |
+
def __init__(self, config: Lfm3VlConfig):
|
| 522 |
+
super().__init__()
|
| 523 |
+
self.downsample_factor = config.downsample_factor
|
| 524 |
+
self.layer_norm = nn.LayerNorm(config.vision_config.hidden_size, eps=config.projector_layer_norm_eps)
|
| 525 |
+
self.linear_1 = nn.Linear(
|
| 526 |
+
config.vision_config.hidden_size * config.downsample_factor**2,
|
| 527 |
+
config.projector_hidden_size,
|
| 528 |
+
bias=config.projector_bias,
|
| 529 |
+
)
|
| 530 |
+
self.act = ACT2FN[config.projector_hidden_act]
|
| 531 |
+
self.linear_2 = nn.Linear(
|
| 532 |
+
config.projector_hidden_size, config.text_config.hidden_size, bias=config.projector_bias
|
| 533 |
+
)
|
| 534 |
+
|
| 535 |
+
def forward(self, image_features: torch.Tensor):
|
| 536 |
+
hidden_states = self.layer_norm(image_features).reshape(-1, self.linear_1.in_features)
|
| 537 |
+
hidden_states = self.linear_1(hidden_states)
|
| 538 |
+
hidden_states = self.act(hidden_states)
|
| 539 |
+
hidden_states = self.linear_2(hidden_states)
|
| 540 |
+
return hidden_states
|
| 541 |
+
|
| 542 |
+
|
| 543 |
+
@auto_docstring(
|
| 544 |
+
custom_intro="""
|
| 545 |
+
Base class for Lfm3Vl outputs, with hidden states and attentions.
|
| 546 |
+
"""
|
| 547 |
+
)
|
| 548 |
+
@dataclass
|
| 549 |
+
class Lfm3VlModelOutputWithPast(BaseModelOutputWithPast):
|
| 550 |
+
r"""
|
| 551 |
+
past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
|
| 552 |
+
It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
|
| 553 |
+
|
| 554 |
+
Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
|
| 555 |
+
`past_key_values` input) to speed up sequential decoding.
|
| 556 |
+
image_hidden_states (`torch.FloatTensor`, *optional*):
|
| 557 |
+
A `torch.FloatTensor` of size `(num_image_tokens, hidden_size)`.
|
| 558 |
+
image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
|
| 559 |
+
video_hidden_states (`torch.FloatTensor`, *optional*):
|
| 560 |
+
A `torch.FloatTensor` of size `(num_video_tokens, hidden_size)`.
|
| 561 |
+
video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
|
| 562 |
+
"""
|
| 563 |
+
|
| 564 |
+
image_hidden_states: torch.FloatTensor | None = None
|
| 565 |
+
video_hidden_states: torch.FloatTensor | None = None
|
| 566 |
+
|
| 567 |
+
|
| 568 |
+
@auto_docstring(
|
| 569 |
+
custom_intro="""
|
| 570 |
+
Base class for Lfm3Vl causal language model (or autoregressive) outputs.
|
| 571 |
+
"""
|
| 572 |
+
)
|
| 573 |
+
@dataclass
|
| 574 |
+
class Lfm3VlCausalLMOutputWithPast(CausalLMOutputWithPast):
|
| 575 |
+
r"""
|
| 576 |
+
image_hidden_states (`torch.FloatTensor`, *optional*):
|
| 577 |
+
A `torch.FloatTensor` of size `(num_image_tokens, hidden_size)`.
|
| 578 |
+
image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
|
| 579 |
+
video_hidden_states (`torch.FloatTensor`, *optional*):
|
| 580 |
+
A `torch.FloatTensor` of size `(num_video_tokens, hidden_size)`.
|
| 581 |
+
video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
|
| 582 |
+
"""
|
| 583 |
+
|
| 584 |
+
image_hidden_states: torch.FloatTensor | None = None
|
| 585 |
+
video_hidden_states: torch.FloatTensor | None = None
|
| 586 |
+
|
| 587 |
+
|
| 588 |
+
@auto_docstring(
|
| 589 |
+
custom_intro="""
|
| 590 |
+
The Lfm3Vl model which consists of a vision backbone and a language model, without a language modeling head.
|
| 591 |
+
"""
|
| 592 |
+
)
|
| 593 |
+
class Lfm3VlModel(Lfm3VlPreTrainedModel):
|
| 594 |
+
def __init__(self, config: Lfm3VlConfig):
|
| 595 |
+
super().__init__(config)
|
| 596 |
+
self.vision_tower = AutoModel.from_config(config.vision_config)
|
| 597 |
+
self.multi_modal_projector = Lfm3VlMultiModalProjector(config)
|
| 598 |
+
self.language_model = AutoModel.from_config(config.text_config)
|
| 599 |
+
self.post_init()
|
| 600 |
+
|
| 601 |
+
@accepts_precomputed_kwargs(modality="image")
|
| 602 |
+
@can_return_tuple
|
| 603 |
+
@auto_docstring(
|
| 604 |
+
custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection."
|
| 605 |
+
)
|
| 606 |
+
def get_image_features(
|
| 607 |
+
self,
|
| 608 |
+
pixel_values: torch.FloatTensor,
|
| 609 |
+
image_grid_thw: torch.LongTensor,
|
| 610 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 611 |
+
) -> tuple | BaseModelOutputWithPooling:
|
| 612 |
+
r"""
|
| 613 |
+
pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 614 |
+
The flattened patches of every image, grouped by merge block.
|
| 615 |
+
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
|
| 616 |
+
The temporal, height and width dimensions of the patch grid of each image.
|
| 617 |
+
"""
|
| 618 |
+
image_outputs = self.vision_tower(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs)
|
| 619 |
+
image_features = self.multi_modal_projector(image_outputs.last_hidden_state)
|
| 620 |
+
split_sizes = (image_grid_thw.prod(-1) // self.vision_tower.spatial_merge_size**2).tolist()
|
| 621 |
+
image_outputs.pooler_output = torch.split(image_features, split_sizes)
|
| 622 |
+
return image_outputs
|
| 623 |
+
|
| 624 |
+
def get_placeholder_mask(
|
| 625 |
+
self,
|
| 626 |
+
input_ids: torch.LongTensor,
|
| 627 |
+
inputs_embeds: torch.FloatTensor,
|
| 628 |
+
image_features: torch.FloatTensor | None = None,
|
| 629 |
+
video_features: torch.FloatTensor | None = None,
|
| 630 |
+
):
|
| 631 |
+
"""Placeholder masks from `input_ids` or `inputs_embeds`, checked against the feature counts."""
|
| 632 |
+
if input_ids is None:
|
| 633 |
+
special_image_mask = inputs_embeds == self.get_input_embeddings()(
|
| 634 |
+
torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
|
| 635 |
+
)
|
| 636 |
+
special_image_mask = special_image_mask.all(-1)
|
| 637 |
+
special_video_mask = inputs_embeds == self.get_input_embeddings()(
|
| 638 |
+
torch.full((), self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
|
| 639 |
+
)
|
| 640 |
+
special_video_mask = special_video_mask.all(-1)
|
| 641 |
+
else:
|
| 642 |
+
special_image_mask = input_ids == self.config.image_token_id
|
| 643 |
+
special_video_mask = input_ids == self.config.video_token_id
|
| 644 |
+
|
| 645 |
+
n_image_tokens = special_image_mask.sum()
|
| 646 |
+
special_image_mask = special_image_mask.unsqueeze(-1).to(inputs_embeds.device)
|
| 647 |
+
if image_features is not None:
|
| 648 |
+
torch_compilable_check(
|
| 649 |
+
n_image_tokens * inputs_embeds.shape[-1] == image_features.numel(),
|
| 650 |
+
f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
|
| 651 |
+
)
|
| 652 |
+
|
| 653 |
+
n_video_tokens = special_video_mask.sum()
|
| 654 |
+
special_video_mask = special_video_mask.unsqueeze(-1).to(inputs_embeds.device)
|
| 655 |
+
if video_features is not None:
|
| 656 |
+
torch_compilable_check(
|
| 657 |
+
n_video_tokens * inputs_embeds.shape[-1] == video_features.numel(),
|
| 658 |
+
f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
|
| 659 |
+
)
|
| 660 |
+
return special_image_mask, special_video_mask
|
| 661 |
+
|
| 662 |
+
@can_return_tuple
|
| 663 |
+
@auto_docstring
|
| 664 |
+
def forward(
|
| 665 |
+
self,
|
| 666 |
+
input_ids: torch.LongTensor | None = None,
|
| 667 |
+
attention_mask: torch.Tensor | None = None,
|
| 668 |
+
position_ids: torch.LongTensor | None = None,
|
| 669 |
+
past_key_values: Cache | None = None,
|
| 670 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 671 |
+
pixel_values: torch.FloatTensor | None = None,
|
| 672 |
+
image_grid_thw: torch.LongTensor | None = None,
|
| 673 |
+
pixel_values_videos: torch.FloatTensor | None = None,
|
| 674 |
+
video_grid_thw: torch.LongTensor | None = None,
|
| 675 |
+
use_cache: bool | None = None,
|
| 676 |
+
mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None,
|
| 677 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 678 |
+
) -> tuple | Lfm3VlModelOutputWithPast:
|
| 679 |
+
r"""
|
| 680 |
+
pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
|
| 681 |
+
The flattened patches of every image, grouped by merge block.
|
| 682 |
+
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
|
| 683 |
+
The temporal, height and width dimensions of the patch grid of each image.
|
| 684 |
+
pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
|
| 685 |
+
The flattened patches of every video, grouped by merge block.
|
| 686 |
+
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
|
| 687 |
+
The temporal, height and width dimensions of the patch grid of each video.
|
| 688 |
+
"""
|
| 689 |
+
if (input_ids is None) ^ (inputs_embeds is not None):
|
| 690 |
+
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
| 691 |
+
|
| 692 |
+
if (pixel_values is not None or pixel_values_videos is not None) and mm_encoder_outputs is not None:
|
| 693 |
+
raise ValueError(
|
| 694 |
+
"You cannot specify both pixel_values/pixel_values_videos and mm_encoder_outputs at the same time"
|
| 695 |
+
)
|
| 696 |
+
|
| 697 |
+
if inputs_embeds is None:
|
| 698 |
+
inputs_embeds = self.get_input_embeddings()(input_ids)
|
| 699 |
+
|
| 700 |
+
mm_encoder_outputs = mm_encoder_outputs if mm_encoder_outputs is not None else {}
|
| 701 |
+
if (pixel_values is None) != (image_grid_thw is None):
|
| 702 |
+
raise ValueError("pixel_values and image_grid_thw must be supplied together")
|
| 703 |
+
if (pixel_values_videos is None) != (video_grid_thw is None):
|
| 704 |
+
raise ValueError("pixel_values_videos and video_grid_thw must be supplied together")
|
| 705 |
+
if mm_encoder_outputs.get("image") is None and pixel_values is not None:
|
| 706 |
+
mm_encoder_outputs["image"] = self.get_image_features(
|
| 707 |
+
pixel_values, image_grid_thw, return_dict=True, **kwargs
|
| 708 |
+
)
|
| 709 |
+
if mm_encoder_outputs.get("video") is None and pixel_values_videos is not None:
|
| 710 |
+
mm_encoder_outputs["video"] = self.get_video_features(
|
| 711 |
+
pixel_values_videos, video_grid_thw, return_dict=True, **kwargs
|
| 712 |
+
)
|
| 713 |
+
|
| 714 |
+
image_features = None
|
| 715 |
+
if mm_encoder_outputs.get("image") is not None:
|
| 716 |
+
image_features = torch.cat(mm_encoder_outputs["image"].pooler_output, dim=0).to(
|
| 717 |
+
inputs_embeds.device, inputs_embeds.dtype
|
| 718 |
+
)
|
| 719 |
+
image_mask, _ = self.get_placeholder_mask(
|
| 720 |
+
input_ids, inputs_embeds=inputs_embeds, image_features=image_features
|
| 721 |
+
)
|
| 722 |
+
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_features)
|
| 723 |
+
|
| 724 |
+
video_features = None
|
| 725 |
+
if mm_encoder_outputs.get("video") is not None:
|
| 726 |
+
video_features = torch.cat(mm_encoder_outputs["video"].pooler_output, dim=0).to(
|
| 727 |
+
inputs_embeds.device, inputs_embeds.dtype
|
| 728 |
+
)
|
| 729 |
+
_, video_mask = self.get_placeholder_mask(
|
| 730 |
+
input_ids, inputs_embeds=inputs_embeds, video_features=video_features
|
| 731 |
+
)
|
| 732 |
+
inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_features)
|
| 733 |
+
|
| 734 |
+
outputs = self.language_model(
|
| 735 |
+
attention_mask=attention_mask,
|
| 736 |
+
position_ids=position_ids,
|
| 737 |
+
past_key_values=past_key_values,
|
| 738 |
+
inputs_embeds=inputs_embeds,
|
| 739 |
+
use_cache=use_cache,
|
| 740 |
+
**kwargs,
|
| 741 |
+
)
|
| 742 |
+
|
| 743 |
+
return Lfm3VlModelOutputWithPast(
|
| 744 |
+
last_hidden_state=outputs.last_hidden_state,
|
| 745 |
+
past_key_values=outputs.past_key_values,
|
| 746 |
+
hidden_states=outputs.hidden_states,
|
| 747 |
+
attentions=outputs.attentions,
|
| 748 |
+
image_hidden_states=image_features,
|
| 749 |
+
video_hidden_states=video_features,
|
| 750 |
+
)
|
| 751 |
+
|
| 752 |
+
@accepts_precomputed_kwargs(modality="video")
|
| 753 |
+
@can_return_tuple
|
| 754 |
+
@auto_docstring(
|
| 755 |
+
custom_intro="Obtains video last hidden states from the vision tower and apply multimodal projection."
|
| 756 |
+
)
|
| 757 |
+
def get_video_features(
|
| 758 |
+
self,
|
| 759 |
+
pixel_values_videos: torch.FloatTensor,
|
| 760 |
+
video_grid_thw: torch.LongTensor,
|
| 761 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 762 |
+
) -> tuple | BaseModelOutputWithPooling:
|
| 763 |
+
r"""
|
| 764 |
+
pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 765 |
+
The flattened patches of every video, grouped by merge block.
|
| 766 |
+
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`):
|
| 767 |
+
The temporal, height and width dimensions of the patch grid of each video.
|
| 768 |
+
"""
|
| 769 |
+
return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs)
|
| 770 |
+
|
| 771 |
+
|
| 772 |
+
def reject_packed_generation(config: PreTrainedConfig, input_ids, generation_config, model_kwargs, is_first_iteration):
|
| 773 |
+
"""`generate` decodes from the cache its prefill fills, and no mask it builds isolates packed documents
|
| 774 |
+
(`position_ids` restarting): the prompt is refused before the prefill, one forward or every chunk, writes to it."""
|
| 775 |
+
position_ids, attention_mask = model_kwargs.get("position_ids"), model_kwargs.get("attention_mask")
|
| 776 |
+
cache = model_kwargs.get("past_key_values")
|
| 777 |
+
if cache is None or position_ids is None or position_ids.ndim != 2:
|
| 778 |
+
return
|
| 779 |
+
if attention_mask is not None and attention_mask.ndim != 2:
|
| 780 |
+
return
|
| 781 |
+
has_attention = any(t != "conv" for t in config.get_text_config().layer_types)
|
| 782 |
+
filled = bool(cache.get_seq_length() > 0 if has_attention else cache.has_previous_state())
|
| 783 |
+
inputs_embeds = model_kwargs.get("inputs_embeds") if is_first_iteration else None
|
| 784 |
+
length = input_ids.shape[1] if inputs_embeds is None else inputs_embeds.shape[1]
|
| 785 |
+
if generation_config.prefill_chunk_size is not None:
|
| 786 |
+
# the chunks run from the start of the prompt
|
| 787 |
+
position_ids = position_ids[:, :length]
|
| 788 |
+
else:
|
| 789 |
+
if filled and has_attention:
|
| 790 |
+
# the slice `_prefill` feeds: a prompt that spans the cache is cut to its uncached part
|
| 791 |
+
mask_length = getattr(generation_config, "_mask_length", -1)
|
| 792 |
+
if attention_mask is not None:
|
| 793 |
+
mask_length = attention_mask.shape[1]
|
| 794 |
+
if inputs_embeds is not None or mask_length == length:
|
| 795 |
+
length -= int(cache.get_seq_length())
|
| 796 |
+
position_ids = position_ids[:, position_ids.shape[1] - length :]
|
| 797 |
+
if attention_mask is None:
|
| 798 |
+
attention_mask = torch.ones_like(position_ids)
|
| 799 |
+
if bool(packed_under_padding(position_ids, attention_mask, filled).any()):
|
| 800 |
+
raise ValueError(
|
| 801 |
+
"`generate` cannot decode packed documents (`position_ids` restarting, or starting at 0 on a cache that "
|
| 802 |
+
"holds tokens): attention would span them. Generate each document in its own row of a left-padded batch."
|
| 803 |
+
)
|
| 804 |
+
|
| 805 |
+
|
| 806 |
+
def packed_under_padding(position_ids: torch.Tensor, padding_mask: torch.Tensor, filled) -> torch.Tensor:
|
| 807 |
+
"""Padding repeats position 0, so packing is a restart between two kept tokens, or a kept position 0 on a filled
|
| 808 |
+
cache whose row holds a kept token; attention under a padding mask cannot isolate it."""
|
| 809 |
+
restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
|
| 810 |
+
valid = padding_mask[:, -restarts.shape[-1] :].bool()
|
| 811 |
+
packed = restarts[:, 1:] & valid[:, 1:] & valid[:, :-1]
|
| 812 |
+
if filled:
|
| 813 |
+
# a row whose cached tokens are all padding (chunked prefill of a left-padded batch) has no document to pack
|
| 814 |
+
# against; a mask covering only the chunk keeps the strict rule
|
| 815 |
+
cached = padding_mask[:, : padding_mask.shape[-1] - restarts.shape[-1]]
|
| 816 |
+
cached_valid = cached.bool().any(-1, keepdim=True) if cached.shape[-1] else torch.ones_like(valid[:, :1])
|
| 817 |
+
packed = packed | (valid[:, :1] & (position_ids[:, :1] == 0) & cached_valid)
|
| 818 |
+
return packed
|
| 819 |
+
|
| 820 |
+
|
| 821 |
+
@auto_docstring(
|
| 822 |
+
custom_intro="""
|
| 823 |
+
The LFM3_VL model which consists of a vision backbone and a language model.
|
| 824 |
+
"""
|
| 825 |
+
)
|
| 826 |
+
class Lfm3VlForConditionalGeneration(Lfm3VlPreTrainedModel, GenerationMixin):
|
| 827 |
+
_tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
|
| 828 |
+
|
| 829 |
+
def __init__(self, config: Lfm3VlConfig):
|
| 830 |
+
super().__init__(config)
|
| 831 |
+
self.model = Lfm3VlModel(config)
|
| 832 |
+
self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
|
| 833 |
+
self.post_init()
|
| 834 |
+
|
| 835 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 836 |
+
return self.lm_head
|
| 837 |
+
|
| 838 |
+
@auto_docstring
|
| 839 |
+
def get_image_features(
|
| 840 |
+
self,
|
| 841 |
+
pixel_values: torch.FloatTensor,
|
| 842 |
+
image_grid_thw: torch.LongTensor,
|
| 843 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 844 |
+
) -> tuple | BaseModelOutputWithPooling:
|
| 845 |
+
r"""
|
| 846 |
+
pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 847 |
+
The flattened patches of every image, grouped by merge block.
|
| 848 |
+
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
|
| 849 |
+
The temporal, height and width dimensions of the patch grid of each image.
|
| 850 |
+
"""
|
| 851 |
+
return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs)
|
| 852 |
+
|
| 853 |
+
@can_return_tuple
|
| 854 |
+
@auto_docstring
|
| 855 |
+
def forward(
|
| 856 |
+
self,
|
| 857 |
+
input_ids: torch.LongTensor | None = None,
|
| 858 |
+
attention_mask: torch.Tensor | None = None,
|
| 859 |
+
position_ids: torch.LongTensor | None = None,
|
| 860 |
+
past_key_values: Cache | None = None,
|
| 861 |
+
inputs_embeds: torch.FloatTensor | None = None,
|
| 862 |
+
pixel_values: torch.FloatTensor | None = None,
|
| 863 |
+
image_grid_thw: torch.LongTensor | None = None,
|
| 864 |
+
pixel_values_videos: torch.FloatTensor | None = None,
|
| 865 |
+
video_grid_thw: torch.LongTensor | None = None,
|
| 866 |
+
labels: torch.LongTensor | None = None,
|
| 867 |
+
use_cache: bool | None = None,
|
| 868 |
+
logits_to_keep: int | torch.Tensor = 0,
|
| 869 |
+
mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None,
|
| 870 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 871 |
+
) -> tuple | Lfm3VlCausalLMOutputWithPast:
|
| 872 |
+
r"""
|
| 873 |
+
pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
|
| 874 |
+
The flattened patches of every image, grouped by merge block.
|
| 875 |
+
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
|
| 876 |
+
The temporal, height and width dimensions of the patch grid of each image.
|
| 877 |
+
pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
|
| 878 |
+
The flattened patches of every video, grouped by merge block.
|
| 879 |
+
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
|
| 880 |
+
The temporal, height and width dimensions of the patch grid of each video.
|
| 881 |
+
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
| 882 |
+
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
|
| 883 |
+
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
|
| 884 |
+
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
|
| 885 |
+
|
| 886 |
+
Example:
|
| 887 |
+
|
| 888 |
+
```python
|
| 889 |
+
>>> from transformers import AutoProcessor, AutoModelForImageTextToText
|
| 890 |
+
|
| 891 |
+
>>> model = AutoModelForImageTextToText.from_pretrained("LiquidAI/LFM3-VL-1.9B")
|
| 892 |
+
>>> processor = AutoProcessor.from_pretrained("LiquidAI/LFM3-VL-1.9B")
|
| 893 |
+
|
| 894 |
+
>>> conversation = [
|
| 895 |
+
... {
|
| 896 |
+
... "role": "user",
|
| 897 |
+
... "content": [
|
| 898 |
+
... {"type": "image", "url": "https://www.ilankelman.org/stopsigns/australia.jpg"},
|
| 899 |
+
... {"type": "text", "text": "What is in this image?"},
|
| 900 |
+
... ],
|
| 901 |
+
... },
|
| 902 |
+
... ]
|
| 903 |
+
|
| 904 |
+
>>> inputs = processor.apply_chat_template(
|
| 905 |
+
... conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"
|
| 906 |
+
... )
|
| 907 |
+
|
| 908 |
+
>>> outputs = model.generate(**inputs, max_new_tokens=45)
|
| 909 |
+
>>> processor.batch_decode(outputs, skip_special_tokens=True)[0]
|
| 910 |
+
```"""
|
| 911 |
+
outputs = self.model(
|
| 912 |
+
input_ids=input_ids,
|
| 913 |
+
attention_mask=attention_mask,
|
| 914 |
+
position_ids=position_ids,
|
| 915 |
+
past_key_values=past_key_values,
|
| 916 |
+
inputs_embeds=inputs_embeds,
|
| 917 |
+
pixel_values=pixel_values,
|
| 918 |
+
image_grid_thw=image_grid_thw,
|
| 919 |
+
pixel_values_videos=pixel_values_videos,
|
| 920 |
+
video_grid_thw=video_grid_thw,
|
| 921 |
+
use_cache=use_cache,
|
| 922 |
+
mm_encoder_outputs=mm_encoder_outputs,
|
| 923 |
+
**kwargs,
|
| 924 |
+
)
|
| 925 |
+
|
| 926 |
+
hidden_states = outputs[0]
|
| 927 |
+
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
| 928 |
+
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
| 929 |
+
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
| 930 |
+
|
| 931 |
+
loss = None
|
| 932 |
+
if labels is not None:
|
| 933 |
+
loss = self.loss_function(
|
| 934 |
+
logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
|
| 935 |
+
)
|
| 936 |
+
|
| 937 |
+
return Lfm3VlCausalLMOutputWithPast(
|
| 938 |
+
loss=loss,
|
| 939 |
+
logits=logits,
|
| 940 |
+
past_key_values=outputs.past_key_values,
|
| 941 |
+
hidden_states=outputs.hidden_states,
|
| 942 |
+
attentions=outputs.attentions,
|
| 943 |
+
image_hidden_states=outputs.image_hidden_states,
|
| 944 |
+
video_hidden_states=outputs.video_hidden_states,
|
| 945 |
+
)
|
| 946 |
+
|
| 947 |
+
def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
|
| 948 |
+
# the text model's refusal, before the prefill writes to the cache
|
| 949 |
+
reject_packed_generation(self.config, input_ids, generation_config, model_kwargs, is_first_iteration)
|
| 950 |
+
return super()._prefill(input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration)
|
| 951 |
+
|
| 952 |
+
@auto_docstring
|
| 953 |
+
def get_video_features(
|
| 954 |
+
self,
|
| 955 |
+
pixel_values_videos: torch.FloatTensor,
|
| 956 |
+
video_grid_thw: torch.LongTensor,
|
| 957 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 958 |
+
) -> tuple | BaseModelOutputWithPooling:
|
| 959 |
+
r"""
|
| 960 |
+
pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
|
| 961 |
+
The flattened patches of every video, grouped by merge block.
|
| 962 |
+
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`):
|
| 963 |
+
The temporal, height and width dimensions of the patch grid of each video.
|
| 964 |
+
"""
|
| 965 |
+
return self.model.get_video_features(pixel_values_videos, video_grid_thw, **kwargs)
|
| 966 |
+
|
| 967 |
+
|
| 968 |
+
__all__ = ["Lfm3VlForConditionalGeneration", "Lfm3VlModel", "Lfm3VlPreTrainedModel", "Lfm3VlVisionModel"]
|
native_lfm3.py
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""LFM3's transformers classes, one module a model type: `native_lfm3` (text), `native_lfm3_vl` (pictures),
|
| 2 |
+
`native_lfm3_omni` (pictures and audio), each importing the one before it, so a Hub repository serves only its
|
| 3 |
+
model's files. transformers' own classes once a release ships them; until then copies of transformers' LFM3 code
|
| 4 |
+
beside these (each names its commit), registered with the Auto classes, so a checkpoint's config, tokenizer and
|
| 5 |
+
processor load and the composite models build their parts.
|
| 6 |
+
|
| 7 |
+
Shipping LFM3, transformers makes the fallbacks dead: delete them and the copies (the files headed "Vendored from"),
|
| 8 |
+
nothing else."""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
try:
|
| 13 |
+
from transformers import Lfm3Config, Lfm3ForCausalLM, Lfm3Model
|
| 14 |
+
except ImportError:
|
| 15 |
+
from transformers import AutoConfig, AutoModel
|
| 16 |
+
|
| 17 |
+
from .configuration_lfm3 import Lfm3Config
|
| 18 |
+
from .modeling_lfm3 import Lfm3ForCausalLM, Lfm3Model
|
| 19 |
+
|
| 20 |
+
# exist_ok: a second copy of these files in one process (two snapshots) registers again
|
| 21 |
+
AutoConfig.register(Lfm3Config.model_type, Lfm3Config, exist_ok=True)
|
| 22 |
+
AutoModel.register(Lfm3Config, Lfm3Model, exist_ok=True) # the composite models build their text model by config
|
| 23 |
+
|
| 24 |
+
__all__ = ["Lfm3Config", "Lfm3ForCausalLM", "Lfm3Model"]
|
native_lfm3_vl.py
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""LFM3-VL's transformers classes (`native_lfm3`): its config, vision tower, picture and video processors."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from .native_lfm3 import Lfm3Model # noqa: F401 (registered: the VL builds its text model by config)
|
| 6 |
+
|
| 7 |
+
try:
|
| 8 |
+
from transformers import Lfm3VlConfig, Lfm3VlForConditionalGeneration, Lfm3VlProcessor
|
| 9 |
+
except ImportError:
|
| 10 |
+
from transformers import AutoConfig, AutoImageProcessor, AutoModel, AutoProcessor, AutoVideoProcessor
|
| 11 |
+
|
| 12 |
+
from .configuration_lfm3_vl import Lfm3VlConfig, Lfm3VlVisionConfig
|
| 13 |
+
from .image_processing_lfm3_vl import Lfm3VlImageProcessor
|
| 14 |
+
from .modeling_lfm3_vl import Lfm3VlForConditionalGeneration, Lfm3VlVisionModel
|
| 15 |
+
from .processing_lfm3_vl import Lfm3VlProcessor
|
| 16 |
+
from .video_processing_lfm3_vl import Lfm3VlVideoProcessor
|
| 17 |
+
|
| 18 |
+
for config in (Lfm3VlConfig, Lfm3VlVisionConfig):
|
| 19 |
+
AutoConfig.register(config.model_type, config, exist_ok=True)
|
| 20 |
+
AutoModel.register(Lfm3VlVisionConfig, Lfm3VlVisionModel, exist_ok=True) # the tower, built by config
|
| 21 |
+
AutoImageProcessor.register(Lfm3VlConfig, fast_image_processor_class=Lfm3VlImageProcessor, exist_ok=True)
|
| 22 |
+
AutoVideoProcessor.register(Lfm3VlConfig, Lfm3VlVideoProcessor, exist_ok=True)
|
| 23 |
+
AutoProcessor.register(Lfm3VlConfig, Lfm3VlProcessor, exist_ok=True)
|
| 24 |
+
|
| 25 |
+
__all__ = ["Lfm3VlConfig", "Lfm3VlForConditionalGeneration", "Lfm3VlProcessor"]
|
parsing.py
ADDED
|
@@ -0,0 +1,297 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Strict readings of written values; no arithmetic, unit conversion or invented date components."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import datetime as dt
|
| 6 |
+
import math
|
| 7 |
+
import re
|
| 8 |
+
import unicodedata
|
| 9 |
+
from decimal import Decimal
|
| 10 |
+
from functools import cache
|
| 11 |
+
|
| 12 |
+
# Month names a date is written with (accents folded; Polish and Russian in the genitive a date uses), and weekdays,
|
| 13 |
+
# which a date may lead with.
|
| 14 |
+
_MONTHS = (
|
| 15 |
+
"january janvier januar enero gennaio janeiro januari stycznia styczen января",
|
| 16 |
+
"february fevrier februar febrero febbraio fevereiro februari lutego luty февраля",
|
| 17 |
+
"march mars marz marzo marco maart marca марта",
|
| 18 |
+
"april avril abril aprile kwietnia kwiecien апреля",
|
| 19 |
+
"may mai mayo maggio maio mei maj maja мая",
|
| 20 |
+
"june juin juni junio giugno junho czerwca czerwiec июня",
|
| 21 |
+
"july juillet juli julio luglio julho lipca lipiec июля",
|
| 22 |
+
"august aout agosto augustus augusti sierpnia sierpien августа",
|
| 23 |
+
"september septembre septiembre settembre setembro wrzesnia wrzesien сентября",
|
| 24 |
+
"october octobre oktober octubre ottobre outubro pazdziernika pazdziernik октября",
|
| 25 |
+
"november novembre noviembre novembro listopada listopad ноября",
|
| 26 |
+
"december decembre dezember diciembre dicembre dezembro grudnia grudzien декабря",
|
| 27 |
+
)
|
| 28 |
+
MONTHS = {name: month for month, names in enumerate(_MONTHS, 1) for name in names.split()}
|
| 29 |
+
_ABBREVIATIONS = {}
|
| 30 |
+
for _name, _month in MONTHS.items():
|
| 31 |
+
_ABBREVIATIONS.setdefault(_name[:3], set()).add(_month)
|
| 32 |
+
MONTHS.update({name: next(iter(months)) for name, months in _ABBREVIATIONS.items() if len(months) == 1})
|
| 33 |
+
MONTHS.update(janv=1, fevr=2, juil=7, sept=9)
|
| 34 |
+
WEEKDAYS = set("""monday tuesday wednesday thursday friday saturday sunday mon tue tues wed thu thur thurs fri sat sun
|
| 35 |
+
lundi mardi mercredi jeudi vendredi samedi dimanche montag dienstag mittwoch donnerstag freitag samstag sonntag lunes
|
| 36 |
+
martes miercoles jueves viernes sabado domingo lunedi martedi mercoledi giovedi venerdi sabato domenica segunda terca
|
| 37 |
+
quarta quinta sexta""".split())
|
| 38 |
+
CJK_DATE = re.compile(r"(\d{4})\s*[年년]\s*(\d{1,2})\s*[月월]\s*(\d{1,2})\s*[日일]") # 2026年10月20日, 2026년 10월 20일
|
| 39 |
+
NUMERIC_DATE = re.compile(r"(?<!\d)(\d{1,4})([/.-])(\d{1,2})\2(\d{2,4})(?!\d)")
|
| 40 |
+
_MONTH_NAME = "|".join(sorted(MONTHS, key=len, reverse=True))
|
| 41 |
+
WORD_DATE = re.compile(rf"\b(?:\d{{1,2}}\w*[ .]+(?:de\s+)?(?:{_MONTH_NAME})[ ,.]+(?:de\s+)?\d{{4}}"
|
| 42 |
+
rf"|(?:{_MONTH_NAME})[ .]+\d{{1,2}}\w*[ ,.]+\d{{4}})\b") # in folded text: 20 octobre 2026
|
| 43 |
+
|
| 44 |
+
# Currencies by the decimal mark their amounts are written with: "" for none (whole units), None for both in use.
|
| 45 |
+
CURRENCIES = {
|
| 46 |
+
".": ("$", "£", "₹", "USD", "GBP", "CAD", "AUD", "NZD", "INR", "Rs.", "Rs", "MXN", "CNY", "RMB", "元", "HKD", "SGD"),
|
| 47 |
+
",": ("R$", "BRL", "₽", "RUB", "zł", "PLN", "kr", "SEK", "NOK", "DKK", "₺", "TRY"),
|
| 48 |
+
"": ("¥", "¥", "円", "₩", "₫", "Rp.", "Rp", "IDR", "JPY", "KRW", "VND", "CLP"),
|
| 49 |
+
None: ("€", "EUR", "CHF"),
|
| 50 |
+
}
|
| 51 |
+
_NAMES = {c.casefold(): c for cs in CURRENCIES.values() for c in cs}
|
| 52 |
+
_MARK = {c.casefold(): mark for mark, cs in CURRENCIES.items() for c in cs}
|
| 53 |
+
_CURRENCY = "(?i:" + "|".join(rf"\b{re.escape(c)}(?![a-z])" if c[0].isascii() and c[0].isalpha() else re.escape(c)
|
| 54 |
+
for c in sorted(_NAMES.values(), key=len, reverse=True)) + ")"
|
| 55 |
+
# Scale words, written apart or attached: unambiguous, so no context needed (German "Billion" is 10^12: left out).
|
| 56 |
+
SCALES = {**dict.fromkeys(("thousand", "mil", "tausend", "тыс", "千"), 10**3), **dict.fromkeys(("万", "萬"), 10**4),
|
| 57 |
+
**dict.fromkeys(("lakh", "lakhs"), 10**5), **dict.fromkeys(("crore", "crores"), 10**7),
|
| 58 |
+
**dict.fromkeys(("million", "millions", "mio", "mln", "millón", "millones", "milhão", "milhões", "milione",
|
| 59 |
+
"milioni", "млн"), 10**6), **dict.fromkeys(("亿", "億"), 10**8),
|
| 60 |
+
**dict.fromkeys(("billion", "billions", "milliard", "milliards", "milliarde", "milliarden", "mrd", "miliardo",
|
| 61 |
+
"miliardi", "млрд"), 10**9), "trillion": 10**12}
|
| 62 |
+
_SCALE = re.compile(r"\s*(" + "|".join(sorted(map(re.escape, SCALES), key=len, reverse=True)) + r")\.?$", re.I)
|
| 63 |
+
_SUFFIX = re.compile(r"(?<=\d)(k|m|bn|b)$", re.I) # attached k, M, B: a magnitude only beside money or a count
|
| 64 |
+
_COUNTED = re.compile(r"\b(?:money|price|invoice|revenue|cost|count|quantity|population|people|items|units|users|"
|
| 65 |
+
r"downloads)\b", re.I)
|
| 66 |
+
_GROUPED = re.compile(r"(?<![\d.,])\d{1,3}(?:([,. \u00a0\u202f'’])\d{3})+(?!\1)([.,])\d+(?![\d.,])") # 1 234,50
|
| 67 |
+
_SHORT = re.compile(r"(?<![\d.,])\d+([.,])\d{1,2}(?![\d.,])") # 12,50: a decimal, never a group
|
| 68 |
+
_REPEATED = re.compile(r"(?<![\d.,])\d{1,3}(?:([.,])\d{3}){2,}(?![\d.,])") # 1.234.567: a group, never a decimal
|
| 69 |
+
_LOCALE = re.compile(r'\b(?:locale|language|lang)["\']?\s*[=:]\s*["\']?(en|fr|de|es|it|pt|nl|pl|ru|sv|ja|zh|ko)'
|
| 70 |
+
r"(?:[-_][a-z]{2})?\b")
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def _one(values: set):
|
| 74 |
+
"""The one value evidence settles on; None when it holds none or disagrees with itself."""
|
| 75 |
+
return next(iter(values)) if len(values) == 1 else None
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def written_unit(text: str) -> str | None:
|
| 79 |
+
"""Only an explicitly written percent or currency marker; a currency sign stays a sign."""
|
| 80 |
+
if text.rstrip().endswith("%"):
|
| 81 |
+
return "%"
|
| 82 |
+
markers = re.findall(_CURRENCY, text)
|
| 83 |
+
return _NAMES[markers[0].casefold()] if len(markers) == 1 else None
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def _currency_marks(text: str) -> set[str]:
|
| 87 |
+
"""The decimal marks of the currencies `text` writes ("" for whole units); a currency of both says nothing."""
|
| 88 |
+
return {_MARK[c.casefold()] for c in re.findall(_CURRENCY, text)} - {None}
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def _decimal_mark(written: str, context: str) -> str | None:
|
| 92 |
+
"""The decimal mark ("" when amounts have none), from the first evidence that speaks: the span's own currency,
|
| 93 |
+
the document's other numbers, the document's currencies, a declared locale. Evidence that disagrees with itself
|
| 94 |
+
decides nothing; a markup declaration's numbers (`<?xml version="1.0"?>`) are not the document's."""
|
| 95 |
+
context = re.sub(r"<\?[^>]*\?>", "", context)
|
| 96 |
+
numbers = {m[2] for m in _GROUPED.finditer(context)} | {m[1] for m in _SHORT.finditer(context)}
|
| 97 |
+
numbers |= {"," if m[1] == "." else "." for m in _REPEATED.finditer(context)}
|
| 98 |
+
locales = {"." if lang in ("en", "ja", "zh", "ko") else "," for lang in _LOCALE.findall(context.casefold())}
|
| 99 |
+
return _one(next((marks for marks in (_currency_marks(written), numbers, _currency_marks(context), locales) if marks),
|
| 100 |
+
set()))
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def parse_number(text: str, context: str = "", *, unit: str | None = None) -> int | float | None:
|
| 104 |
+
"""The written magnitude, percent points included: its currency and percent sign set apart, its scale word
|
| 105 |
+
applied; no implicit unit and no conversion. None when the writing does not settle one value."""
|
| 106 |
+
context += f"\n{unit}" if unit and re.fullmatch(_CURRENCY, unit.strip()) else ""
|
| 107 |
+
written = text
|
| 108 |
+
text = text.strip().replace("\u2212", "-").replace("\u00a0", " ").replace("\u202f", " ").removesuffix("%").rstrip()
|
| 109 |
+
negative = text.startswith("(") and text.endswith(")")
|
| 110 |
+
text = text[1:-1].strip() if negative else text
|
| 111 |
+
sign = currency = None
|
| 112 |
+
for _ in range(2): # one currency and one sign, in either order
|
| 113 |
+
if currency is None and (m := re.fullmatch(rf"({_CURRENCY})\s*(.*)|(.*?)\s*({_CURRENCY})", text)):
|
| 114 |
+
currency, text = m[1] or m[4], (m[2] if m[1] else m[3]).strip()
|
| 115 |
+
if sign is None and text[:1] in ("-", "+"):
|
| 116 |
+
sign, text = text[0], text[1:].strip()
|
| 117 |
+
if negative and sign:
|
| 118 |
+
return None
|
| 119 |
+
sign = sign or ""
|
| 120 |
+
scale = 1
|
| 121 |
+
if word := _SCALE.search(text):
|
| 122 |
+
scale, text = SCALES[word[1].lower()] if word[1] != "Billion" else 0, text[: word.start()]
|
| 123 |
+
elif suffix := _SUFFIX.search(text):
|
| 124 |
+
if not (re.search(_CURRENCY, written + "\n" + context) or _COUNTED.search(context)):
|
| 125 |
+
return None
|
| 126 |
+
scale, text = {"k": 10**3, "m": 10**6, "b": 10**9, "bn": 10**9}[suffix[1].lower()], text[: suffix.start()]
|
| 127 |
+
digits = _digits(text, written, context)
|
| 128 |
+
if digits is None or not scale:
|
| 129 |
+
return None
|
| 130 |
+
value = Decimal(("-" if negative else sign) + digits) * scale
|
| 131 |
+
if value == value.to_integral_value():
|
| 132 |
+
return int(value)
|
| 133 |
+
result = float(value)
|
| 134 |
+
return result if math.isfinite(result) else None
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
def _digits(text: str, written: str, context: str) -> str | None:
|
| 138 |
+
"""`text`'s number as plain digits and a point: the decimal mark from the evidence, the other mark a complete
|
| 139 |
+
grouping (threes, or India's lakh grouping beside a rupee)."""
|
| 140 |
+
if not re.fullmatch(r"\d[\d., '’]*", text):
|
| 141 |
+
return None
|
| 142 |
+
if re.search(r"[ '’]", text): # a space or an apostrophe only groups
|
| 143 |
+
grouping = re.fullmatch(r"(\d{1,3}(?:[ '’]\d{3})+)([.,]\d+)?", text)
|
| 144 |
+
return None if grouping is None else _digits(re.sub(r"[ '’]", "", grouping[1]) + (grouping[2] or ""), written,
|
| 145 |
+
context)
|
| 146 |
+
if re.fullmatch(r"\d{1,2}(?:,\d{2})+,\d{3}(?:\.\d+)?", text) and re.search(r"₹|\bINR\b|\bRs\b|lakh|crore",
|
| 147 |
+
written + "\n" + context, re.I):
|
| 148 |
+
return text.replace(",", "")
|
| 149 |
+
if "." in text and "," in text:
|
| 150 |
+
point, group = (".", ",") if text.rfind(".") > text.rfind(",") else (",", ".")
|
| 151 |
+
if not re.fullmatch(rf"\d{{1,3}}(?:{re.escape(group)}\d{{3}})+{re.escape(point)}\d+", text):
|
| 152 |
+
return None
|
| 153 |
+
return text.replace(group, "").replace(point, ".")
|
| 154 |
+
if "." not in text and "," not in text:
|
| 155 |
+
return text
|
| 156 |
+
mark = "." if "." in text else ","
|
| 157 |
+
parts = text.split(mark)
|
| 158 |
+
if not all(parts):
|
| 159 |
+
return None
|
| 160 |
+
point = _decimal_mark(written, context)
|
| 161 |
+
if len(parts) > 2 or point is not None and point != mark: # the mark groups: complete threes
|
| 162 |
+
return "".join(parts) if 1 <= len(parts[0]) <= 3 and all(len(p) == 3 for p in parts[1:]) else None
|
| 163 |
+
if point is None and len(parts[1]) == 3 and len(parts[0]) <= 3: # 1,250: a group or a decimal, unsettled
|
| 164 |
+
return None
|
| 165 |
+
return ".".join(parts)
|
| 166 |
+
|
| 167 |
+
|
| 168 |
+
def _fold(text: str) -> str:
|
| 169 |
+
return "".join(c for c in unicodedata.normalize("NFKD", text.casefold()) if not unicodedata.combining(c))
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
# a declared numeric date format, in English, French (jj/mm/aaaa), German (tt.mm.jjjj), Spanish or Italian (gg/mm/aaaa)
|
| 173 |
+
_DAY, _MONTH, _YEAR = r"(?:dd?|jj?|tt?|gg?|day|jour|tag|dia|giorno)", r"(?:mm?|month|mois|monat|mes|mese)", \
|
| 174 |
+
r"(?:yy(?:yy)?|aa(?:aa)?|jj(?:jj)?|year|annee|jahr|ano|anno)"
|
| 175 |
+
DATE_FORMATS = ((rf"\b{_DAY}\s*[/.-]\s*{_MONTH}\s*[/.-]\s*{_YEAR}\b", "dmy"),
|
| 176 |
+
(rf"\b{_MONTH}\s*[/.-]\s*{_DAY}\s*[/.-]\s*{_YEAR}\b", "mdy"))
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def date_order(context: str, folded: str | None = None) -> str | None:
|
| 180 |
+
"""Declared format/locale first, then the state's other unambiguous numeric dates; conflicts stay unresolved.
|
| 181 |
+
`folded`: `context` already folded."""
|
| 182 |
+
folded = _fold(context) if folded is None else folded
|
| 183 |
+
explicit = set()
|
| 184 |
+
for pattern, order in DATE_FORMATS:
|
| 185 |
+
if re.search(pattern, folded):
|
| 186 |
+
explicit.add(order)
|
| 187 |
+
if explicit:
|
| 188 |
+
return _one(explicit)
|
| 189 |
+
locales = re.findall(
|
| 190 |
+
r'\b(?:locale|language|lang|langue|sprache|idioma|lingua)["\']?\s*[=:]\s*["\']?'
|
| 191 |
+
r"(en(?:[-_](?:us|gb|au|ca))?|fr(?:[-_]fr)?|de(?:[-_]de)?|es(?:[-_]es)?|it(?:[-_]it)?)\b",
|
| 192 |
+
folded,
|
| 193 |
+
)
|
| 194 |
+
if locales:
|
| 195 |
+
return _one({"mdy" if locale.replace("_", "-") in ("en", "en-us") else "dmy" for locale in locales})
|
| 196 |
+
orders = set()
|
| 197 |
+
for match in NUMERIC_DATE.finditer(context):
|
| 198 |
+
a, b, year = int(match[1]), int(match[3]), int(match[4])
|
| 199 |
+
if len(match[4]) != 4 or a > 31 or b > 31 or a < 1 or b < 1:
|
| 200 |
+
continue
|
| 201 |
+
for day, month, order in ((a, b, "dmy"), (b, a, "mdy")):
|
| 202 |
+
try:
|
| 203 |
+
dt.date(year, month, day)
|
| 204 |
+
except ValueError:
|
| 205 |
+
continue
|
| 206 |
+
if (a > 12) != (b > 12):
|
| 207 |
+
orders.add(order)
|
| 208 |
+
return _one(orders)
|
| 209 |
+
|
| 210 |
+
|
| 211 |
+
def _candidates(year: int, a: int, b: int, order: str | None) -> list[tuple[int, int, int]]:
|
| 212 |
+
"""(year, month, day) for `a`/`b`/year: day first, month first, or the one the context's `order` says."""
|
| 213 |
+
dmy, mdy = (year, b, a), (year, a, b)
|
| 214 |
+
return [dmy, mdy] if order is None else [dmy if order == "dmy" else mdy]
|
| 215 |
+
|
| 216 |
+
|
| 217 |
+
def parse_date(text: str, context: str = "") -> tuple[str | None, bool]:
|
| 218 |
+
"""ISO value and unresolved-ambiguity flag; an invalid or unrecognised span is (None, False)."""
|
| 219 |
+
folded = cache(lambda: _fold(context)) # the context folded and its date order, once, when a date needs them
|
| 220 |
+
return _parse_date(text, context, folded, cache(lambda: date_order(context, folded())))
|
| 221 |
+
|
| 222 |
+
|
| 223 |
+
def _parse_date(text: str, context: str, folded, order) -> tuple[str | None, bool]:
|
| 224 |
+
text = text.strip()
|
| 225 |
+
if not text:
|
| 226 |
+
return None, False
|
| 227 |
+
# Keep an explicitly written time and timezone, with no default time or timezone.
|
| 228 |
+
timed = re.fullmatch(r"(.+?)[T\s]+(\d{1,2}:\d{2}(?::\d{2}(?:\.\d+)?)?(?:Z|[+-]\d{2}:\d{2})?)", text)
|
| 229 |
+
if timed:
|
| 230 |
+
date, ambiguous = _parse_date(timed[1], context, folded, order)
|
| 231 |
+
if date is None:
|
| 232 |
+
return None, ambiguous
|
| 233 |
+
try:
|
| 234 |
+
clock = timed[2]
|
| 235 |
+
if clock[1] == ":":
|
| 236 |
+
clock = "0" + clock
|
| 237 |
+
return dt.datetime.fromisoformat(date + "T" + clock).isoformat(), False
|
| 238 |
+
except ValueError:
|
| 239 |
+
return None, False
|
| 240 |
+
numeric = NUMERIC_DATE.fullmatch(text)
|
| 241 |
+
if numeric:
|
| 242 |
+
a, b, c = int(numeric[1]), int(numeric[3]), int(numeric[4])
|
| 243 |
+
if len(numeric[1]) == 4:
|
| 244 |
+
candidates = [(a, b, c)]
|
| 245 |
+
elif len(numeric[4]) == 4:
|
| 246 |
+
candidates = _candidates(c, a, b, order())
|
| 247 |
+
else:
|
| 248 |
+
# A century is never invented: another unambiguous complete date must establish it.
|
| 249 |
+
if len(numeric[4]) != 2:
|
| 250 |
+
return None, False
|
| 251 |
+
centuries = set()
|
| 252 |
+
complete = [
|
| 253 |
+
match[0] for match in NUMERIC_DATE.finditer(context) if len(match[1]) == 4 or len(match[4]) == 4
|
| 254 |
+
]
|
| 255 |
+
complete.extend(match[0] for match in WORD_DATE.finditer(folded()))
|
| 256 |
+
for candidate in complete:
|
| 257 |
+
value, ambiguous = _parse_date(candidate, context, folded, order)
|
| 258 |
+
if value is not None and not ambiguous:
|
| 259 |
+
centuries.add(int(value[:4]) // 100 * 100)
|
| 260 |
+
if len(centuries) != 1:
|
| 261 |
+
return None, False
|
| 262 |
+
candidates = _candidates(centuries.pop() + c, a, b, order())
|
| 263 |
+
values = set()
|
| 264 |
+
for year, month, day in candidates:
|
| 265 |
+
try:
|
| 266 |
+
values.add(dt.date(year, month, day).isoformat())
|
| 267 |
+
except ValueError:
|
| 268 |
+
pass
|
| 269 |
+
if len(values) > 1:
|
| 270 |
+
return None, True
|
| 271 |
+
return next(iter(values), None), False
|
| 272 |
+
if cjk := CJK_DATE.fullmatch(text):
|
| 273 |
+
try:
|
| 274 |
+
return dt.date(*map(int, cjk.groups())).isoformat(), False
|
| 275 |
+
except ValueError:
|
| 276 |
+
return None, False
|
| 277 |
+
words = _fold(text)
|
| 278 |
+
words = re.sub(r"(?<=\d)(?:st|nd|rd|th|er)\b", "", words)
|
| 279 |
+
words = re.sub(r"(?<=\d)\.?(?:o|a|°)\b", "", words)
|
| 280 |
+
words = re.sub(r"\b(?:de|of|the|le|el|il|den)\b", " ", words)
|
| 281 |
+
words = re.sub(r"[,.]", " ", words).split()
|
| 282 |
+
if words and words[0].removesuffix("-feira") in WEEKDAYS: # Tue, 20 Oct 2026; segunda-feira, 20 de ...
|
| 283 |
+
words = words[1:]
|
| 284 |
+
if len(words) != 3:
|
| 285 |
+
return None, False
|
| 286 |
+
if words[0] in MONTHS:
|
| 287 |
+
month, day, year = MONTHS[words[0]], words[1], words[2]
|
| 288 |
+
elif words[1] in MONTHS:
|
| 289 |
+
day, month, year = words[0], MONTHS[words[1]], words[2]
|
| 290 |
+
else:
|
| 291 |
+
return None, False
|
| 292 |
+
if not day.isdigit() or not year.isdigit() or len(year) != 4:
|
| 293 |
+
return None, False
|
| 294 |
+
try:
|
| 295 |
+
return dt.date(int(year), month, int(day)).isoformat(), False
|
| 296 |
+
except ValueError:
|
| 297 |
+
return None, False
|
pointer.py
ADDED
|
@@ -0,0 +1,166 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Parameter-free readouts over a row's state positions, with BOS0 for nothing."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import math
|
| 6 |
+
from bisect import bisect_right
|
| 7 |
+
from dataclasses import dataclass
|
| 8 |
+
|
| 9 |
+
import torch
|
| 10 |
+
import torch.nn.functional as F
|
| 11 |
+
|
| 12 |
+
from .json_spans import json_text, string_paths # torch-free, so code without torch reads spans too
|
| 13 |
+
|
| 14 |
+
SPAN_CAPS = {"text": 512, "number": 32, "date": 32} # a read's longest span, in tokens
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
@dataclass(frozen=True)
|
| 18 |
+
class StatePositions:
|
| 19 |
+
"""Token positions mapped to the caller's rendered state; None means not a text position."""
|
| 20 |
+
|
| 21 |
+
text: str
|
| 22 |
+
offsets: tuple[tuple[int, int] | None, ...]
|
| 23 |
+
paths: tuple[tuple[int, int, str], ...] = ()
|
| 24 |
+
|
| 25 |
+
@property
|
| 26 |
+
def valid(self) -> list[bool]:
|
| 27 |
+
return [i == 0 or span is not None for i, span in enumerate(self.offsets)]
|
| 28 |
+
|
| 29 |
+
@property
|
| 30 |
+
def pointer_mask(self) -> list[bool]:
|
| 31 |
+
"""`valid` up to its last valid position: the keys a read's pointer scores."""
|
| 32 |
+
valid = self.valid
|
| 33 |
+
return valid[:max(i for i, v in enumerate(valid) if v) + 1]
|
| 34 |
+
|
| 35 |
+
def read(self, first: int, last: int, max_tokens: int = SPAN_CAPS["text"]) -> dict:
|
| 36 |
+
if first == last == 0:
|
| 37 |
+
return {"text": None, "span": None}
|
| 38 |
+
if not 0 < first <= last < len(self.offsets) or last - first >= max_tokens:
|
| 39 |
+
raise ValueError("invalid pointer endpoints")
|
| 40 |
+
a, b = self.offsets[first], self.offsets[last]
|
| 41 |
+
if a is None or b is None:
|
| 42 |
+
raise ValueError("pointer leaves the state text")
|
| 43 |
+
span = [a[0], b[1]]
|
| 44 |
+
# A token can carry what is around a value: the separating space, or a JSON string's quote (` €",`). Trim
|
| 45 |
+
# boundary space and clip to the one string the span starts or ends in; never alter a character inside.
|
| 46 |
+
while span[0] < span[1] and self.text[span[0]].isspace():
|
| 47 |
+
span[0] += 1
|
| 48 |
+
while span[1] > span[0] and self.text[span[1] - 1].isspace():
|
| 49 |
+
span[1] -= 1
|
| 50 |
+
held = [(start, end, p) for start, end, p in self.paths if start < span[1] and span[0] < end]
|
| 51 |
+
if len(held) == 1:
|
| 52 |
+
start, end, path = held[0]
|
| 53 |
+
span = [max(span[0], start), min(span[1], end)]
|
| 54 |
+
# inside one JSON string: the caller's characters, not their escapes
|
| 55 |
+
return {"text": json_text(self.text[slice(*span)]), "span": span, "path": path}
|
| 56 |
+
return {"text": self.text[slice(*span)], "span": span}
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def state_positions(
|
| 60 |
+
tokenizer, state, prefix: str, full_text: str, style: str = "json_only", *, encoded=None
|
| 61 |
+
) -> StatePositions:
|
| 62 |
+
"""Align one unchanged rendered prompt to text tokens; template, picture and question tokens stay invalid."""
|
| 63 |
+
from .prompt import plain_offsets, question_header, state_block, state_text
|
| 64 |
+
|
| 65 |
+
if not full_text.startswith(prefix):
|
| 66 |
+
raise ValueError("the tokenized prompt must start with the state prefix")
|
| 67 |
+
if state is None:
|
| 68 |
+
body, escaped, boundaries = "", "", [0]
|
| 69 |
+
start = len(prefix)
|
| 70 |
+
else:
|
| 71 |
+
body = state_block(state, style)
|
| 72 |
+
escaped, boundaries = plain_offsets(tokenizer, body)
|
| 73 |
+
tail = escaped + "\n" + question_header(tokenizer, state)
|
| 74 |
+
if not prefix.endswith(tail):
|
| 75 |
+
raise ValueError("state rendering differs from the prompt prefix")
|
| 76 |
+
start = len(prefix) - len(tail)
|
| 77 |
+
encoded = (
|
| 78 |
+
tokenizer(full_text, add_special_tokens=False, return_offsets_mapping=True) if encoded is None else encoded
|
| 79 |
+
)
|
| 80 |
+
text = state_text(state, style) if state is not None else "" # the blank line after it is the template's
|
| 81 |
+
stop = start + bisect_right(boundaries, len(text)) - 1
|
| 82 |
+
offsets = []
|
| 83 |
+
for a, b in encoded["offset_mapping"]:
|
| 84 |
+
span = None
|
| 85 |
+
lo, hi = max(a, start), min(b, stop)
|
| 86 |
+
if lo < hi and not full_text[a:lo].strip() and not full_text[hi:b].strip():
|
| 87 |
+
left, right = boundaries[lo - start], boundaries[hi - start]
|
| 88 |
+
if left < right <= len(text) and text[left:right].strip():
|
| 89 |
+
span = (left, right)
|
| 90 |
+
offsets.append(span)
|
| 91 |
+
if not offsets or offsets[0] is not None:
|
| 92 |
+
raise ValueError("position 0 must be a template token for the nothing readout")
|
| 93 |
+
paths = (
|
| 94 |
+
string_paths(text)
|
| 95 |
+
if state is not None and not isinstance(state, str) and style.endswith("_only") and style.startswith("json")
|
| 96 |
+
else []
|
| 97 |
+
)
|
| 98 |
+
return StatePositions(text, tuple(offsets), tuple(paths))
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def position_logits(query: torch.Tensor, state: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
|
| 102 |
+
"""[B, slots, D] against [B, T, D], sharing state keys between questions; valid [B, T], boolean, includes BOS0."""
|
| 103 |
+
dtype = torch.float64 if query.dtype == state.dtype == torch.float64 else torch.float32
|
| 104 |
+
logits = torch.bmm(query.to(dtype), state.to(dtype).transpose(1, 2)) / math.sqrt(query.shape[-1])
|
| 105 |
+
return logits.masked_fill(~valid[:, None], -torch.inf)
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def read_logits(hidden: torch.Tensor, states: torch.Tensor, plan: dict) -> torch.Tensor:
|
| 109 |
+
"""The flat read slots of a mixed batch; queries share each retained state's keys."""
|
| 110 |
+
device = hidden.device
|
| 111 |
+
query = hidden[plan["queries"].to(device)]
|
| 112 |
+
logits = position_logits(query, states, plan["valid"].to(device))
|
| 113 |
+
return logits[plan["query_row"].to(device), plan["query_col"].to(device)]
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
def query_plan(rows: list[int], masks: list[torch.Tensor], slots: list[int], count: int) -> dict:
|
| 117 |
+
"""The read slots of a batch of `count` slots, grouped by their row's state: each state's keys shared by its
|
| 118 |
+
queries; `code_slots`, the other slots, which the LM head reads."""
|
| 119 |
+
unique = sorted(set(rows))
|
| 120 |
+
index = {r: i for i, r in enumerate(unique)}
|
| 121 |
+
width = max(map(len, masks))
|
| 122 |
+
valid = torch.zeros(len(unique), width, dtype=torch.bool)
|
| 123 |
+
groups = [[] for _ in unique]
|
| 124 |
+
row_index, column_index = [], []
|
| 125 |
+
for slot, row, keep in zip(slots, rows, masks):
|
| 126 |
+
r = index[row]
|
| 127 |
+
if groups[r] and not torch.equal(valid[r, : len(keep)], keep):
|
| 128 |
+
raise ValueError("read questions must share their state-position mask")
|
| 129 |
+
valid[r, : len(keep)] = keep
|
| 130 |
+
row_index.append(r)
|
| 131 |
+
column_index.append(len(groups[r]))
|
| 132 |
+
groups[r].append(slot)
|
| 133 |
+
queries = torch.zeros(len(unique), max(map(len, groups)), dtype=torch.long)
|
| 134 |
+
for r, group in enumerate(groups):
|
| 135 |
+
queries[r, : len(group)] = torch.tensor(group)
|
| 136 |
+
return dict(
|
| 137 |
+
rows=torch.tensor(unique),
|
| 138 |
+
width=width,
|
| 139 |
+
queries=queries,
|
| 140 |
+
valid=valid,
|
| 141 |
+
query_row=torch.tensor(row_index),
|
| 142 |
+
query_col=torch.tensor(column_index),
|
| 143 |
+
slots=torch.tensor(slots),
|
| 144 |
+
code_slots=torch.tensor(sorted(set(range(count)) - set(slots)), dtype=torch.long),
|
| 145 |
+
)
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
def spans(logits: torch.Tensor, max_tokens: int = SPAN_CAPS["text"]) -> tuple[torch.Tensor, torch.Tensor]:
|
| 149 |
+
"""Each read's span from its start and end scores [R, 2, T]: found when P(found) >= 0.5, then the likeliest span
|
| 150 |
+
within `max_tokens`; (0, 0) for nothing."""
|
| 151 |
+
logp = logits.log_softmax(-1)
|
| 152 |
+
start, end = logp.unbind(1)
|
| 153 |
+
found = -torch.expm1(start[:, 0])
|
| 154 |
+
if logits.shape[2] == 1:
|
| 155 |
+
return torch.zeros(logits.shape[0], 2, dtype=torch.long, device=logits.device), found
|
| 156 |
+
cap = min(max_tokens, logits.shape[2] - 1)
|
| 157 |
+
windows, indices = F.max_pool1d(
|
| 158 |
+
F.pad(end[:, None, 1:], (0, cap - 1), value=-torch.inf), cap, stride=1, return_indices=True
|
| 159 |
+
)
|
| 160 |
+
scores = start[:, 1:] + windows[:, 0]
|
| 161 |
+
index = scores.argmax(-1)
|
| 162 |
+
first = index + 1
|
| 163 |
+
last = indices[:, 0].gather(1, index[:, None])[:, 0] + 1
|
| 164 |
+
at = torch.stack((first, last), -1)
|
| 165 |
+
nothing = found < 0.5
|
| 166 |
+
return torch.where(nothing[:, None], torch.zeros_like(at), at), found
|
processing_lfm3_vl.py
ADDED
|
@@ -0,0 +1,254 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Vendored from transformers' LFM3 code (not in a release yet) at 69da1b47a11e362fbb15bfe5d5c8da0790c688c5, src/transformers/models/lfm3_vl/processing_lfm3_vl.py:
|
| 2 |
+
# its relative imports made absolute, nothing else changed.
|
| 3 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 4 |
+
# This file was automatically generated from src/transformers/models/lfm3_vl/modular_lfm3_vl.py.
|
| 5 |
+
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
| 6 |
+
# the file from the modular. If any change should be done, please apply the change to the
|
| 7 |
+
# modular_lfm3_vl.py file directly. One of our CI enforces this.
|
| 8 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 9 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 10 |
+
#
|
| 11 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 12 |
+
# you may not use this file except in compliance with the License.
|
| 13 |
+
# You may obtain a copy of the License at
|
| 14 |
+
#
|
| 15 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 16 |
+
#
|
| 17 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 18 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 19 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 20 |
+
# See the License for the specific language governing permissions and
|
| 21 |
+
# limitations under the License.
|
| 22 |
+
import math
|
| 23 |
+
|
| 24 |
+
import numpy as np
|
| 25 |
+
|
| 26 |
+
from transformers.image_utils import is_pil_image, is_valid_image, make_flat_list_of_images
|
| 27 |
+
from transformers.processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin
|
| 28 |
+
from transformers.utils import auto_docstring, is_numpy_array, is_torch_tensor, logging
|
| 29 |
+
from transformers.video_utils import VideoMetadata, is_valid_video, make_batched_videos
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
logger = logging.get_logger(__name__)
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
class Lfm3VlProcessorKwargs(ProcessingKwargs, total=False):
|
| 36 |
+
_defaults = {
|
| 37 |
+
"text_kwargs": {
|
| 38 |
+
"padding": False,
|
| 39 |
+
"return_token_type_ids": False,
|
| 40 |
+
"return_mm_token_type_ids": False,
|
| 41 |
+
},
|
| 42 |
+
"videos_kwargs": {"return_metadata": True},
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def _check_video_fps(metadata: VideoMetadata) -> None:
|
| 47 |
+
if metadata.fps is not None and not (math.isfinite(metadata.fps) and metadata.fps > 0):
|
| 48 |
+
raise ValueError(f"`video_metadata.fps` must be a positive number, got {metadata.fps}")
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def _media_per_prompt(media, is_item) -> list[int] | None:
|
| 52 |
+
# rows name the prompt that owns them while the expansion consumes media in flat order; a row is a list of whole
|
| 53 |
+
# items or one stacked batch, and a list of single frames is one video, not a row
|
| 54 |
+
if not isinstance(media, (list, tuple)):
|
| 55 |
+
return None
|
| 56 |
+
counts = []
|
| 57 |
+
for row in media:
|
| 58 |
+
if isinstance(row, (list, tuple)) and all(is_item(item) for item in row):
|
| 59 |
+
counts.append(len(row))
|
| 60 |
+
elif (is_numpy_array(row) or is_torch_tensor(row)) and row.ndim == 4 and (len(row) == 0 or is_item(row[0])):
|
| 61 |
+
counts.append(len(row))
|
| 62 |
+
else:
|
| 63 |
+
return None
|
| 64 |
+
return counts
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def _check_vision_placeholder_counts(processor, prompts, images, videos) -> None:
|
| 68 |
+
for name, token, media, flatten, is_item in (
|
| 69 |
+
(
|
| 70 |
+
"images",
|
| 71 |
+
processor.image_token,
|
| 72 |
+
images,
|
| 73 |
+
make_flat_list_of_images,
|
| 74 |
+
lambda item: is_valid_image(item) and (is_pil_image(item) or item.ndim == 3),
|
| 75 |
+
),
|
| 76 |
+
(
|
| 77 |
+
"videos",
|
| 78 |
+
processor.video_token,
|
| 79 |
+
videos,
|
| 80 |
+
make_batched_videos,
|
| 81 |
+
lambda item: isinstance(item, str) or (not is_pil_image(item) and is_valid_video(item)),
|
| 82 |
+
),
|
| 83 |
+
):
|
| 84 |
+
num_placeholders = sum(prompt.count(token) for prompt in prompts)
|
| 85 |
+
num_media = 0 if media is None else len(flatten(media))
|
| 86 |
+
if num_placeholders != num_media:
|
| 87 |
+
raise ValueError(f"Found {num_placeholders} {token} tokens in the text but received {num_media} {name}")
|
| 88 |
+
# one prompt cannot take another's media
|
| 89 |
+
per_prompt = _media_per_prompt(media, is_item) if len(prompts) > 1 else None
|
| 90 |
+
in_text = [prompt.count(token) for prompt in prompts]
|
| 91 |
+
if per_prompt is not None and per_prompt != in_text:
|
| 92 |
+
raise ValueError(f"Found {in_text} {token} tokens per prompt but received {per_prompt} {name} per prompt")
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
@auto_docstring
|
| 96 |
+
class Lfm3VlProcessor(ProcessorMixin):
|
| 97 |
+
valid_processor_kwargs = Lfm3VlProcessorKwargs
|
| 98 |
+
|
| 99 |
+
def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs):
|
| 100 |
+
self.image_token = getattr(tokenizer, "image_token", "<|image|>")
|
| 101 |
+
self.video_token = getattr(tokenizer, "video_token", "<|video|>")
|
| 102 |
+
self.vision_start_token = getattr(tokenizer, "vision_start_token", "<|vision_start|>")
|
| 103 |
+
self.vision_end_token = getattr(tokenizer, "vision_end_token", "<|vision_end|>")
|
| 104 |
+
self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
|
| 105 |
+
self.video_token_id = tokenizer.convert_tokens_to_ids(self.video_token)
|
| 106 |
+
self.vision_start_token_id = tokenizer.convert_tokens_to_ids(self.vision_start_token)
|
| 107 |
+
self.vision_end_token_id = tokenizer.convert_tokens_to_ids(self.vision_end_token)
|
| 108 |
+
super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template)
|
| 109 |
+
|
| 110 |
+
def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> str:
|
| 111 |
+
merge_length = self.image_processor.merge_size**2
|
| 112 |
+
num_image_tokens = image_inputs["image_grid_thw"][image_idx].prod() // merge_length
|
| 113 |
+
return self.image_token * num_image_tokens
|
| 114 |
+
|
| 115 |
+
def replace_video_token(self, video_inputs: dict, video_idx: int, **kwargs) -> str:
|
| 116 |
+
metadata = video_inputs["video_metadata"][video_idx]
|
| 117 |
+
if metadata.frames_indices is None:
|
| 118 |
+
raise ValueError("Frame timestamps need the sampled frame indices in `video_metadata`")
|
| 119 |
+
_check_video_fps(metadata)
|
| 120 |
+
merge_length = self.video_processor.merge_size**2
|
| 121 |
+
num_frames = video_inputs["video_grid_thw"][video_idx][0]
|
| 122 |
+
frame_seqlen = video_inputs["video_grid_thw"][video_idx][1:].prod() // merge_length
|
| 123 |
+
metadata = video_inputs["video_metadata"][video_idx]
|
| 124 |
+
video_placeholder = ""
|
| 125 |
+
|
| 126 |
+
if metadata.fps is None:
|
| 127 |
+
logger.warning_once(
|
| 128 |
+
"Lfm3Vl requires frame timestamps to construct prompts, but the `fps` of the input video could not be inferred. "
|
| 129 |
+
"Probably `video_metadata` was missing from inputs and you passed pre-sampled frames. "
|
| 130 |
+
"Defaulting to `fps=24`. Please provide `video_metadata` for more accurate results."
|
| 131 |
+
)
|
| 132 |
+
metadata.fps = 24 if metadata.fps is None else metadata.fps
|
| 133 |
+
|
| 134 |
+
# if timestamps are not provided, calculate them
|
| 135 |
+
curr_timestamp = self._calculate_timestamps(
|
| 136 |
+
metadata.frames_indices,
|
| 137 |
+
metadata.fps,
|
| 138 |
+
self.video_processor.temporal_patch_size,
|
| 139 |
+
)
|
| 140 |
+
|
| 141 |
+
for frame_idx in range(num_frames):
|
| 142 |
+
curr_time = curr_timestamp[frame_idx]
|
| 143 |
+
video_placeholder += f"<{curr_time:.1f} seconds>"
|
| 144 |
+
video_placeholder += self.vision_start_token + self.video_token * frame_seqlen + self.vision_end_token
|
| 145 |
+
return video_placeholder
|
| 146 |
+
|
| 147 |
+
def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs):
|
| 148 |
+
"""
|
| 149 |
+
Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
|
| 150 |
+
|
| 151 |
+
Args:
|
| 152 |
+
image_sizes (`list[list[int]]`, *optional*):
|
| 153 |
+
The input sizes formatted as (height, width) per each image.
|
| 154 |
+
video_sizes (`list[list[int]]`, *optional*):
|
| 155 |
+
The input sizes formatted as (num_frames, height, width) per each video.
|
| 156 |
+
|
| 157 |
+
Returns:
|
| 158 |
+
`MultiModalData`: the number of tokens per provided modality.
|
| 159 |
+
"""
|
| 160 |
+
vision_data = {}
|
| 161 |
+
if image_sizes is not None:
|
| 162 |
+
# copy: `update` on the class-level defaults would leak kwargs into every later call
|
| 163 |
+
images_kwargs = dict(Lfm3VlProcessorKwargs._defaults.get("images_kwargs", {}))
|
| 164 |
+
images_kwargs.update(kwargs)
|
| 165 |
+
merge_size = images_kwargs.get("merge_size") or self.image_processor.merge_size
|
| 166 |
+
|
| 167 |
+
num_image_patches = [
|
| 168 |
+
self.image_processor.get_number_of_image_patches(*image_size, images_kwargs)
|
| 169 |
+
for image_size in image_sizes
|
| 170 |
+
]
|
| 171 |
+
num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches]
|
| 172 |
+
vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
|
| 173 |
+
|
| 174 |
+
if video_sizes is not None:
|
| 175 |
+
videos_kwargs = dict(Lfm3VlProcessorKwargs._defaults.get("videos_kwargs", {}))
|
| 176 |
+
videos_kwargs.update(kwargs)
|
| 177 |
+
merge_size = videos_kwargs.get("merge_size") or self.video_processor.merge_size
|
| 178 |
+
num_video_patches = [
|
| 179 |
+
self.video_processor.get_num_of_video_patches(*video_size, videos_kwargs) for video_size in video_sizes
|
| 180 |
+
]
|
| 181 |
+
num_video_tokens = [(num_patches // merge_size**2) for num_patches in num_video_patches]
|
| 182 |
+
vision_data["num_video_tokens"] = num_video_tokens
|
| 183 |
+
|
| 184 |
+
return MultiModalData(**vision_data)
|
| 185 |
+
|
| 186 |
+
def post_process_image_text_to_text(
|
| 187 |
+
self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
|
| 188 |
+
):
|
| 189 |
+
"""
|
| 190 |
+
Post-process the output of the model to decode the text.
|
| 191 |
+
|
| 192 |
+
Args:
|
| 193 |
+
generated_outputs (`torch.Tensor` or `np.ndarray`):
|
| 194 |
+
The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
|
| 195 |
+
or `(sequence_length,)`.
|
| 196 |
+
skip_special_tokens (`bool`, *optional*, defaults to `True`):
|
| 197 |
+
Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
|
| 198 |
+
clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
|
| 199 |
+
Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
|
| 200 |
+
**kwargs:
|
| 201 |
+
Additional arguments to be passed to the tokenizer's `batch_decode method`.
|
| 202 |
+
|
| 203 |
+
Returns:
|
| 204 |
+
`list[str]`: The decoded text.
|
| 205 |
+
"""
|
| 206 |
+
return self.tokenizer.batch_decode(
|
| 207 |
+
generated_outputs,
|
| 208 |
+
skip_special_tokens=skip_special_tokens,
|
| 209 |
+
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
|
| 210 |
+
**kwargs,
|
| 211 |
+
)
|
| 212 |
+
|
| 213 |
+
@property
|
| 214 |
+
def model_input_names(self):
|
| 215 |
+
return list(
|
| 216 |
+
dict.fromkeys(
|
| 217 |
+
self.tokenizer.model_input_names
|
| 218 |
+
+ self.image_processor.model_input_names
|
| 219 |
+
+ self.video_processor.model_input_names
|
| 220 |
+
)
|
| 221 |
+
)
|
| 222 |
+
|
| 223 |
+
def _calculate_timestamps(self, indices: list[int] | np.ndarray, video_fps: float, merge_size: int = 2):
|
| 224 |
+
if not isinstance(indices, list):
|
| 225 |
+
indices = indices.tolist()
|
| 226 |
+
if len(indices) % merge_size != 0:
|
| 227 |
+
indices.extend(indices[-1] for _ in range(merge_size - len(indices) % merge_size))
|
| 228 |
+
timestamps = [idx / video_fps for idx in indices]
|
| 229 |
+
# @JJJYmmm frames are merged by self.merge_size, \
|
| 230 |
+
# so we need to average the timestamps between the first/last frame within the temporal patch
|
| 231 |
+
timestamps = [
|
| 232 |
+
(timestamps[i] + timestamps[i + merge_size - 1]) / 2 for i in range(0, len(timestamps), merge_size)
|
| 233 |
+
]
|
| 234 |
+
return timestamps
|
| 235 |
+
|
| 236 |
+
def prepare_inputs_layout(self, images=None, text=None, videos=None, audio=None, **kwargs):
|
| 237 |
+
for name, media in (("images", images), ("videos", videos)):
|
| 238 |
+
if isinstance(media, (list, tuple)) and len(media) == 0:
|
| 239 |
+
raise ValueError(f"`{name}` must hold at least one item; pass `{name}=None` for prompts without one")
|
| 240 |
+
return super().prepare_inputs_layout(images=images, text=text, videos=videos, audio=audio, **kwargs)
|
| 241 |
+
|
| 242 |
+
def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs):
|
| 243 |
+
super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs)
|
| 244 |
+
if text is not None:
|
| 245 |
+
_check_vision_placeholder_counts(self, [text] if isinstance(text, str) else text, images, videos)
|
| 246 |
+
|
| 247 |
+
def _process_videos(self, videos, **kwargs):
|
| 248 |
+
# the frame timestamps need the metadata whatever the caller's `return_metadata`; `__call__` drops it from
|
| 249 |
+
# the output unless it was asked for
|
| 250 |
+
kwargs["return_metadata"] = True
|
| 251 |
+
return super()._process_videos(videos, **kwargs)
|
| 252 |
+
|
| 253 |
+
|
| 254 |
+
__all__ = ["Lfm3VlProcessor"]
|
processor_config.json
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"image_processor": {
|
| 3 |
+
"do_convert_rgb": true,
|
| 4 |
+
"do_normalize": true,
|
| 5 |
+
"do_rescale": true,
|
| 6 |
+
"do_resize": true,
|
| 7 |
+
"image_mean": [
|
| 8 |
+
0.5,
|
| 9 |
+
0.5,
|
| 10 |
+
0.5
|
| 11 |
+
],
|
| 12 |
+
"image_processor_type": "Lfm3VlImageProcessor",
|
| 13 |
+
"image_std": [
|
| 14 |
+
0.5,
|
| 15 |
+
0.5,
|
| 16 |
+
0.5
|
| 17 |
+
],
|
| 18 |
+
"merge_size": 2,
|
| 19 |
+
"patch_size": 16,
|
| 20 |
+
"resample": 3,
|
| 21 |
+
"rescale_factor": 0.00392156862745098,
|
| 22 |
+
"size": {
|
| 23 |
+
"longest_edge": 16777216,
|
| 24 |
+
"shortest_edge": 65536
|
| 25 |
+
},
|
| 26 |
+
"temporal_patch_size": 2
|
| 27 |
+
},
|
| 28 |
+
"processor_class": "Lfm3VlProcessor",
|
| 29 |
+
"video_processor": {
|
| 30 |
+
"cap_pixels_per_frame": false,
|
| 31 |
+
"do_convert_rgb": true,
|
| 32 |
+
"do_normalize": true,
|
| 33 |
+
"do_rescale": true,
|
| 34 |
+
"do_resize": true,
|
| 35 |
+
"do_sample_frames": true,
|
| 36 |
+
"fps": 2.0,
|
| 37 |
+
"image_mean": [
|
| 38 |
+
0.5,
|
| 39 |
+
0.5,
|
| 40 |
+
0.5
|
| 41 |
+
],
|
| 42 |
+
"image_std": [
|
| 43 |
+
0.5,
|
| 44 |
+
0.5,
|
| 45 |
+
0.5
|
| 46 |
+
],
|
| 47 |
+
"max_frames": 768,
|
| 48 |
+
"max_video_tokens": 768,
|
| 49 |
+
"merge_size": 2,
|
| 50 |
+
"min_frames": 4,
|
| 51 |
+
"patch_size": 16,
|
| 52 |
+
"resample": 3,
|
| 53 |
+
"rescale_factor": 0.00392156862745098,
|
| 54 |
+
"return_metadata": false,
|
| 55 |
+
"size": {
|
| 56 |
+
"longest_edge": 25165824,
|
| 57 |
+
"shortest_edge": 4096
|
| 58 |
+
},
|
| 59 |
+
"temporal_patch_size": 2,
|
| 60 |
+
"video_processor_type": "Lfm3VlVideoProcessor"
|
| 61 |
+
}
|
| 62 |
+
}
|
prompt.py
ADDED
|
@@ -0,0 +1,761 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The prompt and the readout. A System One decision is a single forward pass that stops at the answer
|
| 2 |
+
slot: everything the model sees is built here, and the answer is a softmax over its options' tokens.
|
| 3 |
+
"""
|
| 4 |
+
|
| 5 |
+
from __future__ import annotations
|
| 6 |
+
|
| 7 |
+
import json
|
| 8 |
+
import math
|
| 9 |
+
import re
|
| 10 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import cached_property, lru_cache
|
| 12 |
+
from typing import Any, Mapping, Sequence
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
|
| 16 |
+
from .media import SPOKEN, Media, as_media
|
| 17 |
+
|
| 18 |
+
# The system turn: none by default.
|
| 19 |
+
SYSTEMS: dict[str, str | None] = {"none": None, "d1": "You are d1, a decision model trained by Liquid AI."}
|
| 20 |
+
DEFAULT_SYSTEM = "none"
|
| 21 |
+
|
| 22 |
+
IM_START = "<|im_start|>"
|
| 23 |
+
IM_END = "<|im_end|>"
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
# --------------------------------------------------------------------------- #
|
| 27 |
+
# question types
|
| 28 |
+
# --------------------------------------------------------------------------- #
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
@dataclass(frozen=True)
|
| 32 |
+
class Choice:
|
| 33 |
+
instructions: str | Media # a spoken question's clip (`media`)
|
| 34 |
+
criteria: Mapping[str, str | None]
|
| 35 |
+
type: str = "choice"
|
| 36 |
+
multiple: bool = False # each option its own yes or no
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
@dataclass(frozen=True)
|
| 40 |
+
class Noul:
|
| 41 |
+
instructions: str | Media
|
| 42 |
+
criteria: Mapping[str, Any] | None = None
|
| 43 |
+
type: str = "noul"
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
@dataclass(frozen=True)
|
| 47 |
+
class Score:
|
| 48 |
+
instructions: str | Media
|
| 49 |
+
criteria: Sequence[str]
|
| 50 |
+
type: str = "score"
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
@dataclass(frozen=True)
|
| 54 |
+
class Read:
|
| 55 |
+
"""A span read from the state: text, or a number or a date written there; `unit`, a number's."""
|
| 56 |
+
|
| 57 |
+
instructions: str | Media
|
| 58 |
+
type: str = "text"
|
| 59 |
+
unit: str | None = None
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
@dataclass(frozen=True)
|
| 63 |
+
class Estimate:
|
| 64 |
+
"""A number estimated on a grid over the declared range: `step` apart, at most 128 values (a finer step
|
| 65 |
+
is read on 128)."""
|
| 66 |
+
|
| 67 |
+
instructions: str | Media
|
| 68 |
+
low: float
|
| 69 |
+
high: float
|
| 70 |
+
step: float | None = None
|
| 71 |
+
unit: str | None = None
|
| 72 |
+
type: str = "number"
|
| 73 |
+
|
| 74 |
+
@cached_property
|
| 75 |
+
def grid(self):
|
| 76 |
+
from .grid import Grid
|
| 77 |
+
|
| 78 |
+
return Grid(self.low, self.high, 128 if self.step is None else min(128, round((self.high - self.low) / self.step) + 1))
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
@dataclass(frozen=True)
|
| 82 |
+
class Box:
|
| 83 |
+
"""Where something is in the picture: its four edges, or nothing."""
|
| 84 |
+
|
| 85 |
+
instructions: str | Media
|
| 86 |
+
type: str = "box"
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
@dataclass(frozen=True)
|
| 90 |
+
class Boxes:
|
| 91 |
+
"""Every instance of something in the picture, in reading order, none if absent: `max` blocks read, else
|
| 92 |
+
`box.INSTANCES`. A readout choice: the question's text never says it."""
|
| 93 |
+
|
| 94 |
+
instructions: str | Media
|
| 95 |
+
type: str = "boxes"
|
| 96 |
+
max: int | None = None
|
| 97 |
+
|
| 98 |
+
@property
|
| 99 |
+
def blocks(self) -> int:
|
| 100 |
+
from .box import INSTANCES
|
| 101 |
+
|
| 102 |
+
return INSTANCES if self.max is None else self.max
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
@dataclass(frozen=True)
|
| 106 |
+
class Position:
|
| 107 |
+
name: str
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
@dataclass(frozen=True)
|
| 111 |
+
class Number:
|
| 112 |
+
"""A slot over a grid's values as the model writes them: value i times `scale`, an integer, is one token of the
|
| 113 |
+
tokenizer's own (`number_ids`). A checkpoint that records `d1_numbers` reads its grids so."""
|
| 114 |
+
|
| 115 |
+
instructions: str
|
| 116 |
+
grid: Any
|
| 117 |
+
scale: int = 1
|
| 118 |
+
type: str = "grid"
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
Question = Choice | Noul | Score | Read | Estimate | Box | Boxes
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
TYPED = ("text", "number", "date", "box", "boxes") # the typed answers' kinds; a choice is one when `multiple`
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
def typed(q: Question) -> bool:
|
| 128 |
+
"""A typed answer, read at its own slots or state positions and answered in its own shape (api.py):
|
| 129 |
+
a read, an estimate, a box, a multiple choice."""
|
| 130 |
+
return q.type in TYPED or isinstance(q, Choice) and q.multiple
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
def cardinality(q: Question) -> int:
|
| 134 |
+
if isinstance(q, (Estimate, Number)):
|
| 135 |
+
return q.grid.size
|
| 136 |
+
return 2 if isinstance(q, Noul) else len(q.criteria)
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
# --------------------------------------------------------------------------- #
|
| 140 |
+
# verbalizer
|
| 141 |
+
# --------------------------------------------------------------------------- #
|
| 142 |
+
|
| 143 |
+
|
| 144 |
+
PROMPTS = (1, 2)
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def version(tokenizer) -> int:
|
| 148 |
+
"""The prompt version the checkpoint behind `tokenizer` trained with (`PROMPTS`)."""
|
| 149 |
+
return getattr(tokenizer, "d1_prompt", 1)
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
def most_blocks(tokenizer) -> int:
|
| 153 |
+
"""The most blocks a `boxes` question's `max` asks of this checkpoint: the 16 prompt 1 trained, else
|
| 154 |
+
`box.MAX_INSTANCES`."""
|
| 155 |
+
from .box import INSTANCES, MAX_INSTANCES
|
| 156 |
+
|
| 157 |
+
return INSTANCES if version(tokenizer) == 1 else MAX_INSTANCES
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
LETTERS = tuple(chr(ord("A") + i) for i in range(26))
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
def option_codes(labels: Sequence[str], scheme: int = 1) -> list[str]:
|
| 164 |
+
"""Native letters when the labels already are letters, else A..Z, else the `scheme`'s codes: 00.. (1), A..Z then
|
| 165 |
+
AA, AB.. (2); one rule for every cardinality."""
|
| 166 |
+
labs = [str(x).strip() for x in labels]
|
| 167 |
+
if labs and all(len(k) == 1 and k.isalpha() for k in labs):
|
| 168 |
+
return labs
|
| 169 |
+
if len(labs) <= 26:
|
| 170 |
+
return list(LETTERS[: len(labs)])
|
| 171 |
+
if scheme == 2:
|
| 172 |
+
return (list(LETTERS) + [a + b for a in LETTERS for b in LETTERS])[: len(labs)]
|
| 173 |
+
return [f"{i:02d}" for i in range(len(labs))]
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
_FALLBACK_POOL = (
|
| 177 |
+
[chr(c) for c in range(ord("A"), ord("Z") + 1)]
|
| 178 |
+
+ [f"{i:02d}" for i in range(100)]
|
| 179 |
+
+ [chr(c) for c in range(ord("a"), ord("z") + 1)]
|
| 180 |
+
+ [f"#{i}" for i in range(200)]
|
| 181 |
+
# Where a tokenizer splits digits, "00".."99" and "#i" are two tokens, and two capital letters are
|
| 182 |
+
# often one. Last in the pool, so a tokenizer with digit pairs never reaches it.
|
| 183 |
+
+ [chr(a) + chr(b) for a in range(ord("A"), ord("Z") + 1) for b in range(ord("A"), ord("Z") + 1)]
|
| 184 |
+
)
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
@lru_cache(maxsize=4096)
|
| 188 |
+
def _aliases_cached(tokenizer_key: int, codes: tuple[str, ...]) -> tuple[tuple[str, int], ...]:
|
| 189 |
+
tokenizer = _TOKENIZERS[tokenizer_key]
|
| 190 |
+
used: set[int] = set()
|
| 191 |
+
out: list[tuple[str, int]] = []
|
| 192 |
+
|
| 193 |
+
def take(raw: str) -> bool:
|
| 194 |
+
enc = tokenizer.encode(raw, add_special_tokens=False)
|
| 195 |
+
if len(enc) != 1 or enc[0] in used:
|
| 196 |
+
return False
|
| 197 |
+
out.append((raw, enc[0]))
|
| 198 |
+
used.add(enc[0])
|
| 199 |
+
return True
|
| 200 |
+
|
| 201 |
+
for code in codes:
|
| 202 |
+
if take(code):
|
| 203 |
+
continue
|
| 204 |
+
if not any(take(raw) for raw in _FALLBACK_POOL):
|
| 205 |
+
raise RuntimeError(f"no single-token alias left for {len(codes)} options")
|
| 206 |
+
return tuple(out)
|
| 207 |
+
|
| 208 |
+
|
| 209 |
+
_TOKENIZERS: dict[int, object] = {}
|
| 210 |
+
|
| 211 |
+
|
| 212 |
+
def aliases(tokenizer, labels: Sequence[str]) -> list[tuple[str, int]]:
|
| 213 |
+
"""Every label's distinct single-token code: [(code, token_id)]. Memoised on the codes, which are positional unless
|
| 214 |
+
the labels are letters, so every ordering of one option list shares an entry."""
|
| 215 |
+
_TOKENIZERS.setdefault(id(tokenizer), tokenizer)
|
| 216 |
+
codes = tuple(option_codes(labels, getattr(tokenizer, "d1_codes", 1)))
|
| 217 |
+
return list(_aliases_cached(id(tokenizer), codes))
|
| 218 |
+
|
| 219 |
+
|
| 220 |
+
@lru_cache(maxsize=8192)
|
| 221 |
+
def _ids_cached(tokenizer_key: int, texts: tuple[str, ...]) -> tuple[int, ...]:
|
| 222 |
+
tokenizer = _TOKENIZERS[tokenizer_key]
|
| 223 |
+
out, seen = [], set()
|
| 224 |
+
for t in texts:
|
| 225 |
+
enc = tokenizer.encode(t, add_special_tokens=False)
|
| 226 |
+
if len(enc) == 1 and enc[0] not in seen:
|
| 227 |
+
out.append(enc[0])
|
| 228 |
+
seen.add(enc[0])
|
| 229 |
+
return tuple(out)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
def _ids(tokenizer, texts: Sequence[str]) -> list[int]:
|
| 233 |
+
_TOKENIZERS.setdefault(id(tokenizer), tokenizer)
|
| 234 |
+
return list(_ids_cached(id(tokenizer), tuple(texts)))
|
| 235 |
+
|
| 236 |
+
|
| 237 |
+
QUESTION_TYPES = ("noul", "choice", "score", "text", "number", "date", "box", "boxes")
|
| 238 |
+
CHOICE = '"criteria" names the options: {"billing": "Charges and refunds", "shipping": "Delivery"}'
|
| 239 |
+
SCORE = '"criteria" lists 2 to 10 levels, lowest first: ["Can wait", "Today", "Now"]'
|
| 240 |
+
|
| 241 |
+
|
| 242 |
+
def check_question(q: Any, tokenizer=None) -> None:
|
| 243 |
+
"""A caller's question in the API's JSON, refused with what to change; a class passes. `tokenizer` carries the
|
| 244 |
+
checkpoint's prompt version, which bounds a `boxes` question's `max` (`most_blocks`; prompt 1 without it)."""
|
| 245 |
+
if isinstance(q, (Choice, Noul, Score, Read, Estimate, Box, Boxes)):
|
| 246 |
+
return
|
| 247 |
+
if not isinstance(q, Mapping):
|
| 248 |
+
raise ValueError(f'a question is a dict such as {{"type": "noul", "instructions": "Is it urgent?"}}, not '
|
| 249 |
+
f"{type(q).__name__}")
|
| 250 |
+
kind, said, criteria = q.get("type"), q.get("instructions"), q.get("criteria")
|
| 251 |
+
if kind not in QUESTION_TYPES:
|
| 252 |
+
raise ValueError(f'"type" is one of {", ".join(QUESTION_TYPES)}, not {kind!r}')
|
| 253 |
+
if not (isinstance(said, str) and said.strip() or isinstance(said, Mapping)):
|
| 254 |
+
raise ValueError('"instructions" is the question as text (or a spoken one, {"type": "audio", "path": ...})')
|
| 255 |
+
if kind == "choice" and not (isinstance(criteria, Mapping) and criteria
|
| 256 |
+
and all(isinstance(k, str) and (v is None or isinstance(v, str))
|
| 257 |
+
for k, v in criteria.items())):
|
| 258 |
+
raise ValueError(f"a choice's {CHOICE}" + ("; a list of levels is a score" if isinstance(criteria, list) else ""))
|
| 259 |
+
if kind == "score" and not (isinstance(criteria, (list, tuple)) and 2 <= len(criteria) <= 10
|
| 260 |
+
and all(isinstance(c, str) for c in criteria)):
|
| 261 |
+
raise ValueError(f"a score's {SCORE}")
|
| 262 |
+
if kind == "noul" and criteria is not None and not isinstance(criteria, Mapping):
|
| 263 |
+
raise ValueError('a noul\'s "criteria", when given, is {"true": "what yes means", "false": "what no means"}')
|
| 264 |
+
if q.get("multiple") and kind != "choice":
|
| 265 |
+
raise ValueError('"multiple" belongs to a choice: {"type": "choice", "multiple": true, "criteria": {...}}')
|
| 266 |
+
if "max" in q and not (kind == "boxes" and isinstance(q["max"], int) and not isinstance(q["max"], bool)
|
| 267 |
+
and 1 <= q["max"] <= most_blocks(tokenizer)):
|
| 268 |
+
raise ValueError(f'"max" belongs to boxes, a whole number from 1 to {most_blocks(tokenizer)}: '
|
| 269 |
+
'{"type": "boxes", "instructions": "Every cup.", "max": 4}')
|
| 270 |
+
if "step" in q and "range" not in q:
|
| 271 |
+
raise ValueError('"step" needs a "range": {"type": "number", "range": [0, 10], "step": 1}')
|
| 272 |
+
if "range" in q:
|
| 273 |
+
bounds = q["range"]
|
| 274 |
+
if kind != "number" or not (isinstance(bounds, (list, tuple)) and len(bounds) == 2
|
| 275 |
+
and all(isinstance(b, (int, float)) and not isinstance(b, bool) for b in bounds)):
|
| 276 |
+
raise ValueError('an estimate is {"type": "number", "range": [low, high]}, with an optional "step"')
|
| 277 |
+
|
| 278 |
+
|
| 279 |
+
def as_question(q: Mapping | Question) -> Question:
|
| 280 |
+
"""A question in the API's JSON (`type`, `instructions`, `criteria`; `multiple`, `range`, `step`, `unit` for the
|
| 281 |
+
typed answers; `instructions` a media dict for a spoken question), as one of the classes above; a class passes
|
| 282 |
+
through."""
|
| 283 |
+
if not isinstance(q, Mapping):
|
| 284 |
+
return q
|
| 285 |
+
if isinstance(q.get("instructions"), Mapping):
|
| 286 |
+
q = {**q, "instructions": as_media(q["instructions"])}
|
| 287 |
+
kind = q.get("type", "choice")
|
| 288 |
+
if kind == "number" and "range" in q:
|
| 289 |
+
low, high = q["range"]
|
| 290 |
+
step = q.get("step")
|
| 291 |
+
if not low < high or step is not None and not (step > 0 and math.isclose((high - low) / step, round((high - low) / step))):
|
| 292 |
+
raise ValueError(f"an estimate's range runs low to high and its step, above 0, divides it: not "
|
| 293 |
+
f"{[low, high]} in steps of {step}")
|
| 294 |
+
return Estimate(q["instructions"], low, high, step, q.get("unit"))
|
| 295 |
+
if kind in ("text", "number", "date"):
|
| 296 |
+
return Read(q["instructions"], kind, q.get("unit"))
|
| 297 |
+
if kind == "box":
|
| 298 |
+
return Box(q["instructions"])
|
| 299 |
+
if kind == "boxes":
|
| 300 |
+
return Boxes(q["instructions"], max=q.get("max"))
|
| 301 |
+
if kind == "noul":
|
| 302 |
+
return Noul(q["instructions"], q.get("criteria"))
|
| 303 |
+
if kind == "score":
|
| 304 |
+
return Score(q["instructions"], list(q["criteria"]))
|
| 305 |
+
return Choice(q["instructions"], q["criteria"], multiple=bool(q.get("multiple", False)))
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
def question_slots(q: Question, tokenizer=None) -> tuple[Question, ...]:
|
| 309 |
+
"""Each answer position's readout, in question order; a grid's through the tokenizer's numbers where the
|
| 310 |
+
checkpoint reads them so (`grid_slot`)."""
|
| 311 |
+
if isinstance(q, Read):
|
| 312 |
+
return Position("start"), Position("end")
|
| 313 |
+
if isinstance(q, Estimate):
|
| 314 |
+
return (grid_slot(tokenizer, q.instructions, q.grid),)
|
| 315 |
+
if isinstance(q, Box):
|
| 316 |
+
from .box import EDGES, edge_grid
|
| 317 |
+
|
| 318 |
+
grid, scale = edge_grid(tokenizer)
|
| 319 |
+
return (Noul(q.instructions), *(grid_slot(tokenizer, q.instructions, grid, scale) for _ in EDGES))
|
| 320 |
+
if isinstance(q, Boxes):
|
| 321 |
+
return question_slots(Box(q.instructions), tokenizer) * q.blocks
|
| 322 |
+
if isinstance(q, Choice) and q.multiple:
|
| 323 |
+
if not q.criteria:
|
| 324 |
+
raise ValueError("choice with no options")
|
| 325 |
+
return tuple(Noul(q.instructions) for _ in q.criteria)
|
| 326 |
+
return (q,)
|
| 327 |
+
|
| 328 |
+
|
| 329 |
+
YES_FORMS = ("yes", "Yes", "YES")
|
| 330 |
+
NO_FORMS = ("no", "No", "NO")
|
| 331 |
+
|
| 332 |
+
|
| 333 |
+
def readout_ids(tokenizer, q: Question) -> list[list[int]]:
|
| 334 |
+
"""Token ids to score, one group per option, max-pooled: the answer is a softmax over these and
|
| 335 |
+
nothing else."""
|
| 336 |
+
if isinstance(q, Estimate):
|
| 337 |
+
return readout_ids(tokenizer, question_slots(q, tokenizer)[0])
|
| 338 |
+
if isinstance(q, Number):
|
| 339 |
+
return [list(group) for group in number_ids(tokenizer, q.grid, q.scale)]
|
| 340 |
+
if isinstance(q, (Read, Position)):
|
| 341 |
+
raise ValueError("a pointer reads state positions, not vocabulary codes")
|
| 342 |
+
if isinstance(q, Choice) and q.multiple:
|
| 343 |
+
raise ValueError("a multiple choice reads one yes/no distribution per option")
|
| 344 |
+
if isinstance(q, Noul):
|
| 345 |
+
yes, no = _ids(tokenizer, YES_FORMS), _ids(tokenizer, NO_FORMS)
|
| 346 |
+
if not yes or not no:
|
| 347 |
+
raise RuntimeError("tokenizer has no single-token yes/no")
|
| 348 |
+
return [yes, no]
|
| 349 |
+
if isinstance(q, Score):
|
| 350 |
+
groups = [_ids(tokenizer, [str(i)]) for i in range(len(q.criteria))]
|
| 351 |
+
if any(not g for g in groups):
|
| 352 |
+
raise RuntimeError(f"score with {len(q.criteria)} levels needs single-token digits; the primitive is defined "
|
| 353 |
+
"for 2 to 10")
|
| 354 |
+
return groups
|
| 355 |
+
groups = []
|
| 356 |
+
for code, tid in aliases(tokenizer, list(q.criteria.keys())):
|
| 357 |
+
extra = _ids(tokenizer, [f" {code}"])
|
| 358 |
+
groups.append([tid] + [i for i in extra if i != tid])
|
| 359 |
+
if not groups:
|
| 360 |
+
raise RuntimeError("choice with no options")
|
| 361 |
+
return groups
|
| 362 |
+
|
| 363 |
+
|
| 364 |
+
def readout(tokenizer, q: Question, logz, calibration=None) -> list[float]:
|
| 365 |
+
"""Option probabilities from the log-probabilities at the answer slot, `logz` (a vocabulary tensor, or a dict
|
| 366 |
+
holding at least the option tokens). Each option scores its best form."""
|
| 367 |
+
if not hasattr(logz, "tolist"):
|
| 368 |
+
return probabilities(q, [max(float(logz[i]) for i in g) for g in readout_ids(tokenizer, q)], calibration)
|
| 369 |
+
tokens, forms, _ = readout_table(tokenizer, q)
|
| 370 |
+
values = np.asarray(logz[list(tokens)].tolist()) # one device transfer
|
| 371 |
+
return probabilities(q, values[forms].max(-1).tolist(), calibration)
|
| 372 |
+
|
| 373 |
+
|
| 374 |
+
def readouts(tokenizer, qs: Sequence[Question], logz, calibrations: Sequence) -> list[list[float]]:
|
| 375 |
+
"""`readout` for each row of `logz` ([slots, vocabulary], a torch tensor) and its question, as one gather for
|
| 376 |
+
the rows sharing a readout table (a box's edges) and one transfer for all."""
|
| 377 |
+
import torch
|
| 378 |
+
|
| 379 |
+
tables = [readout_table(tokenizer, q) for q in qs]
|
| 380 |
+
shared: dict[int, list[int]] = {}
|
| 381 |
+
for j, table in enumerate(tables):
|
| 382 |
+
shared.setdefault(id(table), []).append(j)
|
| 383 |
+
picked = [logz[rows][:, torch.from_numpy(tables[rows[0]][2]).to(logz.device)] for rows in shared.values()]
|
| 384 |
+
host = torch.cat([p.flatten() for p in picked]).cpu().numpy() if picked else None
|
| 385 |
+
out, at = [None] * len(qs), 0
|
| 386 |
+
for rows, part in zip(shared.values(), picked):
|
| 387 |
+
values = host[at:at + part.numel()].reshape(part.shape)
|
| 388 |
+
at += part.numel()
|
| 389 |
+
scores = values[:, tables[rows[0]][1]].max(-1).astype(np.float64)
|
| 390 |
+
scores /= np.array([1.0 if calibrations[j] is None else calibrations[j].temperature(qs[j]) for j in rows])[:, None]
|
| 391 |
+
exps = np.exp(scores - scores.max(-1, keepdims=True)) # the softmax of `probabilities`, a block at once
|
| 392 |
+
for j, p in zip(rows, (exps / exps.sum(-1, keepdims=True)).tolist()):
|
| 393 |
+
out[j] = p
|
| 394 |
+
return out
|
| 395 |
+
|
| 396 |
+
|
| 397 |
+
def probabilities(q: Question, scores: list[float], calibration=None) -> list[float]:
|
| 398 |
+
"""Each option's best score, through the calibration's temperature, as a softmax."""
|
| 399 |
+
if calibration is not None:
|
| 400 |
+
scores = calibration.apply(q, scores)
|
| 401 |
+
m = max(scores)
|
| 402 |
+
exps = [math.exp(s - m) for s in scores]
|
| 403 |
+
total = sum(exps)
|
| 404 |
+
return [e / total for e in exps]
|
| 405 |
+
|
| 406 |
+
|
| 407 |
+
def readout_table(tokenizer, q: Question) -> tuple[tuple[int, ...], np.ndarray, np.ndarray]:
|
| 408 |
+
"""`readout_ids` as one gather: every token once, each group's forms as positions among them (a short group
|
| 409 |
+
repeats its first form, which leaves its best as it is), and the tokens as an array. Cached: a box edge scores
|
| 410 |
+
1,000 numbers."""
|
| 411 |
+
if isinstance(q, Number):
|
| 412 |
+
_TOKENIZERS.setdefault(id(tokenizer), tokenizer)
|
| 413 |
+
return _number_table(id(tokenizer), q.grid.low, q.grid.high, q.grid.size, q.scale)
|
| 414 |
+
return _table(tuple(map(tuple, readout_ids(tokenizer, q))))
|
| 415 |
+
|
| 416 |
+
|
| 417 |
+
@lru_cache(maxsize=4096)
|
| 418 |
+
def _table(groups: tuple[tuple[int, ...], ...]) -> tuple[tuple[int, ...], np.ndarray, np.ndarray]:
|
| 419 |
+
tokens = tuple(dict.fromkeys(i for g in groups for i in g))
|
| 420 |
+
at = {t: k for k, t in enumerate(tokens)}
|
| 421 |
+
width = max(map(len, groups))
|
| 422 |
+
return tokens, np.array([[at[i] for i in g] + [at[g[0]]] * (width - len(g)) for g in groups]), np.array(tokens)
|
| 423 |
+
|
| 424 |
+
|
| 425 |
+
@lru_cache(maxsize=256)
|
| 426 |
+
def _number_table(tokenizer_key: int, low: float, high: float, size: int, scale: int):
|
| 427 |
+
from .grid import Grid
|
| 428 |
+
|
| 429 |
+
return _table(tuple(map(tuple, readout_ids(_TOKENIZERS[tokenizer_key], Number("", Grid(low, high, size), scale)))))
|
| 430 |
+
|
| 431 |
+
|
| 432 |
+
# --------------------------------------------------------------------------- #
|
| 433 |
+
# state and question rendering
|
| 434 |
+
# --------------------------------------------------------------------------- #
|
| 435 |
+
|
| 436 |
+
|
| 437 |
+
# How a state is rendered: `json_only`, the default, writes every state as the object it is; `json` keeps
|
| 438 |
+
# three shortcuts (`Message:`, `Passage:` / `Asked:`, a lone question's text); `sections` writes nested
|
| 439 |
+
# states as labelled blocks.
|
| 440 |
+
DEFAULT_STATE_STYLE = "json_only"
|
| 441 |
+
|
| 442 |
+
|
| 443 |
+
def _is_scalar(v: Any) -> bool:
|
| 444 |
+
return v is None or isinstance(v, (str, int, float, bool)) and "\n" not in str(v)
|
| 445 |
+
|
| 446 |
+
|
| 447 |
+
def _sections(obj: Any, path: str, out: list[str]) -> None:
|
| 448 |
+
"""Flatten a nested state into labelled blocks, keeping the real newlines `json.dumps` would escape."""
|
| 449 |
+
head = f"[{path}]\n" if path else ""
|
| 450 |
+
if isinstance(obj, dict):
|
| 451 |
+
scalars = [(k, v) for k, v in obj.items() if _is_scalar(v)]
|
| 452 |
+
rest = [(k, v) for k, v in obj.items() if not _is_scalar(v)]
|
| 453 |
+
if scalars:
|
| 454 |
+
body = "\n".join(f"{k}: {'' if v is None else v}" for k, v in scalars)
|
| 455 |
+
out.append(f"{head}{body}")
|
| 456 |
+
for k, v in rest:
|
| 457 |
+
_sections(v, f"{path}.{k}" if path else str(k), out)
|
| 458 |
+
return
|
| 459 |
+
if isinstance(obj, (list, tuple)):
|
| 460 |
+
if obj and all(_is_scalar(v) for v in obj):
|
| 461 |
+
body = "\n".join(f"- {'' if v is None else v}" for v in obj)
|
| 462 |
+
out.append(f"{head}{body}")
|
| 463 |
+
return
|
| 464 |
+
for i, v in enumerate(obj, start=1):
|
| 465 |
+
_sections(v, f"{path} {i}/{len(obj)}" if path else f"{i}/{len(obj)}", out)
|
| 466 |
+
return
|
| 467 |
+
out.append(f"{head}{'' if obj is None else obj}")
|
| 468 |
+
|
| 469 |
+
|
| 470 |
+
def render_state(state: Any) -> str:
|
| 471 |
+
if isinstance(state, str):
|
| 472 |
+
return state
|
| 473 |
+
out: list[str] = []
|
| 474 |
+
_sections(state, "", out)
|
| 475 |
+
return "\n\n".join(out)
|
| 476 |
+
|
| 477 |
+
|
| 478 |
+
def state_text(state: Any, style: str = DEFAULT_STATE_STYLE) -> str:
|
| 479 |
+
"""The state as the prompt renders it, without the blank line after it: what a read's span offsets index."""
|
| 480 |
+
return state_block(state, style)[:-2]
|
| 481 |
+
|
| 482 |
+
|
| 483 |
+
def state_block(state: Any, style: str = DEFAULT_STATE_STYLE) -> str:
|
| 484 |
+
"""The state as the block before `QUESTION:`. The `_only` styles (the default) write every state as the object
|
| 485 |
+
it is; the others give three shapes a shortcut (`Message:`, `Passage:` / `Asked:`, a lone question's text)."""
|
| 486 |
+
if isinstance(state, dict) and not style.endswith("_only"):
|
| 487 |
+
keys = set(state.keys())
|
| 488 |
+
if keys == {"text"}:
|
| 489 |
+
return f"Message: {state['text']}\n\n"
|
| 490 |
+
if {"passage", "question"} <= keys and len(keys) == 2:
|
| 491 |
+
return f"Passage: {state['passage']}\n\nAsked: {state['question']}\n\n"
|
| 492 |
+
if keys == {"question"}:
|
| 493 |
+
return f"{state['question']}\n\n"
|
| 494 |
+
if style.startswith("json"):
|
| 495 |
+
if isinstance(state, str):
|
| 496 |
+
return f"{state}\n\n"
|
| 497 |
+
return json.dumps(state, ensure_ascii=False, indent=2) + "\n\n"
|
| 498 |
+
return f"{render_state(state)}\n\n"
|
| 499 |
+
|
| 500 |
+
|
| 501 |
+
# A zero-width space after the first character of a special token's string: it reads the same and is no longer the
|
| 502 |
+
# token. Without it a state holding `<|im_end|>` closed the user turn, and `<image>` became an image token
|
| 503 |
+
# with no picture behind it (the processor then refused the row).
|
| 504 |
+
BREAK = "\u200b"
|
| 505 |
+
|
| 506 |
+
|
| 507 |
+
@lru_cache(maxsize=64)
|
| 508 |
+
def _specials_cached(tokenizer_key: int) -> re.Pattern | None:
|
| 509 |
+
tok = _TOKENIZERS[tokenizer_key]
|
| 510 |
+
tok = getattr(tok, "tokenizer", tok) # a processor's own
|
| 511 |
+
names = set(getattr(tok, "all_special_tokens", ()))
|
| 512 |
+
names |= {t.content for t in getattr(tok, "added_tokens_decoder", {}).values() if t.special}
|
| 513 |
+
names = sorted((n for n in names if len(n) > 1 and BREAK not in n), key=len, reverse=True)
|
| 514 |
+
return re.compile("|".join(map(re.escape, names))) if names else None
|
| 515 |
+
|
| 516 |
+
|
| 517 |
+
def plain(tokenizer, text: str) -> str:
|
| 518 |
+
"""`text` as the caller's text: no special token of `tokenizer` in it once encoded. Text without one is returned
|
| 519 |
+
unchanged."""
|
| 520 |
+
return _plain(tokenizer, text, False)[0]
|
| 521 |
+
|
| 522 |
+
|
| 523 |
+
def plain_offsets(tokenizer, text: str) -> tuple[str, list[int]]:
|
| 524 |
+
"""The same escaped text, with every character boundary mapped to the original text."""
|
| 525 |
+
return _plain(tokenizer, text, True)
|
| 526 |
+
|
| 527 |
+
|
| 528 |
+
def _plain(tokenizer, text: str, offsets: bool) -> tuple[str, list[int] | None]:
|
| 529 |
+
_TOKENIZERS.setdefault(id(tokenizer), tokenizer)
|
| 530 |
+
boundaries = list(range(len(text) + 1)) if offsets else None
|
| 531 |
+
pattern = _specials_cached(id(tokenizer))
|
| 532 |
+
if pattern is None:
|
| 533 |
+
return text, boundaries
|
| 534 |
+
n = 1
|
| 535 |
+
while n: # again for a match that overlapped the one broken before it
|
| 536 |
+
if offsets:
|
| 537 |
+
insertions = [m.start() + 1 for m in pattern.finditer(text)]
|
| 538 |
+
for at in reversed(insertions):
|
| 539 |
+
boundaries.insert(at, boundaries[at])
|
| 540 |
+
text, n = pattern.subn(lambda m: m[0][0] + BREAK + m[0][1:], text)
|
| 541 |
+
return text, boundaries
|
| 542 |
+
|
| 543 |
+
|
| 544 |
+
_PLACEHOLDER = re.compile(r"^opt\d+$")
|
| 545 |
+
|
| 546 |
+
|
| 547 |
+
def _option_line(code: str, label: str, desc: str | None, style: str) -> str:
|
| 548 |
+
text = desc or label.replace("_", " ")
|
| 549 |
+
if style == "name_desc" and not _PLACEHOLDER.match(label) and label != text:
|
| 550 |
+
return f"{code} {label}: {text}"
|
| 551 |
+
return f"{code} {text}"
|
| 552 |
+
|
| 553 |
+
|
| 554 |
+
def said(q: Question) -> str:
|
| 555 |
+
"""A question's instructions as its text writes them: their place (`media.SPOKEN`) when they are spoken."""
|
| 556 |
+
return SPOKEN if isinstance(q.instructions, Media) else q.instructions
|
| 557 |
+
|
| 558 |
+
|
| 559 |
+
def edge_legend(tokenizer) -> str:
|
| 560 |
+
"""How a box's edges are written: number tokens on a checkpoint that reads numbers, else codes."""
|
| 561 |
+
from .box import edge_grid
|
| 562 |
+
|
| 563 |
+
grid, scale = edge_grid(tokenizer)
|
| 564 |
+
if scale > 1:
|
| 565 |
+
return ("Each edge is a number from 0 to 999, in thousandths of the picture: top and bottom from its top edge, "
|
| 566 |
+
"left and right from its left edge.")
|
| 567 |
+
codes = [code for code, _ in aliases(tokenizer, list(grid_question("", grid).criteria))]
|
| 568 |
+
return (f"Each edge is a code on a {grid.size}-value scale: {codes[0]} is 0 (top or left), {codes[-1]} is 1 "
|
| 569 |
+
"(bottom or right). Values are equally spaced.")
|
| 570 |
+
|
| 571 |
+
|
| 572 |
+
def question_block(tokenizer, q: Question, option_style: str = "desc") -> str:
|
| 573 |
+
if isinstance(q, Box):
|
| 574 |
+
one = "one number" if isinstance(question_slots(q, tokenizer)[-1], Number) else "one code"
|
| 575 |
+
return (f"{said(q)}\n\nPresence is yes or no. {edge_legend(tokenizer)} Reply with yes or no at present, "
|
| 576 |
+
f"then {one} at top, bottom, left and right.")
|
| 577 |
+
if isinstance(q, Boxes):
|
| 578 |
+
from .box import INSTANCES
|
| 579 |
+
|
| 580 |
+
most = f"{INSTANCES} at most, " if version(tokenizer) == 1 else ""
|
| 581 |
+
return (f"{said(q)}\n\nEvery one of them, {most}in reading order (left edge, then top). "
|
| 582 |
+
f"Present k is yes when there are at least k. {edge_legend(tokenizer)} Reply with yes or no at each "
|
| 583 |
+
"present, then its top, bottom, left and right.")
|
| 584 |
+
if isinstance(q, Estimate) and isinstance(question_slots(q, tokenizer)[0], Number):
|
| 585 |
+
unit = f" {q.unit}" if q.unit is not None else ""
|
| 586 |
+
return (f"{said(q)}\n\nEstimate from {q.grid.low:g} to {q.grid.high:g}{unit}, in steps of "
|
| 587 |
+
f"{q.grid.step:g}. Reply with the number only.")
|
| 588 |
+
if isinstance(q, Estimate):
|
| 589 |
+
bins = grid_question("", q.grid)
|
| 590 |
+
codes = ([str(i) for i in range(q.grid.size)] if isinstance(bins, Score)
|
| 591 |
+
else [code for code, _ in aliases(tokenizer, list(bins.criteria))])
|
| 592 |
+
unit = f" {q.unit}" if q.unit is not None else ""
|
| 593 |
+
return (f"{said(q)}\n\nEstimate from {q.grid.low:g} to {q.grid.high:g}{unit} on "
|
| 594 |
+
f"{q.grid.size} equally spaced values, step {q.grid.step:g}: "
|
| 595 |
+
f"{codes[0]} is {q.grid.low:g}, {codes[-1]} is {q.grid.high:g}. Reply with the code only.")
|
| 596 |
+
if isinstance(q, Read):
|
| 597 |
+
unit = f" (unit: {q.unit})" if q.type == "number" and q.unit else ""
|
| 598 |
+
return f"{said(q)}\n\nRead the {q.type}{unit} from the state, or nothing if it is absent."
|
| 599 |
+
if isinstance(q, Choice):
|
| 600 |
+
labels = list(q.criteria.keys())
|
| 601 |
+
if q.multiple:
|
| 602 |
+
lines = "\n".join(f"{json.dumps(name, ensure_ascii=False)}: {q.criteria[name] or name}" for name in labels)
|
| 603 |
+
return f"{said(q)}\n\nOptions:\n{lines}\n\nReply with yes or no for each option."
|
| 604 |
+
lines = "\n".join(_option_line(code, lab, q.criteria[lab], option_style)
|
| 605 |
+
for (code, _), lab in zip(aliases(tokenizer, labels), labels))
|
| 606 |
+
return f"{said(q)}\n\nOptions:\n{lines}\n\nReply with the option code only."
|
| 607 |
+
if isinstance(q, Noul):
|
| 608 |
+
extra = ""
|
| 609 |
+
if q.criteria:
|
| 610 |
+
extra = f"\nYes: {q.criteria.get('true')}\nNo: {q.criteria.get('false')}"
|
| 611 |
+
return f"{said(q)}{extra}\n\nReply with yes or no only."
|
| 612 |
+
if isinstance(q, Score):
|
| 613 |
+
legend = "\n".join(f"{i} {name}" for i, name in enumerate(q.criteria))
|
| 614 |
+
return f"{said(q)}\n\n{legend}\n\nReply with a single digit 0-{len(q.criteria) - 1} only."
|
| 615 |
+
raise TypeError(f"unknown question type {type(q)}")
|
| 616 |
+
|
| 617 |
+
|
| 618 |
+
def prefix_text(tokenizer, state: Any, bos: str = "", style: str = DEFAULT_STATE_STYLE, system: str = DEFAULT_SYSTEM,
|
| 619 |
+
media: str = "") -> str:
|
| 620 |
+
"""Everything before the question, shared by all questions on one state: the media's markup (`media`, as
|
| 621 |
+
`media.Reader` writes it) at the head of the user turn, then the state. With no state (`None`) the question
|
| 622 |
+
follows the media directly. The state is the caller's text (`plain`)."""
|
| 623 |
+
text = SYSTEMS[system]
|
| 624 |
+
turn = "" if text is None else f"{IM_START}system\n{text}{IM_END}\n"
|
| 625 |
+
if state is None:
|
| 626 |
+
body = ""
|
| 627 |
+
else:
|
| 628 |
+
body = f"{plain(tokenizer, state_block(state, style))}\n{question_header(tokenizer, state)}"
|
| 629 |
+
return f"{bos}{turn}{IM_START}user\n{media}{body}"
|
| 630 |
+
|
| 631 |
+
|
| 632 |
+
def question_header(tokenizer, state: Any) -> str:
|
| 633 |
+
"""What `prefix_text` ends with to open a question. `set` attention repeats it before each further question
|
| 634 |
+
packed after the state (`attention.pack_set`)."""
|
| 635 |
+
return "" if state is None else "QUESTION:\n"
|
| 636 |
+
|
| 637 |
+
|
| 638 |
+
# What sits between the assistant header and the answer slot. A family module registers its measured one
|
| 639 |
+
# (`LEADS`); any other family's is read off its chat template (`template_lead`).
|
| 640 |
+
DEFAULT_LEAD = ""
|
| 641 |
+
LEADS: dict[str, str] = {}
|
| 642 |
+
|
| 643 |
+
|
| 644 |
+
def default_lead(model_type: str | None, tokenizer=None) -> str:
|
| 645 |
+
"""What a checkpoint's own template writes before a non-thinking answer: `LEADS` for the families measured,
|
| 646 |
+
else read off the tokenizer's chat template when one is given."""
|
| 647 |
+
if model_type in LEADS or tokenizer is None:
|
| 648 |
+
return LEADS.get(model_type, DEFAULT_LEAD)
|
| 649 |
+
return template_lead(tokenizer)
|
| 650 |
+
|
| 651 |
+
|
| 652 |
+
def template_lead(tokenizer) -> str:
|
| 653 |
+
"""What a ChatML template writes after the assistant header for a non-thinking reply. The prompt is ChatML
|
| 654 |
+
(`IM_START`, `IM_END`), so a model with another template is refused rather than mis-prompted."""
|
| 655 |
+
text = tokenizer.apply_chat_template([{"role": "user", "content": "x"}], add_generation_prompt=True,
|
| 656 |
+
tokenize=False, enable_thinking=False)
|
| 657 |
+
head = f"{IM_START}assistant\n"
|
| 658 |
+
if head not in text:
|
| 659 |
+
raise ValueError("not a ChatML chat template: d1 renders ChatML prompts only")
|
| 660 |
+
return text[text.rindex(head) + len(head):]
|
| 661 |
+
|
| 662 |
+
|
| 663 |
+
def suffix_text(tokenizer, q: Question, lead: str = DEFAULT_LEAD, option_style: str = "desc") -> str:
|
| 664 |
+
"""The question and the assistant header, up to the answer slot. The question is the caller's text (`plain`)."""
|
| 665 |
+
return suffix_layout(tokenizer, q, lead, option_style)[0]
|
| 666 |
+
|
| 667 |
+
|
| 668 |
+
def suffix_layout(tokenizer, q: Question, lead: str = DEFAULT_LEAD,
|
| 669 |
+
option_style: str = "desc") -> tuple[str, list[int]]:
|
| 670 |
+
"""The suffix and each slot's token offset from its last token; old questions have one slot at zero."""
|
| 671 |
+
body = question_block(tokenizer, q, option_style)
|
| 672 |
+
if isinstance(q, Read):
|
| 673 |
+
return marked_suffix(tokenizer, body, ["start", "end"], lead)
|
| 674 |
+
if isinstance(q, Box):
|
| 675 |
+
from .box import EDGES
|
| 676 |
+
|
| 677 |
+
return marked_suffix(tokenizer, body, ["present", *EDGES], lead)
|
| 678 |
+
if isinstance(q, Boxes):
|
| 679 |
+
from .box import EDGES, INSTANCES
|
| 680 |
+
|
| 681 |
+
# the last block ends on the newline a next block would follow, where the colon and it are one token, so every
|
| 682 |
+
# slot is the token the whole read has there: prompt 1 when `max` reads fewer than its 16, prompt 2 always
|
| 683 |
+
return marked_suffix(tokenizer, body, [f"{name} {k}" for k in range(1, q.blocks + 1)
|
| 684 |
+
for name in ("present", *EDGES)], lead,
|
| 685 |
+
"\n" if q.blocks < INSTANCES or version(tokenizer) > 1 else "")
|
| 686 |
+
if isinstance(q, Choice) and q.multiple:
|
| 687 |
+
names = [json.dumps(name, ensure_ascii=False) for name in q.criteria]
|
| 688 |
+
return marked_suffix(tokenizer, body, names, lead)
|
| 689 |
+
return f"{plain(tokenizer, body)}{IM_END}\n{IM_START}assistant\n{lead}", [0]
|
| 690 |
+
|
| 691 |
+
|
| 692 |
+
def marked_suffix(tokenizer, body: str, names: Sequence[str], lead: str = DEFAULT_LEAD,
|
| 693 |
+
after: str = "") -> tuple[str, list[int]]:
|
| 694 |
+
"""A question's named answer markers and their token offsets from its last token; `after`, whitespace the last
|
| 695 |
+
marker's colon token takes where the tokenizer writes the two as one token."""
|
| 696 |
+
if not names or len(set(names)) != len(names) or any(not isinstance(n, str) or not n or "\n" in n for n in names):
|
| 697 |
+
raise ValueError("slot names must be distinct nonempty single lines")
|
| 698 |
+
text = f"{plain(tokenizer, body)}{IM_END}\n{IM_START}assistant\n{lead}"
|
| 699 |
+
ends = []
|
| 700 |
+
for i, name in enumerate(names):
|
| 701 |
+
text += ("\n" if i else "") + plain(tokenizer, name) + ":"
|
| 702 |
+
ends.append(len(text))
|
| 703 |
+
if after and tokenizer.decode(tokenizer.encode(text + after, add_special_tokens=False)[-1:]).rstrip().endswith(":"):
|
| 704 |
+
text += after
|
| 705 |
+
encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
|
| 706 |
+
ids = encoded["input_ids"]
|
| 707 |
+
slots = []
|
| 708 |
+
for chars in ends:
|
| 709 |
+
matching = [(i, stop) for i, (start, stop) in enumerate(encoded["offset_mapping"]) if start <= chars - 1 < stop]
|
| 710 |
+
if len(matching) != 1:
|
| 711 |
+
raise ValueError("a slot marker must have an exact colon token position")
|
| 712 |
+
end, stop = matching[0]
|
| 713 |
+
# The colon may share its token with following whitespace, never with the next marker's text.
|
| 714 |
+
if text[chars:stop].strip() or not tokenizer.decode([ids[end]]).rstrip().endswith(":"):
|
| 715 |
+
raise ValueError("a slot marker token must end with its colon and optional whitespace")
|
| 716 |
+
slots.append(end - len(ids) + 1)
|
| 717 |
+
return text, slots
|
| 718 |
+
|
| 719 |
+
|
| 720 |
+
def number_ids(tokenizer, grid, scale: int = 1) -> tuple[tuple[int, ...], ...] | None:
|
| 721 |
+
"""Each grid value times `scale` as the tokenizer writes it, plain and after a space, one token a form: None when
|
| 722 |
+
a value is no integer or takes more than one token (a tokenizer of digits)."""
|
| 723 |
+
_TOKENIZERS.setdefault(id(tokenizer), tokenizer)
|
| 724 |
+
return _number_ids_cached(id(tokenizer), grid.low, grid.high, grid.size, scale)
|
| 725 |
+
|
| 726 |
+
|
| 727 |
+
@lru_cache(maxsize=256)
|
| 728 |
+
def _number_ids_cached(tokenizer_key: int, low: float, high: float, size: int, scale: int):
|
| 729 |
+
from .grid import Grid
|
| 730 |
+
|
| 731 |
+
groups = []
|
| 732 |
+
for value in Grid(low, high, size).values:
|
| 733 |
+
n = round(value * scale)
|
| 734 |
+
if not math.isclose(n, value * scale, abs_tol=1e-6):
|
| 735 |
+
return None
|
| 736 |
+
group = _ids_cached(tokenizer_key, (str(n), f" {n}"))
|
| 737 |
+
if not group or len(_TOKENIZERS[tokenizer_key].encode(str(n), add_special_tokens=False)) != 1:
|
| 738 |
+
return None
|
| 739 |
+
groups.append(group)
|
| 740 |
+
return tuple(groups)
|
| 741 |
+
|
| 742 |
+
|
| 743 |
+
def grid_slot(tokenizer, instructions: str, grid, scale: int = 1) -> Number | Choice | Score:
|
| 744 |
+
"""A grid's slot: the tokenizer's own numbers where the checkpoint reads them so (`d1_numbers`) and every value is
|
| 745 |
+
one (`number_ids`), else the grid's codes (`grid_question`)."""
|
| 746 |
+
if getattr(tokenizer, "d1_numbers", False) and number_ids(tokenizer, grid, scale) is not None:
|
| 747 |
+
return Number(instructions, grid, scale)
|
| 748 |
+
return grid_question(instructions, grid)
|
| 749 |
+
|
| 750 |
+
|
| 751 |
+
def grid_question(instructions: str, grid) -> Choice | Score:
|
| 752 |
+
"""The grid's ordered bins on the existing digit or option-code readout."""
|
| 753 |
+
labels = [f"{v:g}" for v in grid.values]
|
| 754 |
+
return (Score(instructions, labels) if grid.size <= 10
|
| 755 |
+
else Choice(instructions, {str(i): v for i, v in enumerate(labels)}))
|
| 756 |
+
|
| 757 |
+
|
| 758 |
+
def render(tokenizer, state: Any, q: Question, bos: str = "", lead: str = DEFAULT_LEAD,
|
| 759 |
+
style: str = DEFAULT_STATE_STYLE, system: str = DEFAULT_SYSTEM, option_style: str = "desc",
|
| 760 |
+
media: str = "") -> str:
|
| 761 |
+
return prefix_text(tokenizer, state, bos, style, system, media) + suffix_text(tokenizer, q, lead, option_style)
|
runner.py
ADDED
|
@@ -0,0 +1,379 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""One-pass typed decisions on any backbone.
|
| 2 |
+
|
| 3 |
+
The model sees the state and the question once, and the answer is read off the
|
| 4 |
+
logits at the answer slot. Nothing is decoded, so a schema violation is
|
| 5 |
+
impossible and `tokens generated per decision` is zero.
|
| 6 |
+
|
| 7 |
+
- Questions are read at their answer slots, a whole prompt or packed after their state (`SlotReader._padded`):
|
| 8 |
+
every typed answer, and every question on a backbone with no faster path of its own (`SystemOne.run`).
|
| 9 |
+
- A request is `(state, questions, media)`: its pictures, audio clips and videos head the state, a spoken question's
|
| 10 |
+
clip sits in its own text (`media`).
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
import os
|
| 16 |
+
from typing import Any, Sequence
|
| 17 |
+
|
| 18 |
+
import torch
|
| 19 |
+
|
| 20 |
+
from .api import SystemOneApi
|
| 21 |
+
from .attention import SET_SIZE, TRAINED_SET_SIZE, both_ways, language_model, pack_set
|
| 22 |
+
from .box import EDGE_READS
|
| 23 |
+
from .media import Media, Reader, media_list, spoken
|
| 24 |
+
from .pointer import SPAN_CAPS, query_plan, spans, state_positions
|
| 25 |
+
from .prompt import (
|
| 26 |
+
DEFAULT_STATE_STYLE,
|
| 27 |
+
DEFAULT_SYSTEM,
|
| 28 |
+
PROMPTS,
|
| 29 |
+
Box,
|
| 30 |
+
Boxes,
|
| 31 |
+
Choice,
|
| 32 |
+
Position,
|
| 33 |
+
Question,
|
| 34 |
+
Read,
|
| 35 |
+
default_lead,
|
| 36 |
+
prefix_text,
|
| 37 |
+
question_header,
|
| 38 |
+
question_slots,
|
| 39 |
+
readout,
|
| 40 |
+
readouts,
|
| 41 |
+
render,
|
| 42 |
+
suffix_layout,
|
| 43 |
+
suffix_text,
|
| 44 |
+
typed,
|
| 45 |
+
)
|
| 46 |
+
from .slots import slot_logits, text_backbone
|
| 47 |
+
|
| 48 |
+
# Every picture is bounded at this many pixels before the processor (`media.cap_pixels`), in training and serving alike. A checkpoint
|
| 49 |
+
# records its own (`d1_vision_max_pixels` in its config); D1_VISION_MAX_PIXELS, when set, overrides it for a run or a
|
| 50 |
+
# check (reading a checkpoint at another resolution).
|
| 51 |
+
PIXELS_SET = os.environ.get("D1_VISION_MAX_PIXELS")
|
| 52 |
+
VISION_MAX_PIXELS = int(PIXELS_SET or 1024 * 1024)
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def group_size(n: int | None) -> int | None:
|
| 56 |
+
"""A group size (`set_size`): a whole number from 1 (each question alone), or None for the engine's own."""
|
| 57 |
+
if n is None:
|
| 58 |
+
return None
|
| 59 |
+
if isinstance(n, bool) or not isinstance(n, int):
|
| 60 |
+
raise ValueError(f"set_size must be a whole number, not {n!r}")
|
| 61 |
+
if n < 1:
|
| 62 |
+
raise ValueError(f"set_size must be at least 1, not {n}")
|
| 63 |
+
return n
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def remote_ok(model_id) -> bool:
|
| 67 |
+
"""Whether a checkpoint's own code may run: a local folder or a LiquidAI repo.
|
| 68 |
+
Any other Hub id loads with transformers' classes only, and one that needs its own code fails."""
|
| 69 |
+
return os.path.isdir(str(model_id)) or str(model_id).startswith("LiquidAI/")
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def auto_processor(path: str, revision: str | None = None):
|
| 73 |
+
"""The media processor a checkpoint holds."""
|
| 74 |
+
from transformers import AutoProcessor
|
| 75 |
+
|
| 76 |
+
return AutoProcessor.from_pretrained(path, revision=revision, trust_remote_code=remote_ok(path))
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def _under(config: dict | None, compiled):
|
| 80 |
+
"""`compiled` run under inductor's `config` (`SystemOne.compile`)."""
|
| 81 |
+
if not config:
|
| 82 |
+
return compiled
|
| 83 |
+
import torch._inductor.config as inductor
|
| 84 |
+
|
| 85 |
+
def run(*args, **kwargs):
|
| 86 |
+
with inductor.patch(config):
|
| 87 |
+
return compiled(*args, **kwargs)
|
| 88 |
+
|
| 89 |
+
return run
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
class SlotReader:
|
| 93 |
+
"""Prompts read at their answer slots, whatever runs the backbone: a request's sequences (`_sequences`), each
|
| 94 |
+
question's slots, and the answers read off them (`_padded`). An engine sets what `SystemOne.__init__` sets (the
|
| 95 |
+
checkpoint's record through `_serve`, its `processor` or `_load_processor`) and gives `_batches` and
|
| 96 |
+
`_slot_logits`."""
|
| 97 |
+
|
| 98 |
+
def _serve(self, recorded, codes: int | None = None, max_pixels: int | None = None, attention: str | None = None,
|
| 99 |
+
set_size: int | None = None, numbers: bool | None = None, box_edges: str | None = None,
|
| 100 |
+
prompt: int | None = None) -> None:
|
| 101 |
+
"""The option codes, number readout, picture cap, attention mode, set size, box-edge read and prompt version
|
| 102 |
+
`recorded` (a checkpoint's config) holds (`d1_codes`, `d1_numbers`, `d1_vision_max_pixels`, `d1_attention`,
|
| 103 |
+
`d1_set_size`, `d1_box_edges`: the median without it, `d1_prompt`: 1 without it), unless given. A recorded set
|
| 104 |
+
size is 2 to `TRAINED_SET_SIZE`, the default; one given is any whole number from 1 (each question alone):
|
| 105 |
+
larger groups read faster, in packs training never saw."""
|
| 106 |
+
self.tokenizer.d1_codes = codes if codes is not None else getattr(recorded, "d1_codes", 1)
|
| 107 |
+
self.tokenizer.d1_numbers = bool(numbers if numbers is not None else getattr(recorded, "d1_numbers", False))
|
| 108 |
+
self.tokenizer.d1_box_edges = box_edges or getattr(recorded, "d1_box_edges", None) or "median"
|
| 109 |
+
if self.tokenizer.d1_box_edges not in EDGE_READS:
|
| 110 |
+
raise ValueError(f"d1_box_edges is one of {EDGE_READS}, not {self.tokenizer.d1_box_edges!r}")
|
| 111 |
+
self.tokenizer.d1_prompt = prompt or getattr(recorded, "d1_prompt", 1)
|
| 112 |
+
if self.tokenizer.d1_prompt not in PROMPTS:
|
| 113 |
+
raise ValueError(f"d1_prompt is one of {PROMPTS}, not {self.tokenizer.d1_prompt!r}")
|
| 114 |
+
self.max_pixels = (VISION_MAX_PIXELS if PIXELS_SET else
|
| 115 |
+
max_pixels or getattr(recorded, "d1_vision_max_pixels", None) or VISION_MAX_PIXELS)
|
| 116 |
+
self.attention = attention or getattr(recorded, "d1_attention", None) or "causal"
|
| 117 |
+
self.set_size = group_size(getattr(recorded, "d1_set_size", SET_SIZE) if set_size is None else set_size)
|
| 118 |
+
if set_size is None and not 2 <= self.set_size <= TRAINED_SET_SIZE:
|
| 119 |
+
raise ValueError(f"d1_set_size must be between 2 and {TRAINED_SET_SIZE}")
|
| 120 |
+
|
| 121 |
+
def _readout(self, q: Question, logz, tempered: bool = True) -> list[float]:
|
| 122 |
+
"""An option answer through the served temperature; a typed answer's slot raw (`tempered=False`): its own
|
| 123 |
+
table calibrates it (`calibration.json` `tables`, fitted on raw reads)."""
|
| 124 |
+
return readout(self.tokenizer, q, logz, self.calibration if tempered else None)
|
| 125 |
+
|
| 126 |
+
def _readouts(self, qs: Sequence[Question], logz, tempered: Sequence[bool]) -> list[list[float]]:
|
| 127 |
+
"""`_readout` for each code slot of a batch, `logz` as `_slot_logits` gives them."""
|
| 128 |
+
return [self._readout(q, z, t) for q, z, t in zip(qs, logz, tempered)]
|
| 129 |
+
|
| 130 |
+
@property
|
| 131 |
+
def media(self) -> Reader:
|
| 132 |
+
"""The checkpoint's media reader, its processor loaded on first use."""
|
| 133 |
+
if getattr(self, "_reader", None) is None or self._reader.processor is not self.processor:
|
| 134 |
+
if self.processor is None:
|
| 135 |
+
self.processor = self._load_processor()
|
| 136 |
+
self._reader = Reader(self.processor, self.max_pixels)
|
| 137 |
+
return self._reader
|
| 138 |
+
|
| 139 |
+
def _prefix(self, state: Any, media: Sequence[Media]) -> tuple[str, list[dict]]:
|
| 140 |
+
"""`prompt.prefix_text` headed by `media`'s markup, and their processor outputs."""
|
| 141 |
+
markup, parts = self.media.encode(media) if media else ("", [])
|
| 142 |
+
return prefix_text(self.tokenizer, state, self.bos, self.state_style, self.system, markup), parts
|
| 143 |
+
|
| 144 |
+
def _suffixes(self, qs: Sequence[Question]) -> tuple[list[str], list[list[int]], list[list[dict]]]:
|
| 145 |
+
"""Each question's suffix (a spoken one's clip in place), its slots' offsets and its clip's outputs."""
|
| 146 |
+
layouts = [suffix_layout(self.tokenizer, q, self.lead, self.option_style) for q in qs]
|
| 147 |
+
placed = [self.media.place(text, q) if spoken(q) else (text, []) for (text, _), q in zip(layouts, qs)]
|
| 148 |
+
return [text for text, _ in placed], [ends for _, ends in layouts], [parts for _, parts in placed]
|
| 149 |
+
|
| 150 |
+
def _sequences(self, requests: Sequence[tuple[Any, list[Question], Sequence[Media]]], *,
|
| 151 |
+
position_maps=None, set_size: int | None = None) -> list[tuple]:
|
| 152 |
+
"""Every request's sequences, `(request, first question, ids, state length, each token's question or None,
|
| 153 |
+
answer slots, media inputs)`: under `set` up to `set_size` questions after their state (`attention.pack_set`),
|
| 154 |
+
else one whole prompt a question. A request's media head its state, a spoken question's clip sits in its
|
| 155 |
+
suffix, each as its processor expands it, so the text is the sequence. `position_maps` gets each read's state
|
| 156 |
+
positions; `set_size`, this call's group size (the engine's by default)."""
|
| 157 |
+
seqs = []
|
| 158 |
+
for i, (state, qs, media) in enumerate(requests):
|
| 159 |
+
media = media_list(media)
|
| 160 |
+
if any(isinstance(q, (Box, Boxes)) for q in qs) and (n := sum(m.type == "image" for m in media)) != 1:
|
| 161 |
+
raise ValueError(f"a box or boxes question reads exactly one picture in `media`; this request has {n}")
|
| 162 |
+
prefix, state_parts = self._prefix(state, media)
|
| 163 |
+
suffixes, offsets, voices = self._suffixes(qs)
|
| 164 |
+
step = (self.set_size if set_size is None else set_size) if self.attention == "set" else 1
|
| 165 |
+
for k in range(0, len(qs), step):
|
| 166 |
+
text = prefix + suffixes[k]
|
| 167 |
+
encoded = None
|
| 168 |
+
if position_maps is not None and any(isinstance(q, Read) for q in qs[k:k + step]):
|
| 169 |
+
encoded = self.tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
|
| 170 |
+
ids = self.tokenizer.encode(text, add_special_tokens=False) if encoded is None else encoded["input_ids"]
|
| 171 |
+
if self.attention == "set":
|
| 172 |
+
ids, state_len, question, ends = pack_set(self.tokenizer, prefix, suffixes[k:k + step],
|
| 173 |
+
question_header(self.tokenizer, state),
|
| 174 |
+
offsets[k:k + step], first_ids=ids)
|
| 175 |
+
else:
|
| 176 |
+
state_len, question = both_ways(self.attention, self.tokenizer, ids, suffixes[k]), None
|
| 177 |
+
ends = [[len(ids) - 1 + offset for offset in offsets[k]]]
|
| 178 |
+
parts = state_parts + [part for own in voices[k:k + step] for part in own]
|
| 179 |
+
inputs = self.media.inputs(parts, ids) if parts else {}
|
| 180 |
+
if encoded is not None:
|
| 181 |
+
position_maps[len(seqs)] = state_positions(self.tokenizer, state, prefix, text, self.state_style,
|
| 182 |
+
encoded=encoded)
|
| 183 |
+
seqs.append((i, k, ids, state_len, question, ends, inputs))
|
| 184 |
+
return seqs
|
| 185 |
+
|
| 186 |
+
def _padded(self, requests: Sequence[tuple[Any, list[Question], Sequence]],
|
| 187 |
+
set_size: int | None = None) -> list[tuple[list[list[float]], int]]:
|
| 188 |
+
"""`_sequences` in the engine's batches (`_batches`), each read at its answer slots as training reads them
|
| 189 |
+
(`_slot_logits`): option codes through the readout, reads through the pointer."""
|
| 190 |
+
position_maps = {}
|
| 191 |
+
seqs = self._sequences(requests, position_maps=position_maps, set_size=set_size)
|
| 192 |
+
read = [0] * len(requests)
|
| 193 |
+
for i, _, ids, *_ in seqs:
|
| 194 |
+
read[i] += len(ids)
|
| 195 |
+
|
| 196 |
+
def asked(key: tuple[int, int]) -> Question:
|
| 197 |
+
return requests[key[0]][1][key[1]]
|
| 198 |
+
|
| 199 |
+
probs: dict[tuple[int, int], list[float]] = {}
|
| 200 |
+
for batch in self._batches(seqs):
|
| 201 |
+
slots = [] # (row in the batch, (request, question), the slot's readout, its position)
|
| 202 |
+
for r, b in enumerate(batch):
|
| 203 |
+
i, first = seqs[b][:2]
|
| 204 |
+
for k, ends in enumerate(seqs[b][5]):
|
| 205 |
+
key = (i, first + k)
|
| 206 |
+
slots.extend((r, key, q, end) for q, end in zip(question_slots(asked(key), self.tokenizer), ends))
|
| 207 |
+
reads = [j for j, s in enumerate(slots) if isinstance(s[2], Position)] # a read's start then its end
|
| 208 |
+
pointer = None
|
| 209 |
+
if reads:
|
| 210 |
+
masks = [torch.tensor(position_maps[batch[slots[j][0]]].pointer_mask) for j in reads]
|
| 211 |
+
pointer = query_plan([slots[j][0] for j in reads], masks, reads, len(slots))
|
| 212 |
+
logz, positions = self._slot_logits(seqs, batch, slots, pointer)
|
| 213 |
+
if pointer is not None:
|
| 214 |
+
logits = positions.reshape(-1, 2, positions.shape[-1])
|
| 215 |
+
capped = {} # the reads of each span cap
|
| 216 |
+
for index, j in enumerate(reads[::2]):
|
| 217 |
+
capped.setdefault(SPAN_CAPS[asked(slots[j][1]).type], []).append((index, j))
|
| 218 |
+
for cap, own in capped.items():
|
| 219 |
+
at, found = spans(logits[[index for index, _ in own]], max_tokens=cap)
|
| 220 |
+
for (_, j), span, p in zip(own, at.tolist(), found.tolist()):
|
| 221 |
+
r, key = slots[j][:2]
|
| 222 |
+
probs[key] = {**position_maps[batch[r]].read(*span, max_tokens=cap), "found": p}
|
| 223 |
+
codes = pointer["code_slots"].tolist() if pointer else range(len(slots))
|
| 224 |
+
scored = self._readouts([slots[j][2] for j in codes], logz, [not typed(asked(slots[j][1])) for j in codes])
|
| 225 |
+
for j, p in zip(codes, scored):
|
| 226 |
+
key = slots[j][1]
|
| 227 |
+
q = asked(key)
|
| 228 |
+
if isinstance(q, Choice) and q.multiple:
|
| 229 |
+
probs.setdefault(key, []).append(p[0])
|
| 230 |
+
elif isinstance(q, (Box, Boxes)):
|
| 231 |
+
probs.setdefault(key, []).append(p)
|
| 232 |
+
else:
|
| 233 |
+
probs[key] = p
|
| 234 |
+
return [([probs[(i, k)] for k in range(len(qs))], read[i]) for i, (_, qs, _) in enumerate(requests)]
|
| 235 |
+
|
| 236 |
+
|
| 237 |
+
class SystemOne(SlotReader, SystemOneApi):
|
| 238 |
+
"""State in, calibrated distribution out, one forward pass."""
|
| 239 |
+
|
| 240 |
+
def __init__(
|
| 241 |
+
self,
|
| 242 |
+
model=None,
|
| 243 |
+
tokenizer=None,
|
| 244 |
+
device: str | None = None,
|
| 245 |
+
calibration=None,
|
| 246 |
+
lead: str | None = None,
|
| 247 |
+
state_style: str = DEFAULT_STATE_STYLE,
|
| 248 |
+
system: str = DEFAULT_SYSTEM,
|
| 249 |
+
option_style: str = "desc",
|
| 250 |
+
codes: int | None = None,
|
| 251 |
+
max_pixels: int | None = None,
|
| 252 |
+
attention: str | None = None,
|
| 253 |
+
set_size: int | None = None,
|
| 254 |
+
numbers: bool | None = None,
|
| 255 |
+
box_edges: str | None = None,
|
| 256 |
+
prompt: int | None = None,
|
| 257 |
+
compile: bool = False,
|
| 258 |
+
token_budget: int = 65536,
|
| 259 |
+
revision: str | None = None,
|
| 260 |
+
):
|
| 261 |
+
"""`model` and `tokenizer`: a backbone already loaded (the Hub's class passes itself, and the `revision` its
|
| 262 |
+
processor loads at); it stays on its device unless `device` says otherwise. `codes` and `max_pixels` default
|
| 263 |
+
to what the checkpoint's config records (`d1_codes`, `d1_vision_max_pixels`), else scheme 1 and the cap;
|
| 264 |
+
`attention` likewise (`d1_attention`, else causal; `attention`)."""
|
| 265 |
+
if model is None or tokenizer is None or isinstance(model, str):
|
| 266 |
+
raise TypeError("SystemOne reads a loaded backbone: pass `model=` and `tokenizer=` (a checkpoint's own "
|
| 267 |
+
"class serves it: AutoModel.from_pretrained(repo, trust_remote_code=True).system_one)")
|
| 268 |
+
self.model, self.tokenizer = model, tokenizer
|
| 269 |
+
self.model_id, device = model.config._name_or_path, device or next(model.parameters()).device
|
| 270 |
+
self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
|
| 271 |
+
if getattr(self.model, "hf_device_map", None): # spread over GPUs at load: inputs go to the embeddings
|
| 272 |
+
self.device = self.model.get_input_embeddings().weight.device
|
| 273 |
+
else:
|
| 274 |
+
self.model.to(self.device)
|
| 275 |
+
self.model.eval()
|
| 276 |
+
if hasattr(self.model, "fuse_projections"): # one matmul for the linears on one input
|
| 277 |
+
self.model.fuse_projections()
|
| 278 |
+
bos = getattr(self.tokenizer, "bos_token", None)
|
| 279 |
+
self.bos = bos if isinstance(bos, str) else ""
|
| 280 |
+
self.lead = default_lead(model.config.model_type, tokenizer) if lead is None else lead
|
| 281 |
+
self.calibration, self.state_style, self.system, self.option_style = calibration, state_style, system, option_style
|
| 282 |
+
self._serve(self.model.config, codes, max_pixels, attention, set_size, numbers, box_edges, prompt)
|
| 283 |
+
self.token_budget = token_budget
|
| 284 |
+
self.processor, self.revision = None, revision
|
| 285 |
+
if compile:
|
| 286 |
+
self.compile()
|
| 287 |
+
|
| 288 |
+
def compile(self, mode: str | None = None) -> None:
|
| 289 |
+
"""torch.compile what the passes run: the language model's forward, shapes dynamic; not its hooks, so
|
| 290 |
+
`state_attention`'s plan is made outside the graph. Also the media encoders the backbone names
|
| 291 |
+
(`compiled_parts`; 20 to 15 ms on a 384 px picture), under the inductor config it names (`compile_options`),
|
| 292 |
+
set while they run: dynamo reads `freezing` as it traces, where `torch.compile`'s own `options` do not reach
|
| 293 |
+
(d1-2B on an MI325X: 4.5 ms a question with `options`, 3.5 set). CUDA graphs measured no faster there, and
|
| 294 |
+
inductor's graph pool refuses the packed varlen calls."""
|
| 295 |
+
if mode in ("reduce-overhead", "max-autotune") and self.attention != "causal":
|
| 296 |
+
raise ValueError(f"mode={mode!r} records CUDA graphs, which `{self.attention}` attention does not take: "
|
| 297 |
+
"model.compile() fuses without them")
|
| 298 |
+
body, parts = text_backbone(self.model), getattr(self.model, "compiled_parts", ())
|
| 299 |
+
config = getattr(self.model, "compile_options", None)
|
| 300 |
+
for part in (language_model(self.model), *(getattr(body, name) for name in parts)):
|
| 301 |
+
part.forward = _under(config, torch.compile(part.forward, dynamic=True, mode=mode))
|
| 302 |
+
|
| 303 |
+
# ---------------------------------------------------------------- prompt
|
| 304 |
+
|
| 305 |
+
def render(self, state: Any, q: Question) -> str:
|
| 306 |
+
return render(self.tokenizer, state, q, self.bos, self.lead, self.state_style, self.system, self.option_style)
|
| 307 |
+
|
| 308 |
+
# ------------------------------------------------------------------- api
|
| 309 |
+
|
| 310 |
+
@torch.inference_mode()
|
| 311 |
+
def run(self, requests: Sequence[tuple[Any, list[Question], Sequence[Media]]],
|
| 312 |
+
set_size: int | None = None) -> list[tuple[list[list[float]], int]]:
|
| 313 |
+
"""Each `(state, questions, media)` request's probabilities and the tokens it read, at its answer slots
|
| 314 |
+
(`_padded`), every request in the same batches. `set_size`: this call's questions a group under `set` (the
|
| 315 |
+
engine's by default)."""
|
| 316 |
+
set_size = group_size(set_size)
|
| 317 |
+
requests = [(state, qs, media_list(media)) for state, qs, media in requests]
|
| 318 |
+
return self._padded(requests, set_size)
|
| 319 |
+
|
| 320 |
+
def _batches(self, seqs: list[tuple]) -> list[list[int]]:
|
| 321 |
+
"""`_sequences` right-padded into batches of at most `token_budget` padded tokens, whatever their media: a
|
| 322 |
+
request's sequences in one batch (one forward) unless they alone pass the budget, requests of similar length
|
| 323 |
+
together. Every row is whole, its media given with it (`_rows`)."""
|
| 324 |
+
own: dict[int, list[int]] = {}
|
| 325 |
+
for j in sorted(range(len(seqs)), key=lambda j: len(seqs[j][2])):
|
| 326 |
+
own.setdefault(seqs[j][0], []).append(j)
|
| 327 |
+
batches: list[list[int]] = [[]]
|
| 328 |
+
for rows in sorted(own.values(), key=lambda rows: len(seqs[rows[-1]][2])):
|
| 329 |
+
if batches[-1] and len(seqs[rows[-1]][2]) * (len(batches[-1]) + len(rows)) > self.token_budget:
|
| 330 |
+
batches.append([])
|
| 331 |
+
for j in rows: # shortest first, so a batch's width is its last sequence's length
|
| 332 |
+
if batches[-1] and len(seqs[j][2]) * (len(batches[-1]) + 1) > self.token_budget:
|
| 333 |
+
batches.append([])
|
| 334 |
+
batches[-1].append(j)
|
| 335 |
+
return list(filter(None, batches))
|
| 336 |
+
|
| 337 |
+
def _rows(self, seqs: list[tuple], batch: list[int]) -> tuple[torch.Tensor, list[int], dict]:
|
| 338 |
+
"""A batch's right-padded ids, each row's length and its media inputs: every row whole, its own ids and its
|
| 339 |
+
media's inputs, joined in row order (`media.Reader.batch`) for the one forward."""
|
| 340 |
+
pad = self.tokenizer.pad_token_id or 0
|
| 341 |
+
lengths = [len(seqs[b][2]) for b in batch]
|
| 342 |
+
ids = torch.tensor([seqs[b][2] + [pad] * (max(lengths) - n) for b, n in zip(batch, lengths)], device=self.device)
|
| 343 |
+
inputs = [seqs[b][6] for b in batch]
|
| 344 |
+
return ids, lengths, self.media.batch(inputs, ids) if any(inputs) else {}
|
| 345 |
+
|
| 346 |
+
def _slot_logits(self, seqs: list[tuple], batch: list[int], slots: list[tuple], pointer: dict | None):
|
| 347 |
+
"""A batch's code slots' log-probabilities over the vocabulary, and its read slots' position logits
|
| 348 |
+
(`slots.slot_logits`), in one forward (`_rows`)."""
|
| 349 |
+
ids, lengths, media = self._rows(seqs, batch)
|
| 350 |
+
width = ids.shape[1]
|
| 351 |
+
mask = (torch.arange(width, device=self.device) < torch.tensor(lengths, device=self.device)[:, None]).long()
|
| 352 |
+
state = None if self.attention == "causal" else torch.tensor([seqs[b][3] for b in batch])
|
| 353 |
+
question = None if self.attention != "set" else torch.tensor(
|
| 354 |
+
[seqs[b][4] + [-1] * (width - len(seqs[b][4])) for b in batch])
|
| 355 |
+
rows = slot_logits(self.model, ids, mask, torch.tensor([s[3] for s in slots]), state, question,
|
| 356 |
+
torch.tensor([s[0] for s in slots]), pointer=pointer, **self._on_device(media))
|
| 357 |
+
codes = rows.codes.float()
|
| 358 |
+
return codes - torch.logsumexp(codes, dim=-1, keepdim=True), rows.positions
|
| 359 |
+
|
| 360 |
+
def _readouts(self, qs: Sequence[Question], logz: torch.Tensor, tempered: Sequence[bool]) -> list[list[float]]:
|
| 361 |
+
"""Every code slot of a batch read at once (`prompt.readouts`): one transfer, not one a slot."""
|
| 362 |
+
return readouts(self.tokenizer, qs, logz, [self.calibration if t else None for t in tempered])
|
| 363 |
+
|
| 364 |
+
def tokens(self, state: Any, questions: Sequence[Question]) -> int:
|
| 365 |
+
"""The longest prompt one of `questions` makes over a text state: its state's tokens and its own."""
|
| 366 |
+
prefix, _ = self._prefix(state, ())
|
| 367 |
+
return len(self.tokenizer.encode(prefix, add_special_tokens=False)) + max(
|
| 368 |
+
len(self.tokenizer.encode(suffix_text(self.tokenizer, q, self.lead, self.option_style),
|
| 369 |
+
add_special_tokens=False)) for q in questions)
|
| 370 |
+
|
| 371 |
+
# ----------------------------------------------------------------- media
|
| 372 |
+
|
| 373 |
+
def _load_processor(self):
|
| 374 |
+
return auto_processor(self.model.config._name_or_path, self.revision)
|
| 375 |
+
|
| 376 |
+
def _on_device(self, inputs: dict) -> dict:
|
| 377 |
+
"""The media inputs on the model's device, but what the backbone reads on the host (`host_inputs`)."""
|
| 378 |
+
host = getattr(self.model, "host_inputs", ())
|
| 379 |
+
return {k: (v.to(self.device) if hasattr(v, "to") and k not in host else v) for k, v in inputs.items()}
|
slots.py
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Answer positions of a right-padded batch, read through whatever wraps the model."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from dataclasses import dataclass
|
| 6 |
+
|
| 7 |
+
import torch
|
| 8 |
+
|
| 9 |
+
from .attention import state_attention, text_backbone
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
@dataclass(frozen=True)
|
| 13 |
+
class SlotLogits:
|
| 14 |
+
"""A batch's answer slots: the code slots' logits, in slot order, and the read slots' over their rows' state
|
| 15 |
+
positions, in the pointer plan's order (pointer.py `query_plan`)."""
|
| 16 |
+
|
| 17 |
+
codes: torch.Tensor
|
| 18 |
+
positions: torch.Tensor | None = None
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def split(slots: torch.Tensor, states: torch.Tensor | None, plan: dict | None, project) -> SlotLogits:
|
| 22 |
+
"""The hidden states at the slots [S, D]: the code slots through `project` (the LM head, or its option rows),
|
| 23 |
+
the read slots against their rows' state positions `states` [R, T, D]."""
|
| 24 |
+
if plan is None:
|
| 25 |
+
return SlotLogits(project(slots))
|
| 26 |
+
from .pointer import read_logits
|
| 27 |
+
|
| 28 |
+
return SlotLogits(project(slots[plan["code_slots"].to(slots.device)]), read_logits(slots, states, plan))
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def output_head(model):
|
| 32 |
+
"""The LM head, through whatever wrappers are in the way."""
|
| 33 |
+
head = getattr(model, "module", model).get_output_embeddings()
|
| 34 |
+
if head is None:
|
| 35 |
+
raise RuntimeError("no output embedding on this model")
|
| 36 |
+
return head
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def slot_logits(
|
| 40 |
+
model,
|
| 41 |
+
ids: torch.Tensor,
|
| 42 |
+
mask: torch.Tensor,
|
| 43 |
+
last: torch.Tensor,
|
| 44 |
+
state_len: torch.Tensor | None = None,
|
| 45 |
+
question: torch.Tensor | None = None,
|
| 46 |
+
row: torch.Tensor | None = None,
|
| 47 |
+
pointer: dict | None = None,
|
| 48 |
+
**vision,
|
| 49 |
+
):
|
| 50 |
+
"""Logits at one position per row (or at each `row`, `last` pair: a question may have several slots in a row),
|
| 51 |
+
the backbone read directly and only the slots projected, never `B x T x vocab`. `vision`: the processor's picture
|
| 52 |
+
inputs, for a vision-language model."""
|
| 53 |
+
with state_attention(model, state_len, question): # None: the model's own causal mask
|
| 54 |
+
out = text_backbone(model)(input_ids=ids, attention_mask=mask, use_cache=False, **vision)
|
| 55 |
+
hidden = getattr(out, "last_hidden_state", None)
|
| 56 |
+
if hidden is None:
|
| 57 |
+
hidden = out[0]
|
| 58 |
+
row = torch.arange(hidden.size(0)) if row is None else row
|
| 59 |
+
slots = hidden[row.to(hidden.device), last.to(hidden.device)] # a model over GPUs ends elsewhere
|
| 60 |
+
states = hidden[pointer["rows"].to(hidden.device), : pointer["width"]] if pointer is not None else None
|
| 61 |
+
return split(slots, states, pointer, output_head(model))
|
tokenizer.json
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a39eed58008cbcbc09e5b629b81338a654269c64cd13348aaa6090f4db75c7e7
|
| 3 |
+
size 17908258
|
tokenizer_config.json
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"backend": "tokenizers",
|
| 3 |
+
"bos_token": "<|startoftext|>",
|
| 4 |
+
"eos_token": "<|im_end|>",
|
| 5 |
+
"extra_special_tokens": [],
|
| 6 |
+
"image_token": "<|image|>",
|
| 7 |
+
"model_max_length": 1000000000000000019884624838656,
|
| 8 |
+
"model_specific_special_tokens": {
|
| 9 |
+
"image_token": "<|image|>",
|
| 10 |
+
"video_token": "<|video|>",
|
| 11 |
+
"vision_end_token": "<|vision_end|>",
|
| 12 |
+
"vision_start_token": "<|vision_start|>"
|
| 13 |
+
},
|
| 14 |
+
"pad_token": "<|pad|>",
|
| 15 |
+
"processor_class": "Lfm3VlProcessor",
|
| 16 |
+
"tokenizer_class": "TokenizersBackend",
|
| 17 |
+
"video_token": "<|video|>",
|
| 18 |
+
"vision_end_token": "<|vision_end|>",
|
| 19 |
+
"vision_start_token": "<|vision_start|>"
|
| 20 |
+
}
|
video_processing_lfm3_vl.py
ADDED
|
@@ -0,0 +1,404 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Vendored from transformers' LFM3 code (not in a release yet) at 69da1b47a11e362fbb15bfe5d5c8da0790c688c5, src/transformers/models/lfm3_vl/video_processing_lfm3_vl.py:
|
| 2 |
+
# its relative imports made absolute, nothing else changed.
|
| 3 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 4 |
+
# This file was automatically generated from src/transformers/models/lfm3_vl/modular_lfm3_vl.py.
|
| 5 |
+
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
| 6 |
+
# the file from the modular. If any change should be done, please apply the change to the
|
| 7 |
+
# modular_lfm3_vl.py file directly. One of our CI enforces this.
|
| 8 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 9 |
+
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
| 10 |
+
#
|
| 11 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 12 |
+
# you may not use this file except in compliance with the License.
|
| 13 |
+
# You may obtain a copy of the License at
|
| 14 |
+
#
|
| 15 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 16 |
+
#
|
| 17 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 18 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 19 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 20 |
+
# See the License for the specific language governing permissions and
|
| 21 |
+
# limitations under the License.
|
| 22 |
+
import math
|
| 23 |
+
|
| 24 |
+
import numpy as np
|
| 25 |
+
import torch
|
| 26 |
+
from torchvision.transforms.v2 import functional as tvF
|
| 27 |
+
|
| 28 |
+
from transformers.image_processing_utils import BatchFeature
|
| 29 |
+
from transformers.image_utils import PILImageResampling, SizeDict
|
| 30 |
+
from transformers.processing_utils import Unpack, VideosKwargs
|
| 31 |
+
from transformers.utils import TensorType, auto_docstring, logging
|
| 32 |
+
from transformers.video_processing_utils import BaseVideoProcessor
|
| 33 |
+
from transformers.video_utils import (
|
| 34 |
+
VideoMetadata,
|
| 35 |
+
group_videos_by_shape,
|
| 36 |
+
make_batched_metadata,
|
| 37 |
+
make_batched_videos,
|
| 38 |
+
reorder_videos,
|
| 39 |
+
)
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
logger = logging.get_logger(__name__)
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
class Lfm3VlVideoProcessorInitKwargs(VideosKwargs, total=False):
|
| 46 |
+
r"""
|
| 47 |
+
patch_size (`int`, *optional*, defaults to 16):
|
| 48 |
+
The spatial patch size of the vision encoder.
|
| 49 |
+
temporal_patch_size (`int`, *optional*, defaults to 2):
|
| 50 |
+
The temporal patch size of the vision encoder.
|
| 51 |
+
merge_size (`int`, *optional*, defaults to 2):
|
| 52 |
+
The merge size of the vision encoder to llm encoder.
|
| 53 |
+
min_frames (`int`, *optional*, defaults to 4):
|
| 54 |
+
The minimum number of frames to sample from video.
|
| 55 |
+
max_frames (`int`, *optional*, defaults to 768):
|
| 56 |
+
The maximum number of frames to sample from video.
|
| 57 |
+
cap_pixels_per_frame (`bool`, *optional*):
|
| 58 |
+
Whether to cap each frame's pixel cost the way the reference implementation (qwen-vl-utils) does:
|
| 59 |
+
per-frame pixels are limited to `min(max_video_tokens * factor**2, size["longest_edge"] / num_frames)`
|
| 60 |
+
and floored at `1.05 * size["shortest_edge"]`, so token cost scales with clip duration. Without the
|
| 61 |
+
cap, videos that sample few frames spend the whole `size["longest_edge"]` budget on those frames and
|
| 62 |
+
keep near-native per-frame resolution, so a short clip can cost almost as many tokens as a long
|
| 63 |
+
video. If unset, the current uncapped behavior is kept and a warning is emitted: the default will
|
| 64 |
+
change to `True` in v5.22, after which the argument will be removed.
|
| 65 |
+
max_video_tokens (`int`, *optional*, defaults to 768):
|
| 66 |
+
The per-frame token ceiling applied by `cap_pixels_per_frame`, in vision tokens per frame
|
| 67 |
+
(qwen-vl-utils' `VIDEO_MAX_TOKEN_NUM`).
|
| 68 |
+
"""
|
| 69 |
+
|
| 70 |
+
patch_size: int
|
| 71 |
+
temporal_patch_size: int
|
| 72 |
+
merge_size: int
|
| 73 |
+
min_frames: int
|
| 74 |
+
max_frames: int
|
| 75 |
+
cap_pixels_per_frame: bool
|
| 76 |
+
max_video_tokens: int
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def smart_resize(
|
| 80 |
+
num_frames: int,
|
| 81 |
+
height: int,
|
| 82 |
+
width: int,
|
| 83 |
+
temporal_factor: int = 2,
|
| 84 |
+
factor: int = 28,
|
| 85 |
+
min_pixels: int = 112 * 112,
|
| 86 |
+
max_pixels: int = 14 * 14 * 2 * 2 * 2 * 6144,
|
| 87 |
+
):
|
| 88 |
+
if num_frames < temporal_factor:
|
| 89 |
+
raise ValueError(f"t:{num_frames} must be larger than temporal_factor:{temporal_factor}")
|
| 90 |
+
if height < factor or width < factor:
|
| 91 |
+
scale = max(factor / height, factor / width)
|
| 92 |
+
height = int(height * scale)
|
| 93 |
+
width = int(width * scale)
|
| 94 |
+
|
| 95 |
+
if max(height, width) / min(height, width) > 200:
|
| 96 |
+
raise ValueError(
|
| 97 |
+
f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
|
| 98 |
+
)
|
| 99 |
+
h_bar = round(height / factor) * factor
|
| 100 |
+
w_bar = round(width / factor) * factor
|
| 101 |
+
t_bar = round(num_frames / temporal_factor) * temporal_factor
|
| 102 |
+
|
| 103 |
+
if t_bar * h_bar * w_bar > max_pixels:
|
| 104 |
+
beta = math.sqrt((num_frames * height * width) / max_pixels)
|
| 105 |
+
h_bar = max(factor, math.floor(height / beta / factor) * factor)
|
| 106 |
+
w_bar = max(factor, math.floor(width / beta / factor) * factor)
|
| 107 |
+
elif t_bar * h_bar * w_bar < min_pixels:
|
| 108 |
+
beta = math.sqrt(min_pixels / (num_frames * height * width))
|
| 109 |
+
h_bar = math.ceil(height * beta / factor) * factor
|
| 110 |
+
w_bar = math.ceil(width * beta / factor) * factor
|
| 111 |
+
|
| 112 |
+
return h_bar, w_bar
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
def _check_video_fps(metadata: VideoMetadata) -> None:
|
| 116 |
+
if metadata.fps is not None and not (math.isfinite(metadata.fps) and metadata.fps > 0):
|
| 117 |
+
raise ValueError(f"`video_metadata.fps` must be a positive number, got {metadata.fps}")
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def _check_video_metadata_count(videos, video_metadata) -> None:
|
| 121 |
+
# the shared sampler pairs clips with metadata by `zip`, which drops the clips past the last entry
|
| 122 |
+
if video_metadata is None:
|
| 123 |
+
return
|
| 124 |
+
videos = make_batched_videos(videos)
|
| 125 |
+
entries = make_batched_metadata(videos, video_metadata) if video_metadata != [] else []
|
| 126 |
+
if len(entries) != len(videos):
|
| 127 |
+
raise ValueError(
|
| 128 |
+
f"Received {len(videos)} videos but {len(entries)} `video_metadata` entries; pass one entry per video"
|
| 129 |
+
)
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
@auto_docstring
|
| 133 |
+
class Lfm3VlVideoProcessor(BaseVideoProcessor):
|
| 134 |
+
resample = PILImageResampling.BICUBIC
|
| 135 |
+
size = {"shortest_edge": 4_096, "longest_edge": 25_165_824}
|
| 136 |
+
image_mean = [0.5, 0.5, 0.5]
|
| 137 |
+
image_std = [0.5, 0.5, 0.5]
|
| 138 |
+
do_resize = True
|
| 139 |
+
do_rescale = True
|
| 140 |
+
do_normalize = True
|
| 141 |
+
do_convert_rgb = True
|
| 142 |
+
patch_size = 16
|
| 143 |
+
temporal_patch_size = 2
|
| 144 |
+
merge_size = 2
|
| 145 |
+
min_frames = 4
|
| 146 |
+
max_frames = 768
|
| 147 |
+
do_sample_frames = True
|
| 148 |
+
cap_pixels_per_frame = False
|
| 149 |
+
max_video_tokens = 768
|
| 150 |
+
valid_kwargs = Lfm3VlVideoProcessorInitKwargs
|
| 151 |
+
model_input_names = ["pixel_values_videos", "video_grid_thw"]
|
| 152 |
+
max_image_size = {"longest_edge": 28 * 28 * 2 * 30000}
|
| 153 |
+
max_duration = None
|
| 154 |
+
num_frames = None
|
| 155 |
+
fps = 2.0
|
| 156 |
+
|
| 157 |
+
def __init__(self, **kwargs: Unpack[Lfm3VlVideoProcessorInitKwargs]):
|
| 158 |
+
super().__init__(**kwargs)
|
| 159 |
+
|
| 160 |
+
def sample_frames(
|
| 161 |
+
self, metadata: VideoMetadata, num_frames: int | None = None, fps: int | float | None = None, **kwargs
|
| 162 |
+
):
|
| 163 |
+
"""Decoded frames whose metadata carries no fps are sampled and timestamped at the 24 fps the shared sampler
|
| 164 |
+
assumes (with a warning); pass `video_metadata` with the clip's fps for accurate timestamps."""
|
| 165 |
+
_check_video_fps(metadata)
|
| 166 |
+
if fps is not None and not (math.isfinite(fps) and fps > 0):
|
| 167 |
+
raise ValueError(f"`fps` must be a positive number, got {fps}")
|
| 168 |
+
if fps is not None and num_frames is not None:
|
| 169 |
+
raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")
|
| 170 |
+
|
| 171 |
+
total_num_frames = metadata.total_num_frames
|
| 172 |
+
fps = fps if fps is not None else self.fps
|
| 173 |
+
|
| 174 |
+
# If num_frames is not given but fps is, calculate num_frames from fps
|
| 175 |
+
if num_frames is None and fps is not None:
|
| 176 |
+
if metadata.fps is None:
|
| 177 |
+
metadata.fps = 24
|
| 178 |
+
logger.warning_once(
|
| 179 |
+
"Asked to sample `fps` frames per second but no video metadata was provided which is required when sampling with `fps`. "
|
| 180 |
+
"Defaulting to `fps=24`. Please provide `video_metadata` for more accurate results."
|
| 181 |
+
)
|
| 182 |
+
num_frames = int(total_num_frames / metadata.fps * fps)
|
| 183 |
+
num_frames = min(max(num_frames, self.min_frames), self.max_frames, total_num_frames)
|
| 184 |
+
|
| 185 |
+
if num_frames is None:
|
| 186 |
+
num_frames = min(max(total_num_frames, self.min_frames), self.max_frames)
|
| 187 |
+
|
| 188 |
+
indices = np.linspace(0, total_num_frames - 1, num_frames).round().astype(int)
|
| 189 |
+
|
| 190 |
+
return indices
|
| 191 |
+
|
| 192 |
+
def resize(
|
| 193 |
+
self,
|
| 194 |
+
videos: "torch.Tensor",
|
| 195 |
+
size: SizeDict,
|
| 196 |
+
resample: "PILImageResampling | tvF.InterpolationMode | int | None",
|
| 197 |
+
factor: int,
|
| 198 |
+
temporal_factor: int,
|
| 199 |
+
cap_pixels_per_frame: bool | None = None,
|
| 200 |
+
**kwargs,
|
| 201 |
+
) -> "torch.Tensor":
|
| 202 |
+
"""Resize dynamically based on input video aspect ratio."""
|
| 203 |
+
if not size.shortest_edge or not size.longest_edge:
|
| 204 |
+
raise ValueError(f"`size` dict must contain 'shortest_edge' and 'longest_edge' keys but got {size}.")
|
| 205 |
+
|
| 206 |
+
num_frames = videos.shape[1]
|
| 207 |
+
max_pixels = size.longest_edge
|
| 208 |
+
if cap_pixels_per_frame:
|
| 209 |
+
# per-frame pixels are capped at `max_video_tokens` patches or the budget's even share per frame
|
| 210 |
+
frame_cap = self.max_video_tokens * factor * factor
|
| 211 |
+
pixels_per_frame = max(min(frame_cap, size.longest_edge // num_frames), int(size.shortest_edge * 1.05))
|
| 212 |
+
max_pixels = pixels_per_frame * num_frames
|
| 213 |
+
|
| 214 |
+
height, width = videos.shape[-2:]
|
| 215 |
+
resized_height, resized_width = smart_resize(
|
| 216 |
+
height=height,
|
| 217 |
+
width=width,
|
| 218 |
+
num_frames=num_frames,
|
| 219 |
+
factor=factor,
|
| 220 |
+
temporal_factor=temporal_factor,
|
| 221 |
+
min_pixels=size.shortest_edge,
|
| 222 |
+
max_pixels=max_pixels,
|
| 223 |
+
)
|
| 224 |
+
return super().resize(
|
| 225 |
+
image=videos,
|
| 226 |
+
size=SizeDict(height=resized_height, width=resized_width),
|
| 227 |
+
resample=resample,
|
| 228 |
+
)
|
| 229 |
+
|
| 230 |
+
def patchify(
|
| 231 |
+
self,
|
| 232 |
+
videos: "torch.Tensor",
|
| 233 |
+
patch_size: int,
|
| 234 |
+
merge_size: int,
|
| 235 |
+
temporal_patch_size: int,
|
| 236 |
+
) -> tuple["torch.Tensor", int, int]:
|
| 237 |
+
"Patchifies each video into flat layout of shape (`seq_len`, `patch_dim`) so we can concat dynamically shaped pixels."
|
| 238 |
+
batch_size, num_frames, channel, resized_height, resized_width = videos.shape
|
| 239 |
+
|
| 240 |
+
# Check that videos have `num_frames` divisible by `temporal_patch_size`
|
| 241 |
+
if pad := -num_frames % temporal_patch_size:
|
| 242 |
+
repeats = videos[:, -1:].expand(-1, pad, -1, -1, -1)
|
| 243 |
+
videos = torch.cat((videos, repeats), dim=1)
|
| 244 |
+
num_frames += pad
|
| 245 |
+
|
| 246 |
+
grid_t = num_frames // temporal_patch_size
|
| 247 |
+
grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
|
| 248 |
+
|
| 249 |
+
patches = videos.view(
|
| 250 |
+
batch_size,
|
| 251 |
+
grid_t,
|
| 252 |
+
temporal_patch_size,
|
| 253 |
+
channel,
|
| 254 |
+
grid_h // merge_size,
|
| 255 |
+
merge_size,
|
| 256 |
+
patch_size,
|
| 257 |
+
grid_w // merge_size,
|
| 258 |
+
merge_size,
|
| 259 |
+
patch_size,
|
| 260 |
+
)
|
| 261 |
+
patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
|
| 262 |
+
flatten_patches = patches.reshape(
|
| 263 |
+
batch_size,
|
| 264 |
+
grid_t * grid_h * grid_w,
|
| 265 |
+
channel * temporal_patch_size * patch_size * patch_size,
|
| 266 |
+
)
|
| 267 |
+
|
| 268 |
+
return flatten_patches, grid_t, grid_h, grid_w
|
| 269 |
+
|
| 270 |
+
def _preprocess(
|
| 271 |
+
self,
|
| 272 |
+
videos: list["torch.Tensor"],
|
| 273 |
+
do_resize: bool,
|
| 274 |
+
size: SizeDict,
|
| 275 |
+
resample: "PILImageResampling | tvF.InterpolationMode | int | None",
|
| 276 |
+
do_rescale: bool,
|
| 277 |
+
rescale_factor: float,
|
| 278 |
+
do_normalize: bool,
|
| 279 |
+
image_mean: float | list[float] | None,
|
| 280 |
+
image_std: float | list[float] | None,
|
| 281 |
+
patch_size: int | None = None,
|
| 282 |
+
temporal_patch_size: int | None = None,
|
| 283 |
+
merge_size: int | None = None,
|
| 284 |
+
cap_pixels_per_frame: bool | None = None,
|
| 285 |
+
return_tensors: str | TensorType | None = None,
|
| 286 |
+
**kwargs,
|
| 287 |
+
):
|
| 288 |
+
if cap_pixels_per_frame is None:
|
| 289 |
+
logger.warning_once(
|
| 290 |
+
"Lfm3Vl video processing does not apply the per-frame pixel cap the reference "
|
| 291 |
+
"implementation (qwen-vl-utils) applies, so some videos cost far more tokens than they "
|
| 292 |
+
"would there. In v5.22 the capped behavior will become the default and "
|
| 293 |
+
"`cap_pixels_per_frame` will be removed. Pass `cap_pixels_per_frame=True` to adopt the "
|
| 294 |
+
"reference behavior now, or `False` to keep the current behavior and silence this "
|
| 295 |
+
"warning."
|
| 296 |
+
)
|
| 297 |
+
cap_pixels_per_frame = False
|
| 298 |
+
# Group videos by size for batched resizing
|
| 299 |
+
grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
|
| 300 |
+
resized_videos_grouped = {}
|
| 301 |
+
for shape, stacked_videos in grouped_videos.items():
|
| 302 |
+
if do_resize:
|
| 303 |
+
stacked_videos = self.resize(
|
| 304 |
+
videos=stacked_videos,
|
| 305 |
+
size=size,
|
| 306 |
+
resample=resample,
|
| 307 |
+
factor=patch_size * merge_size,
|
| 308 |
+
temporal_factor=temporal_patch_size,
|
| 309 |
+
cap_pixels_per_frame=cap_pixels_per_frame,
|
| 310 |
+
)
|
| 311 |
+
resized_videos_grouped[shape] = stacked_videos
|
| 312 |
+
resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
|
| 313 |
+
|
| 314 |
+
# Group videos by size for further processing
|
| 315 |
+
# Needed in case do_resize is False, or resize returns videos with different sizes
|
| 316 |
+
grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
|
| 317 |
+
processed_videos_grouped = {}
|
| 318 |
+
processed_grids = {}
|
| 319 |
+
for shape, stacked_videos in grouped_videos.items():
|
| 320 |
+
# Fused rescale and normalize
|
| 321 |
+
stacked_videos = self.rescale_and_normalize(
|
| 322 |
+
stacked_videos, do_rescale, rescale_factor, do_normalize, image_mean, image_std
|
| 323 |
+
)
|
| 324 |
+
patches, grid_t, grid_h, grid_w = self.patchify(
|
| 325 |
+
stacked_videos,
|
| 326 |
+
patch_size=patch_size,
|
| 327 |
+
merge_size=merge_size,
|
| 328 |
+
temporal_patch_size=temporal_patch_size,
|
| 329 |
+
)
|
| 330 |
+
|
| 331 |
+
processed_videos_grouped[shape] = patches
|
| 332 |
+
processed_grids[shape] = [[grid_t, grid_h, grid_w]] * len(stacked_videos)
|
| 333 |
+
|
| 334 |
+
processed_videos = reorder_videos(processed_videos_grouped, grouped_videos_index)
|
| 335 |
+
processed_grids = reorder_videos(processed_grids, grouped_videos_index)
|
| 336 |
+
pixel_values_videos = torch.cat(processed_videos, dim=0)
|
| 337 |
+
video_grid_thw = torch.tensor(processed_grids)
|
| 338 |
+
|
| 339 |
+
return BatchFeature(
|
| 340 |
+
data={"pixel_values_videos": pixel_values_videos, "video_grid_thw": video_grid_thw},
|
| 341 |
+
tensor_type=return_tensors,
|
| 342 |
+
)
|
| 343 |
+
|
| 344 |
+
def get_num_of_video_patches(self, num_frames: int, height: int, width: int, videos_kwargs=None):
|
| 345 |
+
"""
|
| 346 |
+
A utility that returns number of video patches a given video size.
|
| 347 |
+
|
| 348 |
+
Args:
|
| 349 |
+
num_frames (`int`):
|
| 350 |
+
Number of frames in the input video.
|
| 351 |
+
height (`int`):
|
| 352 |
+
Height of the input video.
|
| 353 |
+
width (`int`):
|
| 354 |
+
Width of the input video.
|
| 355 |
+
videos_kwargs (`dict`, *optional*)
|
| 356 |
+
Any kwargs to override defaults of the video processor.
|
| 357 |
+
Returns:
|
| 358 |
+
`int`: Number of video patches per video.
|
| 359 |
+
"""
|
| 360 |
+
videos_kwargs = videos_kwargs if videos_kwargs is not None else {}
|
| 361 |
+
size = videos_kwargs.get("size", None) or self.size
|
| 362 |
+
min_pixels = size["shortest_edge"]
|
| 363 |
+
max_pixels = size["longest_edge"]
|
| 364 |
+
patch_size = videos_kwargs.get("patch_size", None) or self.patch_size
|
| 365 |
+
merge_size = videos_kwargs.get("merge_size", None) or self.merge_size
|
| 366 |
+
temporal_patch_size = videos_kwargs.get("temporal_patch_size", None) or self.temporal_patch_size
|
| 367 |
+
cap_pixels_per_frame = videos_kwargs.get("cap_pixels_per_frame", None) or self.cap_pixels_per_frame
|
| 368 |
+
|
| 369 |
+
factor = patch_size * merge_size
|
| 370 |
+
if cap_pixels_per_frame:
|
| 371 |
+
# Keep the count in sync with `resize` when the per-frame cap is active.
|
| 372 |
+
frame_cap = self.max_video_tokens * factor * factor
|
| 373 |
+
pixels_per_frame = max(min(frame_cap, max_pixels // num_frames), int(min_pixels * 1.05))
|
| 374 |
+
max_pixels = pixels_per_frame * num_frames
|
| 375 |
+
|
| 376 |
+
resized_height, resized_width = smart_resize(
|
| 377 |
+
num_frames,
|
| 378 |
+
height,
|
| 379 |
+
width,
|
| 380 |
+
temporal_factor=temporal_patch_size,
|
| 381 |
+
factor=factor,
|
| 382 |
+
min_pixels=min_pixels,
|
| 383 |
+
max_pixels=max_pixels,
|
| 384 |
+
)
|
| 385 |
+
grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
|
| 386 |
+
grid_t = (num_frames + -num_frames % temporal_patch_size) // temporal_patch_size
|
| 387 |
+
return grid_t * grid_h * grid_w
|
| 388 |
+
|
| 389 |
+
def preprocess(self, videos, **kwargs):
|
| 390 |
+
_check_video_metadata_count(videos, kwargs.get("video_metadata"))
|
| 391 |
+
# the shared `preprocess` fills every unset option from the instance before `sample_frames` runs, so the
|
| 392 |
+
# caller's options are told apart from the instance ones here: a caller option beats an instance one, and a
|
| 393 |
+
# configured `num_frames` beats the default fps
|
| 394 |
+
num_frames, fps = kwargs.get("num_frames"), kwargs.get("fps")
|
| 395 |
+
if num_frames is not None and fps is not None:
|
| 396 |
+
raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")
|
| 397 |
+
if num_frames is not None or (fps is None and self.num_frames is not None):
|
| 398 |
+
kwargs["fps"] = None
|
| 399 |
+
elif fps is not None:
|
| 400 |
+
kwargs["num_frames"] = None
|
| 401 |
+
return super().preprocess(videos, **kwargs)
|
| 402 |
+
|
| 403 |
+
|
| 404 |
+
__all__ = ["Lfm3VlVideoProcessor"]
|