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

.gitattributes

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

DEPENDENCIES.json

@@ -0,0 +1,10 @@
+[
+  {
+    "repo": "jhu-clsp/ettin-encoder-400m",
+    "revision": "7662476d60abb071a5bd319c9f3074f3072c062d"
+  },
+  {
+    "repo": "google/siglip2-base-patch16-224",
+    "revision": "75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2"
+  }
+]

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,43 @@
+---
+license: apache-2.0
+base_model: mukaddamzaid/amos-1-mini
+tags: [coreai, core-ai, apple, aimodel, decisions, vision]
+---
+
+# AMOS-1 mini — Apple Core AI
+
+An independently converted AMOS-1 mini decision model. Neural execution (encoder, trained projections, vision tower and projector) runs through Apple Core AI `.aimodel` assets. Upstream Python packing, pooling, probability normalization and decoding run on the host using PyTorch tensors. This is a decision model, not a text generator.
+
+## Use
+
+Requires macOS 27 with the Core AI Python runtime. Tested with `coreai-torch==0.4.3`, PyTorch 2.14.0, NumPy, Pillow, tokenizers and huggingface_hub. The converter warns that this PyTorch version is newer than its validated range.
+
+```python
+import asyncio
+from coreai_mini import AmosMini
+
+async def main():
+    model = await AmosMini.load(".")
+    print(await model.request({
+        "state": "Please refund the duplicate charge.",
+        "questions": {"refund": {
+            "type": "noul", "instructions": "Is a refund requested?"
+        }}
+    }))
+
+asyncio.run(main())
+```
+
+`predict(state, questions, images=())` also accepts upstream internal questions (`text`, `options`) and PIL images, including image options. The request API accepts base64 images in the original AMOS format.
+
+## Capacity and verification
+
+Batch one, **512 combined tokens**, 224×224 images, FP32 weights. Requests exceeding capacity raise an error. This conversion does not provide the original 9,216-token capacity, state-cache acceleration, hosted API server, explain endpoint or benchmark claims. No iPhone or Neural Engine placement/performance has been verified. Core AI uses default specialization.
+
+Nine reproducible smoke/parity cases cover noul, choice, score, multi, span, region, Unicode, 24 options, state images and image options. All tested choices, spans and region cells matched upstream. Maximum measured logit error was 0.000833. These smoke cases do not establish benchmark accuracy. `validation.json` contains exact results. `runtime-smoke.json` records inference with PyTorch Linear/Embedding/Conv2d execution disabled.
+
+## Reproduce and provenance
+
+The original AMOS author is Zaid Mukaddam. AMOS source and weights are Apache-2.0; see `LICENSE`, `SOURCE.json`, and upstream [AMOS-1 mini](https://huggingface.co/mukaddamzaid/amos-1-mini). The encoder derives from jhu-clsp/ettin-encoder-400m and the frozen vision tower from google/siglip2-base-patch16-224; dependency revisions are in `DEPENDENCIES.json`. Conversion adds no restrictions.
+
+Run `convert_mini.py --source PATH_TO_PINNED_AMOS_MINI --output OUTPUT`, then `validate_mini.py` and `smoke_mini.py` from the original workspace layout. Original serving sources are included in `serve/`. `SHA256SUMS` covers release files.

SHA256SUMS

@@ -0,0 +1,32 @@
+630e9474f3de7d236bb679a3ce81e5d2beda5f925cba2b5d19f06268a42ca3a6  DEPENDENCIES.json
+cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30  LICENSE
+2c7ae2534f71c5b654eab75d7a974ebdb21ab137354bb3d0bd59aca2a736b8fe  README.md
+9a27b3030743c234d6971b79592d0c39bf29f5e469ec406071b303a2f917a551  SOURCE.json
+13b2a7499c445d25eb9e6c7be9c242b30f7f88e2f6f1d70bc3b3becd50b0cc55  config.json
+d720cb43ee361fb162de1c5b8fb0460a9fde9f5a5f7f463e415293efd9e954cf  convert_mini.py
+6e91d7729b36183cbbb7df684d1601efcd8ef25858f9c93fa31bf96a1c720bb8  coreai_mini.py
+3a83da88f746a41c45a23f3fcb3306978191bea6d93177a9207c4c479599c456  manifest.json
+755a21be5ba30e8469e6625539a549883ecd974e36ef342a6ab0284f079d9b6f  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
+5ebfa87f58e3eaf3a7f9d6e493a906c1df59b3890f805d87859a52378644646b  smoke_mini.py
+a8d97fff1f3efd130401987198943d69c67202f47e7f029d8887fea06bbcf573  text.aimodel/main.hash
+b330fa3eef9565b6d4d2743ab3f2c71e87bed2b1246f92d7d7702be7da4bff33  text.aimodel/main.mlirb
+4c017ad5ee82a62b803dd371fa4a1b5585839d1896aee649ca1c909480304292  text.aimodel/metadata.json
+9fd55248d51d33976b324fc11592e28071da7d41e0e9401dfb7082e30574b7b1  tokenizer.json
+d733722602caf019fa2243033a62d890ac69339da38af6ddcbd99d79f07d6813  validate_mini.py
+72225fef1b5ebbdf0813c64b2bd63f491a77a36d06452eb77f9a50f7deae23aa  validation.json
+046186ca6d28b322f5bc0d83cf26f9be271d2f1df6070f5c84c92e00ff9bb8c6  vision.aimodel/main.hash
+3c25d133ca87ec965d1e579e5027e9d71bc7aabe9356d98361806dddfd8ff596  vision.aimodel/main.mlirb
+65296e091507f9f0775a538bb5a28d21464522b3b8051cdd07e722a97dfb29ae  vision.aimodel/metadata.json

SOURCE.json

@@ -0,0 +1,112 @@
+{
+  "repo": "mukaddamzaid/amos-1-mini",
+  "revision": "340d8af4bfbf420b2cec31b79d27a528958eaf3a",
+  "bytes": 1629805383,
+  "files": [
+    {
+      "path": ".gitattributes",
+      "size": 1519,
+      "lfs": null
+    },
+    {
+      "path": "LICENSE",
+      "size": 11358,
+      "lfs": null
+    },
+    {
+      "path": "README.md",
+      "size": 4150,
+      "lfs": null
+    },
+    {
+      "path": "assets/amos-1-mini.png",
+      "size": 94372,
+      "lfs": null
+    },
+    {
+      "path": "config.json",
+      "size": 600,
+      "lfs": null
+    },
+    {
+      "path": "model.pt",
+      "size": 1629560074,
+      "lfs": {
+        "size": 1629560074,
+        "sha256": "5e415a743b2d568bad8525b9bd85c73a511136eb6b7ef60f26d51134f83dadec",
+        "pointer_size": 135
+      }
+    },
+    {
+      "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": 1604,
+      "lfs": null
+    }
+  ],
+  "checkpoint_sha256": "5e415a743b2d568bad8525b9bd85c73a511136eb6b7ef60f26d51134f83dadec"
+}

config.json

@@ -0,0 +1,25 @@
+{
+ "repo": "jhu-clsp/ettin-encoder-400m",
+ "max_state": 8192,
+ "digits": false,
+ "score_rank": false,
+ "encoder": {
+  "vocab_size": 50368,
+  "hidden_size": 1024,
+  "intermediate_size": 2624,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "global_attn_every_n_layers": 3,
+  "local_attention": 128,
+  "global_rope_theta": 160000.0,
+  "local_rope_theta": 160000.0,
+  "norm_eps": 1e-05,
+  "pad_token_id": 50283,
+  "cls_token_id": 50281,
+  "sep_token_id": 50282,
+  "model_type": "modernbert",
+  "compare_options": false,
+  "date_facts": false
+ },
+ "vision": "google/siglip2-base-patch16-224"
+}

convert_mini.py

@@ -0,0 +1,45 @@
+"""Reproduce the AMOS mini Core AI export from pinned source files."""
+import argparse,json,sys
+from pathlib import Path
+import torch
+from coreai_torch import TorchConverter,get_decomp_table
+
+class Text(torch.nn.Module):
+ def __init__(self,m):
+  super().__init__(); self.m=m
+ def forward(self,input_ids,position_ids,segment_ids,option_ids,visual_embeddings,visual_mask):
+  m=self.m
+  emb=torch.where(visual_mask,visual_embeddings,m.encoder.tok_embeddings(input_ids))
+  h=m.encoder(input_ids,position_ids,segment_ids,option_ids,inputs_embeds=emb)
+  return m.ptr_q(h),m.ptr_o(h),m.span_q(h),m.span_k(h),m.region_q(h),m.region_k(h),m.multi_bias
+
+class Vision(torch.nn.Module):
+ def __init__(self,m):
+  super().__init__(); self.vision=m.vision; self.projector=m.projector
+ def forward(self,pixel_values):
+  v=self.vision
+  x=v.patch_embedding(pixel_values).flatten(2).transpose(1,2)+v.position_embedding
+  for b in v.blocks:x=b(x)
+  return self.projector(v.post_layernorm(x))
+
+def save(m,args,path):
+ ep=torch.export.export(m,args).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('--source',default='upstream/amos-1-mini');p.add_argument('--output',default='bundles/amos-1-mini-CoreAI');p.add_argument('--length',type=int,default=512);a=p.parse_args()
+ source=Path(a.source).resolve();out=Path(a.output);out.mkdir(parents=True,exist_ok=True)
+ sys.path.insert(0,str(source/'serve'))
+ from amos.model import load_run
+ torch.set_num_threads(4)
+ m,codec=load_run(source);m.eval()
+ args=(torch.full((1,a.length),codec.pad,dtype=torch.int32),torch.arange(a.length,dtype=torch.int32)[None],torch.zeros(1,a.length,dtype=torch.int32),torch.zeros(1,a.length,dtype=torch.int32),torch.zeros(1,a.length,m.encoder.config.hidden_size),torch.zeros(1,a.length,1,dtype=torch.bool))
+ print('Exporting text',flush=True)
+ outputs=save(Text(m).eval(),args,out/'text.aimodel')
+ print('Exporting vision',flush=True)
+ vo=save(Vision(m).eval(),(torch.zeros(1,3,224,224),),out/'vision.aimodel')
+ (out/'manifest.json').write_text(json.dumps({'sequence_length':a.length,'hidden_size':m.encoder.config.hidden_size,'outputs':outputs,'vision_output':vo[0],'precision':'float32'},indent=2))
+ print('Saved',flush=True)
+if __name__=='__main__':main()

coreai_mini.py

@@ -0,0 +1,89 @@
+"""Core AI neural execution; upstream packing/decoding uses PyTorch tensors only."""
+import json,math,sys
+from pathlib import Path
+import numpy as np
+import torch
+from coreai.runtime import AIModel,NDArray,SpecializationOptions
+
+class AmosMini:
+ @classmethod
+ async def load(cls,bundle):
+  self=cls();self.root=Path(bundle).resolve();sys.path.insert(0,str(self.root/'serve'))
+  from amos.packing import Codec
+  cfg=json.loads((self.root/'config.json').read_text())
+  # Codec's loader accepts a locally pinned tokenizer through this scoped patch.
+  import amos.packing as packing
+  old=packing.hf_hub_download
+  packing.hf_hub_download=lambda repo,filename: str(self.root/filename)
+  try:self.codec=Codec(cfg['repo'],max_state=cfg['max_state'],digits=cfg.get('digits',False),score_rank=cfg.get('score_rank',False))
+  finally:packing.hf_hub_download=old
+  self.manifest=json.loads((self.root/'manifest.json').read_text())
+  self.text=await AIModel.load(self.root/'text.aimodel',SpecializationOptions.default());self.tf=self.text.load_function('main')
+  self.vision=await AIModel.load(self.root/'vision.aimodel',SpecializationOptions.default());self.vf=self.vision.load_function('main')
+  return self
+
+ async def raw(self,state,questions,images=()):
+  from amos.packing import pack,collate
+  from amos.vision import preprocess
+  feats=[]
+  async def feature(im):
+   pix=preprocess(im)[None].numpy()
+   out=await self.vf({'pixel_values':NDArray(pix)})
+   return torch.from_numpy(out[self.manifest['vision_output']].numpy()[0].copy())
+  # Pack expects pre-projection 196-cell features. Instead use placeholders;
+  # inject the Core AI projected 49-cell vectors directly into embedding slots.
+  for im in images:feats.append(await feature(im))
+  qs=[]
+  for q in questions:
+   q=dict(q)
+   if 'options' in q:
+    opts=[]
+    for o in q['options']:
+     if isinstance(o,str):opts.append(o)
+     else:opts.append(torch.zeros(196,768));feats.append(await feature(o))
+    q['options']=opts
+   qs.append(q)
+  p=pack(self.codec,state,qs,[torch.zeros(196,768) for _ in images],strict=True)
+  n=len(p);L=self.manifest['sequence_length'];d=self.manifest['hidden_size']
+  if n>L:raise ValueError(f'Request has {n} tokens; Core AI capacity is {L}.')
+  b=collate([p],self.codec.pad)
+  args={}
+  for src,dst,pad in [('ids','input_ids',self.codec.pad),('pos','position_ids',0),('seg','segment_ids',-1),('opt','option_ids',0)]:
+   arr=np.full((1,L),pad,dtype=np.int32);arr[0,:n]=getattr(p,src);args[dst]=NDArray(arr)
+  ve=np.zeros((1,L,d),np.float32);vm=np.zeros((1,L,1),bool)
+  for at,f in zip(p.image_at,feats):ve[0,at:at+49]=f.numpy();vm[0,at:at+49]=True
+  args['visual_embeddings']=NDArray(ve);args['visual_mask']=NDArray(vm)
+  result=await self.tf(args)
+  pq,po,sq,sk,rq,rk,bias=[torch.from_numpy(result[k].numpy().copy()) for k in self.manifest['outputs']]
+  Q=len(p.decide);O=len(b['o_q']);out={'opt':torch.full((Q,max(1,max(p.n_options or [0]))), -torch.inf)}
+  for i in range(O):
+   valid=b['option_of'][0]==i
+   pooled=po[0,:valid.shape[0]][valid].mean(0)
+   qi=int(b['o_q'][i]);rank=int(b['o_rank'][i]);at=p.decide[qi]
+   out['opt'][qi,rank]=(pq[0,at]*pooled).sum()/math.sqrt(d)+(bias if p.types[qi]=='multi' else 0)
+  spans=[i for i,t in enumerate(p.types) if t=='span'];regions=[i for i,t in enumerate(p.types) if t=='region']
+  out['span_q']=torch.tensor(spans);out['region_q']=torch.tensor(regions)
+  if spans:
+   q=sq[0,p.decide].view(Q,2,d)[spans];k=sk[0,:b['ids'].shape[1]].view(-1,2,d)
+   out['span']=torch.einsum('qcd,lcd->qcl',q,k)/math.sqrt(d)
+   out['span'].masked_fill_(~b['state_mask'][:,None,:],-torch.inf)
+  if regions:
+   out['region']=rq[0,[p.decide[i] for i in regions]]@rk[0,:b['ids'].shape[1]].T/math.sqrt(d)
+   out['region'].masked_fill_(~b['region_mask'],-torch.inf)
+  return out,[p]
+
+ async def predict(self,state,questions,images=()):
+  from amos.model import decode
+  out,ps=await self.raw(state,questions,images)
+  return decode(out,ps)[0]
+
+ async def request(self,req):
+  from amos.serve import parse,answer
+  from PIL import Image
+  import base64,io
+  state,ids,qs=parse(req)
+  def image(s):return Image.open(io.BytesIO(base64.b64decode(s.split(',')[-1])))
+  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]
+  results=await self.predict(state,qs,ims)
+  return {'answers':{i:answer(q,r) for i,q,r in zip(ids,qs,results)}}

manifest.json

@@ -0,0 +1,15 @@
+{
+  "sequence_length": 512,
+  "hidden_size": 1024,
+  "outputs": [
+    "view_281",
+    "view_283",
+    "view_285",
+    "view_287",
+    "view_289",
+    "view_291",
+    "p_m_multi_bias"
+  ],
+  "vision_output": "view_198",
+  "precision": "float32"
+}

runtime-smoke.json

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

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_mini.py

@@ -0,0 +1,13 @@
+import asyncio,json
+from pathlib import Path
+import torch
+from coreai_mini import AmosMini
+async def main():
+ def forbidden(*a,**kw):raise AssertionError('PyTorch neural execution is forbidden in runtime smoke')
+ torch.nn.Linear.forward=forbidden;torch.nn.Embedding.forward=forbidden;torch.nn.Conv2d.forward=forbidden
+ root=Path('bundles/amos-1-mini-CoreAI');m=await AmosMini.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__':asyncio.run(main())

text.aimodel/main.hash

@@ -0,0 +1 @@
+�0�>�e���t:�����ұ$o���p+��K�3

text.aimodel/main.mlirb

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

text.aimodel/metadata.json

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

validate_mini.py

@@ -0,0 +1,53 @@
+import asyncio,json,sys,shutil
+from pathlib import Path
+import numpy as np
+import torch
+from PIL import Image
+from huggingface_hub import hf_hub_download,HfApi
+from coreai_mini import AmosMini
+
+async def main():
+ root=Path('bundles/amos-1-mini-CoreAI');source=Path('upstream/amos-1-mini').resolve()
+ shutil.copytree(source/'serve',root/'serve',dirs_exist_ok=True)
+ for f in ['config.json','LICENSE']:shutil.copy2(source/f,root/f)
+ cfg=json.loads((source/'config.json').read_text());api=HfApi();deps=[]
+ for repo in [cfg['repo'],cfg['vision']]:
+  info=api.model_info(repo);deps.append({'repo':repo,'revision':info.sha})
+  if repo==cfg['repo']:shutil.copy2(hf_hub_download(repo,'tokenizer.json',revision=info.sha),root/'tokenizer.json')
+ (root/'DEPENDENCIES.json').write_text(json.dumps(deps,indent=2))
+ sys.path.insert(0,str(source/'serve'))
+ from amos.model import load_run,pack_request,decode
+ from amos.packing import collate
+ from amos.vision import preprocess
+ torch.set_num_threads(4);reference,codec=load_run(source)
+ ai=await AmosMini.load(root)
+ 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)]}],[]),
+ ]
+ pixels=np.zeros((224,224,3),np.uint8);pixels[:,:112]=[255,0,0];pixels[:,112:]=[0,0,255];im=Image.fromarray(pixels)
+ cases += [('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=[]
+ for name,state,qs,ims in cases:
+  refqs=[{**q,'options':[preprocess(o) if isinstance(o,Image.Image) else o for o in q['options']]} if 'options' in q else q for q in qs]
+  with torch.no_grad():
+   ps=pack_request(reference,codec,state,refqs,[preprocess(x) for x in ims]);expected=reference(collate(ps,codec.pad));ed=decode(expected,ps)[0]
+  actual,aps=await ai.raw(state,qs,ims);ad=decode(actual,aps)[0]
+  errors={}
+  for key in ['opt','span','region']:
+   if key not in expected:continue
+   x=expected[key];y=actual[key];mask=torch.isfinite(x)
+   if not torch.equal(mask,torch.isfinite(y)):raise AssertionError((name,key,x.shape,y.shape,x,y))
+   errors[key]=float((x[mask]-y[mask]).abs().max()) if mask.any() else 0
+   if errors[key]>0.005:raise AssertionError((name,key,errors))
+  def decisions(rs):
+   return [{k:v for k,v in r.items() if k in ['type','text','start','end','cell','choice','index']} for r in rs]
+  if decisions(ed)!=decisions(ad):raise AssertionError((name,ed,ad))
+  rec={'case':name,'tokens':len(ps[0]),'max_logit_error':errors,'upstream':ed,'coreai':ad};records.append(rec);print(json.dumps(rec),flush=True)
+ (root/'validation.json').write_text(json.dumps(records,indent=2));print('PASS',flush=True)
+if __name__=='__main__':asyncio.run(main())

validation.json

@@ -0,0 +1,438 @@
+[
+  {
+    "case": "refund",
+    "tokens": 45,
+    "max_logit_error": {
+      "opt": 2.6226043701171875e-06
+    },
+    "upstream": [
+      {
+        "type": "noul",
+        "p_yes": 0.94333815574646
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.7296061515808105,
+          0.016437293961644173,
+          0.2539565861225128
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "noul",
+        "p_yes": 0.94333815574646
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.7296063899993896,
+          0.016437331214547157,
+          0.25395628809928894
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "span",
+    "tokens": 28,
+    "max_logit_error": {
+      "opt": 0,
+      "span": 1.52587890625e-05
+    },
+    "upstream": [
+      {
+        "type": "span",
+        "start": 21,
+        "end": 63,
+        "text": "\u00a3306.00 and the invoice number is INV-2231",
+        "p": 0.16281552612781525
+      }
+    ],
+    "coreai": [
+      {
+        "type": "span",
+        "start": 21,
+        "end": 63,
+        "text": "\u00a3306.00 and the invoice number is INV-2231",
+        "p": 0.16281506419181824
+      }
+    ]
+  },
+  {
+    "case": "score",
+    "tokens": 27,
+    "max_logit_error": {
+      "opt": 3.0994415283203125e-06
+    },
+    "upstream": [
+      {
+        "type": "score",
+        "probs": [
+          0.06895732879638672,
+          0.02768923155963421,
+          0.4086642265319824,
+          0.49468913674354553
+        ],
+        "choice": 3,
+        "expected": 2.329085111618042
+      }
+    ],
+    "coreai": [
+      {
+        "type": "score",
+        "probs": [
+          0.0689573660492897,
+          0.027689190581440926,
+          0.40866488218307495,
+          0.49468863010406494
+        ],
+        "choice": 3,
+        "expected": 2.329084873199463
+      }
+    ]
+  },
+  {
+    "case": "multi",
+    "tokens": 25,
+    "max_logit_error": {
+      "opt": 2.86102294921875e-06
+    },
+    "upstream": [
+      {
+        "type": "multi",
+        "probs": [
+          0.965849757194519,
+          0.9332621097564697,
+          0.07576480507850647
+        ]
+      }
+    ],
+    "coreai": [
+      {
+        "type": "multi",
+        "probs": [
+          0.9658498764038086,
+          0.9332619309425354,
+          0.07576475292444229
+        ]
+      }
+    ]
+  },
+  {
+    "case": "negative",
+    "tokens": 22,
+    "max_logit_error": {
+      "opt": 1.6093254089355469e-06
+    },
+    "upstream": [
+      {
+        "type": "noul",
+        "p_yes": 0.077876016497612
+      }
+    ],
+    "coreai": [
+      {
+        "type": "noul",
+        "p_yes": 0.07787591218948364
+      }
+    ]
+  },
+  {
+    "case": "unicode",
+    "tokens": 32,
+    "max_logit_error": {
+      "opt": 9.775161743164062e-06,
+      "span": 1.1444091796875e-05
+    },
+    "upstream": [
+      {
+        "type": "span",
+        "start": 10,
+        "end": 17,
+        "text": "\u20ac123.45",
+        "p": 0.44328346848487854
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.896871030330658,
+          0.03730686381459236,
+          0.06582211703062057
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "span",
+        "start": 10,
+        "end": 17,
+        "text": "\u20ac123.45",
+        "p": 0.44328081607818604
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.8968717455863953,
+          0.03730650246143341,
+          0.06582175195217133
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "options",
+    "tokens": 85,
+    "max_logit_error": {
+      "opt": 4.410743713378906e-06
+    },
+    "upstream": [
+      {
+        "type": "choice",
+        "probs": [
+          0.021480435505509377,
+          0.012768608517944813,
+          0.014833449386060238,
+          0.02430936135351658,
+          0.015337659977376461,
+          0.014616258442401886,
+          0.01716616190969944,
+          0.019972851499915123,
+          0.018173672258853912,
+          0.02504866011440754,
+          0.026172664016485214,
+          0.02132749743759632,
+          0.024641145020723343,
+          0.023125993087887764,
+          0.0218864344060421,
+          0.020913584157824516,
+          0.0224751066416502,
+          0.5114471316337585,
+          0.029352843761444092,
+          0.027540743350982666,
+          0.024416472762823105,
+          0.020348869264125824,
+          0.019210753962397575,
+          0.023433661088347435
+        ],
+        "choice": 17
+      }
+    ],
+    "coreai": [
+      {
+        "type": "choice",
+        "probs": [
+          0.021480349823832512,
+          0.01276865229010582,
+          0.014833447523415089,
+          0.02430935949087143,
+          0.015337690711021423,
+          0.014616241678595543,
+          0.017166148871183395,
+          0.019972823560237885,
+          0.01817365735769272,
+          0.025048643350601196,
+          0.02617262303829193,
+          0.021327510476112366,
+          0.024641163647174835,
+          0.023125983774662018,
+          0.021886436268687248,
+          0.02091355249285698,
+          0.022474996745586395,
+          0.5114476680755615,
+          0.029352732002735138,
+          0.027540666982531548,
+          0.02441641129553318,
+          0.020348837599158287,
+          0.01921072229743004,
+          0.023433679714798927
+        ],
+        "choice": 17
+      }
+    ]
+  },
+  {
+    "case": "vision",
+    "tokens": 78,
+    "max_logit_error": {
+      "opt": 1.4901161193847656e-05,
+      "region": 0.0008322000503540039
+    },
+    "upstream": [
+      {
+        "type": "region",
+        "cell": 12,
+        "box": [
+          0.7142857313156128,
+          0.1428571492433548,
+          0.8571428656578064,
+          0.2857142984867096
+        ],
+        "p": 0.10726314783096313,
+        "heatmap": [
+          0.03254103288054466,
+          0.040585536509752274,
+          0.01753769814968109,
+          0.008257847279310226,
+          0.011045021004974842,
+          0.00599908921867609,
+          0.004235319793224335,
+          0.0024749385192990303,
+          0.017212174832820892,
+          0.005958267953246832,
+          0.0028916022274643183,
+          0.07630306482315063,
+          0.10726314783096313,
+          0.016866952180862427,
+          0.01265009492635727,
+          0.026430141180753708,
+          0.0291758943349123,
+          0.01452744659036398,
+          0.01801556907594204,
+          0.030956637114286423,
+          0.022234106436371803,
+          0.01916903629899025,
+          0.015629546716809273,
+          0.02098844386637211,
+          0.028073374181985855,
+          0.03741556778550148,
+          0.02767399698495865,
+          0.03067740611732006,
+          0.01079087145626545,
+          0.04043445363640785,
+          0.05705111101269722,
+          0.02413317747414112,
+          0.02471594698727131,
+          0.01658044196665287,
+          0.012068545445799828,
+          0.003539793426170945,
+          0.014276671223342419,
+          0.008514068089425564,
+          0.007855590432882309,
+          0.005705713294446468,
+          0.003973916172981262,
+          0.001432417775504291,
+          0.0023278759326785803,
+          0.010151573456823826,
+          0.027151266112923622,
+          0.019837509840726852,
+          0.010044840164482594,
+          0.0030451989732682705,
+          0.013580089434981346
+        ]
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.6527983546257019,
+          0.3472016155719757
+        ],
+        "choice": 0
+      }
+    ],
+    "coreai": [
+      {
+        "type": "region",
+        "cell": 12,
+        "box": [
+          0.7142857313156128,
+          0.1428571492433548,
+          0.8571428656578064,
+          0.2857142984867096
+        ],
+        "p": 0.10718144476413727,
+        "heatmap": [
+          0.032547224313020706,
+          0.04059035703539848,
+          0.017539363354444504,
+          0.008259006775915623,
+          0.01104627177119255,
+          0.005999898537993431,
+          0.004236394539475441,
+          0.002475050510838628,
+          0.017213277518749237,
+          0.005958556197583675,
+          0.0028917340096086264,
+          0.07630617171525955,
+          0.10718144476413727,
+          0.01685546711087227,
+          0.012653937563300133,
+          0.02643437869846821,
+          0.029178880155086517,
+          0.014528198167681694,
+          0.01801961474120617,
+          0.03096183016896248,
+          0.02223971113562584,
+          0.019169794395565987,
+          0.015628913417458534,
+          0.02098875492811203,
+          0.028077254071831703,
+          0.03741911053657532,
+          0.027675963938236237,
+          0.03067953512072563,
+          0.010791921056807041,
+          0.04043751209974289,
+          0.05705570429563522,
+          0.024136871099472046,
+          0.02471994049847126,
+          0.0165830347687006,
+          0.012070251628756523,
+          0.003540291218087077,
+          0.01427785586565733,
+          0.00851534865796566,
+          0.007856183685362339,
+          0.0057067349553108215,
+          0.003974433988332748,
+          0.001432594726793468,
+          0.0023280533496290445,
+          0.010152487084269524,
+          0.027152718976140022,
+          0.019839033484458923,
+          0.010046061128377914,
+          0.0030458515975624323,
+          0.01358106080442667
+        ]
+      },
+      {
+        "type": "choice",
+        "probs": [
+          0.6527963280677795,
+          0.34720364212989807
+        ],
+        "choice": 0
+      }
+    ]
+  },
+  {
+    "case": "image-option",
+    "tokens": 119,
+    "max_logit_error": {
+      "opt": 3.2901763916015625e-05
+    },
+    "upstream": [
+      {
+        "type": "choice",
+        "probs": [
+          0.2963208258152008,
+          0.7036792039871216
+        ],
+        "choice": 1
+      }
+    ],
+    "coreai": [
+      {
+        "type": "choice",
+        "probs": [
+          0.2963253855705261,
+          0.7036746144294739
+        ],
+        "choice": 1
+      }
+    ]
+  }
+]

vision.aimodel/main.hash

@@ -0,0 +1 @@
+<%�3ʇ�]W�P'��Ǫ��Vكa�m�����

vision.aimodel/main.mlirb

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

vision.aimodel/metadata.json

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