Publish validated AMOS Core AI conversion with source provenance and runtime adapter

.gitattributes

@@ -33,3 +33,7 @@ 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
+heads.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
+model/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+text.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
+vision.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text

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.

README.md

@@ -0,0 +1,45 @@
+---
+license: apache-2.0
+base_model: mukaddamzaid/amos-1-small
+tags: [coreai, core-ai, apple, aimodel, decisions, vision]
+---
+
+# amos-1-small — Apple Core AI
+
+Independent AMOS-1 conversion with **bidirectional 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.17 development, NumPy and Pillow. 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.
+
+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.00001019**. 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 8 additional cases, maximum probability/box error was 0.034142. 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-small](https://huggingface.co/mukaddamzaid/amos-1-small). 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/`.
+
+From the original workspace layout, run `convert_decoder.py amos-1-small`, `validate_decoder.py amos-1-small`, then `smoke_decoder.py amos-1-small`. Release integrity is recorded in `SHA256SUMS`.

SHA256SUMS

@@ -0,0 +1,44 @@
+cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30  LICENSE
+917c1c22fba4cc984a5e95b01cc5e1a386e013fb599782b055b909c91647c06d  README.md
+9f25e439543c4493b66f609c2a4c57a3d97c11194ac247c3c043f4657725911e  SOURCE.json
+19dd2e174b2f7468407243220912947feb04a9260f347223021f880b532439a1  bf16-attempt.json
+5865684cabaa8e1a23243db7d9d32569a57c7b5011bf47ed95814b9813af57dd  bf16-comparison.json
+f9ca221eceeea6dc52774885a8c50ad727bb8c269a7ff4fbb439131d2f5be0f2  calibration.json
+2636cc2d838bc7395655ebb3d2ba933b186a43b13ce134c9a441ee8ba45a9d0d  config.json
+650d3c6c51a4fead342fba413ed8339f6e04d1f48569225e2215e14a837d62c0  convert_decoder.py
+ef80ea9244fb5d713b91e077dde578ebd1e1455ea0788ecec4b3691de1084e41  coreai_decoder.py
+3708c801947fd75c057ff534edc404aacb4cc0acc9e653c12f09207e0d848d37  heads.aimodel/main.hash
+132384c42c3a5c6269a79862d6540a5a6834b532674ba847d83a5b181f85a860  heads.aimodel/main.mlirb
+3284600a3ad7c8e378bbe228ce3eb967d5da56bb40959aba7e56b177a3fc928a  heads.aimodel/metadata.json
+cc3b3cabf8162d51af8764eb9f645fa21c45cde4eba25a31d73caf1534511a29  manifest.json
+273d8e0e683b885071fb17e08d71e5f2a5ddfb5309756181681de4f5a1822d80  model/chat_template.jinja
+90c4a5672369ff2b3010f057c5dd1abafafcd29c9890ad61e7a893f6a1fa2b5e  model/config.json
+62153eb6c69f2e1f426beaa8002b7186437e949c7588167085df14e10e9c0a73  model/generation_config.json
+3a159dfec9978a186a72ba085e0ad6a050f3968d8b364218d7bd13f5c89381f2  model/preprocessor_config.json
+aa7582a6b90dc3711049d293d985e549f2b22c0c4e5ad11b9272a03c6294886e  model/tokenizer.json
+8c54bd45ef2fc821a711502da04af1bff020fe02c2207117c7ffdbf1866e6e25  model/tokenizer_config.json
+7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13  model/video_preprocessor_config.json
+d7910a3390f8073d95908f3301c0ea0e2df4551699546430003bc8226373c2bc  qwen_export_helpers.py
+8d39d10513cef09276fb806d7d16b33c3c9b7d30a212bdd87f0eb16ad8796d2e  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
+c5b0aa4de369327de6d03ba19d322557568e9d9c312b901ab42478dc679199d2  smoke_decoder.py
+bfe4ef23e8e7bb5041f9e10a769b5c2169b26225ffdd7016e9062c689e593932  text.aimodel/main.hash
+1c469415a143a1d9d75486ed48dd9b32c1033f1854a6bfc394ae9eb08caa8c72  text.aimodel/main.mlirb
+a678dc0dcb7ff0a16708ee3e534460d0ba3c0679c9f43fac4c003b4d68d17b35  text.aimodel/metadata.json
+94dc6f950a92250e5711be9af9c1b7893109e1dd207f97f9ae00e5584b71ae0e  validate_decoder.py
+41865477766f6b49dcfa43354734bb6b2158b8395e80f6d9b711a13f1dcae541  validation.json
+7631b43ef5f505089f17d48a3a71ec71bb5fc210c933189ab779c176da39e051  vision.aimodel/main.hash
+1f646ae2677fb9ba2793ec942d2b9ed82690199d7a75685ecdf4a84ea791f38f  vision.aimodel/main.mlirb
+ffd58384957fd6de58cf8bb5ac50261aba88af370aa53499e065b0a635b903f8  vision.aimodel/metadata.json

SOURCE.json

@@ -0,0 +1,179 @@
+{
+  "repo": "mukaddamzaid/amos-1-small",
+  "revision": "09ba3e9d8f506cfcb96ea130845d5d76d9d92028",
+  "bytes": 4555951728,
+  "files": [
+    {
+      "path": ".gitattributes",
+      "size": 1636,
+      "lfs": null
+    },
+    {
+      "path": "LICENSE",
+      "size": 11358,
+      "lfs": null
+    },
+    {
+      "path": "README.md",
+      "size": 4347,
+      "lfs": null
+    },
+    {
+      "path": "assets/amos-1-small.png",
+      "size": 116696,
+      "lfs": {
+        "size": 116696,
+        "sha256": "fa3963e990c3e0e6f1f8f1296027b44ca080a1adab4855c7cae515a122d79089",
+        "pointer_size": 131
+      }
+    },
+    {
+      "path": "calibration.json",
+      "size": 147,
+      "lfs": null
+    },
+    {
+      "path": "config.json",
+      "size": 28,
+      "lfs": null
+    },
+    {
+      "path": "model/chat_template.jinja",
+      "size": 7755,
+      "lfs": null
+    },
+    {
+      "path": "model/config.json",
+      "size": 2855,
+      "lfs": null
+    },
+    {
+      "path": "model/generation_config.json",
+      "size": 116,
+      "lfs": null
+    },
+    {
+      "path": "model/heads.pt",
+      "size": 109123355,
+      "lfs": {
+        "size": 109123355,
+        "sha256": "30ba05e777e4a07e8aec5139eb6700d99b4d294d86ad67f312fc14e3e23575c1",
+        "pointer_size": 134
+      }
+    },
+    {
+      "path": "model/model.safetensors",
+      "size": 4426558864,
+      "lfs": {
+        "size": 4426558864,
+        "sha256": "790c9700fe4a0be2b79498919c04d74a3cbfaf01e451d4c589db601608f4f557",
+        "pointer_size": 135
+      }
+    },
+    {
+      "path": "model/preprocessor_config.json",
+      "size": 443,
+      "lfs": null
+    },
+    {
+      "path": "model/tokenizer.json",
+      "size": 19989510,
+      "lfs": {
+        "size": 19989510,
+        "sha256": "aa7582a6b90dc3711049d293d985e549f2b22c0c4e5ad11b9272a03c6294886e",
+        "pointer_size": 133
+      }
+    },
+    {
+      "path": "model/tokenizer_config.json",
+      "size": 1172,
+      "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": 1355,
+      "lfs": null
+    }
+  ],
+  "local_source_hashes": {
+    "model/model.safetensors": "790c9700fe4a0be2b79498919c04d74a3cbfaf01e451d4c589db601608f4f557",
+    "model/tokenizer_config.json": "8c54bd45ef2fc821a711502da04af1bff020fe02c2207117c7ffdbf1866e6e25",
+    "model/config.json": "90c4a5672369ff2b3010f057c5dd1abafafcd29c9890ad61e7a893f6a1fa2b5e",
+    "model/tokenizer.json": "aa7582a6b90dc3711049d293d985e549f2b22c0c4e5ad11b9272a03c6294886e",
+    "model/heads.pt": "30ba05e777e4a07e8aec5139eb6700d99b4d294d86ad67f312fc14e3e23575c1",
+    "model/generation_config.json": "62153eb6c69f2e1f426beaa8002b7186437e949c7588167085df14e10e9c0a73",
+    "model/chat_template.jinja": "273d8e0e683b885071fb17e08d71e5f2a5ddfb5309756181681de4f5a1822d80",
+    "model/video_preprocessor_config.json": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13",
+    "model/preprocessor_config.json": "3a159dfec9978a186a72ba085e0ad6a050f3968d8b364218d7bd13f5c89381f2"
+  }
+}

bf16-attempt.json

@@ -0,0 +1,298 @@
+[
+  {
+    "case": "refund",
+    "tokens": 122,
+    "max_probability_or_box_error": 0.0008294582366943359,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "noul",
+        "p_yes": 0.9724245667457581
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9623249769210815,
+          0.01487948652356863,
+          0.02279551699757576
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "noul",
+        "p_yes": 0.9723380208015442
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9614955186843872,
+          0.014982474967837334,
+          0.023522058501839638
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "span",
+    "tokens": 123,
+    "max_probability_or_box_error": 0,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "span",
+        "start": 55,
+        "end": 63,
+        "text": "INV-2231",
+        "p": 1.0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "span",
+        "start": 55,
+        "end": 63,
+        "text": "INV-2231",
+        "p": 1.0
+      }
+    ]
+  },
+  {
+    "case": "score",
+    "tokens": 98,
+    "max_probability_or_box_error": 0.001256406307220459,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "score",
+        "probs": [
+          0.011553974822163582,
+          0.012325819581747055,
+          0.05834781378507614,
+          0.9177724123001099
+        ],
+        "choice": 3,
+        "expected": 2.882338684052229
+      }
+    ],
+    "coreai": [
+      {
+        "type": "score",
+        "probs": [
+          0.011898641474545002,
+          0.012468610890209675,
+          0.059116706252098083,
+          0.9165160059928894
+        ],
+        "choice": 3,
+        "expected": 2.880250041373074
+      }
+    ]
+  },
+  {
+    "case": "multi",
+    "tokens": 144,
+    "max_probability_or_box_error": 0.0004277229309082031,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "multi",
+        "probs": [
+          0.9700461626052856,
+          0.9614417552947998,
+          0.022039415314793587
+        ]
+      }
+    ],
+    "coreai": [
+      {
+        "type": "multi",
+        "probs": [
+          0.9697818160057068,
+          0.961869478225708,
+          0.02190219797194004
+        ]
+      }
+    ]
+  },
+  {
+    "case": "negative",
+    "tokens": 83,
+    "max_probability_or_box_error": 0.0006557293236255646,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "noul",
+        "p_yes": 0.049154266715049744
+      }
+    ],
+    "coreai": [
+      {
+        "type": "noul",
+        "p_yes": 0.04980999603867531
+      }
+    ]
+  },
+  {
+    "case": "unicode",
+    "tokens": 128,
+    "max_probability_or_box_error": 0.0008890628814697266,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "span",
+        "start": 10,
+        "end": 17,
+        "text": "\u20ac123.45",
+        "p": 0.9967826008796692
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9754987955093384,
+          0.013173527084290981,
+          0.011327654123306274
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "span",
+        "start": 10,
+        "end": 17,
+        "text": "\u20ac123.45",
+        "p": 0.9976716637611389
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9759361147880554,
+          0.01289563812315464,
+          0.011168276891112328
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "options",
+    "tokens": 230,
+    "max_probability_or_box_error": 0.0005109310150146484,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "choice",
+        "probs": [
+          0.003980829380452633,
+          0.003826719941571355,
+          0.004242099355906248,
+          0.0034746285527944565,
+          0.003326982958242297,
+          0.0027517760172486305,
+          0.0033690459094941616,
+          0.0029055955819785595,
+          0.002698659896850586,
+          0.002124715130776167,
+          0.00298517057672143,
+          0.0031158491037786007,
+          0.003178240265697241,
+          0.00276554049924016,
+          0.0025187907740473747,
+          0.0038957030046731234,
+          0.00339102977886796,
+          0.9302572011947632,
+          0.0037882667966187,
+          0.0032934194896370173,
+          0.0014340636553242803,
+          0.0015281529631465673,
+          0.0019327420741319656,
+          0.0032147751189768314
+        ],
+        "choice": 17
+      }
+    ],
+    "coreai": [
+      {
+        "type": "choice",
+        "probs": [
+          0.004058705177158117,
+          0.003873375477269292,
+          0.004297157749533653,
+          0.0035347286611795425,
+          0.003358113579452038,
+          0.0027418085373938084,
+          0.0034121156204491854,
+          0.002912522992119193,
+          0.0026929345913231373,
+          0.002130611566826701,
+          0.0030158632434904575,
+          0.003093816339969635,
+          0.0032024341635406017,
+          0.002781065646559,
+          0.002530055120587349,
+          0.003932687919586897,
+          0.0034593497402966022,
+          0.9297462701797485,
+          0.003791194409132004,
+          0.003320459509268403,
+          0.0014135388191789389,
+          0.0015202258946374059,
+          0.0019439932657405734,
+          0.0032370148692280054
+        ],
+        "choice": 17
+      }
+    ]
+  },
+  {
+    "case": "vision",
+    "tokens": 155,
+    "max_probability_or_box_error": 0.047274649143218994,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "region",
+        "cell": 9,
+        "box": [
+          0.2031,
+          0.0,
+          0.5674,
+          0.86
+        ],
+        "p": 0.5970777273178101
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9851310849189758,
+          0.014868861995637417
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "region",
+        "cell": 9,
+        "box": [
+          0.2005,
+          0.0,
+          0.5598,
+          0.8525
+        ],
+        "p": 0.5498030781745911
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9848902821540833,
+          0.015109767206013203
+        ],
+        "choice": 0
+      }
+    ]
+  }
+]

bf16-comparison.json

@@ -0,0 +1,282 @@
+[
+  {
+    "case": "refund",
+    "max_probability_or_box_error": 0.0006063971668481827,
+    "upstream_bf16": [
+      {
+        "type": "noul",
+        "p_yes": 0.9724245667457581
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9623249769210815,
+          0.01487948652356863,
+          0.02279551699757576
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai_float32": [
+      {
+        "type": "noul",
+        "p_yes": 0.9725583791732788
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9617581963539124,
+          0.014839922077953815,
+          0.023401914164423943
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "span",
+    "max_probability_or_box_error": 0,
+    "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.0008858442306518555,
+    "upstream_bf16": [
+      {
+        "type": "score",
+        "probs": [
+          0.011553974822163582,
+          0.012325819581747055,
+          0.05834781378507614,
+          0.9177724123001099
+        ],
+        "choice": 3,
+        "expected": 2.882338684052229
+      }
+    ],
+    "coreai_float32": [
+      {
+        "type": "score",
+        "probs": [
+          0.01181461475789547,
+          0.012464579194784164,
+          0.058834243565797806,
+          0.916886568069458
+        ],
+        "choice": 3,
+        "expected": 2.880792770534754
+      }
+    ]
+  },
+  {
+    "case": "multi",
+    "max_probability_or_box_error": 0.00027555227279663086,
+    "upstream_bf16": [
+      {
+        "type": "multi",
+        "probs": [
+          0.9700461626052856,
+          0.9614417552947998,
+          0.022039415314793587
+        ]
+      }
+    ],
+    "coreai_float32": [
+      {
+        "type": "multi",
+        "probs": [
+          0.9698116779327393,
+          0.9617173075675964,
+          0.021801499649882317
+        ]
+      }
+    ]
+  },
+  {
+    "case": "negative",
+    "max_probability_or_box_error": 0.0009536556899547577,
+    "upstream_bf16": [
+      {
+        "type": "noul",
+        "p_yes": 0.049154266715049744
+      }
+    ],
+    "coreai_float32": [
+      {
+        "type": "noul",
+        "p_yes": 0.0501079224050045
+      }
+    ]
+  },
+  {
+    "case": "unicode",
+    "max_probability_or_box_error": 0.0007423162460327148,
+    "upstream_bf16": [
+      {
+        "type": "span",
+        "start": 10,
+        "end": 17,
+        "text": "\u20ac123.45",
+        "p": 0.9967826008796692
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9754987955093384,
+          0.013173527084290981,
+          0.011327654123306274
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai_float32": [
+      {
+        "type": "span",
+        "start": 10,
+        "end": 17,
+        "text": "\u20ac123.45",
+        "p": 0.9975249171257019
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9758961200714111,
+          0.012953538447618484,
+          0.011150323785841465
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "options",
+    "max_probability_or_box_error": 0.000368654727935791,
+    "upstream_bf16": [
+      {
+        "type": "choice",
+        "probs": [
+          0.003980829380452633,
+          0.003826719941571355,
+          0.004242099355906248,
+          0.0034746285527944565,
+          0.003326982958242297,
+          0.0027517760172486305,
+          0.0033690459094941616,
+          0.0029055955819785595,
+          0.002698659896850586,
+          0.002124715130776167,
+          0.00298517057672143,
+          0.0031158491037786007,
+          0.003178240265697241,
+          0.00276554049924016,
+          0.0025187907740473747,
+          0.0038957030046731234,
+          0.00339102977886796,
+          0.9302572011947632,
+          0.0037882667966187,
+          0.0032934194896370173,
+          0.0014340636553242803,
+          0.0015281529631465673,
+          0.0019327420741319656,
+          0.0032147751189768314
+        ],
+        "choice": 17
+      }
+    ],
+    "coreai_float32": [
+      {
+        "type": "choice",
+        "probs": [
+          0.004064517095685005,
+          0.00385798211209476,
+          0.004268613178282976,
+          0.0035152393393218517,
+          0.0033423572313040495,
+          0.0027638806495815516,
+          0.0034003155305981636,
+          0.0028926795348525047,
+          0.0026991423219442368,
+          0.00212881900370121,
+          0.0030107346829026937,
+          0.00310140079818666,
+          0.0031818472780287266,
+          0.002777676796540618,
+          0.002535800449550152,
+          0.003921480383723974,
+          0.003430022159591317,
+          0.9298885464668274,
+          0.003782423911616206,
+          0.0032975096255540848,
+          0.0014300619950518012,
+          0.0015298567013815045,
+          0.0019381335005164146,
+          0.0032409762497991323
+        ],
+        "choice": 17
+      }
+    ]
+  },
+  {
+    "case": "vision",
+    "max_probability_or_box_error": 0.03414243459701538,
+    "upstream_bf16": [
+      {
+        "type": "region",
+        "cell": 9,
+        "box": [
+          0.2031,
+          0.0,
+          0.5674,
+          0.86
+        ],
+        "p": 0.5970777273178101
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9851310849189758,
+          0.014868861995637417
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai_float32": [
+      {
+        "type": "region",
+        "cell": 9,
+        "box": [
+          0.1936,
+          0.0,
+          0.5538,
+          0.8461
+        ],
+        "p": 0.5629352927207947
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9847150444984436,
+          0.015284995548427105
+        ],
+        "choice": 0
+      }
+    ]
+  }
+]

calibration.json

@@ -0,0 +1,7 @@
+{
+ "span": {
+  "cap": 30,
+  "penalty": 0.2,
+  "fitted_on": "held-out quotes: the quote-amounts slice and longquotes clean.jsonl (fit halves)"
+ }
+}

config.json

@@ -0,0 +1 @@
+{"base": "Qwen/Qwen3.5-2B"}

convert_decoder.py

@@ -0,0 +1,46 @@
+"""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.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)
+  h=self.language_model(inputs_embeds=emb,attention_mask={'full_attention':attention_mask[:,None,None,:].bool(),'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},indent=2))
+ print('Saved',flush=True)
+if __name__=='__main__':main()

coreai_decoder.py

@@ -0,0 +1,85 @@
+"""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={}
+  for name in ['text','vision','heads']:
+   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
+  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<n:
+   if cursor in image_at:
+    rows,cols=image_at[cursor]
+    if (rows,cols)!=(7,7):raise ValueError('Only 224x224 image grids are supported')
+    end=cursor+49;args['position_ids'][1,0,cursor:end]=offset;args['position_ids'][2,0,cursor:end]=offset+np.repeat(np.arange(7),7);args['position_ids'][3,0,cursor:end]=offset+np.tile(np.arange(7),7)
+    args['visual_embeddings'][0,cursor:end]=features[fi];args['visual_mask'][0,cursor:end]=True;offset+=7;cursor=end;fi+=1
+   else:
+    args['position_ids'][1:,0,cursor]=offset;offset+=1;cursor+=1
+  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)}}

heads.aimodel/main.hash

@@ -0,0 +1,2 @@
+#��,:\bi��b�T
+Zh4�2gK�G�:[��`

heads.aimodel/main.mlirb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:132384c42c3a5c6269a79862d6540a5a6834b532674ba847d83a5b181f85a860
+size 109141184

heads.aimodel/metadata.json

@@ -0,0 +1,5 @@
+{
+  "creationDate" : "20261014T012237Z",
+  "assetVersion" : "2.0",
+  "producer" : "coreai-core 1.0.0b3"
+}

manifest.json

@@ -0,0 +1,23 @@
+{
+  "sequence_length": 512,
+  "hidden_size": 2048,
+  "outputs": [
+    "mul_1744",
+    "view_3923"
+  ],
+  "vision_output": "addmm_98",
+  "head_outputs": [
+    "div",
+    "div_1",
+    "clamp_2"
+  ],
+  "precision": "float32",
+  "echo": true,
+  "kinds": [
+    "span",
+    "region"
+  ],
+  "keep": 0.25,
+  "has_box": true,
+  "image_side": 224
+}

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<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> 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</IMPORTANT>' }}
+    {%- 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('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- 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 '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\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<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\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</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- 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 true %}
+        {{- '<think>\n' }}
+    {%- else %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

model/config.json

@@ -0,0 +1,110 @@
+{
+  "amos_export_dtype": "bfloat16",
+  "architectures": [
+    "Qwen3_5ForConditionalGeneration"
+  ],
+  "dtype": "bfloat16",
+  "image_token_id": 248056,
+  "model_type": "qwen3_5",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_output_gate": true,
+    "bos_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "full_attention_interval": 4,
+    "head_dim": 256,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 6144,
+    "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_conv_kernel_dim": 4,
+    "linear_key_head_dim": 128,
+    "linear_num_key_heads": 16,
+    "linear_num_value_heads": 16,
+    "linear_value_head_dim": 128,
+    "mamba_ssm_dtype": "float32",
+    "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": 8,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 2,
+    "pad_token_id": null,
+    "partial_rotary_factor": 0.25,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        11,
+        11,
+        10
+      ],
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 10000000,
+      "rope_type": "default"
+    },
+    "tie_word_embeddings": true,
+    "use_cache": false,
+    "vocab_size": 248320,
+    "transformers_version": "5.17.0"
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "use_cache": false,
+  "video_token_id": 248057,
+  "vision_config": {
+    "deepstack_visual_indexes": [],
+    "depth": 24,
+    "dtype": "bfloat16",
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1024,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "model_type": "qwen3_5_vision",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 2048,
+    "patch_size": 16,
+    "rope_parameters": {
+      "rope_theta": 10000.0,
+      "rope_type": "axial"
+    },
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 248054,
+  "vision_start_token_id": 248053,
+  "amos_codes": 255,
+  "amos_noncausal": true
+}

model/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 248044,
+  "transformers_version": "5.17.0",
+  "use_cache": true
+}

model/preprocessor_config.json

@@ -0,0 +1,26 @@
+{
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "Qwen2VLImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "merge_size": 2,
+  "patch_size": 16,
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "longest_edge": 16777216,
+    "shortest_edge": 65536
+  },
+  "temporal_patch_size": 2
+}

model/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa7582a6b90dc3711049d293d985e549f2b22c0c4e5ad11b9272a03c6294886e
+size 19989510

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": [
+    "<answer>"
+  ],
+  "image_token": "<|image_pad|>",
+  "is_local": true,
+  "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|>"
+}

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"
+}

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()

runtime-smoke.json

@@ -0,0 +1,12 @@
+{
+  "passed": true,
+  "pytorch_neural_execution_disabled": true,
+  "result": {
+    "answers": {
+      "refund": {
+        "type": "noul",
+        "noul": 0.9714967012405396
+      }
+    }
+  }
+}

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.

serve/amos/__init__.py

@@ -0,0 +1 @@
+"""Amos 1 serving code. See README.md."""

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
+

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 = "<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, <think>, 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 ("<think>", "</think>"))
+        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, <image|>; ->
+        (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("</think>")[0].replace("<|im_end|>", "").strip()
+        full = codec.tok(prefix + reasoning + "\n</think>\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
+

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()

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)

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]

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

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": ["<base64 image>", ...],                 # optional state images (vision runs)
+   "questions": {"<id>": {"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 {"<option>": description | null |
+{"image": "<base64>"}}; score [level descriptions, lowest first]; span and region none. Objects and
+arrays are rendered as JSON text. Encoded images are cached by content hash.
+
+Response: {"model", "answers": {"<id>": {...}}, "usage": {"input_tokens", "output_tokens": 0}}
+  noul    {"type", "noul": P(yes)}
+  choice  {"type", "choice", "probabilities": {option: p}, "confidence"}
+  score   {"type", "score": expected level, "legend": {"0": ...}, "probabilities": {"0": p}, "confidence"}
+  multi   {"type", "probabilities": {option: p}}
+  span    {"type", "span": text, "start", "end", "confidence"}
+  region  {"type", "box": [x0, y0, x1, y1] normalized, "cell": 0-48 (a decoder: its image's row-major token grid),
+           "confidence"}
+A decoder run answers what its checkpoint supports (/v1/info card.types) and refuses the rest, and abstain, with 422.
+With "abstain", noul/choice/score answers add "set" (options kept at that alpha) and "abstain" (set size != 1).
+"""
+
+import argparse
+import base64
+import hashlib
+import io
+import json
+import math
+import mimetypes
+import os
+import queue
+import re
+import shutil
+import sys
+import threading
+import time
+import traceback
+import urllib.request
+from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from urllib.parse import parse_qs, urlsplit
+
+import torch
+
+from .calibrate import apply
+from .model import decode, pack_request, to_device
+from .packing import MAX_OPTIONS, TYPES, batches, collate, pack
+
+STATE_CACHE_TOKENS = 16384  # cached state tokens per engine, across all cached states
+
+
+class BadRequest(ValueError):
+    pass
+
+
+def render(x) -> str:
+    return x if isinstance(x, str) else json.dumps(x, indent=2, ensure_ascii=False)
+
+
+def image_bytes(x: str) -> bytes:
+    """An image as sent: plain base64, or a data URL ("data:image/png;base64,...")."""
+    return base64.b64decode(x.split(",", 1)[1] if x.startswith("data:") else x)
+
+
+def parse(req: dict):
+    """Validate a request -> (state text, [question id], [internal question])."""
+    if not isinstance(req, dict) or "state" not in req or not isinstance(req.get("questions"), dict) or not req["questions"]:
+        raise BadRequest("need 'state' and a non-empty 'questions' object")
+    ids, qs = [], []
+    for qid, q in req["questions"].items():
+        if not isinstance(q, dict) or q.get("type") not in TYPES or "instructions" not in q:
+            raise BadRequest(f"question {qid!r}: need 'type' in {list(TYPES)} and 'instructions'")
+        t, c = q["type"], q.get("criteria")
+        out = {"type": t, "text": render(q["instructions"])}
+        if t == "noul" and c is not None:
+            if not isinstance(c, dict):
+                raise BadRequest(f"question {qid!r}: noul criteria must be an object")
+            out["criteria"] = {k: render(v) for k, v in c.items() if k in ("true", "false") and v is not None}
+        elif t in ("choice", "multi"):
+            if not isinstance(c, dict) or not 2 <= len(c) <= MAX_OPTIONS:
+                raise BadRequest(f"question {qid!r}: {t} criteria must map 2-{MAX_OPTIONS} options to descriptions")
+            out["names"] = list(c)
+            out["options"] = [k if v is None else v if isinstance(v, dict) and "image" in v else f"{k}: {render(v)}"
+                              for k, v in c.items()]
+        elif t == "score":
+            if not isinstance(c, list) or not 2 <= len(c) <= MAX_OPTIONS:
+                raise BadRequest(f"question {qid!r}: score criteria must be a list of 2-{MAX_OPTIONS} levels")
+            out["options"] = [render(v) for v in c]
+        ids.append(qid)
+        qs.append(out)
+    return render(req["state"]), ids, qs
+
+
+def free_gpu_cache(min_unused_bytes: int = 1 << 30) -> bool:
+    """Hand the CUDA caching allocator's unused blocks back to the device when they pass min_unused_bytes. #3 serves
+    two models on one GPU: a long request leaves the large holding gigabytes it no longer uses, and the medium then
+    runs out of memory on a long request of its own (2026-10-09, JevBench's hard tier). True when it freed."""
+    try:
+        import torch
+    except ImportError:
+        return False
+    if not torch.cuda.is_available() or torch.cuda.memory_reserved() - torch.cuda.memory_allocated() < min_unused_bytes:
+        return False
+    torch.cuda.empty_cache()
+    return True
+
+
+def shift(probs, cut: float) -> list[float]:
+    """Pick-any probabilities moved on the logit scale so 0.5 falls where the fitted cut was: p >= cut <=> p' >= 0.5."""
+    b = math.log(cut / (1 - cut))
+    return [1 / (1 + math.exp(b - math.log(max(p, 1e-9) / max(1 - p, 1e-9)))) for p in probs]
+
+
+def choice_confidence(probs):
+    k = len(probs)
+    return (max(probs) - 1 / k) / (1 - 1 / k)
+
+
+def score_confidence(probs):
+    L = len(probs)
+    mode = max(range(L), key=probs.__getitem__)
+    spread = sum(abs(i - (L - 1) / 2) for i in range(L)) / L
+    return max(0.0, 1 - sum(p * abs(i - mode) for i, p in enumerate(probs)) / spread)
+
+
+def answer(q: dict, r: dict) -> dict:
+    t = q["type"]
+    if t == "noul":
+        a = {"type": t, "noul": r["p_yes"]}
+    elif t == "choice":
+        a = {"type": t, "choice": q["names"][r["choice"]], "probabilities": dict(zip(q["names"], r["probs"])),
+             "confidence": choice_confidence(r["probs"])}
+    elif t == "score":
+        a = {"type": t, "score": r["expected"], "legend": {str(i): o for i, o in enumerate(q["options"])},
+             "probabilities": {str(i): p for i, p in enumerate(r["probs"])}, "confidence": score_confidence(r["probs"])}
+    elif t == "multi":
+        a = {"type": t, "probabilities": dict(zip(q["names"], r["probs"]))}
+    elif t == "region":
+        a = {"type": t, "box": r["box"], "cell": r["cell"], "confidence": r["p"]}
+    else:
+        a = {"type": t, "span": r["text"], "start": r["start"], "end": r["end"], "confidence": r["p"]}
+    if "set" in r:
+        labels = [False, True] if t == "noul" else q.get("names") or [str(i) for i in range(len(q["options"]))]
+        a["set"], a["abstain"] = [labels[i] for i in r["set"]], r["abstain"]
+    return a
+
+
+class Engine:
+    def __init__(self, run: str, device: str, name: str | None = None, *, strict: bool = False,
+                 state_cache: bool = True):
+        from .model import load_run
+
+        self.model, self.codec = load_run(run, device)
+        self.device, self.run = device, run
+        self.vision = self.model.vision is not None
+        self.strict = strict
+        self.use_state_cache = state_cache
+        cal = Path(run) / "calibration.json"
+        report = json.loads(cal.read_text()) if cal.exists() else {}
+        self.cal = report.get("calibration")
+        cfg = Path(run) / "config.json"
+        cfg = json.loads(cfg.read_text()) if cfg.exists() else {}
+        enc, args = cfg.get("encoder", {}), cfg.get("args", {})
+        self.strict = strict or args.get("strict_inputs", False)
+        self.max_len = args.get("max_len", self.codec.max_state + 1024)
+        size = (cfg.get("repo") or "").rsplit("-", 1)[-1]  # jhu-clsp/ettin-encoder-32m -> 32m
+        self.name = name or {"32m": "amos-1-mini", "68m": "amos-1"}.get(size, f"amos-1-{size}")
+        self.card = {  # what /v1/info reports about the model
+            "base": cfg.get("repo"),
+            "parameters": sum(p.numel() for p in self.model.parameters()),
+            "layers": enc.get("num_hidden_layers"),
+            "hidden": enc.get("hidden_size"),
+            "heads": enc.get("num_attention_heads"),
+            "max_state": cfg.get("max_state"),
+            "max_len": args.get("max_len"),
+            "strict_inputs": self.strict,
+            "compare_options": self.model.encoder.config.compare_options,
+            # a run may have trained a region head yet be served without a vision tower: only then can it point
+            "types": [*(t for t in args.get("types", "").split(",") if t and t != "region"),
+                      *(["region"] if self.model.vision else [])],
+            "temperature": (self.cal or {}).get("temperature", {}),
+            "held_out": report.get("test_held_out", {}),
+            "in_distribution": report.get("test_in_distribution", {}),
+            "benchmarks": {},
+        }
+        td = Path(run) / "typed_decisions.json"  # written by `python -m amos.typed_decisions RUN --out ...`
+        if td.exists():
+            day = time.strftime("%Y-%m-%d", time.localtime(td.stat().st_mtime))
+            self.card["benchmarks"]["typed_decisions"] = {"measured": day, **json.loads(td.read_text())["all"]}
+        self.lock = threading.Lock()  # one forward at a time; requests queue on it
+        self.image_cache: OrderedDict[str, torch.Tensor] = OrderedDict()  # content hash -> vision features
+        # state + image hashes -> encoded state; bounded by cached tokens, since one 8k-token state is ~2 GB of keys
+        # and values on a 400m encoder
+        self.state_cache: OrderedDict[str, dict] = OrderedDict()
+
+    def image(self, data: str) -> torch.Tensor:
+        """Base64 image -> frozen-tower features, cached by content hash (LRU, 1024 images)."""
+        from PIL import Image, UnidentifiedImageError
+
+        from .vision import preprocess
+
+        if self.model.vision is None:
+            raise BadRequest("this model has no vision tower")
+        key = hashlib.sha256(data.encode()).hexdigest()
+        if key not in self.image_cache:
+            try:
+                pixels = preprocess(Image.open(io.BytesIO(image_bytes(data))))
+            except (ValueError, UnidentifiedImageError) as e:
+                raise BadRequest(f"could not decode image: {e}") from e
+            self.image_cache[key] = self.model.vision(pixels[None].to(self.device))[0]
+            if len(self.image_cache) > 1024:
+                self.image_cache.popitem(last=False)
+        self.image_cache.move_to_end(key)
+        return self.image_cache[key]
+
+    @torch.no_grad()
+    def __call__(self, req: dict) -> dict:
+        state, ids, qs = parse(req)
+        alpha = req.get("abstain")
+        if alpha is not None and (self.cal is None or str(alpha) not in self.cal["qhat"].get("choice", {})):
+            raise BadRequest(f"abstain={alpha!r} needs a calibration.json fitted at that alpha")
+        with self.lock:
+            images = [self.image(x) for x in req.get("images") or []]
+            qs_in = [{**q, "options": [self.image(o["image"]) if isinstance(o, dict) else o for o in q["options"]]}
+                     if "options" in q else q for q in qs]
+            try:
+                ps = pack_request(self.model, self.codec, state, qs_in, images, strict=self.strict)
+                (p,) = ps
+                if self.strict and len(p) > self.max_len:
+                    raise ValueError(f"maximum context length: request has {len(p)} tokens, limit {self.max_len}")
+            except ValueError as e:
+                raise BadRequest(str(e)) from e
+            key = hashlib.sha256(json.dumps([state, [hashlib.sha256(x.encode()).hexdigest() for x in req.get("images") or []]]).encode()).hexdigest()
+            batch = to_device(collate(ps, self.codec.pad), self.device)
+            if self.use_state_cache:
+                out, self.state_cache[key] = self.model.forward_cached(
+                    batch, p.decide[0], self.state_cache.get(key))
+                self.state_cache.move_to_end(key)
+                while len(self.state_cache) > 1 and (
+                        len(self.state_cache) > 32
+                        or sum(c["pos"].shape[1] for c in self.state_cache.values()) > STATE_CACHE_TOKENS):
+                    self.state_cache.popitem(last=False)
+            else:
+                out = self.model(batch)
+        results = [r for rs in decode(out, ps) for r in rs]
+        if self.cal:
+            results = [apply(r, self.cal, alpha if r["type"] in ("noul", "choice", "score") else None) for r in results]
+        return {"model": self.name, "answers": {i: answer(q, r) for i, q, r in zip(ids, qs, results)},
+                "usage": {"input_tokens": len(p), "output_tokens": 0, "state_tokens": p.decide[0],
+                          "branch_tokens": {i: p.seg.count(k) for k, i in enumerate(ids, start=1)}}}
+
+    @torch.no_grad()
+    def explain(self, req: dict) -> dict:
+        """How much each sentence of the state moves one question's answer.
+
+        Every variant of the state with one sentence removed is packed as its own row and all rows run in a
+        single batched forward pass. Influence = answer score with the full state minus the score without
+        that sentence (positive: the sentence supports the current answer).
+        """
+        state, ids, qs = parse(req)
+        qid = req.get("question")
+        if qid not in ids:
+            raise BadRequest("'question' must name one of the questions")
+        q = qs[ids.index(qid)]
+        if q["type"] == "region":
+            raise BadRequest("explanations cover text states; region questions are not supported")
+        sents = [(m.start(), m.end()) for m in re.finditer(r"[^\n.!?]+[.!?]*[\"'”’)]*", state) if m.group().strip()][:64]
+        variants = [state] + [state[:a] + " " * (b - a) + state[b:] for a, b in sents]  # blank, keep offsets
+        with self.lock:
+            images = [self.image(x) for x in req.get("images") or []]
+            q_in = {**q, "options": [self.image(o["image"]) if isinstance(o, dict) else o for o in q["options"]]} if "options" in q else q
+            packs = [pack(self.codec, s, [q_in], images) for s in variants]
+            res = [None] * len(packs)
+            for idx in batches([len(p) for p in packs], 24000):  # 65 variants of an 8k state do not fit in one batch
+                chunk = [packs[i] for i in idx]
+                for i, r in zip(idx, decode(self.model(to_device(collate(chunk, self.codec.pad), self.device)), chunk)):
+                    res[i] = r[0]
+
+        def score(r):
+            if r["type"] == "noul":
+                return r["p_yes"]
+            if r["type"] in ("choice", "score"):
+                return r["probs"][res[0]["choice"]]
+            if r["type"] == "multi":
+                return sum(r["probs"]) / len(r["probs"])
+            return r["p"] if r["text"] == res[0]["text"] else 0.0
+
+        base = score(res[0])
+        return {"question": qid, "base": base,
+                "sentences": [{"start": a, "end": b, "influence": base - score(r)} for (a, b), r in zip(sents, res[1:])]}
+
+
+class DecoderEngine:
+    """A decoder run (amos.decoder: RUN/model is a Hugging Face checkpoint) behind the same API. Yes/no, choice and
+    score read option letters, pick-any asks one yes/no per option in the same prompt; state images and image options
+    need a checkpoint with a vision tower (decoder vision-merge); quote and point need trained heads."""
+
+    MAX_LEN = 16384  # tokens per request, images included; inputs are never shortened
+
+    def __init__(self, run: str, device: str, name: str | None = None, max_len: int | None = None,
+                 rotations: int = 1, think_below: float | None = None, think_tokens: int = 512):
+        from .decoder import (
+            calibration,
+            load_heads,
+            load_model,
+            multi_cut,
+            span_decoding,
+        )
+
+        self.multi_cut = multi_cut(run)
+        self.think_below, self.think_tokens = think_below, think_tokens  # decoder.results: think when unsure
+        self.MAX_LEN = max_len or self.MAX_LEN
+        path = Path(run) / "model"
+        self.model, self.codec = load_model(str(path), device)
+        self.model.eval()
+        self.heads = load_heads(path, self.model, self.codec)
+        if self.heads is not None:
+            self.heads.span = span_decoding(run)  # the run's own quote cap and penalty, if fitted
+        self.rotations, self.temperature = rotations, calibration(run)  # temperatures from `decoder calibrate`, if fitted
+        self.device, self.run, self.cal = device, run, None
+        self.vision = self.codec.image_processor is not None
+        self.autocast = device == "cuda" and getattr(self.model.config, "amos_export_dtype", None) != "float32"
+        cfg = Path(run) / "config.json"
+        cfg = json.loads(cfg.read_text()) if cfg.exists() else {}
+        text = getattr(self.model.config, "text_config", self.model.config)
+        self.name = name or f"amos-1-{(cfg.get('base') or 'decoder').rsplit('-', 1)[-1].lower()}"
+        heads = self.heads.kinds if self.heads is not None else []
+        self.types = ["noul", "choice", "score", "multi", *(k for k in heads if k == "span" or self.vision)]
+        self.card = {
+            "base": cfg.get("base"),
+            "parameters": sum(p.numel() for p in self.model.parameters()),
+            "layers": text.num_hidden_layers,
+            "hidden": text.hidden_size,
+            "heads": text.num_attention_heads,
+            "max_state": None,
+            "max_len": self.MAX_LEN,
+            "strict_inputs": True,
+            "compare_options": None,
+            "types": self.types,
+            "explain": False,  # no /v1/explain for decoder runs
+            "temperature": self.temperature or {},
+            "rotations": self.rotations,
+            "held_out": {},
+            "in_distribution": {},
+            "benchmarks": {},
+        }
+        td = Path(run) / "typed_decisions.json"
+        if td.exists():
+            day = time.strftime("%Y-%m-%d", time.localtime(td.stat().st_mtime))
+            self.card["benchmarks"]["typed_decisions"] = {"measured": day, **json.loads(td.read_text())["all"]}
+        self.lock = threading.Lock()
+        # thoughts written by a vLLM server serving the same checkpoint (AMOS_THINK_URL, e.g. http://10.0.0.5:8100):
+        # batched across requests there, while the letters are still read here; this model is free meanwhile
+        self.think_url = os.environ.get("AMOS_THINK_URL")
+
+    def remote_thoughts(self, prefixes: list[str], max_tokens: int) -> list[tuple[str, int]]:
+        def one(prefix):
+            body = json.dumps({"model": "amos", "prompt": prefix, "max_tokens": max_tokens, "temperature": 0,
+                               "stop": ["</think>"], "skip_special_tokens": False}).encode()
+            req = urllib.request.Request(f"{self.think_url}/v1/completions", body, {"Content-Type": "application/json"})
+            out = json.loads(urllib.request.urlopen(req, timeout=1800).read())
+            return out["choices"][0]["text"], out["usage"]["completion_tokens"]
+
+        self.lock.release()  # held by __call__; other requests use the model while this one waits for its thoughts
+        try:
+            with ThreadPoolExecutor(len(prefixes)) as pool:
+                return list(pool.map(one, prefixes))
+        finally:
+            self.lock.acquire()
+
+    def image(self, data: str):
+        from PIL import Image, UnidentifiedImageError
+
+        if not self.vision:
+            raise BadRequest("this model has no vision tower")
+        try:
+            return Image.open(io.BytesIO(image_bytes(data))).convert("RGB")
+        except (ValueError, UnidentifiedImageError) as e:
+            raise BadRequest(f"could not decode image: {e}") from e
+
+    @torch.no_grad()
+    def __call__(self, req: dict) -> dict:
+        from .decoder import results
+
+        state, ids, qs = parse(req)
+        if req.get("abstain") is not None:
+            raise BadRequest("abstain is not available for this model yet: it has no calibration.json")
+        for qid, q in zip(ids, qs):
+            if q["type"] not in self.types:
+                raise BadRequest(f"question {qid!r}: {q['type']} is not available for this model yet "
+                                 f"(it has no trained {q['type']} head{'' if self.vision else ' or no vision tower'})")
+        images = [self.image(x) for x in req.get("images") or []]
+        qs_in = [{**q, "options": [self.image(o["image"]) if isinstance(o, dict) else o for o in q["options"]]}
+                 if "options" in q else q for q in qs]
+        with self.lock, torch.autocast("cuda", dtype=torch.bfloat16, enabled=self.autocast):
+            try:
+                rs, n = results(self.model, self.codec, state, qs_in, images, self.heads, self.MAX_LEN, self.rotations,
+                                self.temperature, self.think_below, self.think_tokens,
+                                self.remote_thoughts if self.think_url else None)
+            except ValueError as e:
+                raise BadRequest(str(e)) from e
+        if self.multi_cut:
+            rs = [{**r, "probs": shift(r["probs"], self.multi_cut)} if r["type"] == "multi" else r for r in rs]
+        return {"model": self.name, "answers": {i: answer(q, r) for i, q, r in zip(ids, qs, rs)},
+                "usage": {"input_tokens": n, "output_tokens": 0}}
+
+    def explain(self, req: dict) -> dict:
+        raise BadRequest("explanations are not available for this model yet")
+
+
+def load_engine(run: str, device: str, name: str | None = None, max_len: int | None = None, rotations: int = 1, **kw):
+    """DecoderEngine when the run holds a decoder checkpoint (RUN/model/config.json), else the encoder Engine."""
+    think = {k: kw.pop(k) for k in ("think_below", "think_tokens") if k in kw}
+    if (Path(run) / "model" / "config.json").exists():
+        return DecoderEngine(run, device, name, max_len, rotations, **think)
+    return Engine(run, device, name, **kw)
+
+
+UI = Path(__file__).parent / "ui"  # the playground's static build (cd playground && bun run build)
+MAX_BODY = 4 << 20  # bytes
+
+
+class Server(ThreadingHTTPServer):
+    """socketserver's default backlog of 5 resets connections under a burst; queued requests wait in max_queue instead."""
+
+    request_queue_size = 256
+    daemon_threads = True
+RELEASE_DATE = "2026-09-30"
+
+
+class RequestLog:
+    """Every answered request, kept to improve the models (the site's /privacy page says so). One folder per UTC day:
+    requests.jsonl and, under images/, each photo once by content hash. A record holds the time, model, source
+    (playground or api), the case, questions, photo hashes, answers and server time; never the caller's address.
+    Folders older than `days` are deleted.
+
+    Writes run on a background thread, so the disk never slows an answer; when it falls behind, records are dropped
+    rather than queued without bound."""
+
+    def __init__(self, root, days: int = 90, clock=time.time):
+        self.root, self.days, self.clock = Path(root), days, clock
+        self.root.mkdir(parents=True, exist_ok=True)
+        self.queue: queue.Queue = queue.Queue(maxsize=1000)
+        self.pruned = None
+        threading.Thread(target=self.run, daemon=True).start()
+
+    def add(self, record: dict, images: list) -> bool:
+        try:
+            self.queue.put_nowait((record, images))
+            return True
+        except queue.Full:
+            return False
+
+    def run(self):
+        while True:
+            record, images = self.queue.get()
+            try:
+                self.write(record, images)
+            except Exception:
+                traceback.print_exc()
+            finally:
+                self.queue.task_done()
+
+    def write(self, record: dict, images: list):
+        now = self.clock()
+        day = time.strftime("%Y-%m-%d", time.gmtime(now))
+        folder = self.root / day
+        (folder / "images").mkdir(parents=True, exist_ok=True)
+        refs = []
+        for x in images:
+            raw = image_bytes(x)
+            ext = ".jpg" if raw[:2] == b"\xff\xd8" else ".png" if raw[:4] == b"\x89PNG" else \
+                ".webp" if raw[8:12] == b"WEBP" else ".bin"
+            name = hashlib.sha256(raw).hexdigest() + ext
+            if not (folder / "images" / name).exists():
+                (folder / "images" / name).write_bytes(raw)
+            refs.append(name)
+        line = {"ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(now)), **record, "images": refs}
+        with open(folder / "requests.jsonl", "a") as f:
+            f.write(json.dumps(line, ensure_ascii=False) + "\n")
+        if self.pruned != day:  # once a day: drop folders past retention
+            self.pruned = day
+            cutoff = time.strftime("%Y-%m-%d", time.gmtime(now - self.days * 86400))
+            for old in self.root.iterdir():
+                if old.is_dir() and re.fullmatch(r"\d{4}-\d{2}-\d{2}", old.name) and old.name < cutoff:
+                    shutil.rmtree(old)
+
+
+class RateLimit:
+    """Per-client token bucket: per_min requests a minute on average, bursts of up to burst."""
+
+    def __init__(self, per_min: float, burst: int, clock=time.monotonic):
+        self.rate, self.burst, self.clock = per_min / 60, burst, clock
+        self.buckets, self.lock = {}, threading.Lock()
+
+    def wait(self, key: str) -> float:
+        """0 if this request may go now (and it is counted), else the seconds until one may."""
+        with self.lock:
+            now = self.clock()
+            if len(self.buckets) > 10000:  # forget clients whose bucket has refilled
+                self.buckets = {k: v for k, v in self.buckets.items() if (now - v[1]) * self.rate < self.burst}
+            tokens, last = self.buckets.get(key, (self.burst, now))
+            tokens = min(self.burst, tokens + (now - last) * self.rate)
+            self.buckets[key] = (tokens - 1 if tokens >= 1 else tokens, now)
+            return 0.0 if tokens >= 1 else (1 - tokens) / self.rate
+
+
+PLAYGROUND_ORIGINS = ("https://amos.scira.ai",)
+
+
+def handler(engine: Engine, cors: str | None = None, max_queue: int = 8, limit: RateLimit | None = None,
+            extra: list | None = None, log: RequestLog | None = None, retired: dict[str, str] | None = None):
+    """engine answers by default; extra engines answer requests whose "model" names them. With a log, every answered
+    /v1/systemone request is kept. retired maps a withdrawn model's name to its replacement: asking for it gets 410
+    rather than another model's answers under the old name."""
+    retired = retired or {}
+
+    def gone(name):
+        return {"error": f"{name} has been retired; use {retired[name]} instead"}
+    # Forwards run one at a time (Engine.lock); past max_queue waiting requests, refuse rather than pile up.
+    slots = threading.BoundedSemaphore(max_queue)
+    extra = extra or []
+    models = {e.name: e for e in extra}  # lookups fall back to engine, so it needs no entry
+
+    class H(BaseHTTPRequestHandler):
+        def client(self) -> str:
+            # Behind Vercel: x-vercel-forwarded-for is set by Vercel, while a caller's own x-forwarded-for passes
+            # through untouched, so it is only the fallback when Vercel is not in front. These headers are trusted as
+            # sent: a caller who reaches the server directly can set any of them.
+            forwarded = (self.headers.get("x-vercel-forwarded-for") or self.headers.get("x-real-ip")
+                         or (self.headers.get("x-forwarded-for") or "").split(",")[0])
+            return forwarded.strip() or self.client_address[0]
+
+        def address_string(self):  # the access log names the caller, not the proxy in front
+            return self.client()
+
+        def do_GET(self):
+            path = self.path.split("?")[0]
+            if path == "/v1/info":  # ?model=NAME for a model other than the default
+                name = parse_qs(urlsplit(self.path).query).get("model", [engine.name])[0]
+                if name in retired:
+                    return self.send(410, gone(name))
+                e = engine if name == engine.name else models.get(name)
+                if e is None:
+                    return self.send(404, {"error": f"no model {name!r}; served: {[engine.name, *models]}"})
+                alphas = sorted(e.cal["qhat"].get("choice", {})) if e.cal else []
+                return self.send(200, {"model": e.name, "vision": e.vision,
+                                       "abstain": [float(a) for a in alphas], "card": e.card})
+            if path == "/v1/models":  # what the TypeSafe SDKs' models.list() reads
+                return self.send(200, {"models": [{"name": e.name, "release_date": RELEASE_DATE,
+                                                   "description": f"{e.card['parameters'] / 1e9:.1f}B parameters: "
+                                                                  + ", ".join(e.card["types"])}
+                                                  for e in [engine, *extra]]})
+            if not UI.is_dir():
+                return self.send(404, {"error": "playground not built: run `bun install && bun run build` in playground/"})
+            file = (UI / path.lstrip("/")).resolve()
+            if file.is_dir():
+                file = file / "index.html"
+            if not file.is_relative_to(UI.resolve()) or not file.is_file():
+                return self.send(404, {"error": "not found"})
+            self.send_bytes(200, file.read_bytes(), mimetypes.guess_type(file.name)[0] or "application/octet-stream")
+
+        def do_OPTIONS(self):  # CORS preflight, only when a dev origin was allowed
+            self.send_bytes(204, b"", "text/plain")
+
+        def do_POST(self):
+            if self.path not in ("/v1/systemone", "/v1/explain"):
+                return self.send(404, {"error": "not found"})
+            try:
+                size = int(self.headers.get("Content-Length", 0))
+            except ValueError:
+                return self.send(400, {"error": "bad Content-Length"})
+            if size > MAX_BODY:
+                self.close_connection = True  # the unread body cannot be reused as the next request
+                return self.send(413, {"error": f"request body over {MAX_BODY >> 20} MB"})
+            if limit and (wait := limit.wait(self.client())):
+                self.close_connection = True
+                return self.send(429, {"error": "rate limit reached; retry after the Retry-After header"},
+                                 {"Retry-After": str(math.ceil(wait))})
+            if not slots.acquire(blocking=False):
+                self.close_connection = True
+                return self.send(503, {"error": "busy, retry shortly"}, {"Retry-After": "1"})
+            try:
+                req = json.loads(self.rfile.read(size))
+                if isinstance(req, dict) and req.get("model") in retired:
+                    return self.send(410, gone(req["model"]))
+                # the request's model picks among served models; any other name (the SDKs send jev-latest) gets the first
+                eng = models.get(req.get("model"), engine) if isinstance(req, dict) else engine
+                if self.path != "/v1/systemone":
+                    return self.send(200, eng.explain(req))
+                t = time.perf_counter()
+                try:
+                    res = eng(req)
+                except Exception as e:  # out of GPU memory: give back the cache and try once more
+                    if type(e).__name__ != "OutOfMemoryError" or not free_gpu_cache(0):
+                        raise
+                    res = eng(req)
+                if log:
+                    log.add({"model": res.get("model"), "source": "playground" if self.headers.get("Origin") in
+                             PLAYGROUND_ORIGINS else "api", "state": req["state"], "questions": req["questions"],
+                             "answers": res.get("answers"), "ms": round((time.perf_counter() - t) * 1000, 1)},
+                            list(req.get("images") or []))
+                self.send(200, res)
+            except (BadRequest, json.JSONDecodeError, ValueError) as e:
+                self.send(422, {"error": str(e)})
+            except Exception:  # a bug, not the request: answer instead of dropping the connection
+                traceback.print_exc()
+                self.send(500, {"error": "internal error: the server failed to answer this request"})
+            finally:
+                slots.release()
+                free_gpu_cache()  # so the other model on this GPU can use what this request no longer needs
+
+        def send(self, code, body, headers=None):
+            self.send_bytes(code, json.dumps(body).encode(), "application/json", headers)
+
+        def send_bytes(self, code, data, content_type, headers=None):
+            self.send_response(code)
+            self.send_header("Content-Type", content_type)
+            self.send_header("Content-Length", str(len(data)))
+            for k, v in (headers or {}).items():
+                self.send_header(k, v)
+            if cors:
+                self.send_header("Access-Control-Allow-Origin", cors)
+                self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
+                self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization")
+                self.send_header("Access-Control-Max-Age", "7200")  # a far caller would otherwise pay a preflight each call
+            self.end_headers()
+            self.wfile.write(data)
+
+        def log_message(self, *args):
+            pass
+
+    return H
+
+
+def bench(engine: Engine, runs=20):
+    """Median latency by state length and number of questions: new state (cold) and cached state (warm)."""
+    word = "The shipment of 40 units arrived late and two boxes were damaged in transit. "
+    rows = []
+    for words in (50, 200, 800):
+        for n_q in (1, 8, 32):
+            req = {"state": (word * (words // 13 + 1))[: words * 6],
+                   "questions": {f"q{i}": [{"type": "noul", "instructions": "Was anything damaged?"},
+                                           {"type": "choice", "instructions": "Main issue?", "criteria": {"late": None, "damage": None, "missing": None}},
+                                           {"type": "score", "instructions": "How severe?", "criteria": ["minor", "moderate", "major"]}][i % 3]
+                                 for i in range(n_q)}}
+            engine(req)
+            med = {}
+            for mode in ("cold", "warm"):
+                times = []
+                for _ in range(runs):
+                    if mode == "cold":
+                        getattr(engine, "state_cache", {}).clear()  # a decoder keeps no state cache
+                    t = time.perf_counter()
+                    engine(req)
+                    times.append(time.perf_counter() - t)
+                med[mode] = sorted(times)[runs // 2] * 1000
+            rows.append((words, n_q, med["cold"], med["warm"]))
+            print(f"{engine.device}  state ~{words} words  {n_q:>2} questions  cold {med['cold']:6.1f} ms  warm {med['warm']:6.1f} ms", flush=True)
+    return rows
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser()
+    ap.add_argument("run", help="a run directory (encoder, or decoder with RUN/model)")
+    ap.add_argument("--host", default="127.0.0.1", help="0.0.0.0 inside a container")
+    ap.add_argument("--port", type=int, default=int(os.environ.get("PORT", 8000)))
+    ap.add_argument("--max-queue", type=int, default=8, help="requests allowed to wait for the model before 503")
+    ap.add_argument("--rate-limit", type=float, default=0,
+                    help="requests a minute per caller before 429 (zero: no limit); for a public endpoint")
+    ap.add_argument("--burst", type=int, default=20, help="requests a caller may send at once under --rate-limit")
+    ap.add_argument("--device", default="cpu")
+    ap.add_argument("--name", help="model name to report (default: amos-1 for 68m, amos-1-mini for 32m)")
+    ap.add_argument("--bench", action="store_true")
+    ap.add_argument("--check", action="store_true", help="load the model and exit (an image build caches the tokenizer)")
+    ap.add_argument("--strict", action="store_true", help="preserve question/option text and reject context overflow")
+    ap.add_argument("--max-len", type=int, help="decoder runs: tokens allowed per request (default DecoderEngine.MAX_LEN)")
+    ap.add_argument("--rotations", type=int, default=1, help="decoder runs: average choice reads over this many option orders")
+    ap.add_argument("--think-below", type=float, help="decoder runs: re-ask an answer less sure than this after the "
+                                                      "model reasons in its thinking block")
+    ap.add_argument("--think-tokens", type=int, default=512, help="decoder runs: the most tokens a thought may take")
+    ap.add_argument("--no-state-cache", action="store_true", help="encode each request without retaining state tensors")
+    ap.add_argument("--cors", help="allow this browser origin to call the API (e.g. http://localhost:3000 for `bun dev`)")
+    ap.add_argument("--also", action="append", default=[], metavar="NAME=RUN",
+                    help="serve another run beside the first, answering requests whose \"model\" is NAME")
+    ap.add_argument("--retired", action="append", default=[], metavar="NAME=REPLACEMENT",
+                    help="a withdrawn model: requests naming it get 410 and the replacement's name")
+    ap.add_argument("--typed-decisions", help="typed-decisions metrics JSON for /v1/info (default: the run's "
+                                              "typed_decisions.json)")
+    ap.add_argument("--log-dir", help="keep every answered request here to improve the models (RequestLog). Publish "
+                                      "the notice (/privacy) before turning this on")
+    ap.add_argument("--log-days", type=int, default=90, help="days of logged requests to keep")
+    a = ap.parse_args(argv)
+    engine = load_engine(a.run, a.device, a.name, a.max_len, a.rotations, strict=a.strict,
+                         state_cache=not a.no_state_cache, think_below=a.think_below, think_tokens=a.think_tokens)
+    if a.typed_decisions:
+        td = Path(a.typed_decisions)
+        day = time.strftime("%Y-%m-%d", time.localtime(td.stat().st_mtime))
+        engine.card["benchmarks"]["typed_decisions"] = {"measured": day, **json.loads(td.read_text())["all"]}
+    if a.check:
+        print(f"{engine.name}: {engine.card['parameters']:,} parameters, max_state {engine.card['max_state']}", flush=True)
+        return
+    if a.bench:
+        bench(engine)
+        return
+    extra = [load_engine(run, a.device, name, a.max_len, a.rotations)
+             for name, _, run in (s.partition("=") for s in a.also)]
+    print(f"serving {a.run}{''.join(f', {e.name}' for e in extra)}: playground http://{a.host}:{a.port}/, "
+          f"API POST http://{a.host}:{a.port}/v1/systemone", flush=True)
+    limit = RateLimit(a.rate_limit, a.burst) if a.rate_limit else None
+    log = RequestLog(a.log_dir, a.log_days) if a.log_dir else None
+    Server((a.host, a.port), handler(engine, a.cors, a.max_queue, limit, extra, log,
+                                                         dict(s.split("=", 1) for s in a.retired))).serve_forever()
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

serve/amos/vision.py

@@ -0,0 +1,133 @@
+"""SigLIP-architecture vision tower written from scratch (loads SigLIP 2 weights), plus the projector.
+
+An image becomes 14 x 14 patch vectors; the projector merges each 2 x 2 block (4x fewer tokens) and maps
+it into the text encoder's embedding space with a 2-layer MLP, giving 49 image tokens.
+"""
+
+import json
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from torch import nn
+
+IMAGE_SIZE, PATCH, GRID = 224, 16, 14
+MERGED = (GRID // 2) ** 2  # image tokens per image after 2 x 2 merging
+
+
+def preprocess(image) -> torch.Tensor:
+    """PIL image -> [3, 224, 224] in [-1, 1] (bilinear resize, as SigLIP's processor does)."""
+    from PIL import Image
+
+    image = image.convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE), Image.Resampling.BILINEAR)
+    x = torch.frombuffer(bytearray(image.tobytes()), dtype=torch.uint8).view(IMAGE_SIZE, IMAGE_SIZE, 3)
+    return x.permute(2, 0, 1).float() / 127.5 - 1
+
+
+class Block(nn.Module):
+    def __init__(self, d, heads, ff, eps):
+        super().__init__()
+        self.heads = heads
+        self.layer_norm1, self.layer_norm2 = nn.LayerNorm(d, eps=eps), nn.LayerNorm(d, eps=eps)
+        self.q_proj, self.k_proj, self.v_proj, self.out_proj = (nn.Linear(d, d) for _ in range(4))
+        self.fc1, self.fc2 = nn.Linear(d, ff), nn.Linear(ff, d)
+
+    def forward(self, x):
+        B, N, d = x.shape
+        h = self.layer_norm1(x)
+        q, k, v = (p(h).view(B, N, self.heads, -1).transpose(1, 2) for p in (self.q_proj, self.k_proj, self.v_proj))
+        x = x + self.out_proj(F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(B, N, d))
+        return x + self.fc2(F.gelu(self.fc1(self.layer_norm2(x)), approximate="tanh"))
+
+
+class VisionTower(nn.Module):
+    def __init__(self, d=768, layers=12, heads=12, ff=3072, eps=1e-6):
+        super().__init__()
+        self.patch_embedding = nn.Conv2d(3, d, PATCH, PATCH)
+        self.position_embedding = nn.Parameter(torch.zeros(GRID * GRID, d))
+        self.blocks = nn.ModuleList(Block(d, heads, ff, eps) for _ in range(layers))
+        self.post_layernorm = nn.LayerNorm(d, eps=eps)
+        self.hidden_size = d
+
+    def forward(self, pixels):  # [B, 3, 224, 224] -> [B, 196, d]
+        # cuDNN disabled: its runtime-compiled engines fail to load on some CUDA hosts. The one conv is a cheap patchify.
+        from torch.nn.attention import SDPBackend, sdpa_kernel
+
+        with torch.backends.cudnn.flags(enabled=False), sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]):
+            x = self.patch_embedding(pixels).flatten(2).transpose(1, 2) + self.position_embedding
+            for b in self.blocks:
+                x = b(x)
+            return self.post_layernorm(x)
+
+    @classmethod
+    def from_pretrained(cls, repo: str) -> "VisionTower":
+        with open(hf_hub_download(repo, "config.json")) as f:
+            vc = json.load(f).get("vision_config", {})
+        m = cls(vc.get("hidden_size", 768), vc.get("num_hidden_layers", 12), vc.get("num_attention_heads", 12),
+                vc.get("intermediate_size", 3072), vc.get("layer_norm_eps", 1e-6))
+        sd = {}
+        for k, v in load_file(hf_hub_download(repo, "model.safetensors")).items():
+            if not k.startswith("vision_model.") or ".head." in k:
+                continue
+            k = k.removeprefix("vision_model.").replace("embeddings.", "").replace("encoder.layers.", "blocks.")
+            k = k.replace("self_attn.", "").replace("mlp.", "")
+            sd["position_embedding" if k == "position_embedding.weight" else k] = v
+        m.load_state_dict(sd, strict=True)
+        return m.eval().requires_grad_(False)
+
+
+class Projector(nn.Module):
+    """2 x 2 patch merge (concatenate neighbours) + 2-layer MLP into the text embedding space."""
+
+    def __init__(self, d_vision: int, d_text: int):
+        super().__init__()
+        self.fc1, self.fc2 = nn.Linear(4 * d_vision, d_text), nn.Linear(d_text, d_text)
+
+    def forward(self, feats):  # [N, 196, dv] -> [N, 49, dt]
+        N, _, dv = feats.shape
+        g = GRID // 2
+        x = feats.view(N, g, 2, g, 2, dv).permute(0, 1, 3, 2, 4, 5).reshape(N, g * g, 4 * dv)
+        return self.fc2(F.gelu(self.fc1(x)))
+
+
+def cell_boxes(rows: int = GRID // 2, cols: int = GRID // 2) -> torch.Tensor:
+    """[rows * cols, 4] normalized (x0, y0, x1, y1) of each merged image token, row-major (49 by default)."""
+    i = torch.arange(rows * cols)
+    r, c = i // cols, i % cols
+    return torch.stack([c / cols, r / rows, (c + 1) / cols, (r + 1) / rows], 1).float()
+
+
+def box_target(box, rows: int = GRID // 2, cols: int = GRID // 2) -> torch.Tensor:
+    """Distribution over the cells proportional to overlap with a normalized box (x0, y0, x1, y1)."""
+    cells, (x0, y0, x1, y1) = cell_boxes(rows, cols), box
+    w = (cells[:, 2].clamp(max=x1) - cells[:, 0].clamp(min=x0)).clamp(min=0)
+    h = (cells[:, 3].clamp(max=y1) - cells[:, 1].clamp(min=y0)).clamp(min=0)
+    area = w * h
+    return area / area.sum() if area.sum() > 0 else torch.full((len(cells),), 1 / len(cells))
+
+
+def peak_box(p: torch.Tensor, rows: int, cols: int, keep: float) -> tuple[list[float], float]:
+    """Box around the most likely cell and the cells 4-connected to it that hold at least `keep` of its probability ->
+    (normalized (x0, y0, x1, y1), probability inside the box). keep=1 is (about) the most likely cell alone."""
+    g = p.view(rows, cols)
+    r, c = divmod(int(p.argmax()), cols)
+    on, seen, todo = (g >= keep * g[r, c]).tolist(), {(r, c)}, [(r, c)]
+    while todo:
+        i, j = todo.pop()
+        for y, x in ((i + 1, j), (i - 1, j), (i, j + 1), (i, j - 1)):
+            if 0 <= y < rows and 0 <= x < cols and on[y][x] and (y, x) not in seen:
+                seen.add((y, x))
+                todo.append((y, x))
+    r0, r1 = min(y for y, _ in seen), max(y for y, _ in seen) + 1
+    c0, c1 = min(x for _, x in seen), max(x for _, x in seen) + 1
+    return [c0 / cols, r0 / rows, c1 / cols, r1 / rows], float(g[r0:r1, c0:c1].sum())
+
+
+def box_scores(box, gold) -> tuple[bool, float]:
+    """(the box's centre lies inside the gold box, intersection over union) of normalized (x0, y0, x1, y1) boxes."""
+    cx, cy = (box[0] + box[2]) / 2, (box[1] + box[3]) / 2
+    w = max(0.0, min(box[2], gold[2]) - max(box[0], gold[0]))
+    h = max(0.0, min(box[3], gold[3]) - max(box[1], gold[1]))
+    union = (box[2] - box[0]) * (box[3] - box[1]) + (gold[2] - gold[0]) * (gold[3] - gold[1]) - w * h
+    return gold[0] <= cx <= gold[2] and gold[1] <= cy <= gold[3], w * h / union if union > 0 else 0.0

serve/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "amos-serve"
+version = "1.0.0"
+description = "Serving code for the Amos 1 decision models: one HTTP server, POST /v1/systemone, answers typed questions from one forward pass"
+readme = "README.md"
+requires-python = ">=3.12"
+license = "Apache-2.0"
+# pinned to the versions the published numbers were measured with: with these three the package answers byte for
+# byte like our own servers; transformers 5.19 or the pure-PyTorch linear attention move probabilities by up to 0.02
+dependencies = [
+    "torch==2.14.0",
+    "transformers==5.17.0",
+    "flash-linear-attention==0.5.2",
+    "safetensors>=0.8.0",
+    "tokenizers>=0.23.2",
+    "pillow>=12.3.0",
+    "huggingface-hub>=1.33.0",
+]
+
+[build-system]
+requires = ["setuptools>=75"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+include = ["amos*"]

serve/sources.json

@@ -0,0 +1,59 @@
+{
+ "commit": "cd57f1762f3029565fb3a877dce9e598739b971e",
+ "version": "1.0.0",
+ "files": {
+  "amos/serve.py": {
+   "from": "src/amos/serve.py",
+   "verbatim": true,
+   "sha256": "b0523591c320a074a0ec0751b68ebc205aeedd090ff5e7ca10bda3fb94a07ef8"
+  },
+  "amos/model.py": {
+   "from": "src/amos/model.py",
+   "verbatim": true,
+   "sha256": "d4f4354f3e4d3b475b4c1f832a4a9741da757a368a1d2e0e74544d06009c2c27"
+  },
+  "amos/packing.py": {
+   "from": "src/amos/packing.py",
+   "verbatim": true,
+   "sha256": "14f027c9e15c456abb0f264c07dae06afb7ace7b64270f2374e97eabeb04f993"
+  },
+  "amos/encoder.py": {
+   "from": "src/amos/encoder.py",
+   "verbatim": true,
+   "sha256": "b4a5a606a2dee99eedf86cab50be759eca0908a318e4f3037ba92560d837bd67"
+  },
+  "amos/facts.py": {
+   "from": "src/amos/facts.py",
+   "verbatim": true,
+   "sha256": "f2efd7ef25453f7f2dc204dd247c86733a55697ec48714f30560df8ed2da1f23"
+  },
+  "amos/vision.py": {
+   "from": "src/amos/vision.py",
+   "verbatim": true,
+   "sha256": "133d1e29cc64e371f5df1f0cb0328fd77b8b9093d8465319580cf5ea4af5d081"
+  },
+  "amos/decoder.py": {
+   "from": "src/amos/decoder.py",
+   "verbatim": false,
+   "kept": [
+    "load_model",
+    "load_heads",
+    "results",
+    "calibration",
+    "span_decoding",
+    "multi_cut"
+   ],
+   "source_sha256": "8d76889e60e05606ea2de24ee40648208dac2f49b3a0e0b5277285397ddcbb73",
+   "sha256": "0dcacb27a05d31781f8310cbe6123b0ee351c62456cc6b5f8576b40b8ddad080"
+  },
+  "amos/calibrate.py": {
+   "from": "src/amos/calibrate.py",
+   "verbatim": false,
+   "kept": [
+    "apply"
+   ],
+   "source_sha256": "552322fa793eea43148ed0b7860a5bca1484fec806dfe485a2ddb664f17b8524",
+   "sha256": "ba094d9eaedd229a99b2bdcd35a371ffdeb9aaee258a342524cc19e2f85ec421"
+  }
+ }
+}

smoke_decoder.py

@@ -0,0 +1,13 @@
+import argparse,asyncio,json
+from pathlib import Path
+import torch
+from coreai_decoder import AmosDecoder
+async def main(name):
+ def forbidden(*a,**kw):raise AssertionError('PyTorch neural execution forbidden')
+ torch.nn.Linear.forward=forbidden;torch.nn.Embedding.forward=forbidden;torch.nn.Conv2d.forward=forbidden;torch.nn.Conv3d.forward=forbidden
+ root=Path('bundles')/(name+'-CoreAI');m=await AmosDecoder.load(root)
+ out=await m.request({'state':'Please refund the duplicate charge.','questions':{'refund':{'type':'noul','instructions':'Is a refund requested?'}}})
+ assert out['answers']['refund']['noul']>0.5,out
+ (root/'runtime-smoke.json').write_text(json.dumps({'passed':True,'pytorch_neural_execution_disabled':True,'result':out},indent=2));print(out)
+if __name__=='__main__':
+ p=argparse.ArgumentParser();p.add_argument('name');a=p.parse_args();asyncio.run(main(a.name))

text.aimodel/main.hash

@@ -0,0 +1 @@
+F��C���T��Hݛ2�?T��Ô������r

text.aimodel/main.mlirb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c469415a143a1d9d75486ed48dd9b32c1033f1854a6bfc394ae9eb08caa8c72
+size 7534240163

text.aimodel/metadata.json

@@ -0,0 +1,5 @@
+{
+  "producer" : "coreai-core 1.0.0b3",
+  "assetVersion" : "2.0",
+  "creationDate" : "20261014T012228Z"
+}

validate_decoder.py

@@ -0,0 +1,57 @@
+import argparse,asyncio,json,sys,shutil
+from pathlib import Path
+import numpy as np
+import torch
+from PIL import Image
+from coreai_decoder import AmosDecoder
+
+async def main(name,mode='both'):
+ source=(Path('upstream')/name).resolve();root=Path('bundles')/(name+'-CoreAI')
+ shutil.copytree(source/'serve',root/'serve',dirs_exist_ok=True)
+ (root/'model').mkdir(exist_ok=True)
+ for f in (source/'model').iterdir():
+  if f.suffix not in ['.safetensors','.pt']:shutil.copy2(f,root/'model'/f.name)
+ for f in ['config.json','calibration.json','LICENSE']:shutil.copy2(source/f,root/f)
+ sys.path.insert(0,str(source/'serve'))
+ import amos.decoder as decoder
+ from transformers import Qwen2VLImageProcessorPil
+ decoder.IMAGE_PIXELS=224*224
+ precision=json.loads((root/'manifest.json').read_text())['precision']
+ torch.set_num_threads(4)
+ if mode!='coreai':
+  model,codec=decoder.load_model(str(source/'model'),'cpu',dtype=getattr(torch,precision));model.eval()
+  codec.image_processor=Qwen2VLImageProcessorPil.from_pretrained(source/'model',size={'shortest_edge':224*224,'longest_edge':224*224})
+  heads=decoder.load_heads(source/'model',model,codec);heads.span=decoder.span_decoding(source)
+ if mode!='reference':ai=await AmosDecoder.load(root)
+ pixels=np.zeros((224,224,3),np.uint8);pixels[:,:112]=[255,0,0];pixels[:,112:]=[0,0,255];im=Image.fromarray(pixels)
+ cases=[
+ ('refund','Invoice INV-2231 for £306.00 was charged twice. Please refund one.',[{'type':'noul','text':'Is a refund requested?'},{'type':'choice','text':'Which team?','options':['billing','technical','sales']}],[]),
+ ('span','The invoice total is £306.00 and the invoice number is INV-2231.',[{'type':'span','text':'Quote the invoice number.'}],[]),
+ ('score','The service is down for every customer.',[{'type':'score','text':'How severe is the outage?','options':['none','minor','major','complete']}],[]),
+ ('multi','I need a refund and a password reset.',[{'type':'multi','text':'Which topics apply?','options':['refund','password','delivery']}],[]),
+ ('negative','The customer explicitly says no refund is needed.',[{'type':'noul','text':'Is a refund requested?'}],[]),
+ ('unicode','José paid €123.45 for café supplies.',[{'type':'span','text':'Quote the amount paid.'},{'type':'choice','text':'Who paid?','options':['José','Maria','Bob']}],[]),
+ ('options','This is about category 17.',[{'type':'choice','text':'Which category?','options':[f'category {i}' for i in range(24)]}],[]),
+ ('vision','Two colored halves.',[{'type':'region','text':'Point to the red half.'},{'type':'choice','text':'Which colors?','options':['red and blue','green and yellow']}],[im]),
+ ('image-option','Find the red-blue image.',[{'type':'choice','text':'Which image matches?','options':[im,Image.new('RGB',(224,224),'green')]}],[]),
+ ]
+ records=[];refs=json.loads((root/'reference.json').read_text()) if mode=='coreai' else []
+ for case,state,qs,ims in cases:
+  if mode=='coreai':
+   ref=next(r for r in refs if r['case']==case);expected,tokens=ref['upstream'],ref['tokens']
+  else:
+   with torch.no_grad():expected,tokens=decoder.results(model,codec,state,qs,ims,heads,512,temperature=decoder.calibration(source))
+  if mode=='reference':
+   records.append({'case':case,'tokens':tokens,'reference_precision':precision,'upstream':expected});print(case,tokens,flush=True);continue
+  actual=await ai.predict(state,qs,ims)
+  error=0;match=True
+  for x,y in zip(expected,actual):
+   for key in ['p_yes','probs','p','box']:
+    if key in x:error=max(error,float(np.abs(np.asarray(x[key])-np.asarray(y[key])).max()))
+   for key in ['choice','start','end','text','cell']:
+    if key in x and x[key]!=y[key]:match=False
+  rec={'case':case,'tokens':tokens,'reference_precision':precision,'max_probability_or_box_error':error,'decisions_match':match,'upstream':expected,'coreai':actual};records.append(rec);print(json.dumps(rec),flush=True)
+  if not match or error>0.02:raise AssertionError(rec)
+ (root/('reference.json' if mode=='reference' else 'validation.json')).write_text(json.dumps(records,indent=2));print('PASS',flush=True)
+if __name__=='__main__':
+ p=argparse.ArgumentParser();p.add_argument('name');p.add_argument('--mode',choices=['both','reference','coreai'],default='both');a=p.parse_args();asyncio.run(main(a.name,a.mode))

validation.json

@@ -0,0 +1,333 @@
+[
+  {
+    "case": "refund",
+    "tokens": 122,
+    "reference_precision": "float32",
+    "max_probability_or_box_error": 1.1920928955078125e-07,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "noul",
+        "p_yes": 0.9725584983825684
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9617580771446228,
+          0.014839934185147285,
+          0.02340202033519745
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "noul",
+        "p_yes": 0.9725583791732788
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9617581963539124,
+          0.014839922077953815,
+          0.023401914164423943
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "span",
+    "tokens": 123,
+    "reference_precision": "float32",
+    "max_probability_or_box_error": 0,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "span",
+        "start": 55,
+        "end": 63,
+        "text": "INV-2231",
+        "p": 1.0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "span",
+        "start": 55,
+        "end": 63,
+        "text": "INV-2231",
+        "p": 1.0
+      }
+    ]
+  },
+  {
+    "case": "score",
+    "tokens": 98,
+    "reference_precision": "float32",
+    "max_probability_or_box_error": 1.7285346984863281e-06,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "score",
+        "probs": [
+          0.011814874596893787,
+          0.012464769184589386,
+          0.05883542448282242,
+          0.9168848395347595
+        ],
+        "choice": 3,
+        "expected": 2.880790136754513
+      }
+    ],
+    "coreai": [
+      {
+        "type": "score",
+        "probs": [
+          0.01181461475789547,
+          0.012464579194784164,
+          0.058834243565797806,
+          0.916886568069458
+        ],
+        "choice": 3,
+        "expected": 2.880792770534754
+      }
+    ]
+  },
+  {
+    "case": "multi",
+    "tokens": 144,
+    "reference_precision": "float32",
+    "max_probability_or_box_error": 2.384185791015625e-07,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "multi",
+        "probs": [
+          0.9698117971420288,
+          0.9617175459861755,
+          0.021801624447107315
+        ]
+      }
+    ],
+    "coreai": [
+      {
+        "type": "multi",
+        "probs": [
+          0.9698116779327393,
+          0.9617173075675964,
+          0.021801499649882317
+        ]
+      }
+    ]
+  },
+  {
+    "case": "negative",
+    "tokens": 83,
+    "reference_precision": "float32",
+    "max_probability_or_box_error": 1.825392246246338e-07,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "noul",
+        "p_yes": 0.050108104944229126
+      }
+    ],
+    "coreai": [
+      {
+        "type": "noul",
+        "p_yes": 0.0501079224050045
+      }
+    ]
+  },
+  {
+    "case": "unicode",
+    "tokens": 128,
+    "reference_precision": "float32",
+    "max_probability_or_box_error": 2.384185791015625e-07,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "span",
+        "start": 10,
+        "end": 17,
+        "text": "\u20ac123.45",
+        "p": 0.9975249171257019
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.975895881652832,
+          0.012953696772456169,
+          0.011150466278195381
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "span",
+        "start": 10,
+        "end": 17,
+        "text": "\u20ac123.45",
+        "p": 0.9975249171257019
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9758961200714111,
+          0.012953538447618484,
+          0.011150323785841465
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "options",
+    "tokens": 230,
+    "reference_precision": "float32",
+    "max_probability_or_box_error": 5.960464477539063e-08,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "choice",
+        "probs": [
+          0.004064494278281927,
+          0.0038580188993364573,
+          0.004268646240234375,
+          0.0035152093041688204,
+          0.0033423479180783033,
+          0.0027638780884444714,
+          0.0034003546461462975,
+          0.0028926823288202286,
+          0.002699152799323201,
+          0.002128811087459326,
+          0.0030107321217656136,
+          0.0031014070846140385,
+          0.0031818263232707977,
+          0.0027776346541941166,
+          0.002535774139687419,
+          0.003921506926417351,
+          0.0034300386905670166,
+          0.9298884868621826,
+          0.003782431362196803,
+          0.0032975035719573498,
+          0.0014300656039267778,
+          0.001529851695522666,
+          0.0019381233723834157,
+          0.0032409545965492725
+        ],
+        "choice": 17
+      }
+    ],
+    "coreai": [
+      {
+        "type": "choice",
+        "probs": [
+          0.004064517095685005,
+          0.00385798211209476,
+          0.004268613178282976,
+          0.0035152393393218517,
+          0.0033423572313040495,
+          0.0027638806495815516,
+          0.0034003155305981636,
+          0.0028926795348525047,
+          0.0026991423219442368,
+          0.00212881900370121,
+          0.0030107346829026937,
+          0.00310140079818666,
+          0.0031818472780287266,
+          0.002777676796540618,
+          0.002535800449550152,
+          0.003921480383723974,
+          0.003430022159591317,
+          0.9298885464668274,
+          0.003782423911616206,
+          0.0032975096255540848,
+          0.0014300619950518012,
+          0.0015298567013815045,
+          0.0019381335005164146,
+          0.0032409762497991323
+        ],
+        "choice": 17
+      }
+    ]
+  },
+  {
+    "case": "vision",
+    "tokens": 155,
+    "reference_precision": "float32",
+    "max_probability_or_box_error": 1.0192394256591797e-05,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "region",
+        "cell": 9,
+        "box": [
+          0.1936,
+          0.0,
+          0.5538,
+          0.8461
+        ],
+        "p": 0.5629251003265381
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9847151041030884,
+          0.015284925699234009
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "region",
+        "cell": 9,
+        "box": [
+          0.1936,
+          0.0,
+          0.5538,
+          0.8461
+        ],
+        "p": 0.5629352927207947
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.9847150444984436,
+          0.015284995548427105
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "image-option",
+    "tokens": 181,
+    "reference_precision": "float32",
+    "max_probability_or_box_error": 3.3527612686157227e-08,
+    "decisions_match": true,
+    "upstream": [
+      {
+        "type": "choice",
+        "probs": [
+          0.9767023324966431,
+          0.023297706618905067
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "choice",
+        "probs": [
+          0.9767023324966431,
+          0.02329767309129238
+        ],
+        "choice": 0
+      }
+    ]
+  }
+]

vision.aimodel/main.hash

@@ -0,0 +1 @@
+dj�g��'��-+��&��zuh^���N���

vision.aimodel/main.mlirb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f646ae2677fb9ba2793ec942d2b9ed82690199d7a75685ecdf4a84ea791f38f
+size 1317394842

vision.aimodel/metadata.json

@@ -0,0 +1,5 @@
+{
+  "assetVersion" : "2.0",
+  "creationDate" : "20261014T012236Z",
+  "producer" : "coreai-core 1.0.0b3"
+}