paulpak58 mlabonne Aurelien-Lac commited on
Commit
97c5704
·
0 Parent(s):

Co-authored-by: mlabonne <mlabonne@users.noreply.huggingface.co>
Co-authored-by: Aurelien-Lac <Aurelien-Lac@users.noreply.huggingface.co>

.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>{&quot;type&quot;: &quot;image&quot;, &quot;path&quot;: ...}</code>) and videos as their frames (<code>{&quot;type&quot;: &quot;video&quot;, &quot;video&quot;: [PIL images], &quot;fps&quot;: 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"]