diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..03fcb43d6909b35bf5946c592ec099c1f6c164f5 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,40 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +embedding.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +finish.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +heads.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-00.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-01.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-02.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-03.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-04.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-05.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-06.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-07.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-08.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-09.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-10.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-11.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-12.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-13.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-14.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-15.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-16.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-17.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-18.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-19.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-20.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-21.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-22.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-23.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-24.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-25.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-26.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-27.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-28.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-29.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-30.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +layer-31.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text +model/tokenizer.json filter=lfs diff=lfs merge=lfs -text +vision.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..d645695673349e3947e8e5ae42332d0ac3164cd7 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "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. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "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. + + "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 + (an example is provided in the Appendix below). + + "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. + + "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." + + "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. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, 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. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, 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. + + 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: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (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 + + (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. + + 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. + + 5. 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. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. 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. + + 8. 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. + + 9. 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. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c761922684066a259f4359983b55e1aff3705428 --- /dev/null +++ b/README.md @@ -0,0 +1,57 @@ +--- +license: apache-2.0 +base_model: mukaddamzaid/amos-1-medium +tags: [coreai, core-ai, apple, aimodel, decisions, vision] +--- + +# amos-1-medium — Apple Core AI + +Independent AMOS-1 conversion with its checkpoint's **causal full attention** preserved. Text backbone, 255-code readout, vision tower, span head, region head and trained box head execute through Core AI. Original prompt formatting, temperatures and quote decoding are retained. This is a decision model, not a text generator. + +## Use + +Requires macOS 27 and the Core AI Python runtime. Tested with `coreai-torch==0.4.3`, PyTorch 2.14.0, Transformers 5.19 development, NumPy and Pillow. Exact dependencies are in `requirements.txt`. The converter warns that this PyTorch version exceeds its validated range. Host preprocessing/decoding uses PyTorch tensors; it loads no PyTorch neural weights. + +```python +import asyncio +from coreai_decoder import AmosDecoder + +async def main(): + m = await AmosDecoder.load(".") + print(await m.request({ + "state": "Please refund the duplicate charge.", + "questions": {"refund": { + "type": "noul", "instructions": "Is a refund requested?" + }} + })) + +asyncio.run(main()) +``` + +`predict(state, questions, images=())` takes upstream internal questions (`text`, `options`) and PIL images, including image options. `request()` takes base64 images in the original AMOS format. + +## Verified scope and limits + +Batch one, **512 combined tokens**, 224×224 images (49 merged tokens), float32 backbone/vision weights with FP32 readout and heads. Requests over capacity raise an error. Original longer contexts, variable image resolution, wide-choice two-pass routing, rotations, thinking, hosted server and explanation endpoints are not implemented. Up to 255 option codes are available subject to token capacity. No iPhone, Neural Engine placement, throughput or upstream benchmark accuracy claim has been verified. Core AI uses default specialization. + +The text graph is split into an embedding asset, 32 transformer-layer assets and a final normalization/readout asset. They execute sequentially through Core AI in FP32. Conversion and the independent reference expand only one layer at a time. Gated-delta attention uses functional forward substitution instead of the export-only inverse construction; this matches the original eager triangular solve. The replacement passed 21 independent randomized comparisons and was checked against intermediate outputs from the trained model. Tested on a 64 GiB Mac; no smaller-memory device has been verified. + +Nine smoke/parity cases cover all six typed question kinds, quote extraction, Unicode, 24 options, synthetic state images and image options. Tested discrete decisions matched upstream. Maximum measured probability/box-coordinate error: **0.00001425**. Full results are in `validation.json`; these smoke cases do not establish benchmark accuracy. `runtime-smoke.json` verifies inference with PyTorch neural execution disabled. + +The primary parity reference executes the pinned checkpoint in FP32. Compared with the original BF16 serving path in 9 additional cases, maximum probability/box error was 0.109300, and 1 cases changed a discrete output. Region confidence can change discontinuously when a box crosses a grid boundary. See `bf16-comparison.json`; this is not bit-identical BF16 serving. + +## Attribution and reproducibility + +Original model and serving code: Zaid Mukaddam, [mukaddamzaid/amos-1-medium](https://huggingface.co/mukaddamzaid/amos-1-medium). Weights derive from Qwen3.5. Apache-2.0 applies; see `LICENSE` and `SOURCE.json` for the pinned AMOS revision and source hashes. Conversion adds no restrictions. Upstream serving sources are included under `serve/`. + +Download the pinned source and reproduce from a directory holding these scripts: + +```bash +hf download mukaddamzaid/amos-1-medium --revision 6b48a91ed940e08c87dcdd29c71c3b256d757ec7 --local-dir upstream/amos-1-medium +python convert_medium_sharded.py +python validate_decoder.py amos-1-medium --mode reference +python validate_decoder.py amos-1-medium --mode coreai +python smoke_decoder.py amos-1-medium +``` + +The two validation processes keep the reference model and Core AI runtime out of memory at the same time. Release integrity is recorded in `SHA256SUMS`. diff --git a/SHA256SUMS b/SHA256SUMS new file mode 100644 index 0000000000000000000000000000000000000000..35b9af03361365902e5fced1fde74204826862bb --- /dev/null +++ b/SHA256SUMS @@ -0,0 +1,149 @@ +cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30 LICENSE +a4ba5662de53a69d4aeb035ec818ede0facd61b523040dac025f61c2cf68db83 README.md +59baf0b50bf8d6b1b8f7be5f29827a0649d1d23a6510856a7e85d23ef00086c8 SOURCE.json +e43486a0a8ef7a9b042ba95d83086baad5f5b6ecff194d7ab18f581840fb75bd bf16-attempt.json +e89fae9a7630a27a408f2b48acdaf75aebe017957c815654f934769773f10f5c bf16-comparison.json +4a7251b59a4edd324c1063baf454f2b3a3176a7d678b296f19452cd1d45c46ea calibration.json +b705e46d554298420a89287499d84d93ebb50675dd61a42d8752350754beb2c9 config.json +910e2baa29daa433c9dbf4120c4848da3898d94ff8a0d899e82484308690d347 convert_decoder.py +e457409e42e9a5b7dbe12a53180a045239012170ad4481b79536c839fb83f8d2 convert_medium_sharded.py +94db146d96f60709905396f59e6009c76d39a63f8e2e9796406fa3a5611a4e3a coreai_decoder.py +59d4b07cdcb63aadf6f33369f703e7cd963d9b8ad7c6ad7a7f6f03dc8e2fd2b3 embedding.aimodel/main.hash +6fcec91d2cdc76e76a855ccc7dc61d10697852772f81b6d7a8d94af4ac897407 embedding.aimodel/main.mlirb +1a76f7272568c5a7dfd5a7425023d01ce2b641b090168120c3390e9f29ae3f29 embedding.aimodel/metadata.json +6374a4ca1aa66476a2b447a95ba409a857fffac7c2a2462716e345372d392019 finish.aimodel/main.hash +40ed1efe47f26cc15794ce7616798dd2f948207262fd62041b6961e381f37c15 finish.aimodel/main.mlirb +c57af257edcf0598c4a9c744ad72e067eb44ea638c5fcd374d89e7d1b1ae3773 finish.aimodel/metadata.json +30985aab978011567a406cbb997d76adaa4ca35fea8076de85b3177e721f1f09 heads.aimodel/main.hash +455361e9ff9557669def5344f8db6872e1bdcddc81e41141c5888cce54197f27 heads.aimodel/main.mlirb +447ac46fb7493976509699118d00e61ab2c13f50eb0692797210e4dcd10e410f heads.aimodel/metadata.json +a88d1d17ababbcacbcc09029d691734118b42ae394e19b18034286681430ca2e layer-00.aimodel/main.hash +b506c0c52bdffe5dc796662ea73c92c2612ee4dc775b226ec60110b93594c740 layer-00.aimodel/main.mlirb +c1d58f995410e5c18525f3e7ffd65e37b2f997f9c190efe7ba918be75bec9cbd layer-00.aimodel/metadata.json +f3016c678777a8b94fc725ed16925ce0a52a900b3eab54e258ba93b3f72997f9 layer-01.aimodel/main.hash +3656fc65add337eb97637dd47f3421d082f8d11ba555f753dddd24068ce732c1 layer-01.aimodel/main.mlirb +a121c2080592caf2e9ed3f1225f979d7d638c3fe2349b702970858aca6dc7c60 layer-01.aimodel/metadata.json +a0b5be4c51ba958a9de19be9d87312363ae04f29e8e0de2d5d7dbd096aa21f6d layer-02.aimodel/main.hash +4261515422e07c5f633a4a8436fdd724d768694023941be5f8fed0a6c7ce7ef2 layer-02.aimodel/main.mlirb +cce09e3c5e3173f26940c3330cb09f89c8b148de6910a9005c2ce8cd4188f9ec layer-02.aimodel/metadata.json +d614c642d099d22e18c6621acc654b06d70e6641d11d48fe1a124b459f00e71f layer-03.aimodel/main.hash +40b6b0de27b454661002405a4921d3a8a88e0efabdadc45add9d68abe322a135 layer-03.aimodel/main.mlirb +b7777c03d2a2898d08430335f9edf0729db5c3942ee6464079948a35406a03ea layer-03.aimodel/metadata.json +57d92f10bad48b9c74c7b750b06922abcffdc8ea4b531d18d0b21a95a0448ce8 layer-04.aimodel/main.hash +cbf1ad7ff01a89c1aebf1693e059676cff6a7a3219feee2a49237a1aba83b856 layer-04.aimodel/main.mlirb +d5d59d36e35e40ad826f533dbbba26475586ce8ba6175eb655c1f5fa993086bf layer-04.aimodel/metadata.json +fa1799bf1c3d6e0697a853ba9393da65cdce1ca7ef21b17493a20a63d8ea7697 layer-05.aimodel/main.hash +e8ba384095287e7d9f4bd0aa05b34dad6feca189a7f32fae5a5659b86d1b34e2 layer-05.aimodel/main.mlirb +f516ea4680a1a18a56dcb7ed46fa71c5df32207fbdd0796a20939163f49f796e layer-05.aimodel/metadata.json +4afdf63fe5543438c374444f98cb895c0a3394ba064be75c46323ef53ef04962 layer-06.aimodel/main.hash +a15792d07af6107986f80e3cf495be519d976d4ee64bb110fabb4d6e1ba959d8 layer-06.aimodel/main.mlirb +f044e22a815d45b3dfb4e76e888a23be99c444f9f0f2abad7e4ca674088913a1 layer-06.aimodel/metadata.json +ac6f040a6bbd7762ca462bbe262dd6ca92e2e8d2fb5cab550eca013c77e950cb layer-07.aimodel/main.hash +d1ad7434add3c0c3d4effdd2f98f5a3fef4d14372b2c1fe21e8bbd93b1d78d65 layer-07.aimodel/main.mlirb +b4eadb01163b3b27ae3400b2a6f3b6ee90455e90f32d8b3a509454701fc790dc layer-07.aimodel/metadata.json +8a170b0127d3dcf84fff6ecb4aef3447b003670ae12be547c67d4717d465783f layer-08.aimodel/main.hash +c7bffc69c084785234be0faa61a38d4d4f77c1aa1630bb9be44950160ad8f70e layer-08.aimodel/main.mlirb +bc59d87615bc2623beea3a1a059aba9ee025e9a9d42cb6648bffb3b877879d2e layer-08.aimodel/metadata.json +0291d32e7b16b17c760febda7e4c6156211409aa19e751d7a3f0d7e783e5e879 layer-09.aimodel/main.hash +3f7401a15775cd13ac01a5000bbfb55e2aaba194d2e7315540d09b9c66fcbaea layer-09.aimodel/main.mlirb +1fe1ef0ad0fdd11349d12fcfe5283d1e127f2f259ca9d39843f45e11239b97d6 layer-09.aimodel/metadata.json +9b20226b17789582b12a77891a5e933ce6cd0a65c5c01318c121babd14cd867a layer-10.aimodel/main.hash +be8f740a5edf5284760fe1f17be01c99171015509266209c3f6f2f4d6d566918 layer-10.aimodel/main.mlirb +0a2e591e6bdc41f19e7a5375f742349e1b252457e9224918c816e7ac5ef6d8ff layer-10.aimodel/metadata.json +089c57c5f2770ce3629ff2f21b3cf88370d00afba9fc662992a336e7309c3e02 layer-11.aimodel/main.hash +d97200be31b02ab0c92e9fc1ea7ab834e1377773d4ff589100db9077a60b024f layer-11.aimodel/main.mlirb +9004cd907a3b6a871a48c87dffed53fa6d7af2fe045bdb1cdb68099da267a541 layer-11.aimodel/metadata.json +4e327cd93d97b8112472a7b76811da8fe21357a6a7c3bb4e5e9d3ed2f07b37dd layer-12.aimodel/main.hash +ed3314c6de83daa4f9671bb03dcb11e480556e352e7f0aec26adbed7044711f1 layer-12.aimodel/main.mlirb +793fc185515dd7b0fd6f1cd1b04915cbf58d41d35b5d9fa756ef74de616c53dc layer-12.aimodel/metadata.json +29e2babc319322b48dcf2f93ed02894d2dfbf405c3fbabac052154920930beb3 layer-13.aimodel/main.hash +ae258a478059967a8d7d4321ab6e7e3d940492456f4d3f1213dd76a9cd6642e3 layer-13.aimodel/main.mlirb +a3d53f031431a823840bd5440bd0f9cc63aafda3aa99c415dcaa307a7663b906 layer-13.aimodel/metadata.json +a88fd2368f838be967b1423de44bd694ec945cfd427eddb5a629026558bbcffd layer-14.aimodel/main.hash +16716737712bb9ace2958ee507ee4fc067746ef4c08835d61ff8a2e41519f60a layer-14.aimodel/main.mlirb +c3d9f7797644be0384370bf37f3a90a2f8edb13e86b608cef1434c3085e6bc8d layer-14.aimodel/metadata.json +980baeb810e5e1e87ed7e1f13fc1e8e37c7614913ad2330986ec2cff6e62c39e layer-15.aimodel/main.hash +36656b20507894cfa3ef07e529c9fb9732420985d05e27484d2a91d2b8c1b1eb layer-15.aimodel/main.mlirb +4e2d389e81252fb07306c602ece123dfb0338c5c22b15b1177148494937fc31a layer-15.aimodel/metadata.json +cbf7c01dbd356c9e6f9cc7cb0d31d0106fc50fc609045a72ce921ada6ebf3573 layer-16.aimodel/main.hash +6863d7b56af2cc67a2d492425b4b6e4da3717fec0849b02049a33a6afa93ac4d layer-16.aimodel/main.mlirb +1984518024c16eda5d5da35c2f6d9df2db00214c2884b2c32eda675567249134 layer-16.aimodel/metadata.json +3e95bfe3a4e66790f9068b73dce6f93d6b1f8b26cd6660c82267eccb0efb1e93 layer-17.aimodel/main.hash +80d069228ee3bca6a408568637ebc57ff65a16830087681e492ef4155c000651 layer-17.aimodel/main.mlirb +b144c2f42cf676f93f554ccf5f2366f0efdd449416a1c40a7b2510b504a9c554 layer-17.aimodel/metadata.json +e0a16e1351e0819e6696cf7c6a5a661d14178c224b0e9977c80eca4ddfb5c5d0 layer-18.aimodel/main.hash +94e53f7086b760e0ac860a89b7c53ce161c3990c4bc97517a2de83398ff60507 layer-18.aimodel/main.mlirb +6629be3f16b78c283997d635b90fb3f8985f172bd8d2c62d16408d8852636167 layer-18.aimodel/metadata.json +77d378b15a4bda2ba709e395d30057aed2a1e48a62ea5b2257d88e13a6a6c70b layer-19.aimodel/main.hash +8cec2a63e5ca8e9c3a88929dc4867efa993a6e9b012b8e44aa8eb8fa7fcf148b layer-19.aimodel/main.mlirb +806fc243d9d6adeec45461f0c474346d60616ca7caacbf6b6a4c114f106d2b56 layer-19.aimodel/metadata.json +56cba7775b52b4394db21860f9b04d3caa78bd11fc08c08e41855a734f3967d6 layer-20.aimodel/main.hash +e42fc846ef06126b880ec69d6d25ad83dd5e4813fca4f51695271bf1b142b797 layer-20.aimodel/main.mlirb +79bdaeeaf3b24bcc08e29514e4bd18a57c5f7041cc683bda67260fa14d1fc350 layer-20.aimodel/metadata.json +b44638150943c03bd87a2b2ec3a57334fe43f9981cfca90fd9346a343cb01626 layer-21.aimodel/main.hash +0f1dfc2393dbfcfa0129ad965dd7c2fc4d4ac7ddcf9ce25ea68501fa5aeb7df7 layer-21.aimodel/main.mlirb +f5b3cc44231a7f5a51c1d45b8bb254d92bf37d2e2111b0c8cde5ab9383eeaadd layer-21.aimodel/metadata.json +3773f20fae795a15f4835282945e2a7b1e71cb40aa1ea5f1e7ad39553c596e0e layer-22.aimodel/main.hash +8e35f9ef221fee278b0e7df2ac3e4643bfb77ae3644c786f9ed003ee9c10a30f layer-22.aimodel/main.mlirb +f1b8f1e895dc4c5bf4882c0503367d08af0f40de7e5126bb52b3f13a2c7f6bcc layer-22.aimodel/metadata.json +7c103be39bfb6424157fb8e1029bb3f9020f439355b9f6401356ce5542eb4e94 layer-23.aimodel/main.hash +2b58ef90e64fe06fe3b2127eee0705331f0b307a686f21f99ac9a4e6cf955abe layer-23.aimodel/main.mlirb +5d33a18d46fb436783b5982c5901fc2a43e51dac183cae6a957d7419cdd9aa7d layer-23.aimodel/metadata.json +aa1ba06e01059f78007a29fd70ad83dfe6386276b65330d5b03361fed3b66ab0 layer-24.aimodel/main.hash +dfb66edc24171137f37516602a90c143b966f0a74e7f78ed81a30b9bd0257a89 layer-24.aimodel/main.mlirb +2a840be3ffde9943847dae58a22e8a29a2a101bc5cf96b4f8f588626a3d9ecfa layer-24.aimodel/metadata.json +b47a9310bfd8517228d4dfef1f87abc1fec3c315fb74b0e04522bd7d66efc180 layer-25.aimodel/main.hash +8a376d08635b27688c04e11be678087cb6fe7733362266163c749e717bb6b6e2 layer-25.aimodel/main.mlirb +2e3c6c4ee2812522d020c49f7bd23aac2b1c3ef84a77552c390a87e9cb310083 layer-25.aimodel/metadata.json +d33e996542edc6fde1b11516269a2afcad4c70f97670c2b7722fb598b936be3d layer-26.aimodel/main.hash +f11e92ad47e77a9ce11c30eb50ac2c7a09d3abc0c2c0320cfebb78bcca56d1c3 layer-26.aimodel/main.mlirb +4dc937db892b7b857554843057e7afc8a1c3e2da20d593e4e65e91d861f06f9f layer-26.aimodel/metadata.json +c79db0f7a82074c555bc2d56f9678af3628e1ea73dbe01253f9c6609a52e6998 layer-27.aimodel/main.hash +2bf14a05dbb7cdb6f734b7664a5e2b74a477f9f67ca77ca9d0c4ffb4e00184cd layer-27.aimodel/main.mlirb +70c87704121140dbc66be29d6e24396f2ca2b30557736c5c7fb6d13e5e78c196 layer-27.aimodel/metadata.json +7ac711aa4daaced37b16a4c03c75f55af895c36352b2e779ad4e8f2475bc12ca layer-28.aimodel/main.hash +0f3e200196185e3dcb0125ef0d3524a5ce93143e226f6ea357b790347baf5c0b layer-28.aimodel/main.mlirb +5f376729867494509a7f2dcdb5a035166bf579a11f3b64d7428fdfc6fd52fb5a layer-28.aimodel/metadata.json +2ee93cf261591e0ea8344cdb81aaf05d21490c874b032ed8c557ff3f74ae5600 layer-29.aimodel/main.hash +fbca7a309b0090b05e5df56eab1cd7cff96019f08453d23f5ab476bf2f4b6069 layer-29.aimodel/main.mlirb +ba90bbe1fadf75ea129bd02fea88b0a92ab624693bcc9026f8a4f2454ed14eb4 layer-29.aimodel/metadata.json +9ddd0265026c9be5eac4b69db185bc02a05f088509378402245ce7d5b9739a17 layer-30.aimodel/main.hash +8f0c2d9df134e7721b8a7d819be51c0621803a47eb27a738ad75192f04c4c3aa layer-30.aimodel/main.mlirb +f787b4af5978daf5beb93a4dcd60389078ce52f5e9e65872e061e73230bc55ae layer-30.aimodel/metadata.json +281ffa534596df7a7ba4838eb8de4e57bf60e8a27f2df5a402a4e89b9720e21a layer-31.aimodel/main.hash +bd1e995c957256ffd8ae6e46e045f74446aaad41d92bf2d87b8e6d1c41d21f10 layer-31.aimodel/main.mlirb +c57af257edcf0598c4a9c744ad72e067eb44ea638c5fcd374d89e7d1b1ae3773 layer-31.aimodel/metadata.json +c082cbd02da0ebaba873d44dc15c14305ee3b7c9e6944f01ec75fa913ca9ffae manifest.json +a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 model/chat_template.jinja +b647a20915503035f0e9a906980b891567bc04085fde26759016a05bf4445227 model/config.json +62153eb6c69f2e1f426beaa8002b7186437e949c7588167085df14e10e9c0a73 model/generation_config.json +27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516 model/preprocessor_config.json +aa7582a6b90dc3711049d293d985e549f2b22c0c4e5ad11b9272a03c6294886e model/tokenizer.json +13ef66b39ff28a4bb50d8aaf554dbd01315b4497136b93111e5021ac6396db0f model/tokenizer_config.json +7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13 model/video_preprocessor_config.json +d7910a3390f8073d95908f3301c0ea0e2df4551699546430003bc8226373c2bc qwen_export_helpers.py +52cf5a65b92ffade3534ee62e66a12ca7071bbe6de2429eb64111c95ae861f21 reference.json +d4e1bf42e2598b84d5800323fdec87fbf338613e8e866e6bb5ce11ed628c8cb2 requirements.txt +7337ff08fd8f599bcceef33f7a76fc56c49d23196e1fa06778b9cbb330929ac5 runtime-smoke.json +9e17531efd289cf720a36688d943197f9baf9db7c4f743aa4951f6b6f7dbafda serve/README.md +0579eb845ee7af463ec88cf627d99dd636466f2c12d840a7cbe183d8f17044fd serve/amos/__init__.py +ba094d9eaedd229a99b2bdcd35a371ffdeb9aaee258a342524cc19e2f85ec421 serve/amos/calibrate.py +0dcacb27a05d31781f8310cbe6123b0ee351c62456cc6b5f8576b40b8ddad080 serve/amos/decoder.py +b4a5a606a2dee99eedf86cab50be759eca0908a318e4f3037ba92560d837bd67 serve/amos/encoder.py +f2efd7ef25453f7f2dc204dd247c86733a55697ec48714f30560df8ed2da1f23 serve/amos/facts.py +d4f4354f3e4d3b475b4c1f832a4a9741da757a368a1d2e0e74544d06009c2c27 serve/amos/model.py +14f027c9e15c456abb0f264c07dae06afb7ace7b64270f2374e97eabeb04f993 serve/amos/packing.py +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 serve/amos/py.typed +b0523591c320a074a0ec0751b68ebc205aeedd090ff5e7ca10bda3fb94a07ef8 serve/amos/serve.py +133d1e29cc64e371f5df1f0cb0328fd77b8b9093d8465319580cf5ea4af5d081 serve/amos/vision.py +bea3f98882a5d12c891aae4a5d75eda8c5a5436e88889d94b6a7d3d648d126f2 serve/pyproject.toml +0d43aa4e58fe6131f4ed1f4bbcbcaa4a0815f8b1fe5331c7e96489700681136b serve/sources.json +b9ab562bd0e3f3dd74882b1fe7fd86268c09c75db77a984f56446c21f770b978 smoke_decoder.py +39ff61a4f617156941750bfa8dd8aa024a0becb2ac756ac9aa23eecda6b93c6f stable-delta-validation.json +32e5eb172f5386f78314ff2db700a443ba2ddd5c805487f540fb87e4bd457264 stable_delta.py +439b0b8637dd62f3df5f4f0047ed198048c865beefb075f3374c7dfb9d98da68 validate_decoder.py +b59443aa0b8dacaaa64dc7c8eabf9aee30c01ca78435c0ddb0cda5e15231af54 validate_stable_delta.py +dd8e385cdbeaea173fe02e9cfd0d5149d20be9e217225df722bcbf1f74f2f5a7 validation.json +6eda4eb2dc9e8b70c5f44cf3db16c979102043ac04951d6bec07e42ca93e2351 vision.aimodel/main.hash +86679770f2901e1b4b75d10172924f510621b352fed8db9e8241bd1a2992d068 vision.aimodel/main.mlirb +f1f79a62132c7a424567bded352478dbe872d8bf92b344866474919f8ee88126 vision.aimodel/metadata.json diff --git a/SOURCE.json b/SOURCE.json new file mode 100644 index 0000000000000000000000000000000000000000..745e2794b96b19b57c95ac017a071164f3da2d8b --- /dev/null +++ b/SOURCE.json @@ -0,0 +1,179 @@ +{ + "repo": "mukaddamzaid/amos-1-medium", + "revision": "6b48a91ed940e08c87dcdd29c71c3b256d757ec7", + "bytes": 19259546893, + "files": [ + { + "path": ".gitattributes", + "size": 1637, + "lfs": null + }, + { + "path": "LICENSE", + "size": 11358, + "lfs": null + }, + { + "path": "README.md", + "size": 4427, + "lfs": null + }, + { + "path": "assets/amos-1-medium.png", + "size": 120578, + "lfs": { + "size": 120578, + "sha256": "b361e10452312a27c9e2a3238455961f5060a12a6310dd8f32cc6575e9ec0388", + "pointer_size": 131 + } + }, + { + "path": "calibration.json", + "size": 216, + "lfs": null + }, + { + "path": "config.json", + "size": 28, + "lfs": null + }, + { + "path": "model/chat_template.jinja", + "size": 7756, + "lfs": null + }, + { + "path": "model/config.json", + "size": 2877, + "lfs": null + }, + { + "path": "model/generation_config.json", + "size": 116, + "lfs": null + }, + { + "path": "model/heads.pt", + "size": 419551003, + "lfs": { + "size": 419551003, + "sha256": "5f38629976dc15322e24d3dfb6e48d918d5548c887739084eafef3c4671e3de6", + "pointer_size": 134 + } + }, + { + "path": "model/model.safetensors", + "size": 18819722392, + "lfs": { + "size": 18819722392, + "sha256": "4733fa24fd633dadd4f26e941b129ba7c607fdad278f4e2b2ad231b28a51572a", + "pointer_size": 136 + } + }, + { + "path": "model/preprocessor_config.json", + "size": 390, + "lfs": null + }, + { + "path": "model/tokenizer.json", + "size": 19989510, + "lfs": { + "size": 19989510, + "sha256": "aa7582a6b90dc3711049d293d985e549f2b22c0c4e5ad11b9272a03c6294886e", + "pointer_size": 133 + } + }, + { + "path": "model/tokenizer_config.json", + "size": 1173, + "lfs": null + }, + { + "path": "model/video_preprocessor_config.json", + "size": 385, + "lfs": null + }, + { + "path": "serve/README.md", + "size": 2505, + "lfs": null + }, + { + "path": "serve/amos/__init__.py", + "size": 42, + "lfs": null + }, + { + "path": "serve/amos/calibrate.py", + "size": 1712, + "lfs": null + }, + { + "path": "serve/amos/decoder.py", + "size": 43439, + "lfs": null + }, + { + "path": "serve/amos/encoder.py", + "size": 7816, + "lfs": null + }, + { + "path": "serve/amos/facts.py", + "size": 1356, + "lfs": null + }, + { + "path": "serve/amos/model.py", + "size": 15106, + "lfs": null + }, + { + "path": "serve/amos/packing.py", + "size": 10835, + "lfs": null + }, + { + "path": "serve/amos/py.typed", + "size": 0, + "lfs": null + }, + { + "path": "serve/amos/serve.py", + "size": 39697, + "lfs": null + }, + { + "path": "serve/amos/vision.py", + "size": 6633, + "lfs": null + }, + { + "path": "serve/pyproject.toml", + "size": 842, + "lfs": null + }, + { + "path": "serve/sources.json", + "size": 1723, + "lfs": null + }, + { + "path": "typed_decisions.json", + "size": 1341, + "lfs": null + } + ], + "local_source_hashes": { + "model/model.safetensors": "4733fa24fd633dadd4f26e941b129ba7c607fdad278f4e2b2ad231b28a51572a", + "model/tokenizer_config.json": "13ef66b39ff28a4bb50d8aaf554dbd01315b4497136b93111e5021ac6396db0f", + "model/config.json": "b647a20915503035f0e9a906980b891567bc04085fde26759016a05bf4445227", + "model/tokenizer.json": "aa7582a6b90dc3711049d293d985e549f2b22c0c4e5ad11b9272a03c6294886e", + "model/heads.pt": "5f38629976dc15322e24d3dfb6e48d918d5548c887739084eafef3c4671e3de6", + "model/generation_config.json": "62153eb6c69f2e1f426beaa8002b7186437e949c7588167085df14e10e9c0a73", + "model/chat_template.jinja": "a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715", + "model/video_preprocessor_config.json": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13", + "model/preprocessor_config.json": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516" + } +} \ No newline at end of file diff --git a/bf16-attempt.json b/bf16-attempt.json new file mode 100644 index 0000000000000000000000000000000000000000..986d254e40c55655ed24f26f9224b1514ec02637 --- /dev/null +++ b/bf16-attempt.json @@ -0,0 +1,181 @@ +[ + { + "case": "refund", + "tokens": 122, + "reference_precision": "bfloat16", + "upstream": [ + { + "type": "noul", + "p_yes": 0.9697602987289429 + }, + { + "type": "choice", + "probs": [ + 0.9355213642120361, + 0.02544301003217697, + 0.03903552144765854 + ], + "choice": 0 + } + ] + }, + { + "case": "span", + "tokens": 123, + "reference_precision": "bfloat16", + "upstream": [ + { + "type": "span", + "start": 55, + "end": 63, + "text": "INV-2231", + "p": 1.0 + } + ] + }, + { + "case": "score", + "tokens": 98, + "reference_precision": "bfloat16", + "upstream": [ + { + "type": "score", + "probs": [ + 0.01346411369740963, + 0.021021094173192978, + 0.13800419867038727, + 0.8275105953216553 + ], + "choice": 3, + "expected": 2.7795612774789333 + } + ] + }, + { + "case": "multi", + "tokens": 144, + "reference_precision": "bfloat16", + "upstream": [ + { + "type": "multi", + "probs": [ + 0.961353600025177, + 0.9255949854850769, + 0.05097513645887375 + ] + } + ] + }, + { + "case": "negative", + "tokens": 83, + "reference_precision": "bfloat16", + "upstream": [ + { + "type": "noul", + "p_yes": 0.037488047033548355 + } + ] + }, + { + "case": "unicode", + "tokens": 128, + "reference_precision": "bfloat16", + "upstream": [ + { + "type": "span", + "start": 10, + "end": 17, + "text": "\u20ac123.45", + "p": 1.0 + }, + { + "type": "choice", + "probs": [ + 0.9511587619781494, + 0.027312880381941795, + 0.021528292447328568 + ], + "choice": 0 + } + ] + }, + { + "case": "options", + "tokens": 230, + "reference_precision": "bfloat16", + "upstream": [ + { + "type": "choice", + "probs": [ + 0.003604995086789131, + 0.003263695864006877, + 0.002579145599156618, + 0.0032084090635180473, + 0.0021689790301024914, + 0.002523764269426465, + 0.0024777702055871487, + 0.003174568759277463, + 0.0037178073544055223, + 0.0025443080812692642, + 0.003257984761148691, + 0.0031456719152629375, + 0.0022962531074881554, + 0.003111874219030142, + 0.0029890919104218483, + 0.003832802874967456, + 0.007152783218771219, + 0.922189474105835, + 0.005749549716711044, + 0.004194738809019327, + 0.0036519865971058607, + 0.003285554703325033, + 0.002670723246410489, + 0.003208014415577054 + ], + "choice": 17 + } + ] + }, + { + "case": "vision", + "tokens": 155, + "reference_precision": "bfloat16", + "upstream": [ + { + "type": "region", + "cell": 21, + "box": [ + 0.2459, + 0.0897, + 0.6864, + 0.9191 + ], + "p": 0.6435877084732056 + }, + { + "type": "choice", + "probs": [ + 0.9735397100448608, + 0.026460271328687668 + ], + "choice": 0 + } + ] + }, + { + "case": "image-option", + "tokens": 181, + "reference_precision": "bfloat16", + "upstream": [ + { + "type": "choice", + "probs": [ + 0.9695415496826172, + 0.030458426102995872 + ], + "choice": 0 + } + ] + } +] \ No newline at end of file diff --git a/bf16-comparison.json b/bf16-comparison.json new file mode 100644 index 0000000000000000000000000000000000000000..57e2bae937fe292926c9dec61ea0fbe79454fe07 --- /dev/null +++ b/bf16-comparison.json @@ -0,0 +1,315 @@ +[ + { + "case": "refund", + "max_probability_or_box_error": 0.0011562108993530273, + "discrete_outputs_match": true, + "upstream_bf16": [ + { + "type": "noul", + "p_yes": 0.9697602987289429 + }, + { + "type": "choice", + "probs": [ + 0.9355213642120361, + 0.02544301003217697, + 0.03903552144765854 + ], + "choice": 0 + } + ], + "coreai_float32": [ + { + "type": "noul", + "p_yes": 0.9696579575538635 + }, + { + "type": "choice", + "probs": [ + 0.9343651533126831, + 0.025729088112711906, + 0.03990578278899193 + ], + "choice": 0 + } + ] + }, + { + "case": "span", + "max_probability_or_box_error": 0, + "discrete_outputs_match": true, + "upstream_bf16": [ + { + "type": "span", + "start": 55, + "end": 63, + "text": "INV-2231", + "p": 1.0 + } + ], + "coreai_float32": [ + { + "type": "span", + "start": 55, + "end": 63, + "text": "INV-2231", + "p": 1.0 + } + ] + }, + { + "case": "score", + "max_probability_or_box_error": 0.0032760053873062134, + "discrete_outputs_match": true, + "upstream_bf16": [ + { + "type": "score", + "probs": [ + 0.01346411369740963, + 0.021021094173192978, + 0.13800419867038727, + 0.8275105953216553 + ], + "choice": 3, + "expected": 2.7795612774789333 + } + ], + "coreai_float32": [ + { + "type": "score", + "probs": [ + 0.013492124155163765, + 0.02104916237294674, + 0.13472819328308105, + 0.8307304382324219 + ], + "choice": 3, + "expected": 2.7826968636363745 + } + ] + }, + { + "case": "multi", + "max_probability_or_box_error": 0.0007659792900085449, + "discrete_outputs_match": true, + "upstream_bf16": [ + { + "type": "multi", + "probs": [ + 0.961353600025177, + 0.9255949854850769, + 0.05097513645887375 + ] + } + ], + "coreai_float32": [ + { + "type": "multi", + "probs": [ + 0.9614582657814026, + 0.9263609647750854, + 0.051194678992033005 + ] + } + ] + }, + { + "case": "negative", + "max_probability_or_box_error": 0.0009359195828437805, + "discrete_outputs_match": true, + "upstream_bf16": [ + { + "type": "noul", + "p_yes": 0.037488047033548355 + } + ], + "coreai_float32": [ + { + "type": "noul", + "p_yes": 0.036552127450704575 + } + ] + }, + { + "case": "unicode", + "max_probability_or_box_error": 0.00025212764739990234, + "discrete_outputs_match": true, + "upstream_bf16": [ + { + "type": "span", + "start": 10, + "end": 17, + "text": "\u20ac123.45", + "p": 1.0 + }, + { + "type": "choice", + "probs": [ + 0.9511587619781494, + 0.027312880381941795, + 0.021528292447328568 + ], + "choice": 0 + } + ], + "coreai_float32": [ + { + "type": "span", + "start": 10, + "end": 17, + "text": "\u20ac123.45", + "p": 1.0 + }, + { + "type": "choice", + "probs": [ + 0.9514108896255493, + 0.027261406183242798, + 0.021327782422304153 + ], + "choice": 0 + } + ] + }, + { + "case": "options", + "max_probability_or_box_error": 0.0006901025772094727, + "discrete_outputs_match": true, + "upstream_bf16": [ + { + "type": "choice", + "probs": [ + 0.003604995086789131, + 0.003263695864006877, + 0.002579145599156618, + 0.0032084090635180473, + 0.0021689790301024914, + 0.002523764269426465, + 0.0024777702055871487, + 0.003174568759277463, + 0.0037178073544055223, + 0.0025443080812692642, + 0.003257984761148691, + 0.0031456719152629375, + 0.0022962531074881554, + 0.003111874219030142, + 0.0029890919104218483, + 0.003832802874967456, + 0.007152783218771219, + 0.922189474105835, + 0.005749549716711044, + 0.004194738809019327, + 0.0036519865971058607, + 0.003285554703325033, + 0.002670723246410489, + 0.003208014415577054 + ], + "choice": 17 + } + ], + "coreai_float32": [ + { + "type": "choice", + "probs": [ + 0.0035737575963139534, + 0.0032294215634465218, + 0.0025455907452851534, + 0.0031847632490098476, + 0.002165550133213401, + 0.002492170315235853, + 0.0024529574438929558, + 0.003147532232105732, + 0.0037047737278044224, + 0.0024940771982073784, + 0.0032223134767264128, + 0.003116679610684514, + 0.002279103733599186, + 0.003067179350182414, + 0.0030036454554647207, + 0.0037965381052345037, + 0.007012164685875177, + 0.9228795766830444, + 0.0057253348641097546, + 0.004146119579672813, + 0.00362235214561224, + 0.0032782102935016155, + 0.002652023686096072, + 0.0032081108074635267 + ], + "choice": 17 + } + ] + }, + { + "case": "vision", + "max_probability_or_box_error": 0.10930000000000001, + "discrete_outputs_match": false, + "upstream_bf16": [ + { + "type": "region", + "cell": 21, + "box": [ + 0.2459, + 0.0897, + 0.6864, + 0.9191 + ], + "p": 0.6435877084732056 + }, + { + "type": "choice", + "probs": [ + 0.9735397100448608, + 0.026460271328687668 + ], + "choice": 0 + } + ], + "coreai_float32": [ + { + "type": "region", + "cell": 33, + "box": [ + 0.3552, + 0.0785, + 0.747, + 0.9124 + ], + "p": 0.6444705128669739 + }, + { + "type": "choice", + "probs": [ + 0.974250316619873, + 0.025749672204256058 + ], + "choice": 0 + } + ] + }, + { + "case": "image-option", + "max_probability_or_box_error": 0.000291287899017334, + "discrete_outputs_match": true, + "upstream_bf16": [ + { + "type": "choice", + "probs": [ + 0.9695415496826172, + 0.030458426102995872 + ], + "choice": 0 + } + ], + "coreai_float32": [ + { + "type": "choice", + "probs": [ + 0.9698328375816345, + 0.03016720339655876 + ], + "choice": 0 + } + ] + } +] \ No newline at end of file diff --git a/calibration.json b/calibration.json new file mode 100644 index 0000000000000000000000000000000000000000..ae01a38c3859485083c401fa348d347780c912af --- /dev/null +++ b/calibration.json @@ -0,0 +1,7 @@ +{ + "span": { + "cap": 128, + "penalty": 1.6, + "fitted_on": "held-out quotes: the quote-amounts slice and longquotes clean.jsonl (fit halves), 2026-10-08; cap raised, penalty kept so short quotes are unchanged" + } +} diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..cda7454523650530721b0b78c17f578bdac450c1 --- /dev/null +++ b/config.json @@ -0,0 +1 @@ +{"base": "Qwen/Qwen3.5-9B"} diff --git a/convert_decoder.py b/convert_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..9ea7919f69dbe13858e3e181be1805e4cffcbfed --- /dev/null +++ b/convert_decoder.py @@ -0,0 +1,49 @@ +"""Export AMOS Qwen3.5 bidirectional decisions and trained quote/box heads.""" +import argparse,json,math,sys +from pathlib import Path +import torch +from coreai_torch import TorchConverter,get_decomp_table +from qwen_export_helpers import DecisionVision + +class Text(torch.nn.Module): + def __init__(self,model,letters): + super().__init__();self.language_model=model.model.language_model;self.noncausal=bool(getattr(model.config,'amos_noncausal',False)) + self.register_buffer('readout',model.get_output_embeddings().weight[letters].detach().float().clone()) + def forward(self,input_ids,attention_mask,position_ids,visual_embeddings,visual_mask): + emb=self.language_model.embed_tokens(input_ids) + emb=torch.where(visual_mask,visual_embeddings.to(emb.dtype),emb) + mask=attention_mask[:,None,None,:].bool() + if not self.noncausal: + width=input_ids.shape[1];idx=torch.arange(width,device=input_ids.device);mask=mask&(idx[:,None]>=idx[None,:])[None,None] + h=self.language_model(inputs_embeds=emb,attention_mask={'full_attention':mask,'linear_attention':None},position_ids=position_ids,use_cache=False).last_hidden_state + return h.float(),h.float()@self.readout.T + +class Heads(torch.nn.Module): + def __init__(self,heads):super().__init__();self.heads=heads + def forward(self,question_hidden,state_hidden,image_hidden): + h=self.heads;d=question_hidden.shape[0];dtype=h.span_q.weight.dtype + q=question_hidden.to(dtype);s=state_hidden.to(dtype);im=image_hidden.to(dtype) + span=torch.einsum('cd,lcd->cl',h.span_q(q).view(2,d),h.span_k(s).view(-1,2,d))/math.sqrt(d) + region=h.region_k(im)@h.region_q(q)/math.sqrt(d) + box=h.box(q,im,region,7,7) if h.box is not None else torch.zeros(4) + return span.float(),region.float(),box.float() + +def save(m,args,path): + with torch.no_grad():ep=torch.export.export(m,args,strict=False).run_decompositions(get_decomp_table()) + program=TorchConverter().add_exported_program(ep).to_coreai();program.save_asset(path) + return list(ep.graph_signature.user_outputs) + +def main(): + p=argparse.ArgumentParser();p.add_argument('name');p.add_argument('--length',type=int,default=512);p.add_argument('--precision',default='bfloat16');a=p.parse_args() + source=Path('upstream')/a.name;out=Path('bundles')/(a.name+'-CoreAI');out.mkdir(parents=True,exist_ok=True) + sys.path.insert(0,str((source/'serve').resolve())) + from amos.decoder import load_model,load_heads + torch.set_num_threads(4);model,codec=load_model(str(source/'model'),'cpu',dtype=getattr(torch,a.precision));model.eval();heads=load_heads(source/'model',model,codec) + d=model.config.text_config.hidden_size;L=a.length + inputs=(torch.full((1,L),codec.pad,dtype=torch.int32),torch.ones(1,L,dtype=torch.int32),torch.arange(L,dtype=torch.int32).view(1,1,-1).expand(4,1,-1),torch.zeros(1,L,d),torch.zeros(1,L,1,dtype=torch.bool)) + print('Exporting text',flush=True);outputs=save(Text(model,codec.letters).eval(),inputs,out/'text.aimodel') + print('Exporting vision',flush=True);visual=model.model.visual;vo=save(DecisionVision(visual).eval(),(torch.zeros(196,1536),),out/'vision.aimodel') + print('Exporting heads',flush=True);ho=save(Heads(heads).eval(),(torch.zeros(d),torch.zeros(L,d),torch.zeros(49,d)),out/'heads.aimodel') + (out/'manifest.json').write_text(json.dumps({'sequence_length':L,'hidden_size':d,'outputs':outputs,'vision_output':vo[0],'head_outputs':ho,'precision':a.precision,'echo':codec.echo,'kinds':heads.kinds,'keep':heads.keep,'has_box':heads.box is not None,'image_side':224,'noncausal':bool(getattr(model.config,'amos_noncausal',False))},indent=2)) + print('Saved',flush=True) +if __name__=='__main__':main() diff --git a/convert_medium_sharded.py b/convert_medium_sharded.py new file mode 100644 index 0000000000000000000000000000000000000000..dc017b226a6c40ab2a5c03af2eabd66acc7a4eb4 --- /dev/null +++ b/convert_medium_sharded.py @@ -0,0 +1,54 @@ +"""Export FP32 AMOS medium one component at a time to bound conversion memory.""" +import gc,json,sys +from pathlib import Path +import torch +from convert_decoder import save,Heads +from qwen_export_helpers import DecisionVision + +class Embedding(torch.nn.Module): + def __init__(self,decoder):super().__init__();self.embed_tokens=decoder.embed_tokens;self.rotary=decoder.rotary_emb + def forward(self,input_ids,position_ids,visual_embeddings,visual_mask): + h=torch.where(visual_mask,visual_embeddings,self.embed_tokens(input_ids)) + cos,sin=self.rotary(h,position_ids[1:]);return h,cos,sin + +class Layer(torch.nn.Module): + def __init__(self,layer):super().__init__();self.layer=layer + def forward(self,hidden_states,cos,sin,position_ids,attention_mask): + return self.layer(hidden_states,position_embeddings=(cos,sin),position_ids=position_ids,attention_mask=attention_mask if self.layer.block_type=='full_attention' else None,use_cache=False) + +class Finish(torch.nn.Module): + def __init__(self,model,letters): + super().__init__();self.norm=model.model.language_model.norm + self.register_buffer('readout',model.get_output_embeddings().weight[letters].detach().float().clone()) + def forward(self,hidden_states): + h=self.norm(hidden_states);return h,h@self.readout.T + +def main(): + source=Path('upstream/amos-1-medium');root=Path('bundles/amos-1-medium-CoreAI');root.mkdir(exist_ok=True,parents=True) + sys.path.insert(0,str((source/'serve').resolve())) + from amos.decoder import load_model,load_heads + torch.set_num_threads(4) + m,codec=load_model(str(source/'model'),'cpu',dtype=torch.bfloat16);m.eval();heads=load_heads(source/'model',m,codec) + from stable_delta import install + install() + decoder=m.model.language_model;L=512;d=decoder.config.hidden_size + manifest={'sequence_length':L,'hidden_size':d,'precision':'float32','vision_precision':'float32','echo':codec.echo,'kinds':heads.kinds,'keep':heads.keep,'has_box':heads.box is not None,'image_side':224,'noncausal':False,'text_layout':'layer-sharded','layers':[]} + path=root/'embedding.aimodel' + decoder.embed_tokens.float();emb=Embedding(decoder).eval() + args=(torch.ones(1,L,dtype=torch.int32),torch.arange(L,dtype=torch.int32).view(1,1,-1).expand(4,1,-1),torch.zeros(1,L,d),torch.zeros(1,L,1,dtype=torch.bool)) + with torch.no_grad():_,cos,sin=emb(*args) + print('Embedding',flush=True);manifest['embedding_outputs']=save(emb,args,path) + decoder.embed_tokens.to('meta');del emb;gc.collect() + layer_args=(torch.zeros(1,L,d),cos,sin,torch.arange(L,dtype=torch.int32)[None],torch.ones(1,1,L,L,dtype=torch.bool).tril()) + for i,layer in enumerate(decoder.layers): + name=f'layer-{i:02}.aimodel';print(name,flush=True) + layer.float().eval();outputs=save(Layer(layer).eval(),layer_args,root/name) + manifest['layers'].append({'asset':name,'output':outputs[0],'kind':layer.block_type}) + layer.to('meta');gc.collect() + print('Finish',flush=True);decoder.norm.float();manifest['outputs']=save(Finish(m,codec.letters).eval(),(torch.zeros(1,L,d),),root/'finish.aimodel') + m.get_output_embeddings().to('meta');decoder.norm.to('meta');gc.collect() + print('Vision',flush=True);m.model.visual.float();manifest['vision_output']=save(DecisionVision(m.model.visual).eval(),(torch.zeros(196,1536),),root/'vision.aimodel')[0] + m.model.visual.to('meta');gc.collect() + print('Heads',flush=True);manifest['head_outputs']=save(Heads(heads).eval(),(torch.zeros(d),torch.zeros(L,d),torch.zeros(49,d)),root/'heads.aimodel') + (root/'manifest.json').write_text(json.dumps(manifest,indent=2));print('Saved FP32 sharded model',flush=True) +if __name__=='__main__':main() diff --git a/coreai_decoder.py b/coreai_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..ffbd2f722e40184d534262ee4452a616e4fd2e28 --- /dev/null +++ b/coreai_decoder.py @@ -0,0 +1,100 @@ +"""AMOS Qwen3.5 Core AI adapter. No PyTorch neural model is loaded.""" +import json,sys +from pathlib import Path +import numpy as np +import torch +from coreai.runtime import AIModel,NDArray,SpecializationOptions +from transformers import AutoTokenizer,Qwen2VLImageProcessorPil + +class AmosDecoder: + @classmethod + async def load(cls,bundle): + self=cls();self.root=Path(bundle).resolve();sys.path.insert(0,str(self.root/'serve')) + import amos.decoder as decoder + decoder.IMAGE_PIXELS=224*224 + self.manifest=json.loads((self.root/'manifest.json').read_text()) + self.codec=decoder.Codec(AutoTokenizer.from_pretrained(self.root/'model'),Qwen2VLImageProcessorPil.from_pretrained(self.root/'model',size={'shortest_edge':224*224,'longest_edge':224*224}),255) + self.codec.echo=self.manifest['echo'] + self.temperature=decoder.calibration(self.root) or {};self.span=decoder.span_decoding(self.root) + self.models=[];self.functions={} + names=['vision','heads'] + if self.manifest.get('text_layout')=='layer-sharded':names+=['embedding','finish']+[r['asset'].removesuffix('.aimodel') for r in self.manifest['layers']] + else:names+=['text'] + for name in names: + m=await AIModel.load(self.root/(name+'.aimodel'),SpecializationOptions.default());self.models.append(m);self.functions[name]=m.load_function('main') + return self + + async def raw(self,state,questions,images=()): + from amos.decoder import expand + from PIL import Image + def resize(im):return im.convert('RGB').resize((224,224),Image.Resampling.BICUBIC) + images=[resize(im) for im in images] + questions=[{**q,'options':[resize(o) if isinstance(o,Image.Image) else o for o in q['options']]} if 'options' in q else q for q in questions] + flat,owner=expand(questions) + ids,at,extra=self.codec.encode_case(state,flat,images) + n=len(ids);L=self.manifest['sequence_length'];d=self.manifest['hidden_size'] + if n>L:raise ValueError(f'Request needs {n} tokens; Core AI capacity is {L}.') + args={'input_ids':np.full((1,L),self.codec.pad,np.int32),'attention_mask':np.zeros((1,L),np.int32),'position_ids':np.zeros((4,1,L),np.int32),'visual_embeddings':np.zeros((1,L,d),np.float32),'visual_mask':np.zeros((1,L,1),bool)} + args['input_ids'][0,:n]=ids;args['attention_mask'][0,:n]=1;args['position_ids'][0,0,:n]=np.arange(n) + image_at={s:(rows,cols) for s,rows,cols in extra.get('images',[])} + features=[] + if image_at: + pv=extra['pixel_values'].numpy() + for start in range(0,len(pv),196): + out=await self.functions['vision']({'pixel_values':NDArray(pv[start:start+196].astype(np.float32))});features.append(out[self.manifest['vision_output']].numpy()) + cursor=0;offset=0;fi=0 + while cursor=np.arange(L)[None,:])[None,None]&args['attention_mask'][:,None,None,:].astype(bool) + for layer in self.manifest['layers']: + result=await self.functions[layer['asset'].removesuffix('.aimodel')]({'hidden_states':h,'cos':cos,'sin':sin,'position_ids':NDArray(args['position_ids'][0].copy()),'attention_mask':NDArray(causal)}) + h=result[layer['output']] + output=await self.functions['finish']({'hidden_states':h}) + else:output=await self.functions['text']({k:NDArray(v) for k,v in args.items()}) + hidden,logits=[output[k].numpy() for k in self.manifest['outputs']] + head_results={} + for i,(q,pos) in enumerate(zip(flat,at)): + if q['type'] not in ['span','region']:continue + first,offsets=extra['state'];state_hidden=np.zeros((L,d),np.float32);state_hidden[:len(offsets)]=hidden[0,first:first+len(offsets)] + cells=np.zeros((49,d),np.float32) + if extra.get('images'): + start,rows,cols=extra['images'][0];cells=hidden[0,start:start+49].copy() + if q['type']=='region' and not extra.get('images'):raise ValueError('Region requires a state image') + ho=await self.functions['heads']({'question_hidden':NDArray(hidden[0,pos].copy()),'state_hidden':NDArray(state_hidden),'image_hidden':NDArray(cells)}) + span,region,box=[torch.from_numpy(ho[k].numpy().copy()) for k in self.manifest['head_outputs']] + head_results[i]=span[:,:len(offsets)] if q['type']=='span' else torch.cat([region,box]) if self.manifest['has_box'] else region + return {'flat':flat,'owner':owner,'at':at,'extra':extra,'logits':logits[0],'heads':head_results,'tokens':n,'hidden':hidden[0]} + + async def predict(self,state,questions,images=()): + from amos.decoder import options,decode_head + raw=await self.raw(state,questions,images);out=[None]*len(questions) + for j,(q,k,pos) in enumerate(zip(raw['flat'],raw['owner'],raw['at'])): + typ=questions[k]['type'] + if q['type'] in ['span','region']: + out[k]=decode_head(q,raw['heads'][j],raw['extra'],state,self.manifest['keep'],self.span);continue + logits=torch.tensor(raw['logits'][pos,:len(options(q))]);p=(logits/self.temperature.get(q['type'],1)).softmax(-1).tolist() + if typ=='multi':out[k]=out[k] or {'type':typ,'probs':[]};out[k]['probs'].append(p[0]) + elif typ=='noul':out[k]={'type':typ,'p_yes':p[0]} + else: + out[k]={'type':typ,'probs':p,'choice':int(np.argmax(p))} + if typ=='score':out[k]['expected']=sum(i*v for i,v in enumerate(p)) + return out + + async def request(self,req): + from amos.serve import parse,answer + from PIL import Image + import base64,io + def image(s):return Image.open(io.BytesIO(base64.b64decode(s.split(',')[-1]))).convert('RGB') + state,ids,qs=parse(req);ims=[image(s) for s in req.get('images',[])] + qs=[{**q,'options':[image(o['image']) if isinstance(o,dict) else o for o in q['options']]} if 'options' in q else q for q in qs] + rs=await self.predict(state,qs,ims) + return {'answers':{i:answer(q,r) for i,q,r in zip(ids,qs,rs)}} diff --git a/embedding.aimodel/main.hash b/embedding.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..aa91eaf249f9484a835716202b7baf64f717bf4b --- /dev/null +++ b/embedding.aimodel/main.hash @@ -0,0 +1 @@ +o,vj\}ixRw/רJt \ No newline at end of file diff --git a/embedding.aimodel/main.mlirb b/embedding.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..847162b5b62d4ef51ee7afcaa8aaa726d3486b56 --- /dev/null +++ b/embedding.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6fcec91d2cdc76e76a855ccc7dc61d10697852772f81b6d7a8d94af4ac897407 +size 4068486672 diff --git a/embedding.aimodel/metadata.json b/embedding.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..06ac3a3cea40abf0ec3773f4ee85ef8a08e559a9 --- /dev/null +++ b/embedding.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "creationDate" : "20261014T014939Z", + "producer" : "coreai-core 1.0.0b3", + "assetVersion" : "2.0" +} \ No newline at end of file diff --git a/finish.aimodel/main.hash b/finish.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..50b035909d1f103aa829defb6e746d084aeda78a --- /dev/null +++ b/finish.aimodel/main.hash @@ -0,0 +1 @@ +@GlWvyH rbbia| \ No newline at end of file diff --git a/finish.aimodel/main.mlirb b/finish.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..430376bdbe43743a07170d33fc6d5b4afea6de98 --- /dev/null +++ b/finish.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:40ed1efe47f26cc15794ce7616798dd2f948207262fd62041b6961e381f37c15 +size 4199075 diff --git a/finish.aimodel/metadata.json b/finish.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..c09e16ec9534b5a5d2c3ababc396f77388d2ad1a --- /dev/null +++ b/finish.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "assetVersion" : "2.0", + "producer" : "coreai-core 1.0.0b3", + "creationDate" : "20261014T015037Z" +} \ No newline at end of file diff --git a/heads.aimodel/main.hash b/heads.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..0e0527e3345b1181526b37874720817468a5e0dc --- /dev/null +++ b/heads.aimodel/main.hash @@ -0,0 +1 @@ +ESaWfSDhr܁AňT' \ No newline at end of file diff --git a/heads.aimodel/main.mlirb b/heads.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..ce1a957d7ffa27ac30755b5d92bede80be13e3e7 --- /dev/null +++ b/heads.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:455361e9ff9557669def5344f8db6872e1bdcddc81e41141c5888cce54197f27 +size 419568833 diff --git a/heads.aimodel/metadata.json b/heads.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..d583574a66966603c473dd6b3f1e8fa5943eed8d --- /dev/null +++ b/heads.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "producer" : "coreai-core 1.0.0b3", + "assetVersion" : "2.0", + "creationDate" : "20261014T015042Z" +} \ No newline at end of file diff --git a/layer-00.aimodel/main.hash b/layer-00.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..8cbbfe644cdc3a1e8edf9a4c6b0ff718885d6626 --- /dev/null +++ b/layer-00.aimodel/main.hash @@ -0,0 +1 @@ ++]ǖf.FCzdLxo \ No newline at end of file diff --git a/layer-22.aimodel/main.mlirb b/layer-22.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..7e7f725310cce27a12c28bbbc15fce812d36359a --- /dev/null +++ b/layer-22.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8e35f9ef221fee278b0e7df2ac3e4643bfb77ae3644c786f9ed003ee9c10a30f +size 873853565 diff --git a/layer-22.aimodel/metadata.json b/layer-22.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..613d20c9a0f3d3f6bbe73e19048dd6343b05132b --- /dev/null +++ b/layer-22.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "assetVersion" : "2.0", + "producer" : "coreai-core 1.0.0b3", + "creationDate" : "20261014T015022Z" +} \ No newline at end of file diff --git a/layer-23.aimodel/main.hash b/layer-23.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..c39c9fa153448b6f55bd9f84aea5d68029d671f9 --- /dev/null +++ b/layer-23.aimodel/main.hash @@ -0,0 +1 @@ ++XOo~3 0zho!ɤϕZ \ No newline at end of file diff --git a/layer-23.aimodel/main.mlirb b/layer-23.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..751ee6e8f90efe22f9cad4322cea54445bf5d00c --- /dev/null +++ b/layer-23.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2b58ef90e64fe06fe3b2127eee0705331f0b307a686f21f99ac9a4e6cf955abe +size 838922540 diff --git a/layer-23.aimodel/metadata.json b/layer-23.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..a3e3e224cdd56f520799aa6981ecec91a68fb3e2 --- /dev/null +++ b/layer-23.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "assetVersion" : "2.0", + "producer" : "coreai-core 1.0.0b3", + "creationDate" : "20261014T015023Z" +} \ No newline at end of file diff --git a/layer-24.aimodel/main.hash b/layer-24.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..947a4929c98de410731abf32e11532a24d0beaca --- /dev/null +++ b/layer-24.aimodel/main.hash @@ -0,0 +1 @@ +߶n$7u`*CfNx큣 %z \ No newline at end of file diff --git a/layer-24.aimodel/main.mlirb b/layer-24.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..8838afd603d840f5db5aaa5acd6fd8dedc5171d8 --- /dev/null +++ b/layer-24.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dfb66edc24171137f37516602a90c143b966f0a74e7f78ed81a30b9bd0257a89 +size 873853572 diff --git a/layer-24.aimodel/metadata.json b/layer-24.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..c744149e4a3315638406af7d43b17e9963170be3 --- /dev/null +++ b/layer-24.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "creationDate" : "20261014T015025Z", + "producer" : "coreai-core 1.0.0b3", + "assetVersion" : "2.0" +} \ No newline at end of file diff --git a/layer-25.aimodel/main.hash b/layer-25.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..a47e529b3b9fe50a11e7a620e77149914476099e --- /dev/null +++ b/layer-25.aimodel/main.hash @@ -0,0 +1 @@ +7mc['hx|w36"f ^=% 5$Γ>"onW4{\ \ No newline at end of file diff --git a/layer-28.aimodel/main.mlirb b/layer-28.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..09ea72a2477230c7db11f5aea38168a1cd9141ac --- /dev/null +++ b/layer-28.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0f3e200196185e3dcb0125ef0d3524a5ce93143e226f6ea357b790347baf5c0b +size 873853568 diff --git a/layer-28.aimodel/metadata.json b/layer-28.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..9256d45d58d224bebfee8607e476b8e8f33bc3b1 --- /dev/null +++ b/layer-28.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "assetVersion" : "2.0", + "producer" : "coreai-core 1.0.0b3", + "creationDate" : "20261014T015032Z" +} \ No newline at end of file diff --git a/layer-29.aimodel/main.hash b/layer-29.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..184e57c8ccb6b9299a236717f14a26414bd09a92 Binary files /dev/null and b/layer-29.aimodel/main.hash differ diff --git a/layer-29.aimodel/main.mlirb b/layer-29.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..56045837af58204083337730089764377389bc57 --- /dev/null +++ b/layer-29.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fbca7a309b0090b05e5df56eab1cd7cff96019f08453d23f5ab476bf2f4b6069 +size 873853571 diff --git a/layer-29.aimodel/metadata.json b/layer-29.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..896603bc0df4847b29e690bb2b340826bb191833 --- /dev/null +++ b/layer-29.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "creationDate" : "20261014T015034Z", + "producer" : "coreai-core 1.0.0b3", + "assetVersion" : "2.0" +} \ No newline at end of file diff --git a/layer-30.aimodel/main.hash b/layer-30.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..339bd854abe9e3915f2cd464dab9141ae520269b --- /dev/null +++ b/layer-30.aimodel/main.hash @@ -0,0 +1 @@ + -4r}!:G'8u/ê \ No newline at end of file diff --git a/layer-30.aimodel/main.mlirb b/layer-30.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..acbc257117d3a733c8f29d36163bbc596f5eda91 --- /dev/null +++ b/layer-30.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8f0c2d9df134e7721b8a7d819be51c0621803a47eb27a738ad75192f04c4c3aa +size 873853571 diff --git a/layer-30.aimodel/metadata.json b/layer-30.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..95d6362458f730e6158ba2045f5eeb0c2e57be57 --- /dev/null +++ b/layer-30.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "assetVersion" : "2.0", + "producer" : "coreai-core 1.0.0b3", + "creationDate" : "20261014T015036Z" +} \ No newline at end of file diff --git a/layer-31.aimodel/main.hash b/layer-31.aimodel/main.hash new file mode 100644 index 0000000000000000000000000000000000000000..869935bf7fa0eb094f3d5f6daf5a4b139b306823 --- /dev/null +++ b/layer-31.aimodel/main.hash @@ -0,0 +1 @@ +\rVخnFEDFA+{mA \ No newline at end of file diff --git a/layer-31.aimodel/main.mlirb b/layer-31.aimodel/main.mlirb new file mode 100644 index 0000000000000000000000000000000000000000..f294bf61029c05a84c7273f230a706177dbbaa2d --- /dev/null +++ b/layer-31.aimodel/main.mlirb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bd1e995c957256ffd8ae6e46e045f74446aaad41d92bf2d87b8e6d1c41d21f10 +size 838922541 diff --git a/layer-31.aimodel/metadata.json b/layer-31.aimodel/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..c09e16ec9534b5a5d2c3ababc396f77388d2ad1a --- /dev/null +++ b/layer-31.aimodel/metadata.json @@ -0,0 +1,5 @@ +{ + "assetVersion" : "2.0", + "producer" : "coreai-core 1.0.0b3", + "creationDate" : "20261014T015037Z" +} \ No newline at end of file diff --git a/manifest.json b/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..420fb387a09c2b6792f6d41d36a2409f90478ad0 --- /dev/null +++ b/manifest.json @@ -0,0 +1,193 @@ +{ + "sequence_length": 512, + "hidden_size": 4096, + "precision": "float32", + "vision_precision": "float32", + "echo": true, + "kinds": [ + "span", + "region" + ], + "keep": 0.1, + "has_box": true, + "image_side": 224, + "noncausal": false, + "text_layout": "layer-sharded", + "layers": [ + { + "asset": "layer-00.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-01.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-02.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-03.aimodel", + "output": "add_11", + "kind": "full_attention" + }, + { + "asset": "layer-04.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-05.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-06.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-07.aimodel", + "output": "add_11", + "kind": "full_attention" + }, + { + "asset": "layer-08.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-09.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-10.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-11.aimodel", + "output": "add_11", + "kind": "full_attention" + }, + { + "asset": "layer-12.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-13.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-14.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-15.aimodel", + "output": "add_11", + "kind": "full_attention" + }, + { + "asset": "layer-16.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-17.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-18.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-19.aimodel", + "output": "add_11", + "kind": "full_attention" + }, + { + "asset": "layer-20.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-21.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-22.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-23.aimodel", + "output": "add_11", + "kind": "full_attention" + }, + { + "asset": "layer-24.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-25.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-26.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-27.aimodel", + "output": "add_11", + "kind": "full_attention" + }, + { + "asset": "layer-28.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-29.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-30.aimodel", + "output": "add_25", + "kind": "linear_attention" + }, + { + "asset": "layer-31.aimodel", + "output": "add_11", + "kind": "full_attention" + } + ], + "embedding_outputs": [ + "where", + "cat_1", + "cat" + ], + "outputs": [ + "mul_1", + "view_1" + ], + "vision_output": "addmm_110", + "head_outputs": [ + "div", + "div_1", + "clamp_2" + ] +} \ No newline at end of file diff --git a/model/chat_template.jinja b/model/chat_template.jinja new file mode 100644 index 0000000000000000000000000000000000000000..a585dec894e63da457d9440ec6aa7caa16d20860 --- /dev/null +++ b/model/chat_template.jinja @@ -0,0 +1,154 @@ +{%- set image_count = namespace(value=0) %} +{%- set video_count = namespace(value=0) %} +{%- macro render_content(content, do_vision_count, is_system_content=false) %} + {%- if content is string %} + {{- content }} + {%- elif content is iterable and content is not mapping %} + {%- for item in content %} + {%- if 'image' in item or 'image_url' in item or item.type == 'image' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain images.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set image_count.value = image_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Picture ' ~ image_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|image_pad|><|vision_end|>' }} + {%- elif 'video' in item or item.type == 'video' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain videos.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set video_count.value = video_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Video ' ~ video_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|video_pad|><|vision_end|>' }} + {%- elif 'text' in item %} + {{- item.text }} + {%- else %} + {{- raise_exception('Unexpected item type in content.') }} + {%- endif %} + {%- endfor %} + {%- elif content is none or content is undefined %} + {{- '' }} + {%- else %} + {{- raise_exception('Unexpected content type.') }} + {%- endif %} +{%- endmacro %} +{%- if not messages %} + {{- raise_exception('No messages provided.') }} +{%- endif %} +{%- if tools and tools is iterable and tools is not mapping %} + {{- '<|im_start|>system\n' }} + {{- "# Tools\n\nYou have access to the following functions:\n\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n" }} + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n' }} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {%- if content %} + {{- '\n\n' + content }} + {%- endif %} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- else %} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" %} + {%- set content = render_content(message.content, false)|trim %} + {%- if not(content.startswith('') and content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if ns.multi_step_tool %} + {{- raise_exception('No user query found in messages.') }} +{%- endif %} +{%- for message in messages %} + {%- set content = render_content(message.content, true)|trim %} + {%- if message.role == "system" %} + {%- if not loop.first %} + {{- raise_exception('System message must be at the beginning.') }} + {%- endif %} + {%- elif message.role == "user" %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- set reasoning_content = reasoning_content|trim %} + {%- if loop.index0 > ns.last_query_index %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content + '\n\n\n' + content }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- if loop.first %} + {%- if content|trim %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n\n' }} + {%- endif %} + {%- else %} + {{- '\n\n\n' }} + {%- endif %} + {%- if tool_call.arguments is defined %} + {%- for args_name, args_value in tool_call.arguments|items %} + {{- '\n' }} + {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %} + {{- args_value }} + {{- '\n\n' }} + {%- endfor %} + {%- endif %} + {{- '\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.previtem and loop.previtem.role != "tool" %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if not loop.last and loop.nextitem.role != "tool" %} + {{- '<|im_end|>\n' }} + {%- elif loop.last %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- else %} + {{- raise_exception('Unexpected message role.') }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n' }} + {%- endif %} +{%- endif %} \ No newline at end of file diff --git a/model/config.json b/model/config.json new file mode 100644 index 0000000000000000000000000000000000000000..4dc901a3ba46a0a8a2aa5b36ca65bc9044063b8a --- /dev/null +++ b/model/config.json @@ -0,0 +1,110 @@ +{ + "architectures": [ + "Qwen3_5ForConditionalGeneration" + ], + "image_token_id": 248056, + "model_type": "qwen3_5", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "attn_output_gate": true, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_value_head_dim": 128, + "max_position_embeddings": 262144, + "mlp_only_layers": [], + "model_type": "qwen3_5_text", + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "num_attention_heads": 16, + "num_hidden_layers": 32, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-06, + "use_cache": false, + "vocab_size": 248320, + "mamba_ssm_dtype": "float32", + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default" + }, + "bos_token_id": null, + "pad_token_id": null, + "partial_rotary_factor": 0.25, + "tie_word_embeddings": false, + "transformers_version": "5.17.0" + }, + "tie_word_embeddings": false, + "transformers_version": "4.57.0.dev0", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 27, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1152, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4304, + "model_type": "qwen3_5", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 4096, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053, + "amos_export_dtype": "bfloat16", + "amos_codes": 255 +} diff --git a/model/generation_config.json b/model/generation_config.json new file mode 100644 index 0000000000000000000000000000000000000000..31ec2a53cc9cde720b0a8fbc68b9c4e73ac5a18f --- /dev/null +++ b/model/generation_config.json @@ -0,0 +1,6 @@ +{ + "_from_model_config": true, + "eos_token_id": 248044, + "transformers_version": "5.17.0", + "use_cache": true +} diff --git a/model/preprocessor_config.json b/model/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..2ea84a437d448ff71b08df68fdd949d5cc4ebb64 --- /dev/null +++ b/model/preprocessor_config.json @@ -0,0 +1,21 @@ +{ + "size": { + "longest_edge": 16777216, + "shortest_edge": 65536 + }, + "patch_size": 16, + "temporal_patch_size": 2, + "merge_size": 2, + "image_mean": [ + 0.5, + 0.5, + 0.5 + ], + "image_std": [ + 0.5, + 0.5, + 0.5 + ], + "processor_class": "Qwen3VLProcessor", + "image_processor_type": "Qwen2VLImageProcessorFast" +} \ No newline at end of file diff --git a/model/tokenizer.json b/model/tokenizer.json new file mode 100644 index 0000000000000000000000000000000000000000..cd071df754ad74ef10aced18b877b77d5942dcad --- /dev/null +++ b/model/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aa7582a6b90dc3711049d293d985e549f2b22c0c4e5ad11b9272a03c6294886e +size 19989510 diff --git a/model/tokenizer_config.json b/model/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..328c4c6644071294944181cd2c3b1de069e20c41 --- /dev/null +++ b/model/tokenizer_config.json @@ -0,0 +1,35 @@ +{ + "add_prefix_space": false, + "audio_bos_token": "<|audio_start|>", + "audio_eos_token": "<|audio_end|>", + "audio_token": "<|audio_pad|>", + "backend": "tokenizers", + "bos_token": null, + "clean_up_tokenization_spaces": false, + "eos_token": "<|im_end|>", + "errors": "replace", + "extra_special_tokens": [ + "" + ], + "image_token": "<|image_pad|>", + "is_local": false, + "local_files_only": false, + "model_max_length": 262144, + "model_specific_special_tokens": { + "audio_bos_token": "<|audio_start|>", + "audio_eos_token": "<|audio_end|>", + "audio_token": "<|audio_pad|>", + "image_token": "<|image_pad|>", + "video_token": "<|video_pad|>", + "vision_bos_token": "<|vision_start|>", + "vision_eos_token": "<|vision_end|>" + }, + "pad_token": "<|endoftext|>", + "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", + "split_special_tokens": false, + "tokenizer_class": "Qwen2Tokenizer", + "unk_token": null, + "video_token": "<|video_pad|>", + "vision_bos_token": "<|vision_start|>", + "vision_eos_token": "<|vision_end|>" +} diff --git a/model/video_preprocessor_config.json b/model/video_preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..3ba673a5ad7d4d13f54155ecd38b2a94a6dac8fe --- /dev/null +++ b/model/video_preprocessor_config.json @@ -0,0 +1,21 @@ +{ + "size": { + "longest_edge": 25165824, + "shortest_edge": 4096 + }, + "patch_size": 16, + "temporal_patch_size": 2, + "merge_size": 2, + "image_mean": [ + 0.5, + 0.5, + 0.5 + ], + "image_std": [ + 0.5, + 0.5, + 0.5 + ], + "processor_class": "Qwen3VLProcessor", + "video_processor_type": "Qwen3VLVideoProcessor" +} \ No newline at end of file diff --git a/qwen_export_helpers.py b/qwen_export_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..b13258e905ad834af62a6894794add505a7c3f75 --- /dev/null +++ b/qwen_export_helpers.py @@ -0,0 +1,175 @@ +"""Export pinned Decision 3.0 text inference; vision export is separate.""" + +import argparse +import json +from pathlib import Path + +import torch +import torch.nn.functional as F +from coreai_torch import TorchConverter, get_decomp_table +from safetensors.torch import load_file +from transformers import Qwen3_5Model +from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb_vision +from transformers.vision_utils import ( + get_vision_interpolation_indices_and_weights, + get_vision_position_ids, +) + + +class DecisionText(torch.nn.Module): + def __init__(self, backbone, readout): + super().__init__() + self.language_model = backbone.language_model + self.register_buffer("readout", readout.float()) + + def forward( + self, input_ids, attention_mask, position_ids, visual_embeddings, visual_mask + ): + embeddings = self.language_model.embed_tokens(input_ids) + embeddings = torch.where( + visual_mask, visual_embeddings.to(embeddings.dtype), embeddings + ) + hidden = self.language_model( + inputs_embeds=embeddings, + attention_mask={ + "full_attention": attention_mask[:, None, None, :].bool(), + "linear_attention": attention_mask, + }, + position_ids=position_ids, + use_cache=False, + ).last_hidden_state + return hidden[:, -1].float() @ self.readout.T + + +class DecisionVision(torch.nn.Module): + """The unchanged vision weights at a fixed one-image 224px grid.""" + + def __init__(self, visual, side=224): + super().__init__() + self.visual = visual + grid = torch.tensor( + [[1, side // visual.config.patch_size, side // visual.config.patch_size]] + ) + indices, weights = get_vision_interpolation_indices_and_weights( + grid, + num_grid_per_side=visual.num_grid_per_side, + mode=visual.interpolation_mode, + align_corners=visual.interpolation_align_corners, + spatial_merge_size=visual.config.spatial_merge_size, + kwargs={}, + ) + self.register_buffer("indices", indices) + self.register_buffer("weights", weights) + self.register_buffer( + "positions", + get_vision_position_ids(grid, visual.spatial_merge_size, kwargs={}), + ) + + def forward(self, pixel_values): + v = self.visual + hidden = F.linear( + pixel_values.to(v.patch_embed.proj.weight.dtype), + v.patch_embed.proj.weight.flatten(1), + v.patch_embed.proj.bias, + ) + hidden = hidden + (v.pos_embed(self.indices) * self.weights[:, :, None]).sum( + 1 + ).to(hidden.dtype) + cos, sin = v.rotary_pos_emb(hidden, self.positions) + for block in v.blocks: + x = block.norm1(hidden) + attn = block.attn + q, k, value = ( + attn.qkv(x) + .reshape(x.shape[0], 3, attn.num_heads, -1) + .permute(1, 0, 2, 3) + .unbind(0) + ) + q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) + y = F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + k.transpose(0, 1).unsqueeze(0), + value.transpose(0, 1).unsqueeze(0), + scale=attn.scaling, + ) + hidden = hidden + attn.proj( + y.squeeze(0).transpose(0, 1).reshape(x.shape[0], -1) + ) + hidden = hidden + block.mlp(block.norm2(hidden)) + return v.merger(hidden).float() + + +def save(model, sample, path): + with torch.no_grad(): + exported = torch.export.export(model, sample, strict=False) + print("Decomposing", flush=True) + exported = exported.run_decompositions(get_decomp_table()) + print("Converting", flush=True) + program = TorchConverter().add_exported_program(exported).to_coreai() + program.save_asset(path) + return exported.graph_signature.user_outputs[0] + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--source", default="upstream/d3-edge") + p.add_argument("--output", default="bundles/d3-edge-CoreAI") + p.add_argument("--length", type=int, default=512) + p.add_argument("--precision", choices=["bfloat16", "float32"], default="bfloat16") + a = p.parse_args() + torch.set_num_threads(4) + source = Path(a.source) + output = Path(a.output) + output.mkdir(parents=True, exist_ok=True) + print("Loading backbone", flush=True) + backbone = Qwen3_5Model.from_pretrained( + source, dtype=getattr(torch, a.precision), attn_implementation="sdpa" + ).eval() + model = DecisionText( + backbone, load_file(str(source / "readout.safetensors"))["weight"] + ).eval() + hidden = backbone.config.text_config.hidden_size + sample = ( + torch.ones(1, a.length, dtype=torch.int32), + torch.ones(1, a.length, dtype=torch.int32), + torch.arange(a.length, dtype=torch.int32).view(1, 1, -1).expand(4, 1, -1), + torch.zeros(1, a.length, hidden), + torch.zeros(1, a.length, 1, dtype=torch.bool), + ) + print("Exporting", flush=True) + out = save(model, sample, output / "decision.aimodel") + vision = DecisionVision(backbone.visual).eval() + patches = (224 // backbone.config.vision_config.patch_size) ** 2 + width = ( + 3 + * backbone.config.vision_config.temporal_patch_size + * backbone.config.vision_config.patch_size**2 + ) + vision_out = save(vision, (torch.zeros(patches, width),), output / "vision.aimodel") + (output / "manifest.json").write_text( + json.dumps( + { + "source": source.name, + "sequence_length": a.length, + "hidden_size": hidden, + "inputs": [ + "input_ids", + "attention_mask", + "position_ids", + "visual_embeddings", + "visual_mask", + ], + "output": out, + "vision_output": vision_out, + "image_side": 224, + "precision": a.precision, + }, + indent=2, + ) + + "\n" + ) + print("Saved", flush=True) + + +if __name__ == "__main__": + main() diff --git a/reference.json b/reference.json new file mode 100644 index 0000000000000000000000000000000000000000..1d0a2af96b1562e4d024f7012f429b3a8e5d60f4 --- /dev/null +++ b/reference.json @@ -0,0 +1,181 @@ +[ + { + "case": "refund", + "tokens": 122, + "reference_precision": "float32", + "upstream": [ + { + "type": "noul", + "p_yes": 0.9696575999259949 + }, + { + "type": "choice", + "probs": [ + 0.9343641996383667, + 0.025729455053806305, + 0.039906274527311325 + ], + "choice": 0 + } + ] + }, + { + "case": "span", + "tokens": 123, + "reference_precision": "float32", + "upstream": [ + { + "type": "span", + "start": 55, + "end": 63, + "text": "INV-2231", + "p": 1.0 + } + ] + }, + { + "case": "score", + "tokens": 98, + "reference_precision": "float32", + "upstream": [ + { + "type": "score", + "probs": [ + 0.013492000289261341, + 0.021048933267593384, + 0.13472826778888702, + 0.8307308554649353 + ], + "choice": 3, + "expected": 2.7826980352401733 + } + ] + }, + { + "case": "multi", + "tokens": 144, + "reference_precision": "float32", + "upstream": [ + { + "type": "multi", + "probs": [ + 0.9614583849906921, + 0.9263616800308228, + 0.05119440332055092 + ] + } + ] + }, + { + "case": "negative", + "tokens": 83, + "reference_precision": "float32", + "upstream": [ + { + "type": "noul", + "p_yes": 0.036552559584379196 + } + ] + }, + { + "case": "unicode", + "tokens": 128, + "reference_precision": "float32", + "upstream": [ + { + "type": "span", + "start": 10, + "end": 17, + "text": "\u20ac123.45", + "p": 1.0 + }, + { + "type": "choice", + "probs": [ + 0.9514100551605225, + 0.027261875569820404, + 0.021328067407011986 + ], + "choice": 0 + } + ] + }, + { + "case": "options", + "tokens": 230, + "reference_precision": "float32", + "upstream": [ + { + "type": "choice", + "probs": [ + 0.003573790192604065, + 0.003229466499760747, + 0.002545618684962392, + 0.0031847634818404913, + 0.0021655717864632607, + 0.0024921800941228867, + 0.0024529690854251385, + 0.003147563897073269, + 0.0037048303056508303, + 0.0024940893054008484, + 0.0032223397865891457, + 0.0031166989356279373, + 0.002279138658195734, + 0.0030672175344079733, + 0.003003651276230812, + 0.0037966070231050253, + 0.007012245245277882, + 0.9228788614273071, + 0.005725417286157608, + 0.004146159626543522, + 0.0036223935894668102, + 0.0032782373018562794, + 0.0026520290412008762, + 0.0032081292010843754 + ], + "choice": 17 + } + ] + }, + { + "case": "vision", + "tokens": 155, + "reference_precision": "float32", + "upstream": [ + { + "type": "region", + "cell": 33, + "box": [ + 0.3552, + 0.0785, + 0.747, + 0.9124 + ], + "p": 0.6444847583770752 + }, + { + "type": "choice", + "probs": [ + 0.9742496609687805, + 0.02575041353702545 + ], + "choice": 0 + } + ] + }, + { + "case": "image-option", + "tokens": 181, + "reference_precision": "float32", + "upstream": [ + { + "type": "choice", + "probs": [ + 0.9698329567909241, + 0.03016705997288227 + ], + "choice": 0 + } + ] + } +] \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..7184c765b6140a2123c6541564d9c2643bf038c6 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +coreai-core==1.0.0b3 +coreai-torch==0.4.3 +torch==2.14.0 +transformers @ git+https://github.com/huggingface/transformers.git@536ecc007387a50e77603bb5d92100e9b07514cc +numpy==2.4.6 +Pillow==12.3.0 +huggingface_hub==1.33.0 +tokenizers==0.23.3 diff --git a/runtime-smoke.json b/runtime-smoke.json new file mode 100644 index 0000000000000000000000000000000000000000..0cb039742670bed27368493516de64d34779c143 --- /dev/null +++ b/runtime-smoke.json @@ -0,0 +1,26 @@ +{ + "passed": true, + "pytorch_neural_execution_disabled": true, + "capacity_guard": true, + "result": { + "answers": { + "refund": { + "type": "noul", + "noul": 0.9637374877929688 + } + } + }, + "image_result": [ + { + "type": "region", + "cell": 21, + "box": [ + 0.0, + 0.0, + 0.9819, + 0.8978 + ], + "p": 0.9999999403953552 + } + ] +} \ No newline at end of file diff --git a/serve/README.md b/serve/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d36828d99ace6c4f6aa54117690e1150bacdc335 --- /dev/null +++ b/serve/README.md @@ -0,0 +1,46 @@ +# amos-serve + +The serving code of the Amos 1 decision models (`amos-1-large`, `amos-1-medium`, `amos-1-small`, `amos-1-mini`): one +HTTP server that takes a case and typed questions and answers every question from one forward pass, each with a +probability. No text is generated. This is the code behind api.amos.scira.ai and the code the Decision Index runs. + +Install into a fresh environment (Python 3.12 or newer, an NVIDIA GPU): + +```bash +pip install torch==2.14.0 --index-url https://download.pytorch.org/whl/cu130 +pip install . +``` + +The dependencies are pinned (torch 2.14.0, transformers 5.17.0, flash-linear-attention 0.5.2) because they are the +versions every published number was measured with: with them the package answers byte for byte like our own servers. +Other versions run, but probabilities can move by up to 0.02. + +Serve a model directory (a Hugging Face checkout with `model/`, `config.json` and `calibration.json`): + +```bash +python -m amos.serve amos-1-large --device cuda --port 8000 --rotations 1 --max-len 131072 --strict --no-state-cache +``` + +- `--strict` keeps question and option text exactly as sent and refuses requests over `--max-len` tokens instead of + truncating them. +- `--rotations N` averages each choice question over N option orders (1 is what the Index latency rule allows). +- `--no-state-cache` encodes every request from scratch. +- `GET /v1/info` reports the model name; `POST /v1/systemone` takes the System One request shape: + +```bash +curl -s localhost:8000/v1/systemone -H 'content-type: application/json' -d '{ + "state": "Invoice INV-2231 for 306.00 was charged twice this morning. Please refund one of them today.", + "questions": { + "refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}, + "team": {"type": "choice", "instructions": "Which team should handle this?", + "criteria": {"billing": null, "technical": null, "sales": null}}, + "urgency": {"type": "score", "instructions": "How urgent is this?", "criteria": ["Can wait", "This week", "Today"]} + } +}' +``` + +Choice questions take 2 to 255 options and are read in one pass. Photos go in `images` as base64 data URLs; every +question sees them. `span` (quote the case) and `region` (point at the photo) need the model's `heads.pt`. + +`sources.json` lists every file with the repository commit and sha256 it was taken from. `decoder.py` and +`calibrate.py` are the repository's modules cut to the definitions the server reaches; the other modules are verbatim. diff --git a/serve/amos/__init__.py b/serve/amos/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f29936777f01ec2993e064267b8686f9e89521b5 --- /dev/null +++ b/serve/amos/__init__.py @@ -0,0 +1 @@ +"""Amos 1 serving code. See README.md.""" diff --git a/serve/amos/calibrate.py b/serve/amos/calibrate.py new file mode 100644 index 0000000000000000000000000000000000000000..2dc3a429115bde75949c36ab00b5e3fec8eb5dcf --- /dev/null +++ b/serve/amos/calibrate.py @@ -0,0 +1,49 @@ +"""Cut from amos/calibrate.py: only the definitions amos.serve reaches (see sources.json).""" +import argparse +import json +import math +import sys +from collections import defaultdict +from pathlib import Path + +SET_TYPES = ("noul", "choice", "score") + +def _log(p): + return math.log(min(max(p, 1e-9), 1 - 1e-9)) + +def distribution(pred) -> list[float]: + return [1 - pred["p_yes"], pred["p_yes"]] if pred["type"] == "noul" else pred["probs"] + +def rescale(probs, T): + """Softmax(log p / T): temperature applied to the model's logits (log-probs up to a constant).""" + z = [_log(p) / T for p in probs] + m = max(z) + e = [math.exp(x - m) for x in z] + s = sum(e) + return [x / s for x in e] + +def rescale_binary(p, T): + return 1 / (1 + math.exp(-(_log(p) - _log(1 - p)) / T)) + +def apply(pred: dict, cal: dict, alpha: float | None = None) -> dict: + """Recalibrated copy of a decoded prediction; with alpha, adds the conformal set and abstain flag.""" + typ, T = pred["type"], cal["temperature"].get(pred["type"], 1.0) + out = dict(pred) + if typ in SET_TYPES: + probs = rescale(distribution(pred), T) + if typ == "noul": + out["p_yes"] = probs[1] + else: + out["probs"] = probs + if typ == "score": + out["expected"] = sum(i * p for i, p in enumerate(probs)) + if alpha is not None: + qhat = cal["qhat"][typ][str(alpha)] + out["set"] = [i for i, p in enumerate(probs) if p >= 1 - qhat] + out["abstain"] = len(out["set"]) != 1 + elif typ == "multi": + out["probs"] = [rescale_binary(p, T) for p in pred["probs"]] + else: + out["p"] = rescale_binary(pred["p"], T) + return out + diff --git a/serve/amos/decoder.py b/serve/amos/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..4ab0876a18cd91a2d9f50c57485099685c57d47d --- /dev/null +++ b/serve/amos/decoder.py @@ -0,0 +1,727 @@ +"""Cut from amos/decoder.py: only the definitions amos.serve reaches (see sources.json).""" +import argparse +import glob +import hashlib +import itertools +import json +import math +import os +import random +import shutil +import string +import time +from pathlib import Path +import torch + +ANSWER = "" + +IMAGE_PAD = "<|image_pad|>" + +LETTERS = string.ascii_uppercase + string.ascii_lowercase # the first 52 option codes; runs before amos_codes use these alone + +KINDS = ("noul", "choice", "score") # answered by option letters + +HEADED = ("span", "region") # answered by Heads + +SPAN_LENGTH_PENALTY = 1.6 + +SPAN_CAP = 30 + +HINTS = {"score": " (levels from lowest to highest)", "span": " (quote the case)", "region": " (point to it in the image)"} + +IMAGE_PIXELS = int(os.environ.get("AMOS_IMAGE_PIXELS", 1 << 20)) + +HEADS = "heads.pt" + +REGION_KEEP = 0.5 # a point's box: cells joined to the peak holding this share of its probability (vision.peak_box) + +SYSTEM = ("You make decisions about a case. Read the case, then answer every question with the letter of the option " + "that is most likely right.") + +def options(q: dict) -> list[str]: + """Option lines in answer order; yes/no is yes first, so its target is [P(yes), P(no)].""" + if q["type"] == "noul": + c = q.get("criteria") or {} + return [f"yes: {c['true']}" if c.get("true") else "yes", f"no: {c['false']}" if c.get("false") else "no"] + return list(q["options"]) + +def render(state: str, questions: list[dict], images=(), codes=LETTERS) -> list[dict]: + """Chat messages for a case. Images (the state's after "Case:", image options where they are listed, a pick-any + option's image after its question) make the user message a list of text and image items; without any it is plain + text.""" + parts = ["Case:\n", *images, f"{state}\n\nQuestions:\n"] + for i, q in enumerate(questions, start=1): + parts += ["\n" if i > 1 else "", f"q{i}. {q['text']}" + HINTS.get(q["type"], "")] + if "subject" in q: + parts += [" ", q["subject"]] + for j, o in enumerate(options(q) if q["type"] in KINDS else ()): + parts += [f"\n{codes[j]}) ", o] + merged = [] + for x in parts: + if isinstance(x, str) and merged and isinstance(merged[-1], str): + merged[-1] += x + else: + merged.append(x) + user = merged[0] if len(merged) == 1 else [{"type": "text", "text": x} if isinstance(x, str) else {"type": "image", "image": x} + for x in merged if x != ""] + answers = "\n".join(f"q{i}: {ANSWER}" for i in range(1, len(questions) + 1)) + reply = {"role": "assistant", "content": answers} + if len(questions) == 1 and questions[0].get("reasoning"): # a trace goes in the thinking block, as think() reads + reply["reasoning_content"] = questions[0]["reasoning"] + return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}, reply] + +def expand(questions: list[dict]) -> tuple[list[dict], list[int]]: + """Pick-any questions become one yes/no per option ("subject") in the same prompt, each with its own answer -> + (questions, index of each one's original question).""" + flat, owner = [], [] + for k, q in enumerate(questions): + if q["type"] != "multi": + flat.append(q) + owner.append(k) + continue + for j, o in enumerate(q["options"]): + flat.append({"type": "noul", "text": q["text"] + " Does this option apply?", "subject": o, + **({"target": q["target"][j]} if "target" in q else {})}) + owner.append(k) + return flat, owner + +PLAIN_TEMPLATE = ("{% for m in messages %}{{ m['role'] }}:\n{% if m['content'] is string %}{{ m['content'] }}{% else %}" + "{% for x in m['content'] %}{{ x['text'] if x['type'] == 'text' else '' }}{% endfor %}{% endif %}" + "\n\n{% endfor %}") + +def option_codes(tokenizer, n: int) -> list[str]: + """The first n option codes: the 52 letters, then two-letter codes the tokenizer keeps as one token.""" + two = ("".join(p) for p in itertools.product(string.ascii_uppercase, repeat=2)) + codes = list(LETTERS)[:n] + [c for c in two if len(tokenizer.encode(c, add_special_tokens=False)) == 1][:max(0, n - 52)] + if len(codes) != n: + raise ValueError(f"the tokenizer has only {len(codes)} one-token option codes, not {n}") + return codes + +class Codec: + def __init__(self, tokenizer, image_processor=None, n_codes: int = len(LETTERS)): + self.tok, self.image_processor = tokenizer, image_processor + if getattr(tokenizer, "chat_template", "") is None: + tokenizer.chat_template = PLAIN_TEMPLATE + if ANSWER not in tokenizer.get_added_vocab(): + tokenizer.add_special_tokens({"additional_special_tokens": [ANSWER]}) + self.marker = tokenizer.convert_tokens_to_ids(ANSWER) + self.codes = option_codes(tokenizer, n_codes) + ids = [tokenizer.encode(x, add_special_tokens=False) for x in self.codes] + if any(len(i) != 1 for i in ids): + raise ValueError("every option letter must be one token") + self.letters = [i[0] for i in ids] + self.pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + # A unified model (Gemma 4) comes with its whole processor, which turns each image placeholder into its own run + # of tokens; Qwen comes with an image processor alone, and the placeholder is expanded here + self.unified = image_processor is not None and hasattr(image_processor, "image_token") + pad = image_processor.image_token if self.unified else IMAGE_PAD + self.image_pad = tokenizer.convert_tokens_to_ids(pad) if image_processor is not None else None + self.echo = False # set by load_heads when the quote head was trained on an echoed case (see encode_case) + # the tokenizer's added tokens (the answer marker, chat roles, , image placeholders) as user text would + # tokenize as the tokens themselves and break the prompt; defuse() swaps their "<" for a look-alike + self.specials = sorted((t for t in tokenizer.get_added_vocab() if t.startswith("<")), key=len, reverse=True) + + def defuse(self, text): + """User text with any added-token string made ordinary text; same length, so character offsets still hold.""" + if not isinstance(text, str) or "<" not in text: + return text + for t in self.specials: + text = text.replace(t, "‹" + t[1:]) + return text + + def safe(self, state: str, questions: list[dict]) -> tuple[str, list[dict]]: + """state and questions with every user-written string defused (images and other values pass through).""" + d = self.defuse + qs = [{**q, **{k: d(q[k]) for k in ("text", "subject") if k in q}, + **({"options": [d(o) for o in q["options"]]} if "options" in q else {}), + **({"criteria": {k: d(v) for k, v in q["criteria"].items()}} if isinstance(q.get("criteria"), dict) else {})} + for q in questions] + return d(state), qs + + def encode(self, state: str, questions: list[dict]) -> tuple[list[int], list[int]]: + """-> (token ids, the position whose next-token logits answer each question).""" + state, questions = self.safe(state, questions) + text = self.tok.apply_chat_template(render(state, questions, codes=self.codes), tokenize=False, + enable_thinking=False) + ids = self.tok(text, add_special_tokens=False)["input_ids"] + at = [i - 1 for i, t in enumerate(ids) if t == self.marker] + if len(at) != len(questions): + raise ValueError("one answer marker per question") + return ids, at + + def reason_span(self, ids: list[int]) -> tuple[int, int]: + """The tokens inside the thinking block, (first, end): the trace a reasoning record teaches the model to write.""" + start, end = (self.tok.convert_tokens_to_ids(t) for t in ("", "")) + return ids.index(start) + 1, ids.index(end) + + def encode_case(self, state: str, questions: list[dict], images=()) -> tuple[list[int], list[int], dict]: + """-> (token ids, answer positions, extra); without images the ids are exactly encode()'s. extra["state"] is + (first state token, each state token's char offsets in the state). With images extra also holds the + image processor's pixel_values and image_grid_thw, and "images": (first token, rows, cols) of each image's + merged token grid; each image placeholder is expanded to its token count, as the model's processor does.""" + state, questions = self.safe(state, questions) + messages = render(state, questions, images, self.codes) + user = messages[1]["content"] + pics = [] if isinstance(user, str) else [x["image"] for x in user if x["type"] == "image"] + text = self.tok.apply_chat_template(messages, tokenize=False, enable_thinking=False) + enc = self.tok(text, add_special_tokens=False, return_offsets_mapping=True) + ids, spans = enc["input_ids"], enc["offset_mapping"] + # an empty state has no text to find (and Gemma's template trims the newlines that would surround it) + s0 = text.index(state + "\n\nQuestions:\n") if state else 0 + s1 = s0 + len(state) + toks = [i for i, (a, b) in enumerate(spans) if a < s1 and b > s0] + offsets = [(max(spans[i][0], s0) - s0, min(spans[i][1], s1) - s0) for i in toks] + where, extra = list(range(len(ids))), {} + if pics and self.unified: + ids, where, extra = self._unified_images(text, ids, pics) + elif pics: + if self.image_processor is None: + raise ValueError("this model has no vision tower") + ip = self.image_processor + v = ip(images=[p.convert("RGB") for p in pics], max_pixels=IMAGE_PIXELS, return_tensors="pt") + grid, m = v["image_grid_thw"], ip.merge_size + counts, out, where, starts = iter((grid.prod(-1) // m**2).tolist()), [], [], [] + for t in ids: + where.append(len(out)) + if t == self.image_pad: + starts.append(len(out)) + out += [t] * next(counts) + else: + out.append(t) + if len(starts) != len(pics): + raise ValueError("one image placeholder per image") + ids = out + extra = {"pixel_values": v["pixel_values"], "image_grid_thw": grid, + "images": [(s, int(g[1]) // m, int(g[2]) // m) for s, g in zip(starts, grid)]} + extra["state"] = (where[toks[0]] if toks else 0, offsets) + at = [i - 1 for i, t in enumerate(ids) if t == self.marker] + if len(at) != len(questions): + raise ValueError("one answer marker per question") + if self.echo and toks and any(q["type"] == "span" for q in questions): + # Echo: the case again after every answer marker, and quotes are read from this copy. Reading left to + # right, a case token cannot see what follows it ("Total due" after a line item), but every copied token + # has seen the whole case and each question. Attention only looks back, so no earlier answer changes. + sep = self.tok("\n\nThe case again:\n", add_special_tokens=False)["input_ids"] + copy = self.tok(state, add_special_tokens=False, return_offsets_mapping=True) + extra["state"] = (len(ids) + len(sep), [tuple(o) for o in copy["offset_mapping"]]) + ids = ids + sep + copy["input_ids"] + return ids, at, extra + + def _unified_images(self, text: str, ids: list[int], pics) -> tuple[list[int], list[int], dict]: + """A unified model's processor expands each image placeholder into <|image>, its image tokens, ; -> + (expanded ids, each original token's new position, the model's image inputs). The images have no separate + grid of merged tokens here, so "images" is empty and the point and box heads do not apply.""" + ip = self.image_processor + v = ip(text=text, images=[p.convert("RGB") for p in pics], return_tensors="pt", add_special_tokens=False) + new = v["input_ids"][0].tolist() + run = {self.image_pad, self.tok.convert_tokens_to_ids(ip.boi_token), self.tok.convert_tokens_to_ids(ip.eoi_token)} + where, j = [], 0 + for t in ids: + where.append(j) + if t == self.image_pad: + while j < len(new) and new[j] in run: + j += 1 + elif j < len(new) and new[j] == t: + j += 1 + else: + raise ValueError("the processor changed a text token") + if j != len(new): + raise ValueError("the processor added tokens after the text") + return new, where, {"pixel_values": v["pixel_values"], "image_position_ids": v["image_position_ids"], + "mm_token_type_ids": v["mm_token_type_ids"][0], "images": []} + +def load_model(path: str, device: str, *, dtype=None, vision=None): + """-> (model, codec). vision=None loads the image-text model when the checkpoint has a vision tower and an image + processor config beside it, else the text model alone. Qwen's image processor runs on PIL, never torchvision + (not installed, and the Auto classes demand it), so training and serving resize alike on every machine.""" + from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoModelForImageTextToText, + AutoTokenizer, + Qwen2VLImageProcessorPil, + ) + + if device == "cuda": # as in amos.train: cuDNN attention fails to load in bf16 with the pip CUDA wheels + torch.backends.cuda.enable_cudnn_sdp(False) + + tok = AutoTokenizer.from_pretrained(path) + config = AutoConfig.from_pretrained(path) + unified = config.model_type.startswith("gemma4") # images through its own processor (which needs torchvision) + if vision is None: + vision = hasattr(config, "vision_config") and (Path(path) / ("processor_config.json" if unified else + "preprocessor_config.json")).exists() + processor = None + if vision: + from transformers import AutoProcessor + + processor = AutoProcessor.from_pretrained(path) if unified else Qwen2VLImageProcessorPil.from_pretrained(path) + codec = Codec(tok, processor, getattr(config, "amos_codes", len(LETTERS))) + if dtype is None: + dtype = torch.float32 if getattr(config, "amos_export_dtype", None) == "float32" else torch.bfloat16 + from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES + + # some image-text models (Muse Glimmer) have no causal-LM class: load the whole model, its vision tower unused + causal = config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES + model = (AutoModelForImageTextToText if vision or not causal else AutoModelForCausalLM).from_pretrained( + path, config=config, dtype=dtype) + if vision and not unified: + patchify_as_matmul(model.model.visual.patch_embed) + if model.get_input_embeddings().weight.shape[0] < len(tok): + model.resize_token_embeddings(len(tok), mean_resizing=True) + return model.to(device), codec + +def patchify_as_matmul(patch_embed): + """Qwen's patch embedding is a conv3d whose kernel is the whole flattened patch, so it is one matmul with the same + weights; run it as one. cuDNN's conv3d fails to load on some hosts and PyTorch's own takes ~5 s per photo.""" + proj = patch_embed.proj + patch_embed.forward = lambda x: torch.nn.functional.linear(x.to(proj.weight.dtype), proj.weight.flatten(1), proj.bias) + +def noncausal(model) -> bool: + """Whether the run was trained with bidirectional full attention (train --noncausal), so it must be read that way.""" + return bool(getattr(model.config, "amos_noncausal", False)) + +def bidirectional_mask(lengths: list[int], width: int, device) -> torch.Tensor: + """[B, 1, L, L], True where the key is one of its row's real tokens: every query sees the whole case, never pads.""" + keys = torch.arange(width, device=device) < torch.tensor(lengths, device=device)[:, None] + return keys[:, None, None, :].expand(-1, 1, width, -1) + +PREFILL_CHUNK = 16384 + +def chunked_hidden(decoder, x: torch.Tensor) -> torch.Tensor: + """Last hidden states [1, L, d] for one causal text row, read PREFILL_CHUNK tokens at a time.""" + from transformers import DynamicCache + + cache, out = DynamicCache(config=decoder.config), [] + for s in range(0, x.shape[1], PREFILL_CHUNK): + out.append(decoder(input_ids=x[:, s : s + PREFILL_CHUNK], past_key_values=cache, use_cache=True).last_hidden_state) + del cache + return torch.cat(out, 1) + +def layerwise_hidden(decoder, x: torch.Tensor) -> torch.Tensor: + """Last hidden states [1, L, d] for one noncausal text row, read a layer at a time. + + A noncausal run's full-attention layers let every token see the whole case, so it cannot be read in chunks the way + a causal one can. Instead the whole case goes through each layer before the next, PREFILL_CHUNK tokens at a time: + a linear-attention layer is a forward recurrence, so its chunks carry its conv and recurrent state; a full-attention + layer first takes keys and values for every token, then each chunk of queries attends to all of them. The same + computation as one pass, holding one layer's activations for one chunk instead of every layer's for the whole case.""" + import torch.nn.functional as F + from transformers import DynamicCache + from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb + + n, C = x.shape[1], PREFILL_CHUNK + h = decoder.embed_tokens(x) + pos = torch.arange(n, device=x.device).view(1, 1, -1).expand(4, 1, -1) + cos, sin = decoder.rotary_emb(h, pos[1:]) + text_pos, cache = pos[0], DynamicCache(config=decoder.config) + for i, layer in enumerate(decoder.layers[: decoder.config.num_hidden_layers]): + if decoder.config.layer_types[i] == "linear_attention": + for s in range(0, n, C): # each chunk is overwritten after it is read; later chunks are still this layer's input + h[:, s : s + C] = layer(h[:, s : s + C], position_embeddings=(cos[:, s : s + C], sin[:, s : s + C]), + attention_mask=None, position_ids=text_pos[:, s : s + C], past_key_values=cache) + continue + at = layer.self_attn + keys, values = [], [] + for s in range(0, n, C): + hn = layer.input_layernorm(h[:, s : s + C]) + shape = (*hn.shape[:-1], -1, at.head_dim) + k = at.k_norm(at.k_proj(hn).view(shape)).transpose(1, 2) + _, k = apply_rotary_pos_emb(k, k, cos[:, s : s + C], sin[:, s : s + C]) + keys.append(k) + values.append(at.v_proj(hn).view(shape).transpose(1, 2)) + keys, values = torch.cat(keys, 2), torch.cat(values, 2) + for s in range(0, n, C): + res = h[:, s : s + C] + hn = layer.input_layernorm(res) + shape = hn.shape[:-1] + q, gate = torch.chunk(at.q_proj(hn).view(*shape, -1, at.head_dim * 2), 2, dim=-1) + q = at.q_norm(q).transpose(1, 2) + q, _ = apply_rotary_pos_emb(q, q, cos[:, s : s + C], sin[:, s : s + C]) + o = F.scaled_dot_product_attention(q, keys, values, scale=at.scaling, enable_gqa=True) + o = at.o_proj(o.transpose(1, 2).reshape(*shape, -1) * torch.sigmoid(gate.reshape(*shape, -1))) + mid = res + o + h[:, s : s + C] = mid + layer.mlp(layer.post_attention_layernorm(mid)) + del keys, values + del cache + for s in range(0, n, C): + h[:, s : s + C] = decoder.norm(h[:, s : s + C]) + return h + +def hidden(model, codec, batch: list[tuple]) -> torch.Tensor: + """Last hidden states [B, L, d] for a batch of (ids, answer positions, questions[, extra]). + + Sequences are right-padded without a mask: the model is causal (its linear-attention layers are recurrences that + only run forward), so padding after a sequence cannot change anything read from inside it. A noncausal run's + full-attention layers see the whole sequence, so they get a mask hiding the pads; the recurrences stay causal. + Entries with images go through the image-text model, which fills image tokens from the vision tower and gives + them 2-D positions.""" + width = max(len(e[0]) for e in batch) + x = torch.full((len(batch), width), codec.pad, dtype=torch.long, device=model.device) + for b, e in enumerate(batch): + x[b, : len(e[0])] = torch.tensor(e[0], device=model.device) + vis = [e[3] for e in batch if len(e) > 3 and "pixel_values" in e[3]] + kw = {} if not noncausal(model) else {"attention_mask": { + "full_attention": bidirectional_mask([len(e[0]) for e in batch], width, x.device), "linear_attention": None}} + if not vis and len(batch) == 1 and width > PREFILL_CHUNK: + return layerwise_hidden(model.get_decoder(), x) if kw else chunked_hidden(model.get_decoder(), x) + if not vis: + return model.get_decoder()(input_ids=x, **kw).last_hidden_state + if codec.unified and kw: + raise ValueError("noncausal attention is built for Qwen3.5's layer types, not Gemma 4's") + if codec.unified: # Gemma 4: image patches go through the same layers, placed by their token types and positions + types = torch.zeros_like(x) + for b, e in enumerate(batch): + if len(e) > 3 and "mm_token_type_ids" in e[3]: + types[b, : len(e[3]["mm_token_type_ids"])] = e[3]["mm_token_type_ids"].to(x.device) + return model.model(input_ids=x, mm_token_type_ids=types, + pixel_values=torch.cat([v["pixel_values"] for v in vis]).to(model.device, model.dtype), + image_position_ids=torch.cat([v["image_position_ids"] for v in vis]).to(model.device) + ).last_hidden_state + grid, types = torch.cat([v["image_grid_thw"] for v in vis]).to(model.device), (x == codec.image_pad).int() + if kw: # its rope index would read a dict mask as per-row padding, so the positions come unmasked, as causal + kw["position_ids"] = model.model.get_rope_index(x, types, image_grid_thw=grid)[0] + return model.model(input_ids=x, pixel_values=torch.cat([v["pixel_values"] for v in vis]).to(model.device), + image_grid_thw=grid, mm_token_type_ids=types, **kw).last_hidden_state + +def finish_logits(model, logits: torch.Tensor) -> torch.Tensor: + """The model's own last step from raw output-layer logits, as its forward would apply it: Granite divides by + logits_scaling, Gemma caps with final_logit_softcapping. Qwen has neither, so its logits pass unchanged.""" + cfg = getattr(model.config, "text_config", None) or model.config + if scaling := getattr(cfg, "logits_scaling", None): + logits = logits / scaling + if cap := getattr(cfg, "final_logit_softcapping", None): + logits = cap * torch.tanh(logits / cap) + return logits + +def distributions(model, codec, batch: list[tuple], h: torch.Tensor | None = None) -> list[torch.Tensor]: + """Log-probabilities over the option letters of each yes/no, choice and score question in a batch of (ids, answer + positions, questions[, extra]), in order. Only the 52 letter logits are computed, from the tied output + embeddings.""" + h = hidden(model, codec, batch) if h is None else h + picks = [(b, p, q) for b, e in enumerate(batch) for p, q in zip(e[1], e[2]) if q["type"] in KINDS] + if not picks: + return [] + rows = torch.tensor([b for b, _, _ in picks], device=h.device) + cols = torch.tensor([p for _, p, _ in picks], device=h.device) + logits = finish_logits(model, h[rows, cols].float() @ model.get_output_embeddings().weight[codec.letters].float().T) + return [logits[k, : len(options(q))].log_softmax(-1) for k, (_, _, q) in enumerate(picks)] + +class BoxHead(torch.nn.Module): + """A point's box, predicted rather than grown from cells: reads the question's vector, the image features pooled + by the point head's own cell distribution, and that distribution's centre and spread, and moves the centre and + sets the size. Cells alone draw a box around only the few cells the point head is sure of (a car's tail light + rather than the car); this head is trained on overlap with the whole object.""" + + def __init__(self, d: int, width: int = 512): + super().__init__() + self.net = torch.nn.Sequential(torch.nn.Linear(2 * d + 4, width), torch.nn.GELU(), torch.nn.Linear(width, 4)) + torch.nn.init.zeros_(self.net[-1].weight) # starts as a half-image box on the cells' centre + torch.nn.init.zeros_(self.net[-1].bias) + + def forward(self, hq: torch.Tensor, cells: torch.Tensor, logits: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + p = logits.detach().float().softmax(-1) # the point head stays as it is; only the box learns + ys, xs = torch.meshgrid((torch.arange(rows, device=p.device) + 0.5) / rows, + (torch.arange(cols, device=p.device) + 0.5) / cols, indexing="ij") + xs, ys = xs.flatten(), ys.flatten() + mx, my = (p * xs).sum(), (p * ys).sum() + sx, sy = (p * (xs - mx) ** 2).sum().sqrt(), (p * (ys - my) ** 2).sum().sqrt() + pooled = (p[:, None] * cells.float()).sum(0) + o = self.net(torch.cat([hq.float(), pooled, torch.stack([mx, my, sx, sy])]).to(self.net[0].weight.dtype)).float() + cx, cy = (mx + 0.5 * torch.tanh(o[0])).clamp(0, 1), (my + 0.5 * torch.tanh(o[1])).clamp(0, 1) + w, hgt = torch.sigmoid(o[2]), torch.sigmoid(o[3]) + return torch.stack([cx - w / 2, cy - hgt / 2, cx + w / 2, cy + hgt / 2]).clamp(0, 1) + +class Heads(torch.nn.Module): + """Quote and point heads, as amos.model has them: a question's answer-position vector scored against the state's + tokens (start and end of a quote) or against the first state image's tokens (one cell each); `keep` is the + region readout's share (decode_head). With a BoxHead, a region's logits carry four more values: its box.""" + + def __init__(self, d: int, kinds=HEADED, keep: float = REGION_KEEP, box: bool = False): + super().__init__() + self.kinds, self.keep = list(kinds), keep + self.span = (SPAN_CAP, SPAN_LENGTH_PENALTY) # the quote decoder's cap and penalty; serving sets the run's own + self.span_q, self.span_k = torch.nn.Linear(d, 2 * d), torch.nn.Linear(d, 2 * d) + self.region_q, self.region_k = torch.nn.Linear(d, d), torch.nn.Linear(d, d) + for m in (self.span_q, self.span_k, self.region_q, self.region_k): + torch.nn.init.normal_(m.weight, std=0.02) + torch.nn.init.zeros_(m.bias) + self.box = BoxHead(d) if box else None + + def forward(self, h: torch.Tensor, batch: list[tuple]) -> list[torch.Tensor]: + """Logits of each span ([2, state tokens]) and region ([cells], then the 4 box edges when there is a box head) + question in the batch, in order.""" + h, d, out = h.to(self.span_q.weight.dtype), h.shape[-1], [] + for b, e in enumerate(batch): + for p, q in zip(e[1], e[2]): + if q["type"] == "span": + first, offsets = e[3]["state"] + k = self.span_k(h[b, first : first + len(offsets)]).view(-1, 2, d) + out.append(torch.einsum("cd,lcd->cl", self.span_q(h[b, p]).view(2, d), k) / math.sqrt(d)) + elif q["type"] == "region": + if not e[3].get("images"): + raise ValueError("a region question needs a state image") + start, rows, cols = e[3]["images"][0] + cells = h[b, start : start + rows * cols] + x = self.region_k(cells) @ self.region_q(h[b, p]) / math.sqrt(d) + if self.box is not None: + x = torch.cat([x, self.box(h[b, p], cells, x, rows, cols).to(x.dtype)]) + out.append(x) + return out + +def decode_head(q: dict, x: torch.Tensor, extra: dict, state: str, keep: float = REGION_KEEP, + span: tuple[int, float] = (SPAN_CAP, SPAN_LENGTH_PENALTY)) -> dict: + """A span or region result in amos.model.decode's format; a quote is always cut from the state, at most span[0] + tokens long, each token past the start scaled by exp(-span[1]). A point is the most likely cell and the box of the + cells around it holding `keep` of its probability, with the box's mass.""" + from types import SimpleNamespace + + from .model import decode_span + from .vision import peak_box + + x = x.float().cpu() + if q["type"] == "span": + x = torch.cat([x.new_full((2, 1), -torch.inf), x], 1) # decode_span skips a leading [CLS] slot + a, z, _, p = decode_span(x[0], x[1], SimpleNamespace(offsets=extra["state"][1], state=state), span[1], span[0]) + # Qwen fuses a JSON string's closing quote with the last word ('detail."', 'approves.",'), so a quote cut at + # token edges can end (or start) on an unmatched quote mark; drop it and what follows it + while state[a:z].count('"') % 2: + if state[z - 1] in '",}]': + z -= 1 + elif state[a] in '"{[': + a += 1 + else: + break + t = state[a:z] + a, z = a + len(t) - len(t.lstrip()), z - (len(t) - len(t.rstrip())) + return {"type": "span", "start": a, "end": max(a, z), "text": state[a : max(a, z)], "p": p} + _, rows, cols = extra["images"][0] + cells = x[: rows * cols].softmax(-1) + if len(x) > rows * cols: # a box head's box; its mass is the cells it covers + box = [round(v, 4) for v in x[rows * cols :].tolist()] + r0, r1 = int(box[1] * rows), max(int(box[1] * rows) + 1, math.ceil(box[3] * rows)) + c0, c1 = int(box[0] * cols), max(int(box[0] * cols) + 1, math.ceil(box[2] * cols)) + return {"type": "region", "cell": int(cells.argmax()), "box": box, + "p": float(cells.view(rows, cols)[r0:r1, c0:c1].sum())} + box, p = peak_box(cells, rows, cols, keep) + return {"type": "region", "cell": int(cells.argmax()), "box": box, "p": p} + +def load_heads(path, model, codec: Codec | None = None) -> Heads | None: + """The trained heads saved beside a checkpoint, or None. A quote head trained on echoed cases switches the codec's + echo on, since it reads nothing else.""" + f = Path(path) / HEADS + if not f.exists(): + return None + saved = torch.load(f, map_location="cpu", weights_only=True) + heads = Heads(model.get_input_embeddings().weight.shape[1], saved["kinds"], saved.get("keep", REGION_KEEP), + box=any(k.startswith("box.") for k in saved["state"])) + heads.load_state_dict(saved["state"]) + if codec is not None: + codec.echo = bool(saved.get("echo", False)) + return heads.to(model.device).eval() + +def tempered(lp: torch.Tensor, t: float | None = None) -> list[float]: + """Probabilities from log-probabilities, sharpened (t < 1) or softened (t > 1) by a fitted temperature.""" + return (lp / t).softmax(-1).tolist() if t else lp.exp().tolist() + +def rotate(q: dict, r: int) -> dict: + """A choice question with its options rotated r places; anything else as it is (a scale's order is its meaning).""" + if q["type"] != "choice" or not r: + return q + s = r % len(q["options"]) + return {**q, "options": q["options"][s:] + q["options"][:s]} + +def read(model, codec, state: str, questions: list[dict], images=(), max_len: int | None = None, + temperature: dict | None = None): + """-> (option distributions for letter-answered questions, input tokens), one forward pass.""" + ids, at, extra = codec.encode_case(state, questions, images) + if max_len and len(ids) > max_len: + raise ValueError(f"maximum context length: request has {len(ids)} tokens, limit {max_len}") + entry = [(ids, at, questions, extra)] + t = temperature or {} + picked = [q for q in questions if q["type"] in KINDS] + return [tempered(lp, t.get(q["type"])) for lp, q in + zip(distributions(model, codec, entry, hidden(model, codec, entry)), picked)], len(ids) + +def wide_choice(model, codec, state: str, q: dict, images=(), max_len: int | None = None, + temperature: dict | None = None): + """A choice over more options than there are letters, in two passes -> (distribution, most input tokens). + + First the options are split into chunks, each asked as its own question; then the leading options of every chunk + are asked again together. Finalists share their first-pass mass (each chunk weighted equally) in proportion to the + second pass; every other option keeps its first-pass share, so the distribution sums to one and its top answer is + the second pass's.""" + n, per = len(q["options"]), len(codec.codes) + size = -(-n // -(-n // per)) # balanced chunks of at most `per` + chunks = [list(range(s, min(s + size, n))) for s in range(0, n, size)] + first, t1 = read(model, codec, state, [{**q, "options": [q["options"][i] for i in c]} for c in chunks], images, max_len, + temperature) + prior = [0.0] * n + for c, p in zip(chunks, first): + for i, x in zip(c, p): + prior[i] = x / len(chunks) + keep = per // len(chunks) + finalists = sorted(i for c, p in zip(chunks, first) for i, _ in sorted(zip(c, p), key=lambda t: -t[1])[:keep]) + (second,), t2 = read(model, codec, state, [{**q, "options": [q["options"][i] for i in finalists]}], images, max_len, + temperature) + mass, out = sum(prior[i] for i in finalists), prior[:] + for i, x in zip(finalists, second): + out[i] = mass * x + return out, max(t1, t2) + +@torch.no_grad() +def sureness(r: dict) -> float: + """How sure a yes/no, choice or score result is: the probability of its answer.""" + return max(r["p_yes"], 1 - r["p_yes"]) if r["type"] == "noul" else max(r["probs"]) + +@torch.no_grad() +def think_many(model, codec, state: str, qs: list[dict], max_tokens: int, temperature: dict | None = None, + generate=None): + """think for several questions about one case at once: their thoughts are generated as one left-padded batch, so a + case with eight unsure questions costs about one generation instead of eight. generate(prefixes, max_tokens) -> + [(text, tokens)] writes the thoughts elsewhere instead (a vLLM server: serve's AMOS_THINK_URL); the letters are + still read here, so the answers are the same either way.""" + if noncausal(model): + raise ValueError("a noncausal run cannot think: generation reads a causal model's cache") + state, qs = codec.safe(state, qs) + prefixes = [codec.tok.apply_chat_template(render(state, [q], codes=codec.codes)[:2], tokenize=False, + add_generation_prompt=True, enable_thinking=True) for q in qs] + if generate is not None: + thoughts = generate(prefixes, max_tokens) + else: + ids = [codec.tok(p, add_special_tokens=False)["input_ids"] for p in prefixes] + width, pad = max(map(len, ids)), codec.tok.pad_token_id + batch = torch.tensor([[pad] * (width - len(x)) + x for x in ids], device=model.device) + mask = torch.tensor([[0] * (width - len(x)) + [1] * len(x) for x in ids], device=model.device) + out = model.generate(batch, attention_mask=mask, max_new_tokens=max_tokens, do_sample=False, use_cache=True, + pad_token_id=pad) + thoughts = [(codec.tok.decode(row, skip_special_tokens=False).replace(codec.tok.pad_token or "\0", ""), + int((row != pad).sum())) for row in out[:, width:]] + results = [] + for prefix, q, (text, n) in zip(prefixes, qs, thoughts): + reasoning = text.split("")[0].replace("<|im_end|>", "").strip() + full = codec.tok(prefix + reasoning + "\n\n\n" + f"q1: {ANSWER}", add_special_tokens=False)["input_ids"] + at = [i - 1 for i, t in enumerate(full) if t == codec.marker] + lp = distributions(model, codec, [(full, at, [q])])[0] + results.append((tempered(lp, (temperature or {}).get(q["type"])), n)) + return results + +def results(model, codec, state: str, questions: list[dict], images=(), heads: Heads | None = None, + max_len: int | None = None, rotations: int = 1, temperature: dict | None = None, + think_below: float | None = None, think_tokens: int = 512, generate=None) -> tuple[list[dict], int]: + """One typed result per question in amos.model.decode's format (noul p_yes; choice/score probs, choice, expected; + multi independent probs; span start/end/text/p; region cell/box/p), and the number of input tokens. A choice over + more options than there are letters is answered on its own (wide_choice). rotations and temperature: see + one_pass. think_below: a yes/no, choice or score answer less sure than this is asked again after the model + thinks (think), as GLiDE spends compute only where a decision is uncertain; the result then carries "thought".""" + if think_below and noncausal(model): # on every request, not only the unsure ones + raise ValueError("a noncausal run cannot think: generation reads a causal model's cache") + out, tokens = _results(model, codec, state, questions, images, heads, max_len, rotations, temperature) + if think_below and not images: + unsure = [k for k, (q, r) in enumerate(zip(questions, out)) + if q["type"] in KINDS and len(options(q)) <= len(codec.codes) and sureness(r) < think_below] + for k, (p, n) in zip(unsure, think_many(model, codec, state, [questions[k] for k in unsure], think_tokens, + temperature, generate) if unsure else []): + q = questions[k] + if q["type"] == "noul": + out[k] = {"type": "noul", "p_yes": p[0], "thought": n} + else: + out[k] = {"type": q["type"], "probs": p, "choice": max(range(len(p)), key=p.__getitem__), "thought": n} + if q["type"] == "score": + out[k]["expected"] = sum(i * x for i, x in enumerate(p)) + return out, tokens + +def _results(model, codec, state: str, questions: list[dict], images=(), heads: Heads | None = None, + max_len: int | None = None, rotations: int = 1, temperature: dict | None = None) -> tuple[list[dict], int]: + wide = {k for k, q in enumerate(questions) if q["type"] == "choice" and len(q["options"]) > len(codec.codes)} + out, tokens = [None] * len(questions), 0 + for k in sorted(wide): + p, n = wide_choice(model, codec, state, questions[k], images, max_len, temperature) + out[k], tokens = {"type": "choice", "probs": p, "choice": max(range(len(p)), key=p.__getitem__)}, max(tokens, n) + rest = [k for k in range(len(questions)) if k not in wide] + if rest: + sub, n = one_pass(model, codec, state, [questions[k] for k in rest], images, heads, max_len, rotations, + temperature) + for k, r in zip(rest, sub): + out[k] = r + tokens = max(tokens, n) + return out, tokens + +ROTATION_BATCH_TOKENS = 65536 + +def one_pass(model, codec, state: str, questions: list[dict], images=(), heads: Heads | None = None, + max_len: int | None = None, rotations: int = 1, temperature: dict | None = None) -> tuple[list[dict], int]: + """results() for questions that fit the letters, in one forward pass. + + rotations > 1 (and a choice question present): the case is also read with every choice question's options rotated, + each rotation one more row of the same batched forward, and choice distributions are averaged in the original order. + temperature: {question type: t} fitted by `decoder calibrate` (pick-any's per-option yes/no uses noul's).""" + flat, owner = expand(questions) + rotations = rotations if any(q["type"] == "choice" for q in flat) else 1 + entries = [] + for r in range(rotations): + qs = [rotate(q, r) for q in flat] + ids, at, extra = codec.encode_case(state, qs, images) + if max_len and len(ids) > max_len: + raise ValueError(f"maximum context length: request has {len(ids)} tokens, limit {max_len}") + entries.append((ids, at, qs, extra)) + ids, extra = entries[0][0], entries[0][3] + if heads is None and any(q["type"] in HEADED for q in flat): + raise ValueError("this model has no span or region head") + if sum(len(e[0]) for e in entries) <= ROTATION_BATCH_TOKENS: + h = hidden(model, codec, entries) + lps = distributions(model, codec, entries, h) + else: # a long case: one rotation at a time, so memory stays that of a single row + h, lps = None, [] + for e in entries: + he = hidden(model, codec, [e]) + lps += distributions(model, codec, [e], he) + h = he if h is None else h + head = iter(heads(h[:1], entries[:1]) if heads is not None else []) + picked = [j for j, q in enumerate(flat) if q["type"] in KINDS] + temps = temperature or {} + mean = {} + for r in range(rotations): + for m, j in enumerate(picked): + p, n = tempered(lps[r * len(picked) + m], temps.get(flat[j]["type"])), len(options(flat[j])) + s = r % n if flat[j]["type"] == "choice" else 0 + acc = mean.setdefault(j, [0.0] * n) + for i in range(n): # rotated position (i - s) mod n holds original option i + acc[i] += p[(i - s) % n] / rotations + lps = iter(mean[j] for j in picked) + out = [None] * len(questions) + for q, k in zip(flat, owner): + t = questions[k]["type"] + if q["type"] in HEADED: + out[k] = decode_head(q, next(head), extra, state, heads.keep, heads.span) + continue + p = next(lps) + if t == "multi": + out[k] = out[k] or {"type": t, "probs": []} + out[k]["probs"].append(p[0]) + elif t == "noul": + out[k] = {"type": t, "p_yes": p[0]} + else: + out[k] = {"type": t, "probs": p, "choice": max(range(len(p)), key=p.__getitem__)} + if t == "score": + out[k]["expected"] = sum(i * x for i, x in enumerate(p)) + return out, len(ids) + +CALIBRATION = "calibration.json" + +def calibration(run) -> dict | None: + """The per-type temperatures `decoder calibrate` fitted for a run, or None.""" + f = Path(run) / CALIBRATION + return json.loads(f.read_text()).get("temperature") if f.exists() else None + +def span_decoding(run) -> tuple[int, float]: + """The quote decoder's (cap, penalty) fitted for a run on held-out quote records, or the defaults.""" + f = Path(run) / CALIBRATION + s = (json.loads(f.read_text()).get("span") or {}) if f.exists() else {} + return int(s.get("cap", SPAN_CAP)), float(s.get("penalty", SPAN_LENGTH_PENALTY)) + +def multi_cut(run) -> float | None: + """The pick-any cut fitted on our own held-out pick-any questions (an option is picked at p >= cut), or None.""" + f = Path(run) / CALIBRATION + return json.loads(f.read_text()).get("multi_cut") if f.exists() else None + diff --git a/serve/amos/encoder.py b/serve/amos/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..6326909724afc38600ad5359e600f37a4d12903e --- /dev/null +++ b/serve/amos/encoder.py @@ -0,0 +1,161 @@ +"""ModernBERT-architecture encoder (Ettin weights), written from scratch. + +Attention is controlled by per-token vectors instead of a precomputed mask: + position_ids: RoPE positions (branches restart after the state). + segment_ids: 0 = state, k > 0 = question branch k, -1 = padding. + option_ids: 0 = not an option, m > 0 = option m of its question. +A token attends to state tokens, its question's own tokens, and (for option tokens) +its own option. Local layers additionally restrict attention to +|pos_i - pos_j| <= window // 2, measured in position ids (not sequence index) so +packed and separate requests match. +Comparison mode lets each question and its options attend to the whole question branch. +""" + +import json +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from huggingface_hub import hf_hub_download +from torch import nn + + +@dataclass +class Config: + vocab_size: int + hidden_size: int + intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + global_attn_every_n_layers: int = 3 + local_attention: int = 128 + global_rope_theta: float = 160000.0 + local_rope_theta: float = 10000.0 + norm_eps: float = 1e-5 + pad_token_id: int = 50283 + cls_token_id: int = 50281 + sep_token_id: int = 50282 + model_type: str = "modernbert" + compare_options: bool = False + date_facts: bool = False # an input-format setting: facts.date_facts runs on every state (model.pack_request) + + @classmethod + def from_hf(cls, d: dict) -> "Config": + """Reads a Hugging Face config or a run's saved config, ignoring fields this class does not define.""" + return cls(**{k: d[k] for k in cls.__dataclass_fields__ if k in d}) + + +def rope(x: torch.Tensor, pos: torch.Tensor, theta: float) -> torch.Tensor: + """x: [B, H, L, D], pos: [B, L]. Half-split rotation (rotate_half convention).""" + d = x.shape[-1] + inv_freq = 1.0 / (theta ** (torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d)) + freqs = pos[:, None, :, None].float() * inv_freq # [B, 1, L, D/2] + cos, sin = freqs.cos().to(x.dtype), freqs.sin().to(x.dtype) + x1, x2 = x[..., : d // 2], x[..., d // 2 :] + return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1) + + +class Layer(nn.Module): + def __init__(self, c: Config, idx: int): + super().__init__() + self.is_global = idx % c.global_attn_every_n_layers == 0 + self.theta = c.global_rope_theta if self.is_global else c.local_rope_theta + self.heads = c.num_attention_heads + self.attn_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False) if idx else nn.Identity() + self.Wqkv = nn.Linear(c.hidden_size, 3 * c.hidden_size, bias=False) + self.Wo = nn.Linear(c.hidden_size, c.hidden_size, bias=False) + self.mlp_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False) + self.Wi = nn.Linear(c.hidden_size, 2 * c.intermediate_size, bias=False) + self.Wo_mlp = nn.Linear(c.intermediate_size, c.hidden_size, bias=False) + + def forward(self, h, pos, global_mask, local_mask, prefix=None): + """prefix: cached (k, v) of earlier tokens, prepended to this call's keys and values.""" + B, L, _ = h.shape + q, k, v = self.Wqkv(self.attn_norm(h)).view(B, L, 3, self.heads, -1).permute(2, 0, 3, 1, 4) + q, k = rope(q, pos, self.theta), rope(k, pos, self.theta) + kv = (k, v) + if prefix is not None: + k, v = torch.cat([prefix[0], k], 2), torch.cat([prefix[1], v], 2) + mask = global_mask if self.is_global else local_mask + a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + h = h + self.Wo(a.transpose(1, 2).reshape(B, L, -1)) + x, gate = self.Wi(self.mlp_norm(h)).chunk(2, dim=-1) + return h + self.Wo_mlp(F.gelu(x) * gate), kv + + +class Encoder(nn.Module): + def __init__(self, c: Config): + super().__init__() + self.config = c + self.tok_embeddings = nn.Embedding(c.vocab_size, c.hidden_size) + self.emb_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False) + self.layers = nn.ModuleList(Layer(c, i) for i in range(c.num_hidden_layers)) + self.final_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False) + + @staticmethod + def masks(pos, seg, opt, window, pos_k=None, seg_k=None, opt_k=None, *, compare_options=False): + """Boolean [B, 1, Lq, Lk] masks, True = may attend. Global and local variants. + + Keys default to the queries; pass *_k to attend over a longer key sequence (cached state + queries). + The masks are dense, so their memory grows with Lq * Lk. + """ + pos_k, seg_k, opt_k = (x if x is not None else y for x, y in ((pos_k, pos), (seg_k, seg), (opt_k, opt))) + si, sj, oi, oj = seg[:, :, None], seg_k[:, None, :], opt[:, :, None], opt_k[:, None, :] + within = (sj == si) & ((si > 0) if compare_options else ((oj == 0) | (oj == oi))) + allowed = (sj != -1) & ((sj == 0) | within) + local = allowed & ((pos[:, :, None] - pos_k[:, None, :]).abs() <= window // 2) + return allowed[:, None], local[:, None] + + def forward(self, input_ids, position_ids=None, segment_ids=None, option_ids=None, inputs_embeds=None): + B, L = input_ids.shape + if position_ids is None: + position_ids = torch.arange(L, device=input_ids.device).expand(B, L) + if segment_ids is None: + segment_ids = torch.zeros_like(input_ids) + if option_ids is None: + option_ids = torch.zeros_like(input_ids) + emb = self.tok_embeddings(input_ids) if inputs_embeds is None else inputs_embeds + return self.run(emb, position_ids, segment_ids, option_ids)[0] + + def run(self, emb, pos, seg, opt, cache=None): + """Encode embeddings, optionally on top of a cached prefix. Returns (hidden, cache for these tokens). + + A cache holds each layer's rotated keys and values plus the pos/seg/opt ids of its tokens. The state + never attends to questions, so a state encoded once can serve any number of later questions. + """ + if cache is None: + g, l = self.masks(pos, seg, opt, self.config.local_attention, compare_options=self.config.compare_options) + else: + g, l = self.masks(pos, seg, opt, self.config.local_attention, + torch.cat([cache["pos"], pos], 1), torch.cat([cache["seg"], seg], 1), torch.cat([cache["opt"], opt], 1), + compare_options=self.config.compare_options) + h, kvs = self.emb_norm(emb), [] + for i, layer in enumerate(self.layers): + h, kv = layer(h, pos, g, l, None if cache is None else cache["kv"][i]) + kvs.append(kv) + return self.final_norm(h), {"kv": kvs, "pos": pos, "seg": seg, "opt": opt} + + @classmethod + def from_pretrained(cls, repo: str) -> "Encoder": + with open(hf_hub_download(repo, "config.json")) as f: + model = cls(Config.from_hf(json.load(f))) + sd = torch.load(hf_hub_download(repo, "pytorch_model.bin"), map_location="cpu", weights_only=True) + rename = { + "model.embeddings.tok_embeddings.": "tok_embeddings.", + "model.embeddings.norm.": "emb_norm.", + "model.final_norm.": "final_norm.", + ".attn.Wqkv.": ".Wqkv.", + ".attn.Wo.": ".Wo.", + ".mlp.Wi.": ".Wi.", + ".mlp.Wo.": ".Wo_mlp.", + "model.layers.": "layers.", + } + out = {} + for k, v in sd.items(): + if not k.startswith("model."): # MLM head, unused + continue + for a, b in rename.items(): + k = k.replace(a, b) + out[k] = v + model.load_state_dict(out, strict=True) + return model.eval() diff --git a/serve/amos/facts.py b/serve/amos/facts.py new file mode 100644 index 0000000000000000000000000000000000000000..161346816f64cfb906f27542e7239617c46f5b60 --- /dev/null +++ b/serve/amos/facts.py @@ -0,0 +1,32 @@ +"""Computed facts appended to a state, so the model reads date gaps instead of doing calendar arithmetic. + +The same function runs in training (train.py) and at inference (model.pack_request), switched by the run's +encoder config (date_facts), so a model always sees the input format it was trained on. +""" + +import re +from datetime import date +from itertools import combinations + +# An ISO date, optionally the value of a JSON key and optionally followed by a time ("2026-03-18T09:00:00Z"). +DATE = re.compile(r'(?:"([^"\n]{1,60})":\s*")?(\d{4})-(\d{2})-(\d{2})(?=[T "\s,.;)\]}]|$)') +MAX_DATES = 6 # 15 pairs at most, so long records gain a few lines, not a page + + +def date_facts(state: str) -> str: + """The state plus one line per pair of distinct dates it mentions, or the state unchanged if it has fewer than 2.""" + seen = {} + for m in DATE.finditer(state): + try: + d = date(int(m[2]), int(m[3]), int(m[4])) + except ValueError: + continue + seen.setdefault(d, m[1] or "a date in the text") + if len(seen) == MAX_DATES: + break + if len(seen) < 2: + return state + lines = [] + for (a, ka), (b, kb) in combinations(sorted(seen.items()), 2): + lines.append(f"- {ka} ({a}) is {(b - a).days} days before {kb} ({b}).") + return f"{state}\n\nDate facts:\n" + "\n".join(lines) diff --git a/serve/amos/model.py b/serve/amos/model.py new file mode 100644 index 0000000000000000000000000000000000000000..2ad0ea8dac16f61172a1d9961656a069aa4dd7d4 --- /dev/null +++ b/serve/amos/model.py @@ -0,0 +1,321 @@ +"""Amos: encoder + typed decision heads. Every output is an option, a span of the state, or a P(yes). + +All heads read the question's [DECIDE_t] vector: + noul softmax over two implicit options ["no: ...", "yes: ..."]; P(yes) is the second + choice/score softmax over bilinear(h_decide, mean-pooled option tokens) + multi independent sigmoid over the same pointer logits (+ bias) + span start/end softmax, bilinear(h_decide, state tokens); decoded to an exact substring + region softmax over the first state image's 49 cells, bilinear(h_decide, image tokens) + +With a vision tower, image slots are filled with projected vision features before the first layer; +the tower itself stays frozen and is not part of saved checkpoints. +""" + +import json +import math +from pathlib import Path + +import torch +import torch.nn.functional as F +from torch import nn + +from .encoder import Config, Encoder +from .facts import date_facts +from .packing import ( + IMAGE_TOKENS, + TYPES, + Codec, + Packed, + char_span_to_tokens, + collate, + pack, +) +from .vision import Projector, VisionTower, box_target, cell_boxes + +NOUL, CHOICE, SCORE, MULTI, SPAN, REGION = range(len(TYPES)) +MAX_SPAN_TOKENS = 30 + + +def _init(*mods): + for m in mods: + nn.init.normal_(m.weight, std=0.02) + nn.init.zeros_(m.bias) + + +class Amos(nn.Module): + def __init__(self, encoder: Encoder, vision: VisionTower | None = None): + super().__init__() + d = encoder.config.hidden_size + self.encoder = encoder + self.ptr_q, self.ptr_o = nn.Linear(d, d), nn.Linear(d, d) + self.multi_bias = nn.Parameter(torch.zeros(())) + self.span_q, self.span_k = nn.Linear(d, 2 * d), nn.Linear(d, 2 * d) + _init(self.ptr_q, self.ptr_o, self.span_q, self.span_k) + self.vision = vision + if vision is not None: + self.projector = Projector(vision.hidden_size, d) + self.region_q, self.region_k = nn.Linear(d, d), nn.Linear(d, d) + _init(self.region_q, self.region_k) + + @classmethod + def from_encoder(cls, repo: str, codec: Codec, vision_repo: str | None = None) -> "Amos": + enc = Encoder.from_pretrained(repo) + init_markers(enc, codec) + return cls(enc, VisionTower.from_pretrained(vision_repo) if vision_repo else None) + + def saved_state(self) -> dict: + """Weights to checkpoint: everything except the frozen, separately downloadable vision tower.""" + return {k: v for k, v in self.state_dict().items() if not k.startswith("vision.")} + + def embed(self, b: dict) -> torch.Tensor: + emb = self.encoder.tok_embeddings(b["ids"]) + if "images" not in b: + return emb + x = b["images"] + if x.dim() == 4: # pixels -> frozen tower; precomputed features skip it + with torch.no_grad(): + x = self.vision(x) + tokens = self.projector(x.to(emb.dtype)) + cols = b["img_at"][:, None] + torch.arange(IMAGE_TOKENS, device=emb.device) + rows = b["img_row"][:, None].expand_as(cols) + return emb.index_put((rows.reshape(-1), cols.reshape(-1)), tokens.reshape(-1, tokens.shape[-1])) + + def forward(self, b: dict) -> dict: + return self.heads(self.encoder(b["ids"], b["pos"], b["seg"], b["opt"], inputs_embeds=self.embed(b)), b) + + def forward_cached(self, b: dict, state_len: int, cache: dict | None): + """One pack (B = 1) with a reusable state. Returns (outputs, state cache); same result as forward(). + + The state never attends to the questions, so its keys, values and hidden states in a full pass are + exactly what a state-only pass would produce: a miss runs one normal pass and slices the cache out + of it; a hit encodes only the question tokens on top of the cached state. + """ + emb, s = self.embed(b), state_len + if cache is None: + h, kv = self.encoder.run(emb, b["pos"], b["seg"], b["opt"]) + cache = {"kv": [(k[:, :, :s], v[:, :, :s]) for k, v in kv["kv"]], "hidden": h[:, :s], + **{k: b[k][:, :s] for k in ("pos", "seg", "opt")}} + return self.heads(h, b), cache + h_q, _ = self.encoder.run(emb[:, s:], b["pos"][:, s:], b["seg"][:, s:], b["opt"][:, s:], cache) + return self.heads(torch.cat([cache["hidden"], h_q], 1), b), cache + + def heads(self, h: torch.Tensor, b: dict) -> dict: + d = h.shape[-1] + hq = h[b["q_row"], b["q_tok"]] + out = {} + + Q, O = len(hq), len(b["o_q"]) + opt_logits = h.new_full((Q, opt_width(b)), -torch.inf) + if O: + of = b["option_of"] + valid = of >= 0 + pooled = h.new_zeros(O, d).index_add_(0, of[valid], h[valid]) + pooled = pooled / torch.bincount(of[valid], minlength=O).clamp(min=1)[:, None] + ptr = (self.ptr_q(hq)[b["o_q"]] * self.ptr_o(pooled)).sum(-1) / math.sqrt(d) + ptr = ptr + (b["q_type"][b["o_q"]] == MULTI) * self.multi_bias + opt_logits = opt_logits.index_put((b["o_q"], b["o_rank"]), ptr) + out["opt"] = opt_logits + + sq = (b["q_type"] == SPAN).nonzero().squeeze(-1) + out["span_q"] = sq + if len(sq): + rows = b["q_row"][sq] + qs = self.span_q(hq[sq]).view(len(sq), 2, d) + ks = self.span_k(h[rows]).view(len(sq), -1, 2, d) + logits = torch.einsum("qcd,qlcd->qcl", qs, ks) / math.sqrt(d) + out["span"] = logits.masked_fill(~b["state_mask"][rows][:, None, :], -torch.inf) + + rq = (b["q_type"] == REGION).nonzero().squeeze(-1) + out["region_q"] = rq + if len(rq): + rows = b["q_row"][rq] + logits = torch.einsum("qd,qld->ql", self.region_q(hq[rq]), self.region_k(h[rows])) / math.sqrt(d) + out["region"] = logits.masked_fill(~b["region_mask"][rows], -torch.inf) + return out + + +def init_markers(enc: nn.Module, codec: Codec, markers=None): + """Reserved markers start as copies of [CLS] (their pretrained rows were never trained).""" + with torch.no_grad(): + for i in markers or (*codec.decide.values(), codec.opt, codec.img, codec.img_end, codec.patch): + enc.tok_embeddings.weight[i] = enc.tok_embeddings.weight[codec.cls] + + +def add_targets(b: dict, packs: list[Packed], questions: list[list[dict]]) -> dict: + """Attach target tensors to a collated batch. Targets are probabilities (hard labels are one-hot). + + noul: float P(yes). choice/score: distribution over options. multi: independent P per option. + span: list of accepted [start, end) char spans in the state. region: normalized (x0, y0, x1, y1) box + in the first state image, spread over the 49 cells by overlap. + """ + Q, L = len(b["q_row"]), b["ids"].shape[1] + t_opt = torch.zeros(Q, opt_width(b)) + t_span, t_region, has = [], [], torch.ones(Q, dtype=torch.bool) + qi = 0 + for row, (p, qs) in enumerate(zip(packs, questions)): + ok = b["state_mask"][row] + for q in qs: + tgt, typ = q["target"], q["type"] + if typ == "noul": # implicit options are ["no", "yes"] + t_opt[qi, :2] = torch.tensor([1 - float(tgt), float(tgt)]) + elif typ == "region": + dist = torch.zeros(L) + dist[p.image_at[0] : p.image_at[0] + IMAGE_TOKENS] = box_target(tgt) + t_region.append(dist) + elif typ == "span": + dist = torch.zeros(2, L) + for s, e in tgt: + tok = char_span_to_tokens(p.offsets, s, e) + if tok and ok[tok[0]] and ok[tok[1]]: # boundaries must be decodable (not truncated/blank) + dist[0, tok[0]] += 1 + dist[1, tok[1]] += 1 + has[qi] = bool(dist.sum()) + t_span.append(dist / dist.sum(-1, keepdim=True).clamp(min=1)) + else: + t_opt[qi, : len(tgt)] = torch.tensor(tgt, dtype=torch.float) + qi += 1 + b["t_opt"], b["has_target"] = t_opt, has + b["t_span"] = torch.stack(t_span) if t_span else torch.zeros(0, 2, L) + b["t_region"] = torch.stack(t_region) if t_region else torch.zeros(0, L) + return b + + +def opt_width(b: dict) -> int: + return max(1, int(b["q_nopt"].max())) if len(b["q_nopt"]) else 1 + + +def emd2(p: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Squared earth mover's distance between ordinal distributions [N, K] (padded with zeros).""" + return (p.cumsum(-1) - t.cumsum(-1)).pow(2).sum(-1) + + +def soft_ce(logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Cross-entropy against a (soft) target distribution; masked (-inf) logits must have zero target.""" + logp = logits.log_softmax(-1).masked_fill(target == 0, 0) + return -(target * logp).sum(-1) + + +def brier(p: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Half the squared distance between distributions [N, K]; padded options have p = t = 0.""" + return 0.5 * (p - t).pow(2).sum(-1) + + +def loss_fn(out: dict, b: dict, ordinal_weight: float = 1.0, brier_weight: float = 0.0, + type_weights: dict | None = None) -> tuple[torch.Tensor, dict]: + qt, has = b["q_type"], b["has_target"] + parts = {} + for typ in (NOUL, CHOICE, SCORE): + m = (qt == typ) & has + if m.any(): + ce = soft_ce(out["opt"][m], b["t_opt"][m]) + if brier_weight: + ce = ce + brier_weight * brier(out["opt"][m].softmax(-1), b["t_opt"][m]) + if typ == SCORE: + ce = ce + ordinal_weight * emd2(out["opt"][m].softmax(-1), b["t_opt"][m]) + parts[TYPES[typ]] = ce + m = (qt == MULTI) & has + if m.any(): + logits, valid = out["opt"][m], out["opt"][m].isfinite() + bce = F.binary_cross_entropy_with_logits(logits.masked_fill(~valid, 0), b["t_opt"][m], reduction="none") + parts["multi"] = (bce * valid).sum(-1) / valid.sum(-1) + if len(out["span_q"]): + m = has[out["span_q"]] + if m.any(): + parts["span"] = soft_ce(out["span"][m], b["t_span"][m]).mean(-1) + if len(out["region_q"]): + parts["region"] = soft_ce(out["region"], b["t_region"]) + w = type_weights or {} + total = torch.cat([v * w.get(k, 1.0) for k, v in parts.items()]).mean() + return total, {k: v.mean().item() for k, v in parts.items()} + + +def to_device(b: dict, device) -> dict: + return {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in b.items()} + + +def decode_span(start: torch.Tensor, end: torch.Tensor, p: Packed, length_penalty: float = 0.0, + max_tokens: int = MAX_SPAN_TOKENS): + """Best valid span under independent start/end softmax. Returns (char_start, char_end, text, prob). + length_penalty: each token past the start scales a span's score by exp(-length_penalty). max_tokens: the longest + span considered. + + The answer is always p.state[char_start:char_end]: it is cut from the input, never generated. + """ + n = len(p.offsets) + if not n or not start.isfinite().any(): + return 0, 0, "", 0.0 + ps, pe = start.softmax(-1)[1 : n + 1], end.softmax(-1)[1 : n + 1] # state text tokens only + K = min(max_tokens, n) + pe_pad = F.pad(pe, (0, K)) + joint = torch.stack([ps * pe_pad[k : k + n] for k in range(K)], 1) # [start, length - 1] + joint = joint * torch.exp(-length_penalty * torch.arange(K, dtype=joint.dtype)) + joint = joint / joint.sum().clamp(min=1e-12) + s, k = divmod(int(joint.argmax()), K) + prob = float(joint[s, k]) + a, z = p.offsets[s][0], p.offsets[s + k][1] + raw = p.state[a:z] + a, z = a + len(raw) - len(raw.lstrip()), z - (len(raw) - len(raw.rstrip())) + return a, max(a, z), p.state[a : max(a, z)], prob + + +def decode(out: dict, packs: list[Packed]) -> list[list[dict]]: + """Turn a batch forward into one typed result per question, grouped per pack.""" + opt = out["opt"].float().cpu() + span = out["span"].float().cpu() if "span" in out else None + region = out["region"].float().cpu() if "region" in out else None + res, qi, si, ri = [], 0, 0, 0 + for p in packs: + rs = [] + for t, n in zip(p.types, p.n_options): + if t == "noul": + rs.append({"type": t, "p_yes": opt[qi, :2].softmax(-1)[1].item()}) + elif t in ("choice", "score"): + probs = opt[qi, :n].softmax(-1) + r = {"type": t, "probs": probs.tolist(), "choice": int(probs.argmax())} + if t == "score": + r["expected"] = float((probs * torch.arange(n)).sum()) + rs.append(r) + elif t == "multi": + rs.append({"type": t, "probs": torch.sigmoid(opt[qi, :n]).tolist()}) + elif t == "region": + cells = region[ri].softmax(-1)[p.image_at[0] : p.image_at[0] + IMAGE_TOKENS] + k = int(cells.argmax()) + rs.append({"type": t, "cell": k, "box": cell_boxes()[k].tolist(), "p": float(cells[k]), "heatmap": cells.tolist()}) + ri += 1 + else: + a, z, text, prob = decode_span(span[si, 0], span[si, 1], p) + rs.append({"type": t, "start": a, "end": z, "text": text, "p": prob}) + si += 1 + qi += 1 + res.append(rs) + return res + + +def load_run(run, device="cpu"): + """A trained run directory (config.json + model.pt) -> (model in eval mode, its codec).""" + cfg = json.loads((Path(run) / "config.json").read_text()) + enc = cfg["encoder"] + if cfg.get("readout", "pointer") != "pointer" or enc.get("question_reads") or enc.get("model_type", "modernbert") != "modernbert": + raise ValueError(f"{run}: trained with a removed mode (MLM readout, question reads or DeBERTa)") + codec = Codec(cfg["repo"], max_state=cfg["max_state"], digits=cfg.get("digits", False), + score_rank=cfg.get("score_rank", False)) + vision = VisionTower.from_pretrained(cfg["vision"]) if cfg.get("vision") else None + model = Amos(Encoder(Config.from_hf(enc)), vision) + missing, unexpected = model.load_state_dict(torch.load(Path(run) / "model.pt", map_location="cpu", weights_only=True), strict=False) + if unexpected or any(not k.startswith("vision.") for k in missing): + raise ValueError(f"checkpoint mismatch: missing {missing}, unexpected {unexpected}") + return model.to(device).eval(), codec + + +def pack_request(model: Amos, codec: Codec, state: str, questions: list[dict], images=(), *, strict=False) -> list[Packed]: + """One pack holding every question (a list, so callers batch requests as rows).""" + if model.encoder.config.date_facts: + state = date_facts(state) + return [pack(codec, state, questions, images, strict=strict)] + + +@torch.no_grad() +def predict(model: Amos, codec: Codec, state: str, questions: list[dict], device="cpu", images=()) -> list[dict]: + """Run one state with its questions; returns one typed result per question.""" + ps = pack_request(model, codec, state, questions, images) + return [r for rs in decode(model(to_device(collate(ps, codec.pad), device)), ps) for r in rs] diff --git a/serve/amos/packing.py b/serve/amos/packing.py new file mode 100644 index 0000000000000000000000000000000000000000..1bade9ba1d331291cbfd81790cabeabaf3ce48dd --- /dev/null +++ b/serve/amos/packing.py @@ -0,0 +1,224 @@ +"""Pack one state and N typed questions into a single sequence. + +Layout: [CLS] state text [IMG] image tokens [/IMG] [SEP] | [DECIDE_t] q1 [SEP] [OPT] o1 [OPT] o2 ... | ... + +Attention tree (see Encoder.masks): state sees state; a question sees state + itself; +an option sees state + its question + itself. Positions restart after the state for +every question and after the question for every option, so results do not depend on +question order, option order, or which siblings share the pack. + +Images (state images or image options) occupy IMAGE_TOKENS placeholder slots each; the model +fills those slots with projected vision features before the first encoder layer. + +Structural tokens are inserted by id. User text is tokenized with special-token matching +disabled, and the reserved [unusedN] markers are made special, so no text can forge them. +""" + +import json +from dataclasses import dataclass, field + +import torch +from huggingface_hub import hf_hub_download +from tokenizers import Tokenizer, pre_tokenizers + +TYPES =("noul", "choice", "score", "multi", "span", "region") +OPTION_TYPES = ("choice", "score", "multi") +MAX_OPTIONS = 255 +IMAGE_TOKENS = 49 # merged patch tokens per image (vision.MERGED) +# Fixed [unusedN] slots, so adding a marker never moves one a trained checkpoint already uses. +MARKERS = {"noul": 0, "choice": 1, "score": 2, "multi": 3, "span": 4, "opt": 5, "region": 6, "img": 7, "img_end": 8, "patch": 9} + + +class Codec: + def __init__(self, repo: str, max_state: int = 512, max_question: int = 128, max_option: int = 48, + digits: bool = False, score_rank: bool = False): + with open(hf_hub_download(repo, "tokenizer.json")) as f: + spec = json.load(f) + for t in spec["added_tokens"]: + if t["content"].startswith("[unused"): + t["special"] = True + self.tok = Tokenizer.from_str(json.dumps(spec)) + if digits: + # BPE cuts numbers into arbitrary chunks (300020 -> "3000" "20"), so two numbers of the same size + # rarely line up. One token per digit lets the model compare them place by place. + self.tok.pre_tokenizer = pre_tokenizers.Sequence( + [pre_tokenizers.Digits(individual_digits=True), self.tok.pre_tokenizer]) + self.digits = digits + self.score_rank = score_rank + self.tok.encode_special_tokens = True # literal "[SEP]" / "[unused0]" in text stays plain text + self.tok.no_truncation() + self.tok.no_padding() + tid = self.tok.token_to_id + self.cls, self.sep, self.pad = tid("[CLS]"), tid("[SEP]"), tid("[PAD]") + marker = {k: tid(f"[unused{i}]") for k, i in MARKERS.items()} + assert None not in (self.cls, self.sep, self.pad, *marker.values()), f"{repo}: special tokens missing" + self.decide = {t: marker[t] for t in TYPES} + self.opt, self.img, self.img_end, self.patch = marker["opt"], marker["img"], marker["img_end"], marker["patch"] + self.reserved = {self.cls, self.sep, self.pad, *marker.values()} + self.max_state, self.max_question, self.max_option = max_state, max_question, max_option + + def encode(self, text: str) -> list[int]: + return self.tok.encode(text, add_special_tokens=False).ids + + +@dataclass +class Packed: + ids: list[int] = field(default_factory=list) + pos: list[int] = field(default_factory=list) + seg: list[int] = field(default_factory=list) # 0 state, k question k + opt: list[int] = field(default_factory=list) # 0 not an option, m option m of its question + decide: list[int] = field(default_factory=list) # token index of each question's DECIDE token + types: list[str] = field(default_factory=list) + n_options: list[int] = field(default_factory=list) + offsets: list[tuple[int, int]] = field(default_factory=list) # char offsets of state text tokens 1..n + state: str = "" + images: list = field(default_factory=list) # pixels [3, H, W] or vision features [patches, d], in slot order + image_at: list[int] = field(default_factory=list) # token index of each image's first placeholder + + def __len__(self): + return len(self.ids) + + +def noul_options(q: dict) -> list[str]: + """noul is scored as two implicit options, carrying the optional criteria descriptions.""" + c = q.get("criteria") or {} + return [f"no: {c['false']}" if c.get("false") else "no", f"yes: {c['true']}" if c.get("true") else "yes"] + + +def _image_tokens(codec: Codec, p: Packed, image, start: int) -> list[int]: + p.images.append(image) + p.image_at.append(start + 1) + return [codec.img, *[codec.patch] * IMAGE_TOKENS, codec.img_end] + + +def pack(codec: Codec, state: str, questions: list[dict], images=(), *, strict: bool = False) -> Packed: + """questions: [{"type": ..., "text": ..., "options": [...], "criteria": {...}}]. + + options: required for choice/score/multi (score: ordered low to high); an option may be an image + tensor instead of text. criteria: optional {"true": ..., "false": ...} for noul. span and region + take neither. images: state images, placed after the state text. + """ + enc = codec.tok.encode(state, add_special_tokens=False) + if strict and len(enc.ids) > codec.max_state: + raise ValueError(f"maximum context length: state has {len(enc.ids)} tokens, limit {codec.max_state}") + keep = min(len(enc.ids), codec.max_state) + p = Packed(state=state, offsets=enc.offsets[:keep]) + p.ids = [codec.cls, *enc.ids[:keep]] + for im in images: + p.ids += _image_tokens(codec, p, im, len(p.ids)) + p.ids.append(codec.sep) + n = len(p.ids) + p.pos, p.seg, p.opt = list(range(n)), [0] * n, [0] * n + for k, q in enumerate(questions, start=1): + t = q["type"] + if t not in TYPES: + raise ValueError(f"unknown question type {t!r}") + if t == "region" and not images: + raise ValueError("region questions need a state image") + options = noul_options(q) if t == "noul" else q.get("options") or [] + if t != "noul" and ((t in OPTION_TYPES) != bool(options) or len(options) > MAX_OPTIONS): + raise ValueError(f"{t} question needs {'1-255' if t in OPTION_TYPES else 'no'} options") + text = codec.encode(q["text"]) + if not strict: + text = text[: codec.max_question] + branch = [codec.decide[t], *text, codec.sep] + p.decide.append(len(p.ids)) + p.types.append(t) + p.n_options.append(len(options)) + p.ids += branch + p.pos += range(n, n + len(branch)) + p.seg += [k] * len(branch) + p.opt += [0] * len(branch) + o_start = n + len(branch) + for m, o in enumerate(options, start=1): + if isinstance(o, str): + if t == "score" and codec.score_rank: + o = f"Level {m - 1} of {len(options) - 1}: {o}" + text = codec.encode(o) + toks = [codec.opt, *(text if strict else text[: codec.max_option])] + else: + toks = [codec.opt, *_image_tokens(codec, p, o, len(p.ids))] + p.ids += toks + p.pos += range(o_start, o_start + len(toks)) + p.seg += [k] * len(toks) + p.opt += [m] * len(toks) + return p + + +def char_span_to_tokens(offsets, start: int, end: int): + """Map a char span [start, end) of the state to inclusive state-token indices (1-based), or None.""" + toks = [i for i, (a, b) in enumerate(offsets, start=1) if a < end and b > start] + return (toks[0], toks[-1]) if toks else None + + +PAD_TO = 64 # rows are padded to a multiple of this: few distinct shapes, so few cached kernels/graphs + + +def padded(n: int) -> int: + return -(-n // PAD_TO) * PAD_TO + + +def batches(lengths: list[int], budget: int, rng=None, chunk: int = 4096) -> list[list[int]]: + """Group indices into batches with rows * L <= budget and rows * L^2 <= budget * 512 (L = padded longest row). + + The second bound caps attention memory, which grows with L^2. Sorting by length limits padding; + with an rng, sorting happens within random chunks and the batch order is shuffled, so batches + differ every epoch. + """ + idx = list(range(len(lengths))) + if rng: + rng.shuffle(idx) + else: + chunk = len(idx) or 1 + + def fits(rows, longest): + L = padded(longest) + return rows * L <= budget and rows * L * L <= budget * 512 + + out = [] + for c in range(0, len(idx), chunk): + cur, longest = [], 0 + for i in sorted(idx[c : c + chunk], key=lengths.__getitem__): + if cur and not fits(len(cur) + 1, max(longest, lengths[i])): + out.append(cur) + cur, longest = [], 0 + cur.append(i) + longest = max(longest, lengths[i]) + if cur: + out.append(cur) + if rng: + rng.shuffle(out) + return out + + +def collate(packs: list[Packed], pad_id: int) -> dict: + """Right-pad packs into [B, L] tensors and index every question, option, and image in the batch.""" + B, L = len(packs), padded(max(len(p) for p in packs)) + t = {k: torch.full((B, L), v, dtype=torch.long) for k, v in [("ids", pad_id), ("pos", 0), ("seg", -1), ("opt", 0)]} + t["option_of"] = torch.full((B, L), -1, dtype=torch.long) # global option index per token + t["state_mask"] = torch.zeros(B, L, dtype=torch.bool) # span boundary candidates: non-blank state tokens + t["region_mask"] = torch.zeros(B, L, dtype=torch.bool) # region candidates: first state image's tokens + q_row, q_tok, q_type, q_nopt, o_q, o_rank, images, img_row, img_at = [], [], [], [], [], [], [], [], [] + for b, p in enumerate(packs): + n = len(p) + for k in ("ids", "pos", "seg", "opt"): + t[k][b, :n] = torch.tensor(getattr(p, k)) + t["state_mask"][b, 1 : 1 + len(p.offsets)] = torch.tensor([bool(p.state[a:z].strip()) for a, z in p.offsets]) + if p.image_at and p.seg[p.image_at[0]] == 0: + t["region_mask"][b, p.image_at[0] : p.image_at[0] + IMAGE_TOKENS] = True + images += p.images + img_row += [b] * len(p.images) + img_at += p.image_at + base = [] + for tok, typ, n_opt in zip(p.decide, p.types, p.n_options): + base.append(len(o_q)) + o_q += [len(q_row)] * n_opt + o_rank += range(n_opt) + q_row.append(b), q_tok.append(tok), q_type.append(TYPES.index(typ)), q_nopt.append(n_opt) + t["option_of"][b, :n] = torch.tensor([base[s - 1] + o - 1 if o else -1 for s, o in zip(p.seg, p.opt)]) + for k, v in [("q_row", q_row), ("q_tok", q_tok), ("q_type", q_type), ("q_nopt", q_nopt), ("o_q", o_q), + ("o_rank", o_rank), ("img_row", img_row), ("img_at", img_at)]: + t[k] = torch.tensor(v, dtype=torch.long) + if images: + t["images"] = torch.stack([im.float() for im in images]) + return t diff --git a/serve/amos/py.typed b/serve/amos/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/serve/amos/serve.py b/serve/amos/serve.py new file mode 100644 index 0000000000000000000000000000000000000000..7b7a9cc040172f6b7059902f1c21c30036187fcb --- /dev/null +++ b/serve/amos/serve.py @@ -0,0 +1,728 @@ +"""Local HTTP server: POST /v1/systemone. + + uv run python -m amos.serve runs/distilled --port 8000 + uv run python -m amos.serve runs/distilled --bench + uv run python -m amos.serve runs/dec --device cuda # a decoder run (RUN/model/config.json) is detected + uv run python -m amos.serve runs/dec --also amos-1-reason=runs/reason # a second model, picked by "model" + +Request: + {"state": str | object | array, + "images": ["", ...], # optional state images (vision runs) + "questions": {"": {"type": "noul" | "choice" | "score" | "multi" | "span" | "region", + "instructions": str | object | array, + "criteria": ...}}, + "abstain": 0.1} # optional: conformal sets at this alpha + +criteria: noul {"true": ..., "false": ...} (optional); choice and multi {"