Arc Tagger v0.1.0: AT10Mv2 step88000 + refit thresholds + arctag package

LICENSE-arc_audio

@@ -0,0 +1,201 @@
+                                 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.

MANIFEST.sha256

@@ -0,0 +1,38 @@
+432e1c0db00dca1695c84c2e0b05f9b1f6f23a45acbc1bc4173a2669f9fef9a2  ./arctag/_core/arc_audio/__init__.py
+fabe47fdbed606eb9bb17cd1f8cb809d4b0e765b3bc45c8aa0b7ad30f14a6026  ./arctag/_core/arc_audio/mamba3/__init__.py
+640a5d529a810de1cddc0919b4e8744bfc5642257ce1abf6e4e65d74a54fc7c2  ./arctag/_core/arc_audio/mamba3/layer.py
+96fb2a6fba2b0aed4fed121a95bb9307adf319afdff9a897d93fc1007f897de4  ./arctag/_core/arc_audio/mamba3/scan.py
+e5a5b7f94c059f01250a5e8af1063d7aaf98f338b8596048884d018485c65a90  ./arctag/_core/arc_audio/mamba3/ssd.py
+d86b6d7b6ee1653bfd91dff4c4170b30244650ccf992cf8e8f67ea3ae9c9009f  ./arctag/_core/arc_audio/model.py
+365ae9423a5fe89484178cb2823ab594ce58fc85f4df85a1a2019bb1c3568724  ./arctag/_core/arc_audio/ontology.py
+aced3f28bd6f821e2edf522f37c127a8370ca54f508a9ed05d6206447993bcd6  ./arctag/_core/arc_audio/wordch.py
+99497a4f0ae7ed568f593c472b1edd1cf95fcd09a729c7ae666936bf868e13d1  ./arctag/_core/arc_text.py
+427234a679197cd846d13e4e9b293d34492d06814aafbbd842c527fddc162450  ./arctag/_core/engine.py
+abad828a99d8e961277fdc3ce3f9c8af9fa09b61289f4882dde257efef832e17  ./arctag/_core/fast_arc.py
+02bd73a10018f86d788f8784cd12fd10e19836e95975362fbcb0d9e9dfc90381  ./arctag/_core/__init__.py
+3d3fb0551279f43413801a2e03148121f960097cab72124132ebe6a8395d4f9f  ./arctag/_core/kb.py
+71352e7c576d76458431a24ee6df6ebaa95919f79fa3ad9ba66f1f7e37103b0d  ./arctag/_core/kb_strfeat.py
+0ff792d9926e12a99954a0f7a1705ae2dacb2f95827f9e5ba2e334f4755881d7  ./arctag/_core/losses.py
+355e635e6935d60906ab7ee51b7921b28509bb99b4dccc89b69f9a450d3f584f  ./arctag/_core/speaker_link.py
+7799029e19b871b08949af23b2bd1c25da4c720e807fbdfa6c9776b95e43ebae  ./arctag/_core/stream_arc.py
+6804a7d18aa7faa6b632537e9478a95a03bfcc8c7625d044c40acea96c770284  ./arctag/_core/tagtable.py
+552c4ace1562ba31a532a86afc0fe57f7c80ce95cc72b539e2ea6c5b2736c6b6  ./arctag/_core/v2.py
+0b24e8f43145c878966f722ebfdc55e1e937c3005301d4d284306bc1eedc8f50  ./arctag/driver.py
+09b3d76601a0756ef761628112f26886cff287113024605d34eef3d1936a4a01  ./arctag/__init__.py
+505e9b93f1895cf4e29374c11041fe5316b7c44e226f63000a1945910c0d9109  ./config.json
+d11cc9c244b81636f087e2165bfea9e7c3fb6d8c09a1975cb17ea3293f8d7370  ./demo_settings.json
+d28c8d84e9d3829cbb2dd5b2d324812aeab248691d771d5b4b82cb9cd5ce3648  ./example.py
+c45bd62ce519bcc479834b7d94396f1f29dee0b606aeff7e29adadacf6342383  ./label_norm.json
+c71d239df91726fc519c6eb72d318ec65820627232b2f796219e87dcf35d0ab4  ./LICENSE-arc_audio
+e80f2c903d3f087c045c0f9482718533c7cf2d9d42817cfea5943fad72935ed8  ./model.safetensors
+503a2a7e125b7311b3198a1d8b4a37901377093c85b281def26100d2b6b56939  ./ontology/ext_cc_roles_E.npy
+9040786b540fcbfbc9137abf81a05e972e9380ae0c208bb55fae9c459aa4b97f  ./ontology/ext_cc_roles.json
+05cc12d2bdafa7ef463797ca2997ecda2dd914eaf4718a86959f45f0a8840f96  ./ontology/head_E_r8_bge_small.npy
+2a2d312fcea8a6082fe8035e1e35d44ff425ed26bb14a7841205bb6054c845ea  ./ontology/tags_v1_r8.json
+ecabaf702a9ad51efc739a44d6c5b5ecc92403c62d5ea365f1970cc466f2752a  ./ontology/text_tags.json
+90da9da1a2da0eac47fa6ea441b29179e7100ccbe222beb3d979257414a01e91  ./README.md
+174b3daab120ee5ba78a21f811f732f923b40fdfa95b5b12fd6adcedd4ab2369  ./requirements.txt
+20037d4fa059a7efe37acd4ba1c334a13c9c930c53ee116c0dcde185f21d5537  ./thresholds/events_step88000.json
+9297c10a69e3b17a360c4b1f6f62050b247fc74ca28e3cd182804fdea9756300  ./thresholds/step88000.json
+6dc720dd382f90cc514b18bbda9b1c8e083b0cba606bf3ab5129d98bc9389d3b  ./tok/bpe8k/tokenizer_config.json
+2ae4623fc283d9b25b83e9efb9bf376f24c8974a01a2bbbe6cf8b0b29ac6fc35  ./tok/bpe8k/tokenizer.json

README.md

@@ -0,0 +1,194 @@
+---
+language:
+- en
+license: other
+library_name: pytorch
+pipeline_tag: token-classification
+tags:
+- streaming
+- asr
+- ner
+- pii
+- redaction
+- intent
+- dialogue-act
+- mamba
+---
+
+# Arc Tagger (10M), v0.1.0
+
+Arc Tagger is a **text-only, streaming tagger** with about 9.8 M parameters. It reads the word stream of **any ASR** (Deepgram, Nemotron,
+Whisper, our own ASR, or plain text). As each word arrives, it emits entity and number spans, PII spans for redaction, profanity and
+toxicity, discourse acts, intents, safety flags, speaker roles, events (speaker change, turn yield, sports, legal and meeting events),
+content state and emotion (experimental).
+
+* Checkpoint: `AT10Mv2` step 88000. It scored best on all gold sets of its run.
+* Architecture: a word channel (8k byte-level BPE, 2 causal local-attention layers) feeds a 6-layer Mamba-3 trunk. Global KV-memory
+  attention layers, domain FiLM and adapters condition the trunk. The read-outs are definition-keyed tag heads plus a v2 span head (BIO
+  and type per family), a span verifier and pooled utterance, speaker and item heads. The model is causal and runs one recurrent step
+  per word, so a revision re-runs only the revised words.
+* Runs on CPU. Measured on one core (busy shared box, 1 thread): **engine about 16 ms per word p50, 25 ms p95**. A full
+  `push_words()` call, including span decoding and read-outs, takes about 24 ms p50 and 45 ms p95.
+
+## Tag families
+
+276 tags across 10 read-outs. Every tag name comes from the Whissle ontology r8 plus the extension `ext_cc_roles1`.
+
+| family | tags (examples) | read-out |
+|---|---|---|
+| Entities | `entity.mention.{person,organization,team,place,event_name,legal_ref,term,...}` (10) | word spans |
+| Numbers | `entity.number.{date,time_of_day,duration,count,money,percent,measure,score,clock,ordinal,id_pii,...}` (16), `entity.numeral` | word spans |
+| PII / redaction | `safety.pii.{private_name,phone_number,email,street_address,card_account_number,gov_id,date_of_birth,url_handle,public_contact}` spans; `safety.pii.{request,leak}` per utterance | spans and utterance |
+| Profanity / toxicity | `safety.profanity.{mild,strong}`, `safety.toxicity.*` (6) | spans and utterance |
+| Safety | `safety.scam.*` (6), `safety.crisis.self_harm`, `safety.threat.escalation`, `safety.child.*`, `safety.solicitation.dating_romantic` | utterance |
+| Discourse act | `discourse.act.*` (statement, question, command_request, ...), `discourse.answer`, `discourse.backchannel`, `discourse.assistant.*` | utterance |
+| Intent | 26 cross-domain `intent.*` (request_information, schedule_arrange, complain_criticize, ...), 60 assistant `discourse.intent.*` (alarm_set, weather_query, ...) | utterance |
+| Speaker role | `speaker.role.*` (19, including contact-centre agent and customer) | per speaker (needs `spk` ids on the words) |
+| Events | `speaker.change`, `discourse.turn_yield`, `discourse.interruption`, `event.sports.*` (18), `event.legal.*`, `event.meeting.{decision,action_item}` | per word |
+| Content state | `content.mode`, `content.subject`, `content.game_state`, `content.legal_phase`, `content.meeting_phase` | per word (frame) |
+| Emotion (experimental) | `affect.emotion.*` (5), `affect.arousal.*` (3) | utterance; low support, mostly "neutral" |
+
+`ontology/` holds the full definitions. `config.json["columns"]` lists every `tag@readout` column.
+
+## Quickstart
+
+```bash
+pip install huggingface_hub
+huggingface-cli download WhissleAI/arc-tagger --local-dir arc-tagger     # private repo: huggingface-cli login first
+cd arc-tagger
+pip install -r requirements.txt     # CPU torch is fine: pip install torch --index-url https://download.pytorch.org/whl/cpu
+python example.py
+```
+
+```python
+import sys; sys.path.insert(0, "arc-tagger")          # the package ships in the repo (arctag/)
+from arctag import ArcTagger
+
+tg = ArcTagger.from_pretrained("arc-tagger")           # a local dir, or a repo id: ArcTagger.from_pretrained("WhissleAI/arc-tagger")
+out = tg.tag("Hi, this is Maria from Northwind Bank. Can I get your phone number?", domain="talk")
+print(out["redacted_text"]); print(out["spans"]); print(out["utterances"])
+
+# streaming (words as the ASR emits them; interim hypotheses can be revised)
+st = tg.stream(domain="assistant")
+st.push_words(["set", "an", "alarm", "for"], final=False)                       # interim
+up = st.push_words(["set", "an", "alarm", "for", "seven", "thirty."], revise_from=0, final=True)
+up = st.endpoint()                         # end of utterance -> utterance read-outs in up["utterances"]
+res = st.finalize()                        # full result (same schema as tag())
+```
+
+### Input
+
+`tag(x, domain="auto", tier="T0")`
+* `x` can be a `str`, a list of words, or a list of word dicts `{"w": str, "t0": s, "t1": s, "conf": float, "spk": str|int, "final": bool}`.
+  Every field except `w` is optional.
+* Utterances end after words ending in `. ? !`. In streaming, call `endpoint()` at the ASR's end of utterance or end of turn.
+
+`domain` is one of `auto | talk | meeting | court | assistant | sports | news`.
+* It conditions the model through domain FiLM and adapters, and picks the per-domain thresholds.
+* `talk` covers contact centre, podcasts and chat. Inside `sports` and `court`, the specific league or court type is detected from
+  lexicon cues.
+* `auto` runs the "unknown" domain, which is weaker. **Set the domain when you know it.**
+
+`tier`
+* `T0` (default; 0 words of lookahead; the thresholds were fitted on it)
+* `T1` (at most 3 words of delay)
+* `T2` (utterance end)
+
+Streaming: `push_words(words, final=True, revise_from=None, endpoint=False)`
+* `revise_from=k` replaces `words[k:]`. Use it for interim hypotheses such as Deepgram's `is_final=false` or a growing RNNT hypothesis.
+
+### Output (JSON)
+
+```jsonc
+{
+ "meta": {"model": "arc-tagger", "version": "0.1.0", "step": 88000, "domain": "talk", "fine_domain": "talk", "tier": "T0", "thresholds": "..."},
+ "words": ["Hi,", "this", "is", "Maria", ...],
+ "spans": [                                     // all decoded spans (word indices [start, end))
+   {"tag": "safety.pii.private_name", "family": "pii", "start": 3, "end": 4, "text": "Maria", "p": 0.97, "thr": 0.42,
+    "alt": {"tag": "...", "p": 0.01},           // runner-up type
+    "ver": 0.99,                                // span verifier probability
+    "src": "head"}                              // decoded by the v2 span head
+ ],
+ "entities": [...], "pii": [...], "profanity_toxicity": [...],   // spans split by family
+ "redacted_text": "Hi, this is [NAME] from ...",                  // PII / profanity / toxicity spans masked (public_contact is not masked)
+ "utterances": [
+   {"start": 0, "end": 8, "text": "...",
+    "intent": ["intent.request_information"], "act": ["discourse.act.question"], "safety": ["safety.pii.request"],
+    "role": [], "role_note": "speaker role needs diarization (no speaker id on words)",
+    "item": [], "state": {"content.mode": ["content.mode.call_in"]},
+    "emotion": [{"tag": "affect.emotion.neutral", "p": 0.93, "thr": 0.6, "on": true}, ...],
+    "scores": {"intent": [{"tag": ..., "p": ..., "thr": ..., "on": ...}], "act": [...], ...}}   // top-k with probabilities
+ ],
+ "events": [{"tag": "speaker.change", "at": 15, "word": "Yeah,", "p": 0.94, "thr": 0.3, "tier": "T0", "read": "ev_rec"}]
+}
+```
+
+The label lists (`intent`, `act`, `safety`, ...) contain the tags at or above their threshold (`on`). `scores` keeps the top-k with
+their probabilities.
+
+Each streaming update (`push_words` and `endpoint`) returns:
+* `n_words` and `changed_from`
+* `spans`: the re-decoded window; replace your spans that end after its start
+* `events` (new ones only)
+* `utterances` (the ones that just ended)
+* `running`: the live intent, act, safety and state of the open utterance
+
+## Latency tiers
+
+| tier | lookahead | use |
+|---|---|---|
+| T0 | 0 words | live display and zero-lag redaction (default; thresholds fitted here) |
+| T1 | up to 3 words | revisions a few words later |
+| T2 | to utterance end (up to 8 words) | final utterance labels |
+
+All three tiers are delayed read-outs of one causal pass, so they add no extra compute. Utterance-level labels (intent, act, safety,
+role) are read at the endpoint.
+
+## Evaluation
+
+Internal multi-domain dev and test sets cover contact centre, assistant, meetings, news, talk, sports and courts, using real and
+simulated ASR (Deepgram and RNNT word streams).
+
+| metric | old gold | cleaned gold |
+|---|---|---|
+| light eval, primary score (step 88000) | 0.693 | 0.764 |
+| test, T0 macro AP | 0.719 | 0.793 |
+
+**Caveat:** "cleaned gold" re-adjudicated the eval labels after a confident-learning label audit. The model selected the suspect
+labels and Claude adjudicated them. The same audit also cleaned part of the training data, so cleaned-gold numbers are partly
+favourable to this model family. Treat **old-gold numbers as the conservative estimate**. These are internal benchmarks, not public
+leaderboards.
+
+## Files
+
+| file | what |
+|---|---|
+| `model.safetensors` | weights (fp32, 45 MB) |
+| `config.json` | model, ontology and read-out columns; training args with data paths removed |
+| `tok/bpe8k/` | 8k byte-level BPE tokenizer (words are fed as `" " + word`) |
+| `ontology/` | r8 ontology, text-tag table, `ext_cc_roles` extension, frozen tag-definition embeddings (bge-small) |
+| `thresholds/step88000.json` | display thresholds: a precision-first T0 refit on the cleaned dev split, per tag, read-out and domain |
+| `thresholds/events_step88000.json` | display tier, thresholds and hidden list for each event |
+| `demo_settings.json` | the live demo's engine floors and settings (span threshold of at least .3; toxicity of at least .5 outside safety verticals) |
+| `label_norm.json` | label-normalisation conventions behind the tags (training-time; pii1sn3) |
+| `arctag/` | Python package: `ArcTagger`, `driver.py` (session to JSON), `_core/` (the demo's streaming engine and model code, vendored with only file paths changed) |
+| `onnx/` | ONNX export for the C API, if present (published separately) |
+| `MANIFEST.sha256` | checksums |
+
+## Limitations
+
+* **English only.** Non-English input gives garbage-like output.
+* Text only: no audio cues. The model cannot detect speaker identity, emotion from voice or vocal events. **Speaker roles need speaker
+  ids** (`spk`) on the words, for example from ASR diarization.
+* Utterance-level safety and scam flags can fire on short or out-of-domain utterances; use `scores` (p) and your own policy for
+  anything consequential. Emotion is low-support and mostly predicts neutral.
+* The `auto` domain is weaker than setting the domain. Some sports events fire loosely outside sports.
+* Redaction here masks the decoded spans after the fact. The live demo's zero-lag redactor (carrier-phrase pre-arm, regex rules, hold
+  modes) is not included.
+* Thresholds are fitted for precision first on internal dev data. Refit them for your traffic if you can.
+* CPU inference only in this package, which is fp32. The ONNX export, if present, is the path for C and embedded use.
+
+## Licence
+
+Weights: WhissleAI, private release (`license: other`). The vendored `arctag/_core/arc_audio` subset is from WhissleAI Arc 1.4.2,
+licensed under Apache-2.0 (see `LICENSE-arc_audio`).

arctag/__init__.py

@@ -0,0 +1,110 @@
+"""Arc Tagger: text-only streaming tagger (~9.8 M params) for ASR word streams.
+
+    from arctag import ArcTagger
+    tg = ArcTagger.from_pretrained(".")              # or a Hugging Face repo id
+    out = tg.tag("Hi, this is Dana from Acme. Can you move my appointment to Friday?")
+
+    st = tg.stream(domain="talk")                    # streaming: words in as the ASR emits them
+    st.push_words(["hi", "this", "is"], final=False) # interim
+    st.push_words(["hi,", "this", "is", "Dana."], revise_from=0, final=True)
+    st.endpoint()
+    res = st.finalize()
+
+CPU only. The model code is the demo's streaming engine (arctag/_core, vendored unchanged except for file paths).
+"""
+import importlib.util, json, os, sys
+
+__version__ = "0.1.0"
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_CORE = os.path.join(_HERE, "_core")
+_EN = None
+
+
+def _engine(home):
+    """import the vendored engine once (module-level path constants resolve against `home`)"""
+    global _EN
+    if _EN is None:
+        os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+        os.environ.setdefault("ARCTAG_HOME", home)
+        os.environ["ARCTAG_ROLE"] = "offline"          # engine.py: no prod checkpoint-dir guard
+        if _CORE not in sys.path:
+            sys.path.insert(0, _CORE)
+        spec = importlib.util.spec_from_file_location("arctag_engine", os.path.join(_CORE, "engine.py"))
+        _EN = importlib.util.module_from_spec(spec); sys.modules["arctag_engine"] = _EN
+        spec.loader.exec_module(_EN)
+    return _EN
+
+
+def _driver():
+    from . import driver
+    return driver
+
+
+class ArcTagger:
+    def __init__(self, home, EN, loaded, cfg):
+        self.home, self.EN, self.loaded, self.config = home, EN, loaded, cfg
+
+    @classmethod
+    def from_pretrained(cls, path=".", revision=None, thresholds="refit"):
+        """path: release dir (model.safetensors, config.json, tok/, ontology/, thresholds/) or a Hugging Face repo id.
+        thresholds: "refit" (default: thresholds/step88000.json + events_step88000.json, precision-first T0 refit on cleaned dev)
+        or "none" (0.5 everywhere, engine display floors still apply)."""
+        if not os.path.isdir(path):
+            from huggingface_hub import snapshot_download
+            path = snapshot_download(path, revision=revision, allow_patterns=["*.json", "*.safetensors", "*.npy", "*.py", "tok/*",
+                                                                               "ontology/*", "thresholds/*", "arctag/*"])
+        home = os.path.abspath(path)
+        EN = _engine(home)
+        AT = sys.modules["arc_text"]; TT = sys.modules["tagtable"]
+        AT.BPE8K = os.path.join(home, "tok", "bpe8k", "tokenizer.json")
+        cfg = json.load(open(os.path.join(home, "config.json")))
+        ext = os.path.join(home, cfg.get("onto_ext") or "ontology/ext_cc_roles.json")
+        T = TT.TextTable(r8=os.path.join(home, "ontology", "tags_v1_r8.json"),
+                         text_tags=os.path.join(home, "ontology", "text_tags.json"), ext=ext)
+        T.load_E(os.path.join(home, "ontology", "head_E_r8_bge_small.npy"))
+        if T.sha != cfg.get("ontology_sha"):
+            raise RuntimeError(f"ontology sha {T.sha} != config {cfg.get('ontology_sha')}")
+        loaded = _load_model(EN, home, cfg, T, thresholds)
+        return cls(home, EN, loaded, cfg)
+
+    def stream(self, domain="auto", tier="T0"):
+        """a streaming session (driver.Stream): push_words(words, final=True, revise_from=None, endpoint=False), endpoint(),
+        finalize() -> result dict"""
+        return _driver().Stream(self.EN, self.loaded, domain=domain, tier=tier, meta=self._meta())
+
+    def tag(self, x, domain="auto", tier="T0"):
+        """x: text (str), a list of words (str) or word dicts {"w", "t0", "t1", "conf", "spk", "final"}. Utterances end after
+        words ending in . ? ! Returns the structured result (see README "Output")."""
+        return _driver().tag(self.EN, self.loaded, x, domain=domain, tier=tier, meta=self._meta())
+
+    def _meta(self):
+        return {"model": "arc-tagger", "version": __version__, "step": self.loaded.step}
+
+
+def _load_model(EN, home, cfg, T, thresholds="refit"):
+    """= demo/engine.py ModelManager._load + _thresholds for a single release dir"""
+    from safetensors.torch import load_file
+    AT = sys.modules["arc_text"]; SA = sys.modules["stream_arc"]
+    step = int(cfg.get("step", 0))
+    m = AT.ArcText(T, scratch=cfg.get("scratch") or AT.SCRATCH_10M)
+    sd = load_file(os.path.join(home, "model.safetensors"))
+    SA.add_modules_for(m, sd)
+    missing, unexpected = m.load_state_dict(sd, strict=False)
+    if unexpected:
+        raise RuntimeError(f"checkpoint keys without a module: {unexpected[:6]}")
+    bad = [k for k in missing if not k.startswith(("head.E",))]
+    if bad:
+        import warnings
+        warnings.warn(f"{len(bad)} missing keys (e.g. {bad[:3]})")
+    m.eval()
+    own = cfg.get("thresholds") or {}
+    pref, pstep, ev, evstep = {}, None, {}, None
+    if thresholds == "refit":
+        d = os.path.join(home, "thresholds")
+        pref, pstep = EN.load_prod_thresholds(step, d)
+        ev, evstep = EN.load_event_tiers(step, d)
+    elif thresholds != "none":
+        raise ValueError("thresholds must be 'refit' or 'none'")
+    thr = EN.Thresholds(own, {}, None, pref=pref, pref_step=pstep)
+    thr.set_events(ev, evstep)
+    return EN.Loaded(step, home, m, T, thr, cfg)

arctag/_core/__init__.py

@@ -0,0 +1 @@
+# vendored engine modules (copies of the Arc Tagger demo/code modules; see arctag/__init__.py)

arctag/_core/arc_audio/__init__.py

@@ -0,0 +1,2 @@
+"""Subset of arc_audio 1.4.2 (WhissleAI Arc, Apache-2.0, see LICENSE-arc_audio) needed by the Arc Tagger: model.py (Trunk / TagHead /
+EntityHeads), wordch.py (word channel), ontology.py, mamba3/ (pure-PyTorch Mamba-3). Files copied unchanged; the audio API is not included."""

arctag/_core/arc_audio/mamba3/__init__.py

@@ -0,0 +1,5 @@
+"""Pure-PyTorch Mamba-3 layer (Lahoti et al., arXiv 2603.15569; WhissleAI in-house implementation, no CUDA kernels):
+layer.py (recurrent path + step()), scan.py (chunked scan), ssd.py (chunked-SSD parallel path used for offline inference).
+Identical to the copy the model was trained with."""
+from .layer import Mamba3Layer  # noqa: F401
+from .ssd import mamba3_layer_ssd, _f  # noqa: F401

arctag/_core/arc_audio/mamba3/layer.py

@@ -0,0 +1,332 @@
+"""Mamba-3 (Lahoti et al., arXiv 2603.15569), implemented to the paper.
+
+THREE CHANGES FROM MAMBA-2, implemented here in the order the paper derives
+them. Nothing is improvised: where this file departs from the paper it says
+so, and the equivalence tests at the bottom are how that claim is checked.
+
+1. EXPONENTIAL-TRAPEZOIDAL DISCRETIZATION (Sec 3.1). Mamba-1/2 approximate
+
+        integral_{t_{k-1}}^{t_k} e^{A(t_k - tau)} B(tau) x(tau) d tau
+
+   by holding the endpoint (Euler). Mamba-3 averages both endpoints
+   (trapezoidal), which makes the recurrence depend on x_{t-1} as well as
+   x_t:
+
+        h_t = alpha_t h_{t-1} + beta_t B_{t-1} x_{t-1} + gamma_t B_t x_t
+
+   with alpha_t = exp(Delta_t A), beta_t = (Delta_t/2) alpha_t,
+   gamma_t = Delta_t/2. That two-tap dependence IS the implicit convolution
+   the paper describes - it is why Mamba-3 can drop the short causal conv
+   that Mamba-1/2 carry as a separate component (Sec 3.1, and eq 7's
+   factorisation of L into a 1-semiseparable matrix times a 2-band matrix).
+
+2. COMPLEX STATE via the RoPE trick (Sec 3.2, Prop 2 and 3). A complex SSM
+   of state N/2 is equivalent to a REAL SSM of state N whose transition is a
+   scalar decay times a block-diagonal of 2x2 rotations:
+
+        R_t = blockdiag({ R(Delta_t theta_t[i]) })
+        h_t = alpha_t R_t h_{t-1} + ...
+
+   Real non-negative eigenvalues cannot express rotation, which is why
+   Mamba-2 provably fails state tracking (parity). This file implements the
+   real-equivalent form, not complex arithmetic, exactly as Prop 2 licenses.
+
+3. MIMO (Sec 3.3). x_t: (P,) -> (P,R), B_t: (N,) -> (N,R), C_t likewise;
+   the state stays (N,P) and the update becomes a matmul B_t x_t^T rather
+   than an outer product. Memory traffic is dominated by h_t and barely
+   moves; FLOPs rise by R. Arithmetic intensity goes Theta(1) -> Theta(R),
+   which is the whole point during memory-bound decode.
+
+WHAT THIS FILE DOES NOT DO. No fused CUDA scan - the paper's speed comes
+from hardware-aware kernels and this runs on MPS/CPU with an explicit loop
+for the recurrent path and a materialised mask for the parallel path. The
+MODELLING claims are reproducible here; the LATENCY claims are not.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def rotate_pairs(h: torch.Tensor, ang: torch.Tensor) -> torch.Tensor:
+    """Apply blockdiag 2x2 rotations to the state's (even, odd) channel pairs.
+
+    h:   (..., N)   state, N even
+    ang: (..., N/2) rotation angle per pair
+    """
+    a, b = h[..., 0::2], h[..., 1::2]
+    c, s = torch.cos(ang), torch.sin(ang)
+    out = torch.empty_like(h)
+    out[..., 0::2] = c * a - s * b
+    out[..., 1::2] = s * a + c * b
+    return out
+
+
+def ssm_scan(alpha, beta, gamma, theta, S, Cm, P, state=None, prev_S=None):
+    """The Mamba-3 recurrence (paper eq 11), shared by Mamba3Layer and MMSE.
+
+    This is the SINGLE SOURCE OF TRUTH for the recurrence - both the single-
+    stream layer and the multi-stream MMSE core call it, so a paper-faithfulness
+    fix here (e.g. the beta-term rotation) applies to both and cannot drift.
+
+        h_t = R_t(alpha_t h_{t-1} + beta_t S_{t-1}) + gamma_t S_t
+        y_t = C_t^T h_t
+
+    alpha,beta,gamma: (B,T,H)      per-step scalars
+    theta:            (B,T,H,N/2)  per-step rotation angles (already * dt)
+    S:                (B,T,H,N,P)  per-step write B_t x_t^T (or summed streams)
+    Cm:               (B,T,H,N,R)  readout
+    P:                head dim; state/prev_S carry across chunks for streaming.
+    Returns (y_flat (B,T,H*R*P), final state h, last write S[:,-1]).
+    """
+    B, T, H, N, _ = S.shape
+    R = Cm.shape[-1]
+    h = torch.zeros(B, H, N, P, device=S.device, dtype=S.dtype) \
+        if state is None else state
+    ys = []
+    for t in range(T):
+        prev = S[:, t - 1] if t > 0 else prev_S
+        pre = alpha[:, t, :, None, None] * h
+        if prev is not None:
+            pre = pre + beta[:, t, :, None, None] * prev
+        h = rotate_pairs(pre.transpose(-1, -2), theta[:, t].unsqueeze(-2)
+                         .expand(-1, -1, P, -1)).transpose(-1, -2)
+        h = h + gamma[:, t, :, None, None] * S[:, t]
+        ys.append((Cm[:, t].transpose(-1, -2) @ h).reshape(B, H * R * P))
+    return torch.stack(ys, 1), h, S[:, -1]
+
+
+class Mamba3Layer(nn.Module):
+    """One Mamba-3 layer. Recurrent and parallel paths, provably equal.
+
+    d_model  input/output width
+    d_state  N, the SSM state size (even; pairs are rotated)
+    rank     R of the MIMO lift. rank=1 is SISO.
+    head_dim P, channels sharing one (alpha, Delta, B) recurrence.
+    """
+
+    def __init__(self, d_model: int, d_state: int = 64, rank: int = 1,
+                 head_dim: int = 32, dt_min: float = 1e-3, dt_max: float = 0.5):
+        super().__init__()
+        assert d_state % 2 == 0, "d_state must be even: the state rotates in pairs"
+        assert d_model % head_dim == 0, "d_model must divide into heads of head_dim"
+        self.d_model, self.N, self.R, self.P = d_model, d_state, rank, head_dim
+        self.H = d_model // head_dim
+
+        # Gate branch (Mamba-2/3 block, Fig 2): the SSM output is multiplied by
+        # SiLU(in_proj(input)) before out_proj. Dimension d_model*R to match the
+        # pre-out_proj SSM output (H*R*P). Was defined-but-unused before.
+        self.in_proj = nn.Linear(d_model, d_model * rank)
+        self.B_proj = nn.Linear(d_model, self.H * d_state * rank)
+        self.C_proj = nn.Linear(d_model, self.H * d_state * rank)
+        self.x_proj = nn.Linear(d_model, d_model * rank)
+        self.dt_proj = nn.Linear(d_model, self.H)
+        self.theta_proj = nn.Linear(d_model, self.H * (d_state // 2))
+        # lambda_t: the DATA-DEPENDENT trapezoidal weight (paper Prop 1, Remark
+        # 3). gamma_t = lambda_t * dt, beta_t = (1-lambda_t) * dt * alpha_t.
+        # lambda=1/2 is the classical trapezoid; the paper found NOT enforcing
+        # that (letting the model choose per step) is empirically better. Bias 0
+        # -> sigmoid 0.5 at init, so training starts from the classical rule.
+        self.lam_proj = nn.Linear(d_model, self.H)
+        nn.init.zeros_(self.lam_proj.weight)
+        nn.init.zeros_(self.lam_proj.bias)
+        self.out_proj = nn.Linear(d_model * rank, d_model)
+        self.norm = nn.RMSNorm(d_model) if hasattr(nn, "RMSNorm") else nn.LayerNorm(d_model)
+
+        # BCNorm / QKNorm (paper section 3.2). RMS normalisation AFTER the B and C
+        # projections, mirroring QKNorm in Transformers. The paper reports this
+        # is what stabilises large-scale runs, to the point that they remove
+        # Mamba-2's post-gate RMSNorm on the strength of it. It was missing from
+        # the first version of this file and E17's ssm arm then diverged
+        # (0.330 -> 19.313) into a constant predictor, which was patched with
+        # gradient clipping - a generic fix for a problem the paper had already
+        # solved specifically. Normalising over the N state dimensions, which is
+        # the axis B and C index into.
+        rms = nn.RMSNorm if hasattr(nn, "RMSNorm") else nn.LayerNorm
+        self.B_norm = rms(d_state)
+        self.C_norm = rms(d_state)
+        # B, C biases (paper appendix): learnable, head-specific, channel-wise,
+        # data-INDEPENDENT, added after the norm, initialised to all ones. The
+        # paper finds models insensitive to the init as long as it is positive,
+        # and that B and C biases together beat either alone.
+        self.B_bias = nn.Parameter(torch.ones(self.H, d_state, rank))
+        self.C_bias = nn.Parameter(torch.ones(self.H, d_state, rank))
+
+        # A is data-INDEPENDENT here. Remark 1: the paper found data-dependent
+        # A_t performs similarly and chose it for consistency; the simpler
+        # form is used here and the difference is recorded rather than hidden.
+        self.A_log = nn.Parameter(torch.log(torch.rand(self.H) * 0.9 + 0.05))
+        # dt bias initialised so Delta spans [dt_min, dt_max] - the standard
+        # Mamba init. Delta sets the memory horizon, so this is the prior on
+        # how long the layer remembers before it learns anything.
+        dt = torch.exp(torch.rand(self.H) * (math.log(dt_max) - math.log(dt_min))
+                       + math.log(dt_min))
+        self.dt_proj.bias.data = dt + torch.log(-torch.expm1(-dt))
+
+    def _coeffs(self, u: torch.Tensor):
+        """alpha, beta, gamma, B, C, x, theta for every timestep."""
+        Bsz, T, _ = u.shape
+        H, N, R, P = self.H, self.N, self.R, self.P
+        dt = F.softplus(self.dt_proj(u))                      # (B,T,H) > 0
+        A = -torch.exp(self.A_log)                            # (H,) < 0
+        alpha = torch.exp(dt * A)                             # (B,T,H) in (0,1)
+        # Exponential-trapezoidal with a DATA-DEPENDENT weight lambda_t (paper
+        # Prop 1 / eq 5-6): gamma_t = lambda_t*dt (current input), beta_t =
+        # (1-lambda_t)*dt*alpha_t (previous input). lambda=1/2 recovers the
+        # classical trapezoid; lambda=1 recovers Mamba-2's exponential-Euler.
+        lam = torch.sigmoid(self.lam_proj(u))                 # (B,T,H) in (0,1)
+        gamma = lam * dt
+        beta = (1.0 - lam) * dt * alpha
+        # BCNorm then bias, in that order - the paper adds the bias AFTER the
+        # norm, so the bias is not rescaled by it.
+        Bm = self.B_proj(u).view(Bsz, T, H, R, N)
+        Bm = self.B_norm(Bm).permute(0, 1, 2, 4, 3) + self.B_bias
+        Cm = self.C_proj(u).view(Bsz, T, H, R, N)
+        Cm = self.C_norm(Cm).permute(0, 1, 2, 4, 3) + self.C_bias
+        x = self.x_proj(u).view(Bsz, T, H, P, R)
+        theta = self.theta_proj(u).view(Bsz, T, H, N // 2) * dt.unsqueeze(-1)
+        return alpha, beta, gamma, Bm, Cm, x, theta
+
+    def gate_and_proj(self, y_flat: torch.Tensor, nu: torch.Tensor):
+        """Apply the Mamba gate then the output projection.
+
+        y_flat: (..., d_model*R) SSM output before out_proj.
+        nu:     (..., d_model)   the NORMALISED layer input (in_proj sees the
+                                 same normalised input the SSM path does).
+        out = out_proj( y_flat * SiLU(in_proj(nu)) ).
+        """
+        return self.out_proj(y_flat * F.silu(self.in_proj(nu)))
+
+    def forward(self, u: torch.Tensor, state=None, return_state=False):
+        """Recurrent path. u: (B,T,d_model). This is the streaming form.
+
+        The recurrence itself lives in the shared `ssm_scan` (paper eq 11), so
+        this and MMSECore.forward cannot drift. Here the per-step write is the
+        single-stream outer product S_t = B_t x_t^T.
+        """
+        nu = self.norm(u)
+        alpha, beta, gamma, Bm, Cm, x, theta = self._coeffs(nu)
+        S = Bm @ x.transpose(-1, -2)                    # (B,T,H,N,P)
+        y_flat, h, _ = ssm_scan(alpha, beta, gamma, theta, S, Cm, self.P,
+                                state=state)
+        out = self.gate_and_proj(y_flat, nu)
+        return (out, h) if return_state else out
+
+    def step(self, u_t: torch.Tensor, state, prev_bx=None):
+        """One timestep, constant memory. The deployment path.
+
+        state: (B,H,N,P) or None. prev_bx: B_{t-1}x_{t-1}^T for the beta term,
+        which the trapezoidal rule requires and Mamba-2 does not have.
+        """
+        Bsz = u_t.shape[0]
+        nu = self.norm(u_t.unsqueeze(1))
+        alpha, beta, gamma, Bm, Cm, x, theta = self._coeffs(nu)
+        H, N, P = self.H, self.N, self.P
+        h = torch.zeros(Bsz, H, N, P, device=u_t.device, dtype=u_t.dtype) \
+            if state is None else state
+        # Match forward/eq(11): rotate (alpha*h + beta*prev_bx) by R_t, then add
+        # the current (unrotated) write. prev_bx = B_{t-1}x_{t-1}^T from the last
+        # step; bx = B_t x_t^T is returned for the next step's beta term.
+        pre = alpha[:, 0, :, None, None] * h
+        if prev_bx is not None:
+            pre = pre + beta[:, 0, :, None, None] * prev_bx
+        h = rotate_pairs(pre.transpose(-1, -2), theta[:, 0].unsqueeze(-2)
+                         .expand(-1, -1, P, -1)).transpose(-1, -2)
+        bx = Bm[:, 0] @ x[:, 0].transpose(-1, -2)
+        h = h + gamma[:, 0, :, None, None] * bx
+        y = (Cm[:, 0].transpose(-1, -2) @ h).reshape(Bsz, -1)
+        return self.gate_and_proj(y, nu[:, 0]), h, bx
+
+
+# ---------------------------------------------------------------------------
+def selftest(verbose: bool = True) -> dict:
+    """The four invariants, at float64.
+
+    These were an ad-hoc scratch script until BCNorm was added and there was
+    nothing to re-run. A verified file whose verification lives somewhere else
+    is an unverified file the moment anyone edits it.
+
+        python nir/mamba3.py
+    """
+    torch.manual_seed(0)
+    torch.set_default_dtype(torch.float64)
+    net = Mamba3Layer(d_model=32, d_state=16, rank=2, head_dim=8).eval()
+    B, T = 2, 20
+    u = torch.randn(B, T, 32)
+    out = {}
+
+    with torch.no_grad():
+        y = net(u)
+
+        # 1. CAUSALITY, as a SCALE test rather than an equality one. Appending
+        #    future timesteps moves the past by ~1e-15 here: the permute in
+        #    _coeffs leaves B and C non-contiguous, so matmul takes a different
+        #    kernel path at a different T and the last bits differ. An equality
+        #    assertion would call that a leak, and would also fail spuriously in
+        #    float32. A real leak scales with the size of the future; rounding
+        #    does not. Drive the future through four orders of magnitude and
+        #    require the drift to stay put.
+        drifts = []
+        for mag in (1.0, 1e2, 1e4):
+            u2 = torch.cat([u, torch.randn(B, 6, 32) * mag], 1)
+            drifts.append(float((net(u2)[:, :T] - y).abs().max()))
+        out["causality"] = max(drifts)
+        out["causality_scaling"] = max(drifts) / max(min(drifts), 1e-300)
+
+        # 2. STREAMING == BATCHED. The deployment path must equal training.
+        st = bx = None
+        steps = []
+        for t in range(T):
+            o, st, bx = net.step(u[:, t], st, bx)
+            steps.append(o)
+        out["streaming"] = float((torch.stack(steps, 1) - y).abs().max())
+
+        # 3. ROTATION IS ORTHOGONAL. State norm preserved, so it cannot blow up.
+        h = torch.randn(B, 4, 16, 8)
+        ang = torch.randn(B, 4, 1, 8)
+        hr = rotate_pairs(h.transpose(-1, -2), ang).transpose(-1, -2)
+        out["orthogonal"] = float(
+            (hr.norm(dim=-2) - h.norm(dim=-2)).abs().max())
+
+        # 4. TRAPEZOIDAL TWO-TAP. Perturbing x_{t-1} moves y_t, not y_{t-2}.
+        up = u.clone(); up[:, 3] += 1.0
+        yp = net(up)
+        out["twotap_y4"] = float((yp[:, 4] - y[:, 4]).abs().max())
+        out["twotap_y2"] = float((yp[:, 2] - y[:, 2]).abs().max())
+
+        # 5. BCNORM IS ACTIVE. B's pre-bias rows should have unit RMS. Without
+        #    the norm this is whatever the projection happens to produce, and
+        #    the first version of this file had no norm at all.
+        a, b, g, Bm, Cm, x, th = net._coeffs(net.norm(u))
+        #    Bm is (B,T,H,N,R) after the permute, so the normalised axis is -2.
+        rms = (Bm - net.B_bias).pow(2).mean(-2).sqrt()
+        out["bcnorm_rms"] = float((rms - 1.0).abs().max())
+
+    out["params"] = sum(p.numel() for p in net.parameters())
+    torch.set_default_dtype(torch.float32)
+
+    if verbose:
+        ok = lambda c: "OK" if c else "FAIL"
+        print("Mamba-3 selftest  (float64)")
+        print(f"  1 causality (future->past)   {out['causality']:.3e}"
+              f"   {ok(out['causality'] < 1e-12)}")
+        print(f"    drift x1 vs x10000 ratio   {out['causality_scaling']:.3f}"
+              f"   {ok(out['causality_scaling'] < 1.5)}  (a leak would scale)")
+        print(f"  2 step() == forward()        {out['streaming']:.3e}"
+              f"   {ok(out['streaming'] < 1e-10)}")
+        print(f"  3 rotation orthogonal        {out['orthogonal']:.3e}"
+              f"   {ok(out['orthogonal'] < 1e-10)}")
+        print(f"  4 trapezoidal two-tap        y4 {out['twotap_y4']:.3e}"
+              f"  y2 {out['twotap_y2']:.3e}"
+              f"   {ok(out['twotap_y4'] > 1e-6 and out['twotap_y2'] == 0.0)}")
+        print(f"  5 BCNorm rows unit RMS       {out['bcnorm_rms']:.3e}"
+              f"   {ok(out['bcnorm_rms'] < 1e-10)}")
+        print(f"  {out['params']:,} params")
+    return out
+
+
+if __name__ == "__main__":
+    selftest()

arctag/_core/arc_audio/mamba3/scan.py

@@ -0,0 +1,195 @@
+"""Chunked parallel scan for the Mamba-3 recurrence.
+
+WHY. The recurrent path is a Python loop of T sequential steps. Measured on
+this laptop: d_model 256 / 4 layers / N 32 runs at 3,670 tok/s, while
+384 / 6 / 64 runs at 179 tok/s - 20x slower for ~4x the FLOPs. The gap is not
+arithmetic, it is kernel-launch overhead: 6 layers x 128 timesteps is 768
+sequential iterations each firing a handful of tiny kernels, so DEPTH
+MULTIPLIES LAUNCHES. That is what blocks scaling, not parameters.
+
+THE TRICK, and it is the paper's own (Proposition 3). The recurrence
+
+    h_t = alpha_t R(theta_t) h_{t-1} + beta_t S_{t-1} + gamma_t S_t
+
+has a data-dependent rotation inside it, which a scan cannot factor. But
+rotations compose: define the CUMULATIVE angle Phi_t = sum_{s<=t} theta_s and
+change basis by R(-Phi_t). In that frame the rotation disappears and what is
+left is a plain first-order linear recurrence
+
+    g_t = alpha_t g_{t-1} + w_t          w_t = R(-Phi_t) (beta_t S_{t-1} + gamma_t S_t)
+
+with h_t = R(Phi_t) g_t. The paper states this as applying data-dependent
+RoPE to B and C; the same identity in state space is what makes the scan
+possible at all.
+
+THE SCAN. g_t = alpha_t g_{t-1} + w_t has the closed form
+
+    g_t = A_t * sum_{s<=t} w_s / A_s     A_t = prod_{r<=t} alpha_r
+
+which is numerically hopeless directly - A_s underflows for long sequences
+and the division explodes. So this chunks: within a chunk of length L the sum
+is materialised as an L x L lower-triangular decay matrix (no division, decay
+is always relative to the chunk start), and chunk states are carried
+recurrently. T/L sequential steps instead of T. The intra-chunk work becomes
+one batched matmul, which is what GPUs are for.
+
+CORRECTNESS IS CHECKED AGAINST THE RECURRENCE, NOT AGAINST A LOSS CURVE.
+`verify()` compares this to Mamba3Layer.forward at float64. A fast scan that
+computes something slightly different would show up as a slightly worse loss
+curve and be indistinguishable from a hyperparameter difference - which is
+exactly how this project has lost time before.
+
+    python nir/scan.py
+"""
+from __future__ import annotations
+
+import math
+import torch
+
+
+def cumulative_angles(theta: torch.Tensor) -> torch.Tensor:
+    """Phi_t = sum_{s<=t} theta_s.  theta: (B,T,H,N/2)."""
+    return torch.cumsum(theta, dim=1)
+
+
+def rotate_by(v: torch.Tensor, ang: torch.Tensor) -> torch.Tensor:
+    """Rotate pairs of the N axis. v: (...,N,P), ang: (...,N/2) broadcastable.
+
+    Built with stack+reshape rather than writing into an empty tensor. The
+    first version used `out = torch.empty_like(v)` with strided assignment,
+    which is correct in the forward direction and BREAKS AUTOGRAD: in-place
+    writes bump the version counter of a tensor the graph still needs. The
+    forward equivalence check passed anyway, because it ran under no_grad.
+    """
+    a, b = v[..., 0::2, :], v[..., 1::2, :]
+    c, s = torch.cos(ang).unsqueeze(-1), torch.sin(ang).unsqueeze(-1)
+    even, odd = c * a - s * b, s * a + c * b
+    return torch.stack((even, odd), dim=-2).flatten(-3, -2)
+
+
+def chunked_scan(alpha: torch.Tensor, w: torch.Tensor, chunk: int = 64,
+                 state=None):
+    """g_t = alpha_t g_{t-1} + w_t, computed in chunks.
+
+    alpha: (B,T,H)          per-step decay, in (0,1)
+    w:     (B,T,H,N,P)      per-step input, already in the rotated frame
+    returns g: (B,T,H,N,P) and the final state.
+
+    Within a chunk, g_t = d_t * g_0 + sum_{s<=t} (d_t/d_s) w_s where
+    d_t = prod_{r<=t} alpha_r RELATIVE TO THE CHUNK START. Both d_t and d_t/d_s
+    stay bounded because the product never runs longer than `chunk` steps -
+    that is the whole reason for chunking rather than a global cumprod.
+    """
+    B, T, H = alpha.shape
+    N, P = w.shape[-2], w.shape[-1]
+    # Accumulate chunks in a list and cat once. Writing into a preallocated
+    # g_out and then reading g_out[:, e-1] back as the next chunk's state is
+    # an in-place write into a tensor autograd still needs.
+    pieces = []
+    g = torch.zeros(B, H, N, P, device=w.device, dtype=w.dtype) \
+        if state is None else state
+    la = torch.log(alpha.clamp_min(1e-30))                 # (B,T,H)
+    for s in range(0, T, chunk):
+        e = min(s + chunk, T)
+        L = e - s
+        cl = torch.cumsum(la[:, s:e], dim=1)               # (B,L,H) log d_t
+        d = torch.exp(cl)                                  # decay from chunk start
+        # (d_t / d_s) for s<=t, lower triangular, exp of a difference of logs
+        rel = cl.unsqueeze(2) - cl.unsqueeze(1)            # (B,L,L,H)  log d_t - log d_s
+        mask = torch.tril(torch.ones(L, L, device=w.device, dtype=torch.bool))
+        M = torch.where(mask.view(1, L, L, 1), torch.exp(rel),
+                        torch.zeros((), device=w.device, dtype=w.dtype))
+        # intra-chunk: sum_s M[t,s] * w_s  -> one batched matmul over (L,L) x (L, N*P)
+        wf = w[:, s:e].reshape(B, L, H, N * P).permute(0, 2, 1, 3)   # (B,H,L,NP)
+        Mp = M.permute(0, 3, 1, 2)                                   # (B,H,L,L)
+        acc = (Mp @ wf).permute(0, 2, 1, 3).reshape(B, L, H, N, P)
+        # carry the incoming chunk state, decayed
+        piece = acc + d.unsqueeze(-1).unsqueeze(-1) * g.unsqueeze(1)
+        pieces.append(piece)
+        g = piece[:, -1]
+    return torch.cat(pieces, dim=1), g
+
+
+def scan_forward(layer, u: torch.Tensor, chunk: int = 64):
+    """Mamba3Layer.forward computed with the chunked scan instead of the loop."""
+    Bsz, T, _ = u.shape
+    H, N, R, P = layer.H, layer.N, layer.R, layer.P
+    nu = layer.norm(u)
+    alpha, beta, gamma, Bm, Cm, x, theta = layer._coeffs(nu)
+    S = Bm @ x.transpose(-1, -2)                       # (B,T,H,N,P)
+    # trapezoidal: beta carries the PREVIOUS step's write. Built by cat, not
+    # by assigning into a zeros tensor, for the same autograd reason.
+    Sprev = torch.cat([torch.zeros_like(S[:, :1]), S[:, :-1]], dim=1)
+    Phi = cumulative_angles(theta)                     # (B,T,H,N/2) = Phi_s
+    Phi_prev = torch.cat([torch.zeros_like(Phi[:, :1]), Phi[:, :-1]], dim=1)
+    # Paper eq (11): the gamma (current) write rotates by R(-Phi_s), the beta
+    # (previous) write by R(-Phi_{s-1}) - one step less. Rotate the two terms
+    # separately into the global frame, then scan.
+    w = (gamma[..., None, None] * rotate_by(S, -Phi)
+         + beta[..., None, None] * rotate_by(Sprev, -Phi_prev))
+    g, _ = chunked_scan(alpha, w, chunk)
+    h = rotate_by(g, Phi)                              # back out
+    y = (Cm.transpose(-1, -2) @ h).reshape(Bsz, T, H * R * P)
+    return layer.gate_and_proj(y, nu)
+
+
+def verify(chunks=(8, 16, 64), verbose: bool = True) -> dict:
+    """Against the recurrence, at float64 - FORWARD AND BACKWARD.
+
+    The backward half is not optional. The first version of this file passed
+    every forward check and could not backpropagate at all: it built outputs
+    with in-place strided writes, which autograd rejects. The forward checks
+    ran under no_grad and saw nothing. A scan that cannot train is useless,
+    and one that trains with subtly wrong gradients is worse than useless -
+    it would surface as a slightly worse loss curve, indistinguishable from a
+    hyperparameter difference."""
+    import sys
+    from pathlib import Path
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+    from mamba3.layer import Mamba3Layer
+
+    torch.manual_seed(0)
+    torch.set_default_dtype(torch.float64)
+    out = {}
+    lay = Mamba3Layer(d_model=32, d_state=16, rank=2, head_dim=8).eval()
+    u = torch.randn(3, 40, 32)
+    with torch.no_grad():
+        ref = lay(u)
+        for c in chunks:
+            got = scan_forward(lay, u, chunk=c)
+            out[f"chunk{c}"] = float((got - ref).abs().max())
+        # a length that is not a multiple of the chunk
+        u2 = torch.randn(2, 37, 32)
+        r2 = lay(u2)
+        out["ragged"] = float((scan_forward(lay, u2, chunk=16) - r2).abs().max())
+        # causality of the scan path itself
+        u3 = torch.cat([u, torch.randn(3, 9, 32) * 50], 1)
+        out["causality"] = float(
+            (scan_forward(lay, u3, chunk=16)[:, :40]
+             - scan_forward(lay, u, chunk=16)).abs().max())
+
+    # ---- BACKWARD ----
+    ug = torch.randn(3, 24, 32, requires_grad=True)
+    lay.zero_grad()
+    lay(ug).pow(2).sum().backward()
+    gu_rec = ug.grad.clone()
+    gp_rec = {n: p.grad.clone() for n, p in lay.named_parameters()
+              if p.grad is not None}
+    ug.grad = None
+    lay.zero_grad()
+    scan_forward(lay, ug, chunk=8).pow(2).sum().backward()
+    out["grad_input"] = float((gu_rec - ug.grad).abs().max())
+    out["grad_params"] = max(
+        float((gp_rec[n] - p.grad).abs().max())
+        for n, p in lay.named_parameters() if p.grad is not None)
+    torch.set_default_dtype(torch.float32)
+    if verbose:
+        ok = lambda c: "OK  " if c else "FAIL"
+        print("chunked scan vs recurrence  (float64)")
+        for k, v in out.items():
+            print(f"  {k:<12}{v:.3e}   {ok(v < 1e-9)}")
+    return out
+
+
+if __name__ == "__main__":
+    verify()

arctag/_core/arc_audio/mamba3/ssd.py

@@ -0,0 +1,106 @@
+"""Chunked SSD evaluation of Mamba3Layer (copied verbatim from WhissleAI's internal Mamba-3 code)."""
+import torch
+from .scan import rotate_by
+from .layer import Mamba3Layer
+
+def _f(t):
+    return t.float() if t is not None and t.dtype in (torch.float16, torch.bfloat16) else t
+
+
+def _shift(x):
+    return torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1)
+
+
+def ssd_safe(alpha, Beff, xeff, Ctil, chunk: int = 64):
+    """Chunked SSD (as vendor ssd_pytorch) with the decay masked before exp. Returns y (B,T,H,R,P)."""
+    B, T, H, N, K = Beff.shape
+    P = xeff.shape[-2]
+    la = torch.log(alpha.clamp_min(1e-30))
+    g = torch.zeros(B, H, N, P, dtype=Beff.dtype, device=Beff.device)
+    ys = []
+    for s0 in range(0, T, chunk):
+        e = min(s0 + chunk, T)
+        L = e - s0
+        Bc, xc, Cc = Beff[:, s0:e], xeff[:, s0:e], Ctil[:, s0:e]
+        cl = torch.cumsum(la[:, s0:e], dim=1)
+        d = torch.exp(cl)
+        rel = cl.unsqueeze(2) - cl.unsqueeze(1)                       # (B,L,L,H)
+        tril = torch.tril(torch.ones(L, L, dtype=torch.bool, device=Beff.device)).view(1, L, L, 1)
+        Ldec = torch.exp(rel.masked_fill(~tril, float("-inf")))
+        G = torch.einsum('blhnr,bmhnk->bhlmrk', Cc, Bc)
+        y_intra = torch.einsum('blmh,bhlmrk,bmhpk->blhrp', Ldec, G, xc)
+        y_inter = torch.einsum('blh,blhnr,bhnp->blhrp', d, Cc, g)
+        ys.append(y_intra + y_inter)
+        coef = torch.exp(cl[:, -1:, :] - cl)
+        g = d[:, -1].unsqueeze(-1).unsqueeze(-1) * g + torch.einsum('blh,blhnk,blhpk->bhnp', coef, Bc, xc)
+    return torch.cat(ys, dim=1)
+
+
+def _ssd_from_coeffs(alpha, beta, gamma, Bc, xc, Cm, theta, chunk):
+    """Bc (B,T,H,N,K) xc (B,T,H,P,K) Cm (B,T,H,N,R) -> y (B,T,H,R,P); the vendor ssd_inputs reduction."""
+    Phi = torch.cumsum(theta, dim=1)
+    Phi_prev = _shift(Phi)
+    Btil = rotate_by(Bc, -Phi)
+    Ctil = rotate_by(Cm, -Phi)
+    Btil_prev = rotate_by(_shift(Bc), -Phi_prev)
+    Beff = torch.cat([gamma[..., None, None] * Btil, beta[..., None, None] * Btil_prev], dim=-1)
+    xeff = torch.cat([xc, _shift(xc)], dim=-1)
+    return ssd_safe(alpha, Beff, xeff, Ctil, chunk)
+
+
+def mamba3_layer_ssd(layer: Mamba3Layer, u: torch.Tensor, chunk: int = 64):
+    """Same output as Mamba3Layer.forward(u) via SSD."""
+    nu = layer.norm(u)
+    alpha, beta, gamma, Bm, Cm, x, theta = layer._coeffs(nu)
+    y = _ssd_from_coeffs(alpha, beta, gamma, Bm, x, Cm, theta, chunk)
+    Bsz, T = u.shape[:2]
+    return layer.gate_and_proj(y.reshape(Bsz, T, -1), nu)
+
+
+
+
+# ---- stateful SSD (carried state across calls; truncated-BPTT training and reset-free inference) -------------------
+def ssd_safe_state(alpha, Beff, xeff, Ctil, g0=None, chunk: int = 64):
+    """ssd_safe with an initial inter-chunk state g0 (B,H,N,P) (rotated frame of t = -1); returns (y, g_final)."""
+    B, T, H, N, K = Beff.shape
+    P = xeff.shape[-2]
+    la = torch.log(alpha.clamp_min(1e-30))
+    g = torch.zeros(B, H, N, P, dtype=Beff.dtype, device=Beff.device) if g0 is None else g0.to(Beff.dtype)
+    ys = []
+    for s0 in range(0, T, chunk):
+        e = min(s0 + chunk, T)
+        L = e - s0
+        Bc, xc, Cc = Beff[:, s0:e], xeff[:, s0:e], Ctil[:, s0:e]
+        cl = torch.cumsum(la[:, s0:e], dim=1)
+        d = torch.exp(cl)
+        rel = cl.unsqueeze(2) - cl.unsqueeze(1)
+        tril = torch.tril(torch.ones(L, L, dtype=torch.bool, device=Beff.device)).view(1, L, L, 1)
+        Ldec = torch.exp(rel.masked_fill(~tril, float("-inf")))
+        G = torch.einsum('blhnr,bmhnk->bhlmrk', Cc, Bc)
+        y_intra = torch.einsum('blmh,bhlmrk,bmhpk->blhrp', Ldec, G, xc)
+        y_inter = torch.einsum('blh,blhnr,bhnp->blhrp', d, Cc, g)
+        ys.append(y_intra + y_inter)
+        coef = torch.exp(cl[:, -1:, :] - cl)
+        g = d[:, -1].unsqueeze(-1).unsqueeze(-1) * g + torch.einsum('blh,blhnk,blhpk->bhnp', coef, Bc, xc)
+    return torch.cat(ys, dim=1), g
+
+
+def mamba3_layer_ssd_state(layer: Mamba3Layer, u: torch.Tensor, state=None, chunk: int = 64):
+    """SSD with carried state. state = (h (B,H,N,P) absolute frame, B_last (B,H,N,K), x_last (B,H,P,K)) or None (zeros).
+    Returns (out, new_state). Equal to running mamba3_layer_ssd over the concatenated sequence; step()'s prev_bx = B_last @ x_last^T."""
+    nu = layer.norm(u)
+    alpha, beta, gamma, Bm, Cm, x, theta = layer._coeffs(nu)
+    Phi = torch.cumsum(theta, dim=1)
+    Phi_prev = _shift(Phi)
+    if state is None:
+        h0 = None; Bp = torch.zeros_like(Bm[:, :1]); xp = torch.zeros_like(x[:, :1])
+    else:
+        h0, Bl, xl = state; Bp = Bl[:, None].to(Bm.dtype); xp = xl[:, None].to(x.dtype)
+    Bsh = torch.cat([Bp, Bm[:, :-1]], 1); xsh = torch.cat([xp, x[:, :-1]], 1)
+    Btil = rotate_by(Bm, -Phi); Ctil = rotate_by(Cm, -Phi); Btil_prev = rotate_by(Bsh, -Phi_prev)
+    Beff = torch.cat([gamma[..., None, None] * Btil, beta[..., None, None] * Btil_prev], dim=-1)
+    xeff = torch.cat([x, xsh], dim=-1)
+    y, g = ssd_safe_state(alpha, Beff, xeff, Ctil, h0, chunk)
+    hT = rotate_by(g.unsqueeze(1), Phi[:, -1:])[:, 0]
+    Bsz, T = u.shape[:2]
+    return layer.gate_and_proj(y.reshape(Bsz, T, -1), nu), (hT, Bm[:, -1], x[:, -1])

arctag/_core/arc_audio/model.py

@@ -0,0 +1,363 @@
+"""Arc v1.4 network (inference only; 1.4.1 adds the word channel, arc_audio/wordch.py).
+
+trunk:  x (B,T,640) 10 Hz [mean,max]-pooled Moonshine Streaming features -> LayerNorm -> Linear 512 (+ domain emb + delay emb)
+        -> 5 x [FiLM(domain + delay) -> Mamba-3 + GELU MLP] -> per-domain rank-64 adapter -> RMSNorm -> h (B,T,512).
+head:   one flat tag-prompt head over the tag table: key_t = P e_t + r_t (e_t = frozen text embedding of the tag description,
+        shipped as a buffer); z = <Q_q(pool_q(h)), key_t> + b_t + b_t^q for each read-out q (frame, span_in, span_end, utt,
+        speaker, ev_at, ev_imm, ev_rec, item). Pooled read-outs carry (sum, n) state: utt (since the utterance start), speaker
+        (per speaker index), item (since the stream start).
+extra:  name pointer + numeral (NWT) value head (EntityHeads), diarization slots + voice embedding (DiarHeads), non-speech
+        vocal events (BioHead).
+
+Two equivalent trunk paths:
+  Trunk.forward(x, ..., state)   chunked SSD with carried state (offline / large chunks)
+  Trunk.step(x_t, ..., state)    recurrent Mamba3Layer.step, one 100 ms frame at a time (true streaming; constant cost)
+Both carry the same per-layer state; tests/test_release.py checks that they agree.
+"""
+import math
+import torch, torch.nn as nn, torch.nn.functional as F
+
+from .mamba3 import Mamba3Layer
+from .mamba3.layer import rotate_pairs
+from .mamba3.ssd import mamba3_layer_ssd_state, _f
+from .ontology import QUERIES, POOLED, FINE_DOMAINS
+
+NEG = -1e4
+S_MAX = 64   # speaker indices pooled per stream
+
+
+class Trunk(nn.Module):
+    def __init__(self, d_in=640, d_model=512, n_layers=5, d_state=64, head_dim=64, dropout=0.2, in_drop=0.2, chunk=128,
+                 n_dom=len(FINE_DOMAINS), delays=(0, 1, 2, 3, 5, 8), adapters=64, long_mem=True):
+        super().__init__()
+        self.cfg = dict(d_in=d_in, d_model=d_model, n_layers=n_layers, d_state=d_state, head_dim=head_dim, dropout=dropout, in_drop=in_drop,
+                        chunk=chunk, n_dom=n_dom, delays=list(delays), adapters=adapters, long_mem=long_mem)
+        self.chunk = chunk
+        self.in_norm = nn.LayerNorm(d_in); self.inp = nn.Linear(d_in, d_model)
+        self.dom_emb = nn.Embedding(n_dom, d_model)
+        self.film = nn.ModuleList([nn.Linear(d_model, 2 * d_model) for _ in range(n_layers)])
+        self.delays = list(delays)
+        self.delay_emb = nn.Embedding(len(delays), d_model)
+        self.delay_film = nn.Embedding(len(delays), d_model)
+        lut = torch.full((max(delays) + 1,), -1, dtype=torch.long)
+        for i, d in enumerate(delays):
+            lut[d] = i
+        self.register_buffer("delay_lut", lut, persistent=False)
+        self.adapt_down = nn.Parameter(torch.zeros(n_dom, d_model, adapters))
+        self.adapt_up = nn.Parameter(torch.zeros(n_dom, adapters, d_model))
+        self.layers = nn.ModuleList([Mamba3Layer(d_model, d_state=d_state, rank=1, head_dim=head_dim) for _ in range(n_layers)])
+        self.mlps = nn.ModuleList([nn.Sequential(nn.RMSNorm(d_model), nn.Linear(d_model, 2 * d_model), nn.GELU(), nn.Dropout(dropout),
+                                                 nn.Linear(2 * d_model, d_model)) for _ in range(n_layers)])
+        self.out_norm = nn.RMSNorm(d_model)
+
+    def _cond(self, x, dom, delay):
+        di = self.delay_lut[delay.clamp(0, len(self.delay_lut) - 1)]
+        if (di < 0).any():
+            raise ValueError(f"delay {delay.tolist()} frames not in the trained set {self.delays}")
+        h = self.inp(self.in_norm(x))
+        e = self.dom_emb(dom) + self.delay_film(di)
+        return h + (self.dom_emb(dom) + self.delay_emb(di))[:, None], e
+
+    def _adapt(self, h, dom):
+        Dn, Up = self.adapt_down[dom].to(h.dtype), self.adapt_up[dom].to(h.dtype)
+        return h + torch.einsum("btr,brd->btd", F.gelu(torch.einsum("btd,bdr->btr", h, Dn)), Up)
+
+    def forward(self, x, dom, delay, state=None):
+        """chunked SSD. x (B,T,d_in); dom (B,) long; delay (B,) long frames; state: list per layer of (h, B_last, x_last) or None."""
+        h, e = self._cond(x, dom, delay); new = []
+        for i, (lay, mlp) in enumerate(zip(self.layers, self.mlps)):
+            gb = self.film[i](e)[:, None]; h = h * (1 + gb[..., :h.shape[-1]]) + gb[..., h.shape[-1]:]
+            st = None if state is None else tuple(_f(t) for t in state[i][:3])
+            s, ns = mamba3_layer_ssd_state(lay, _f(h), st, self.chunk)
+            h = h + s
+            h = h + mlp(h)
+            new.append(ns)
+        return self.out_norm(self._adapt(h, dom)), new
+
+    def step(self, x, dom, delay, state=None):
+        """recurrent path over T frames (python loop of the Mamba-3 recurrence; constant cost per frame).
+        Same arguments and result as forward(); the state is (h, B_last, x_last, B_last x_last^T) per layer and can be
+        continued by forward() (which uses the first three entries) and vice versa."""
+        h_all, e = self._cond(x, dom, delay)
+        B, T, D = h_all.shape
+        st = [None] * len(self.layers) if state is None else list(state)
+        gbs = [self.film[i](e) for i in range(len(self.layers))]
+        outs = []
+        for t in range(T):
+            h = h_all[:, t]
+            for i, (lay, mlp) in enumerate(zip(self.layers, self.mlps)):
+                gb = gbs[i]; h = h * (1 + gb[..., :D]) + gb[..., D:]
+                y, st[i] = layer_step(lay, _f(h), st[i])
+                h = h + y
+                h = h + mlp(h)
+            outs.append(h)
+        h = torch.stack(outs, 1)
+        return self.out_norm(self._adapt(h, dom)), st
+
+
+def layer_step(lay, u, state):
+    """one frame of Mamba3Layer (= Mamba3Layer.step with the coefficients computed once). u (B, D) -> (y (B, D), state)."""
+    Bsz = u.shape[0]
+    nu = lay.norm(u.unsqueeze(1))
+    alpha, beta, gamma, Bm, Cm, x, theta = lay._coeffs(nu)
+    H, N, P = lay.H, lay.N, lay.P
+    bx = Bm[:, 0] @ x[:, 0].transpose(-1, -2)                                   # (B,H,N,P)
+    if state is None:
+        h = gamma[:, 0, :, None, None] * bx
+    else:
+        h0, Bl, xl = (_f(t) for t in state[:3])
+        prev_bx = _f(state[3]) if len(state) > 3 else Bl @ xl.transpose(-1, -2)
+        pre = alpha[:, 0, :, None, None] * h0 + beta[:, 0, :, None, None] * prev_bx
+        h = rotate_pairs(pre.transpose(-1, -2), theta[:, 0].unsqueeze(-2).expand(-1, -1, P, -1)).transpose(-1, -2)
+        h = h + gamma[:, 0, :, None, None] * bx
+    y = (Cm[:, 0].transpose(-1, -2) @ h).reshape(Bsz, -1)
+    return lay.gate_and_proj(y, nu[:, 0]), (h, Bm[:, 0], x[:, 0], bx)
+
+
+# ---------------------------------------------------------------- pooling with carried state
+def pool_speaker(z, spk, carry):
+    """z (B,T,D); spk (B,T) long speaker index (-1 none); carry (sum (B,S,D) fp64, n (B,S) fp64) -> running mean over that
+    speaker's frames (incl. earlier calls), new carry."""
+    B, T, D = z.shape
+    sp = torch.where((spk >= 0) & (spk < S_MAX), spk, torch.full_like(spk, -1))
+    key = (torch.arange(B, device=z.device)[:, None] * (S_MAX + 1) + (sp + 1)) * T + torch.arange(T, device=z.device)[None]
+    order = torch.argsort(key.reshape(-1), stable=True)
+    zs = z.reshape(B * T, D)[order].double(); gs = key.reshape(-1)[order] // T
+    cs = torch.cumsum(zs, 0)
+    start = torch.ones_like(gs, dtype=torch.bool); start[1:] = gs[1:] != gs[:-1]
+    pos = torch.arange(B * T, device=z.device)
+    first = torch.cummax(torch.where(start, pos, torch.zeros_like(pos)), 0).values
+    base = torch.where(first[:, None] > 0, cs[(first - 1).clamp(min=0)], torch.zeros_like(cs))
+    num = torch.empty_like(cs); num[order] = cs - base
+    cnt = torch.empty(B * T, dtype=torch.float64, device=z.device); cnt[order] = (pos - first + 1).double()
+    num = num.reshape(B, T, D); cnt = cnt.reshape(B, T)
+    cs_, cn_ = carry
+    sidx = sp.clamp(min=0)
+    num = num + torch.gather(cs_, 1, sidx[..., None].expand(B, T, D)); cnt = cnt + torch.gather(cn_, 1, sidx)
+    out = (num / cnt.clamp(min=1)[..., None]).float() * (sp >= 0)[..., None].float()
+    oh = F.one_hot(sidx, S_MAX).double() * (sp >= 0)[..., None].double()
+    new = (cs_ + torch.einsum("bts,btd->bsd", oh, z.double()), cn_ + oh.sum(1))
+    return out, new
+
+
+def pool_segment(z, starts, carry):
+    """causal mean since the last start (True at a unit's first frame); carry (sum (B,D), n (B,)) continues the open unit.
+    starts=None -> one unit for the whole stream (item pooling)."""
+    B, T, D = z.shape
+    st = torch.zeros(B, T, dtype=torch.bool, device=z.device) if starts is None else starts
+    cs = torch.cumsum(z.double(), 1); cs0 = F.pad(cs, (0, 0, 1, 0))[:, :-1]
+    t = torch.arange(T, device=z.device)[None].expand(B, T)
+    first = torch.cummax(torch.where(st, t, torch.full_like(t, -1)), 1).values
+    base = torch.gather(cs0, 1, first.clamp(min=0)[..., None].expand(B, T, D))
+    n = (t - first.clamp(min=0) + 1).double()
+    open_ = (first < 0)
+    num = torch.where(open_[..., None], cs + carry[0][:, None], cs - base)
+    n = torch.where(open_, (t + 1).double() + carry[1][:, None], n)
+    out = (num / n.clamp(min=1)[..., None]).float()
+    last = first[:, -1]
+    new_s = torch.where((last < 0)[:, None], carry[0] + cs[:, -1], cs[:, -1] - torch.gather(cs0, 1, last.clamp(min=0)[:, None, None].expand(B, 1, D))[:, 0])
+    new_n = torch.where(last < 0, carry[1] + T, (T - last).double())
+    return out, (new_s, new_n)
+
+
+def zero_carry(B, D, device):
+    z = lambda *s: torch.zeros(*s, dtype=torch.float64, device=device)
+    return dict(utt=(z(B, D), z(B)), speaker=(z(B, S_MAX, D), z(B, S_MAX)), item=(z(B, D), z(B)))
+
+
+class TagHead(nn.Module):
+    def __init__(self, table, n_text=384, d_model=512, d_k=512):
+        super().__init__()
+        self.table = table
+        self.register_buffer("E", torch.zeros(table.n, n_text))                 # frozen tag-text embeddings (from the checkpoint)
+        self.register_buffer("valid", torch.as_tensor(table.valid), persistent=False)
+        self.P = nn.Linear(n_text, d_k, bias=False)
+        self.r = nn.Parameter(torch.zeros(table.n, d_k)); self.b = nn.Parameter(torch.zeros(table.n))
+        self.bq = nn.ParameterDict({q: nn.Parameter(torch.zeros(table.n)) for q in QUERIES})
+        self.Q = nn.ModuleDict({q: nn.Linear(d_model, d_k) for q in QUERIES})
+        self.pool_in = nn.ModuleDict({q: nn.Sequential(nn.Linear(d_model, d_model), nn.GELU(), nn.Linear(d_model, d_model)) for q in POOLED})
+        for q in QUERIES:
+            self.register_buffer(f"cols_{q}", torch.tensor(table.cols[q], dtype=torch.long), persistent=False)
+
+    def cols(self, q):
+        return getattr(self, f"cols_{q}")
+
+    def keys(self):
+        return self.P(self.E) + self.r
+
+    def pooled(self, h, spk, ustart, carry, h_utt=None):
+        """h_utt: the word-fused h' for the utterance pool (word channel); speaker / item pools read the audio-only h."""
+        zin, new = {}, {}
+        for q in POOLED:
+            hq = h_utt if (q == "utt" and h_utt is not None) else h
+            z = hq + self.pool_in[q](hq)
+            if q == "speaker":
+                zin[q], new[q] = pool_speaker(z.float(), spk, carry[q])
+            elif q == "utt":
+                zin[q], new[q] = pool_segment(z.float(), ustart, carry[q])
+            else:
+                zin[q], new[q] = pool_segment(z.float(), None, carry[q])
+        return zin, new
+
+    def logits_q(self, q, zq, dom, K=None):
+        """raw (prior-free) logits for read-out q; domain-invalid columns -> NEG."""
+        K = self.keys() if K is None else K
+        c = self.cols(q)
+        z = self.Q[q](zq.float()) @ K[c].T + (self.b + self.bq[q])[c]
+        return z.masked_fill(~self.valid[dom][:, c][:, None, :], NEG)
+
+
+# ---------------------------------------------------------------- entities: name pointer + numeral value head
+NAME_MAX_BYTES = 40
+C_KINDS = ["person", "team", "organization", "place", "work_product", "event_name", "legal_ref", "game_entity", "game_map", "other"]
+MENTION_KINDS = {"entity.mention.person": [0, 9], "entity.mention.team": [1, 2], "entity.mention.organization": [2, 1, 9],
+                 "entity.mention.place": [3, 9], "entity.mention.work_product": [4, 9], "entity.mention.event_name": [5, 9],
+                 "entity.mention.legal_ref": [6, 9], "entity.mention.game_entity": [7, 9], "entity.mention.game_map": [8, 9]}
+K_SLOTS = 4
+CLOCK = 3
+SPAN_G_MULT = 20.0
+
+
+def encode_names(names, max_len=NAME_MAX_BYTES, keep="head"):
+    import numpy as np
+    ids = np.zeros((len(names), max_len), np.int64)
+    for i, n in enumerate(names):
+        b = " ".join(str(n).lower().split()).encode("utf-8")
+        b = b[:max_len] if keep == "head" else b[-max_len:]
+        ids[i, :len(b)] = np.frombuffer(b, np.uint8).astype(np.int64) + 1
+    return ids
+
+
+class NameEncoder(nn.Module):
+    """byte embedding -> 3 residual dilated Conv1d -> [masked mean, max] -> MLP (+ kind embedding), unit norm."""
+    def __init__(self, d=128, char_dim=48):
+        super().__init__()
+        self.d = d
+        self.emb = nn.Embedding(257, char_dim, padding_idx=0)
+        self.conv_in = nn.Conv1d(char_dim, d, 3, padding=1)
+        self.convs = nn.ModuleList([nn.Conv1d(d, d, 3, padding=dl, dilation=dl) for dl in (1, 2, 4)])
+        self.kind = nn.Embedding(len(C_KINDS), d)
+        self.out = nn.Sequential(nn.Linear(2 * d, d), nn.GELU(), nn.Linear(d, d))
+
+    def forward(self, ids, kinds):
+        sh = ids.shape[:-1]; ids = ids.reshape(-1, ids.shape[-1]); m = (ids > 0).float()[:, None]
+        x = self.emb(ids).float().transpose(1, 2); h = F.gelu(self.conv_in(x)) * m
+        for cv in self.convs:
+            h = h + F.gelu(cv(h)) * m
+        mean = h.sum(-1) / m.sum(-1).clamp(min=1.0); mx = h.masked_fill(m == 0, -1e4).max(-1).values * (m.sum(-1) > 0).float()
+        e = self.out(torch.cat([mean, mx], -1)) + self.kind(kinds.reshape(-1))
+        return F.normalize(e, dim=-1).reshape(*sh, self.d)
+
+
+class EntityHeads(nn.Module):
+    def __init__(self, d_model=512, d_key=512, d=128, n_units=64, span=True, span_new=True, sim=False):
+        super().__init__()
+        self.names = NameEncoder(d)
+        self.query = nn.Sequential(nn.Linear(d_model, d), nn.GELU(), nn.Linear(d, d))
+        self.key_in = nn.Linear(d_key, d_model, bias=False)
+        self.special = nn.Parameter(torch.zeros(2, d))        # new, generic
+        self.special_b = nn.Parameter(torch.zeros(2))
+        self.log_scale = nn.Parameter(torch.tensor([math.log(10.0)]))
+        self.nwt = nn.Linear(d_model, K_SLOTS * 43)
+        self.unit = nn.Linear(d_model, n_units + 1)
+        self.clock = nn.Linear(d_model, CLOCK)
+        self.approx = nn.Linear(d_model, 2)
+        self.n_units = n_units
+        self.span_q = nn.Linear(d, d) if span else None
+        if span:
+            self.span_g = nn.Parameter(torch.zeros(1))
+        if span and span_new:
+            self.span_new = nn.Parameter(torch.zeros(1))
+        self.sim_w = None
+        if sim:   # 1.4.1: heard-span <-> candidate string / phonetic similarity features (namesim.py), per heard-text source row
+            self.sim_w = nn.Parameter(torch.zeros(3, 6)); self.sim_new = nn.Parameter(torch.zeros(3, 3))
+            if span:
+                self.span_g_src = nn.Parameter(torch.zeros(3))
+
+    def ptr_logits(self, h, cls_key, E, cmask, Es=None, Sf=None, src=0):
+        """h (B,T,D); cls_key (B,T,Dk); E (B,R,d) candidate name embeddings; cmask (B,T,R) bool; Es (B,T,d) heard-text embedding
+        (zeros where no heard text) or None; Sf (B,T,R,6) heard-text similarity features (namesim.sim_feats) or None; src = heard-text
+        source row (0 machine-ASR words, 1 the built-in decoder's words) -> (B,T,2+R) logits [new, generic, candidates]."""
+        qr = self.query(h.float() + self.key_in(cls_key.float()))
+        use_span = Es is not None and self.span_q is not None
+        if use_span:
+            qr = qr + self.span_q(Es.float())
+        q = F.normalize(qr, dim=-1)
+        sc = self.log_scale.clamp(max=math.log(50.0)).exp()
+        s = torch.einsum("btd,brd->btr", q, E.float()) * sc
+        if use_span:
+            g_ = self.span_g if getattr(self, "span_g_src", None) is None else self.span_g + self.span_g_src[int(src)]
+            s = s + SPAN_G_MULT * g_ * sc * torch.einsum("btd,brd->btr", Es.float(), E.float())
+        use_sim = Sf is not None and self.sim_w is not None
+        if use_sim:
+            Sf = Sf.float()
+            s = s + SPAN_G_MULT * sc * torch.einsum("btrf,f->btr", Sf, self.sim_w[int(src)])
+        s = s.masked_fill(~cmask, NEG)
+        sp = torch.einsum("btd,kd->btk", q, F.normalize(self.special, dim=-1)) * sc + self.special_b
+        if use_sim:
+            hs = (Sf.abs().sum((-1, -2)) > 0).float()
+            mb = Sf[..., 4].masked_fill(~cmask, 0.0).max(-1).values; msh = Sf[..., 5].masked_fill(~cmask, 0.0).max(-1).values
+            sn = self.sim_new[int(src)]
+            tn = sn[0] + sn[1] * (1.0 - mb) + sn[2] * (1.0 - msh)
+            sp = torch.cat([sp[..., :1] + (SPAN_G_MULT * sc * hs * tn)[..., None], sp[..., 1:]], -1)
+        if use_span and getattr(self, "span_new", None) is not None:
+            ms = torch.einsum("btd,brd->btr", Es.float(), E.float()).masked_fill(~cmask, -1.0).max(-1).values.clamp(min=-1.0)
+            has = (Es.float().abs().sum(-1) > 0).float()
+            sp = torch.cat([sp[..., :1] + SPAN_G_MULT * self.span_new * sc * has[..., None] * (1.0 - ms[..., None]), sp[..., 1:]], -1)
+        return torch.cat([sp, s], -1)
+
+    def span_embed(self, text):
+        """heard-span text (string) -> (d,) unit embedding through the shared name encoder (kind 0); keeps the most recent bytes."""
+        ids = torch.as_tensor(encode_names([text], keep="tail"), device=self.special.device)
+        return self.names(ids, torch.zeros(1, dtype=torch.long, device=ids.device))[0]
+
+    def value_logits(self, h):
+        return dict(nwt=self.nwt(h).float().reshape(*h.shape[:-1], K_SLOTS, -1), unit=self.unit(h).float(),
+                    clock=self.clock(h).float(), approx=self.approx(h).float())
+
+
+class BioHead(nn.Module):
+    """non-speech vocal events on the frame read-out (sneeze was never trained: masked in the output)."""
+    EV = ("cough", "throat_clear", "breath", "sneeze")
+
+    def __init__(self, d_model=512):
+        super().__init__()
+        self.ev = nn.Linear(d_model, len(self.EV))
+
+
+class DiarHeads(nn.Module):
+    """arrival-ordered speaker slots ([none, slot 1..K]) + unit-norm voice embedding from h_t and the input frame x_t."""
+    def __init__(self, d_model=512, d_in=640, k=8, d=128, hidden=256, n_extra=0):
+        super().__init__()
+        self.k = k + n_extra
+        self.slot = nn.Linear(d_model, 1 + k)
+        self.h_proj = nn.Linear(d_model, hidden); self.x_norm = nn.LayerNorm(d_in); self.x_proj = nn.Linear(d_in, hidden)
+        self.out = nn.Sequential(nn.GELU(), nn.Linear(hidden, hidden), nn.GELU(), nn.Linear(hidden, d))
+        self.slot_ext = nn.Linear(d_model, n_extra) if n_extra else None
+
+    def slot_logits(self, h):
+        z = self.slot(h)
+        return torch.cat([z, self.slot_ext(h)], -1) if self.slot_ext is not None else z
+
+    def embed(self, h, x):
+        return F.normalize(self.out(self.h_proj(h.float()) + self.x_proj(self.x_norm(x.float()))), dim=-1)
+
+
+class ArcNet(nn.Module):
+    def __init__(self, table, cfg):
+        super().__init__()
+        self.trunk = Trunk(**cfg["trunk"])
+        d = self.trunk.cfg["d_model"]
+        self.head = TagHead(table, n_text=int(cfg.get("n_text", 384)), d_model=d)
+        self.ent = EntityHeads(d_model=d, d_key=d, n_units=int(cfg.get("n_units", 64)), span=bool(cfg.get("ptr_span_text")),
+                               span_new=bool(cfg.get("ptr_span_new")), sim=bool(cfg.get("ptr_sim_feat")))
+        self.diar = DiarHeads(d_model=d, d_in=self.trunk.cfg["d_in"], n_extra=max(0, int(cfg.get("diar_slots", 8)) - 8))
+        self.bio = BioHead(d) if cfg.get("bio_head") else None
+        self.word = None
+        if cfg.get("word_ch"):   # 1.4.1 word channel (wordch.py)
+            from .wordch import WordChannel
+            self.word = WordChannel(d_model=d, d_k=d, n_in=len(table.cols["span_in"]), k_slots=K_SLOTS, n_nwt=int(cfg.get("n_nwt", 43)))
+        self.cfg = cfg
+
+    def trunk_run(self, x, dom, delay, state=None, recurrent=False):
+        return (self.trunk.step if recurrent else self.trunk)(x, dom, delay, state)

arctag/_core/arc_audio/ontology.py

@@ -0,0 +1,79 @@
+"""Tag table (tags.json in the model repo): stable tag ids and order, exclusive groups, read-outs, domain validity.
+
+Read-out per tag unit: F/G -> frame, W -> span_in + span_end, U/T -> utt, S -> speaker, E -> ev_at | ev_rec (.recent) |
+ev_imm (.imminent), D -> item.
+"""
+import hashlib, json
+import numpy as np
+
+FINE_DOMAINS = ["nba", "nhl", "nfl", "cricket", "soccer", "esports", "news", "talk", "stream_chat", "meeting", "court_appellate",
+                "court_trial", "assistant", "assistant_noisy", "short_clips", "unknown"]
+QUERIES = ("frame", "span_in", "span_end", "utt", "speaker", "ev_at", "ev_imm", "ev_rec", "item")
+POOLED = ("utt", "speaker", "item")
+SPORTS = ("nba", "nhl", "nfl", "cricket", "soccer")
+
+
+def readouts(t):
+    u = t["unit"]
+    if u in ("F", "G"):
+        return ("frame",)
+    if u == "W":
+        return ("span_in", "span_end")
+    if u in ("U", "T"):
+        return ("utt",)
+    if u == "S":
+        return ("speaker",)
+    if u == "D":
+        return ("item",)
+    if u == "E":
+        return ("ev_rec",) if t["id"].endswith(".recent") else ("ev_imm",) if t["id"].endswith(".imminent") else ("ev_at",)
+    raise ValueError(f"unknown unit {u} for {t['id']}")
+
+
+def domain_index(domain):
+    d = (domain or "unknown").lower()
+    if d not in FINE_DOMAINS:
+        raise ValueError(f"unknown domain {domain!r}; use one of {FINE_DOMAINS}")
+    return FINE_DOMAINS.index(d)
+
+
+class Table:
+    def __init__(self, path):
+        raw = open(path, "rb").read(); self.sha = hashlib.sha256(raw).hexdigest()
+        j = json.loads(raw); self.version = j["version"]; self.j = j
+        tags = sorted(j["tags"], key=lambda t: t["index"])
+        self.n = max(t["index"] for t in tags) + 1
+        self.tag = [None] * self.n
+        for t in tags:
+            self.tag[t["index"]] = t
+        self.ids = [t["id"] if t else f"_unused{i}" for i, t in enumerate(self.tag)]
+        self.idx = {t: i for i, t in enumerate(self.ids)}
+        rg = {g["id"]: g for g in j["groups"]}
+        self.groups, self.members = [], {}
+        for t in tags:
+            if t.get("exclusivity") not in (None, "independent"):
+                if t["group"] not in self.members:
+                    self.groups.append(t["group"]); self.members[t["group"]] = []
+                self.members[t["group"]].append(t["index"])
+        self.group_of = {i: g for g, m in self.members.items() for i in m}
+        self.gate = {g: rg.get(g, {}).get("gate") for g in self.groups}
+        self.group_info = rg
+        self.trained = [i for i, t in enumerate(self.tag) if t and t.get("status") not in ("reserved", "derived")]
+        dg = j["domains"]
+        V = np.zeros((len(FINE_DOMAINS), self.n), bool)
+        for i in self.trained:
+            keys = self.tag[i]["domains"]
+            for d, fd in enumerate(FINE_DOMAINS):
+                V[d, i] = any(k == "*" or k == fd or fd in dg.get(k, []) for k in keys)
+        # 'unknown' (no domain given): every tag that is valid in some domain (the trunk still gets the 'unknown' domain tag)
+        V[FINE_DOMAINS.index("unknown")] |= V.any(0)
+        self.valid = V
+        self.ro = {i: readouts(self.tag[i]) for i in self.trained}
+        self.cols = {q: [i for i in self.trained if q in self.ro[i]] for q in QUERIES}
+        self.facet = {i: self.tag[i]["facet"] for i in self.trained}
+        self.status = {i: self.tag[i]["status"] for i in self.trained}
+
+    def describe(self, tag_id):
+        t = self.tag[self.idx[tag_id]]
+        return dict(id=t["id"], facet=t["facet"], unit=t["unit"], group=t.get("group"), description=t.get("description"),
+                    paraphrases=t.get("paraphrases") or [], status=t.get("status"))

arctag/_core/arc_audio/wordch.py

@@ -0,0 +1,332 @@
+"""Arc 1.4.1 word channel (inference): heard transcript words feed the network next to the audio.
+
+Each heard word j (any ASR: the built-in decoder, or external words) becomes a vector u_j:
+  e_j = LN(W1 mean_p(E[piece] + pos_p) + W2 E[last piece]) + flag / source embeddings     (E = the shipped decoder's frozen
+        piece-embedding table; pieces of the package tokenizer, first 7 + last)
+  x_j = e_j + Wa [mean of the trunk state h over the word's own frames, h at its last own frame] + has-audio embedding
+  u_j = 2-layer causal LOCAL transformer over the words (window 32 words, RoPE)  -> receptive field 62 words back.
+Each label frame t then reads a window of up to 8 words around the newest available word (6 back, 2 ahead when available):
+  c_t = Wo attn(q = h_t, k = u_j + rel(t - f0_j) + rel(t - oe_j), v = u_j)      (0 when no word is available)
+  h'_t = h_t + c_t   feeds the frame / span / event read-outs, the utterance pool and the name / number heads;
+  the speaker / item pools, diarization and vocal-event heads read the audio-only h_t (as trained).
+With no words, h' == h exactly (the audio-only network).
+Word-aligned read-out (war_*): per-word B / I probabilities of every span tag (right context k words) + per-word numeral tokens.
+"""
+import math, re
+import numpy as np
+import torch, torch.nn as nn, torch.nn.functional as F
+
+P_MAX = 8          # pieces per word (first 7 + last when longer)
+K_CTX = 8          # words per frame window
+BACK = 6
+ENC_WIN = 32       # WordEnc causal window (words)
+N_ENC = 2          # WordEnc layers
+RF = (ENC_WIN - 1) * N_ENC    # receptive field in words (62)
+OFF_CLIP = 32
+SRC_ROW = dict(npz=0, ms=1, corrupt=2, none=3)   # 0 = machine-ASR words (external ASR), 1 = the shipped (built-in) decoder's words
+_PUNCT = re.compile(r"^[^\w]+|[^\w]+$")
+
+
+def norm_word(w):
+    return _PUNCT.sub("", " ".join(str(w).lower().split()))
+
+
+def own_ends(f0, f1):
+    f0 = np.asarray(f0, np.int64); f1 = np.asarray(f1, np.int64); oe = f1.copy()
+    nxt = np.r_[f0[1:], np.iinfo(np.int64).max]
+    m = (f0 < nxt) & (nxt <= f1)
+    oe[m] = nxt[m] - 1
+    return np.maximum(oe, f0)
+
+
+def utt_ends(f0, oe, gap=4):
+    W = len(f0)
+    if not W:
+        return np.zeros(0, np.int64)
+    brk = np.r_[f0[1:] - oe[:-1] > gap, True]
+    ue = np.where(brk, oe, np.iinfo(np.int64).max)
+    return np.minimum.accumulate(ue[::-1])[::-1]
+
+
+def rule_arrival(f0, oe, d, mode="final", Lf=15, L=8):
+    """offline word arrival (the trained / evaluated rule): a final transcript reaches Arc Lf frames after its utterance ends
+    (0.4 s pause rule), never before the word's own end + d; monotone. -> first label frame at which word j is readable."""
+    if not len(f0):
+        return np.zeros(0, np.int64)
+    a = (oe + L) if mode == "partial" else (utt_ends(f0, oe) + Lf)
+    a = np.maximum(a, oe + d)
+    return np.maximum.accumulate(a) - d
+
+
+class Pieces:
+    """word -> piece ids of the package tokenizer (cached)."""
+
+    def __init__(self, tok_path):
+        from tokenizers import Tokenizer
+        self.t = Tokenizer.from_file(tok_path); self.c = {}
+
+    def __call__(self, w):
+        if w not in self.c:
+            ids = self.t.encode(w, add_special_tokens=False).ids if w else []
+            if len(ids) > P_MAX:
+                ids = ids[:P_MAX - 1] + ids[-1:]
+            if len(self.c) > 200000:
+                self.c.clear()
+            self.c[w] = ids
+        return self.c[w]
+
+
+def word_rows(texts, f0, oe, src_row, pieces, prev_oe=None):
+    """-> tok (W,P), npc (W,), flags (W,5) int64 (training row_words features). prev_oe: own end of the word before the first
+    one (incremental encoding), for the gap flag."""
+    W = len(texts); tok = np.zeros((W, P_MAX), np.int64); npc = np.zeros(W, np.int64); flags = np.zeros((W, 5), np.int64)
+    for i, w in enumerate(texts):
+        p = pieces(w); npc[i] = len(p); tok[i, :len(p)] = p
+        hd = any(ch.isdigit() for ch in w)
+        flags[i, 0] = int(hd); flags[i, 1] = int(hd and w.isdigit())
+    flags[:, 2] = np.clip(npc, 0, P_MAX)
+    if W:
+        flags[:, 3] = np.clip(oe - f0 + 1, 1, 20)
+        flags[:, 4] = np.clip(np.r_[20 if prev_oe is None else f0[0] - prev_oe - 1, f0[1:] - oe[:-1] - 1], 0, 20)
+    return tok, npc, flags
+
+
+def frame_window(f0, avail, t0, t1, K=K_CTX, back=BACK):
+    """(t1-t0, K) word indices (-1 none) read by frames t0..t1-1; avail monotone (available words are a prefix)."""
+    t = np.arange(t0, t1)
+    n = np.searchsorted(avail, t, side="right")
+    p = np.searchsorted(f0, t, side="right")
+    q = np.minimum(p, n)
+    idx = q[:, None] - back + np.arange(K)[None]
+    ok = (idx >= 0) & (idx < n[:, None])
+    return np.where(ok, idx, -1)
+
+
+def _rope(x, pos):
+    dh = x.shape[-1]; inv = 1.0 / (10000 ** (torch.arange(0, dh, 2, device=x.device, dtype=torch.float32) / dh))
+    ang = pos.float()[..., None] * inv
+    while ang.dim() < x.dim():
+        ang = ang.unsqueeze(1)
+    c, s_ = torch.cos(ang), torch.sin(ang); x1, x2 = x[..., 0::2], x[..., 1::2]
+    return torch.stack([x1 * c - x2 * s_, x1 * s_ + x2 * c], -1).flatten(-2)
+
+
+class EncLayer(nn.Module):
+    """pre-norm causal local self-attention (window ENC_WIN words incl. self) + FFN (blocked, as trained)."""
+
+    def __init__(self, d=256, h=4, ff=512, win=ENC_WIN, dropout=0.1):
+        super().__init__()
+        self.h, self.win = h, win
+        self.n1 = nn.LayerNorm(d); self.qkv = nn.Linear(d, 3 * d); self.o = nn.Linear(d, d)
+        self.n2 = nn.LayerNorm(d); self.ff = nn.Sequential(nn.Linear(d, ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(ff, d)); self.drop = nn.Dropout(dropout)
+
+    def forward(self, x, valid, pos0=0):
+        Bn, W, D = x.shape; h, w = self.h, self.win; dh = D // h
+        q, k, v = self.qkv(self.n1(x)).float().view(Bn, W, 3, h, dh).unbind(2)
+        pos = (torch.arange(W, device=x.device) + pos0)[None].expand(Bn, W)
+        q = _rope(q.transpose(1, 2), pos); k = _rope(k.transpose(1, 2), pos); v = v.transpose(1, 2)
+        nb = (W + w - 1) // w; Wp = nb * w; pad = Wp - W
+        if pad:
+            q, k, v = (F.pad(t_, (0, 0, 0, pad)) for t_ in (q, k, v))
+        vm = F.pad(valid, (0, pad)) if pad else valid
+        kp = F.pad(k, (0, 0, w, 0)).view(Bn, h, nb + 1, w, dh); vp = F.pad(v, (0, 0, w, 0)).view(Bn, h, nb + 1, w, dh)
+        kb = torch.cat([kp[:, :, :-1], kp[:, :, 1:]], 3); vb = torch.cat([vp[:, :, :-1], vp[:, :, 1:]], 3)
+        qb = q.view(Bn, h, nb, w, dh)
+        qi = (torch.arange(nb, device=x.device)[:, None] * w + torch.arange(w, device=x.device)[None])
+        ki = (torch.arange(nb, device=x.device)[:, None] * w - w + torch.arange(2 * w, device=x.device)[None])
+        mp = F.pad(vm, (w, 0)).view(Bn, nb + 1, w); km = torch.cat([mp[:, :-1], mp[:, 1:]], 2)
+        m = (ki[:, None, :] <= qi[:, :, None]) & (ki[:, None, :] > qi[:, :, None] - w) & (ki[:, None, :] >= 0)
+        m = m[None] & km[:, :, None, :]
+        m = m | (ki[:, None, :] == qi[:, :, None])[None]
+        a = F.scaled_dot_product_attention(qb.permute(0, 2, 1, 3, 4).reshape(Bn * nb, h, w, dh),
+                                           kb.permute(0, 2, 1, 3, 4).reshape(Bn * nb, h, 2 * w, dh), vb.permute(0, 2, 1, 3, 4).reshape(Bn * nb, h, 2 * w, dh),
+                                           attn_mask=m.reshape(Bn * nb, 1, w, 2 * w))
+        a = a.view(Bn, nb, h, w, dh).permute(0, 2, 1, 3, 4).reshape(Bn, h, Wp, dh)[:, :, :W].transpose(1, 2).reshape(Bn, W, D)
+        x = x + self.drop(self.o(a.to(x.dtype)))
+        return x + self.drop(self.ff(self.n2(x)))
+
+
+class WordChannel(nn.Module):
+    def __init__(self, d_model=512, d_k=512, n_in=41, k_slots=4, n_nwt=43, d_w=256, d_c=128, heads_c=4, n_tok=32768, d_e=320):
+        super().__init__()
+        self.register_buffer("tok_E", torch.zeros(n_tok, d_e, dtype=torch.float16))    # frozen decoder piece table (from the weights)
+        self.ppos = nn.Parameter(torch.zeros(P_MAX, d_e))
+        self.W1 = nn.Linear(d_e, d_w); self.W2 = nn.Linear(d_e, d_w, bias=False); self.ln = nn.LayerNorm(d_w)
+        self.fl = nn.ModuleList([nn.Embedding(n, d_w) for n in (2, 2, P_MAX + 1, 21, 21)]); self.src = nn.Embedding(len(SRC_ROW) + 1, d_w)
+        self.Wa = nn.Linear(2 * d_model, d_w)
+        self.hav = nn.Embedding(2, d_w)
+        self.enc = nn.ModuleList([EncLayer(d_w) for _ in range(N_ENC)]); self.enc_norm = nn.LayerNorm(d_w)
+        self.hc = heads_c; self.fq = nn.Linear(d_model, d_c); self.fk = nn.Linear(d_w, d_c); self.fv = nn.Linear(d_w, d_c)
+        self.R0 = nn.Embedding(2 * OFF_CLIP + 1, d_c); self.R1 = nn.Embedding(2 * OFF_CLIP + 1, d_c)
+        self.fo = nn.Linear(d_c, d_model)
+        self.war_r = nn.Linear(d_w, d_w); self.war_k = nn.Embedding(3, d_w); self.war_ln = nn.LayerNorm(d_w)
+        self.QB = nn.Linear(d_w, d_k); self.QI = nn.Linear(d_w, d_k)
+        self.bB = nn.Parameter(torch.zeros(n_in)); self.bI = nn.Parameter(torch.zeros(n_in))
+        self.nwt = nn.Linear(d_w, k_slots * n_nwt); self.k_slots, self.n_nwt = k_slots, n_nwt
+
+    def embed_in(self, tok, npc, flags, src, hm, he, have):
+        """per-word encoder input x_j (B,W,d_w). hm / he (B,W,D) own-frame mean / last own-frame trunk state, have (B,W) bool."""
+        Et = self.tok_E[tok].float()
+        pm = (torch.arange(P_MAX, device=tok.device)[None, None] < npc[..., None]).float()[..., None]
+        mean = ((Et + self.ppos) * pm).sum(2) / npc.clamp(min=1)[..., None].float()
+        last = torch.gather(Et, 2, (npc - 1).clamp(min=0)[..., None, None].expand(-1, -1, 1, Et.shape[-1]))[:, :, 0]
+        e = self.ln(self.W1(mean) + self.W2(last))
+        for i, m in enumerate(self.fl):
+            e = e + m(flags[..., i].clamp(0, m.num_embeddings - 1))
+        e = e + self.src(src.clamp(0, self.src.num_embeddings - 1))
+        return e + self.Wa(torch.cat([hm, he], -1)) + self.hav(have.long())
+
+    def run_enc(self, x, valid, pos0=0):
+        for lay in self.enc:
+            x = lay(x, valid, pos0)
+        return self.enc_norm(x) * valid[..., None].float()
+
+    def frame_ctx(self, h, u, f0, oe, wctx, t0=0):
+        """h (B,T,D) trunk states of frames t0..t0+T-1; u (B,W,d_w); f0 / oe (B,W) word frames; wctx (B,T,K) word indices
+        (-1 none) -> c (B,T,D) (0 where the window is empty)."""
+        Bn, T, Kc = wctx.shape; ok = wctx >= 0; idx = wctx.clamp(min=0)
+        dc = self.fk.out_features; hh = self.hc; dh = dc // hh
+        flat = idx.reshape(Bn, T * Kc)
+        gk = torch.gather(self.fk(u), 1, flat[..., None].expand(-1, -1, dc)).view(Bn, T, Kc, dc)
+        gv = torch.gather(self.fv(u), 1, flat[..., None].expand(-1, -1, dc)).view(Bn, T, Kc, dc)
+        t = torch.arange(t0, t0 + T, device=h.device)[None, :, None]
+        o0 = (t - torch.gather(f0, 1, flat).view(Bn, T, Kc)).clamp(-OFF_CLIP, OFF_CLIP) + OFF_CLIP
+        o1 = (t - torch.gather(oe, 1, flat).view(Bn, T, Kc)).clamp(-OFF_CLIP, OFF_CLIP) + OFF_CLIP
+        gk = gk + self.R0(o0) + self.R1(o1)
+        q = self.fq(h).view(Bn, T, 1, hh, dh)
+        sc = (q * gk.view(Bn, T, Kc, hh, dh)).sum(-1).float() / math.sqrt(dh)
+        sc = sc.masked_fill(~ok[..., None], float("-inf"))
+        any_ = ok.any(-1)
+        att = torch.softmax(sc, 2).nan_to_num(0.0)
+        ctx = (att[..., None] * gv.view(Bn, T, Kc, hh, dh).float()).sum(2).reshape(Bn, T, dc)
+        return self.fo(ctx.to(h.dtype)) * any_[..., None].to(h.dtype)
+
+    def war_logits(self, u, valid, K_in, k):
+        Bn, W, D = u.shape
+        if k:
+            us = F.pad(u[:, k:], (0, 0, 0, k)); ok = F.pad(valid[:, k:], (0, k)) & valid
+        else:
+            us = torch.zeros_like(u); ok = valid
+        v = self.war_ln(u + self.war_r(us) + self.war_k.weight[k])
+        return self.QB(v).float() @ K_in.T + self.bB, self.QI(v).float() @ K_in.T + self.bI, ok, v
+
+    @torch.no_grad()
+    def war_probs(self, u, valid, K_in, k=1):
+        """-> P(B), P(I) (B,W,n_in) at right context k (falls back to k-1 / 0 where the right words do not exist yet), v."""
+        zB, zI, ok, v = self.war_logits(u, valid, K_in, k)
+        if k:
+            zB0, zI0, _, _ = self.war_logits(u, valid, K_in, k - 1) if k > 1 else self.war_logits(u, valid, K_in, 0)
+            zB = torch.where(ok[..., None], zB, zB0); zI = torch.where(ok[..., None], zI, zI0)
+        p = torch.softmax(torch.stack([torch.zeros_like(zB), zB, zI], -1), -1)
+        return p[..., 1], p[..., 2], v
+
+    def nwt_logits(self, v):
+        return self.nwt(v).float().reshape(*v.shape[:-1], self.k_slots, self.n_nwt)
+
+
+class WordFeed:
+    """per-stream word-channel state: the word list (append-only; a revision of earlier words re-encodes from there), each word's
+    encoder input / output, its frames and the first label frame at which it is readable."""
+
+    def __init__(self, wch, pieces, dev, d):
+        self.w, self.P, self.dev, self.d = wch, pieces, dev, int(d)
+        self.keys = []                         # (start, end, text) identity of each word in the list
+        self.txt, self.f0, self.oe, self.f1, self.src, self.av = [], [], [], [], [], []
+        self.X = None; self.U = None           # (n_x, d_w) inputs / (n_u, d_w) outputs (prefix of the list)
+        self.n_x = 0
+        self.ver = None
+
+    def __len__(self):
+        return len(self.keys)
+
+    def sync(self, finals, f_now, arrival=None, src_of=None):
+        """bring the list in line with the store's final words. New words become readable at max(f_now, own end) (live) or at
+        `arrival[j]` (offline rule). -> True if the list changed."""
+        ws = [w for w in finals if norm_word(w["w"])]
+        keys = [(w["start"], w["end"], w["w"]) for w in ws]
+        n0 = 0; m = min(len(keys), len(self.keys))
+        while n0 < m and keys[n0] == self.keys[n0]:
+            n0 += 1
+        if n0 == len(keys) == len(self.keys):
+            return False
+        # truncate at the first difference, then append
+        for L in (self.keys, self.txt, self.f0, self.oe, self.f1, self.src, self.av):
+            del L[n0:]
+        self.n_x = min(self.n_x, n0)
+        if self.X is not None:
+            self.X = self.X[:self.n_x]
+        if self.U is not None:
+            self.U = self.U[:min(len(self.U), n0)]
+        for j in range(n0, len(ws)):
+            w = ws[j]
+            self.keys.append(keys[j]); self.txt.append(norm_word(w["w"])); self.f0.append(int(w["f0"])); self.f1.append(int(w["f1"]))
+            self.src.append(src_of(w) if src_of else SRC_ROW["ms"]); self.av.append(None)
+        f0 = np.asarray(self.f0, np.int64); oe = own_ends(f0, np.asarray(self.f1, np.int64))
+        # a word already encoded keeps the own end it was encoded with (a later word sharing its last frame does not re-open it)
+        old = self.oe[:min(n0, self.n_x)]
+        self.oe = [int(x) for x in oe]
+        self.oe[:len(old)] = old
+        if arrival is not None:
+            self.av = [int(x) for x in arrival]
+        else:
+            for j in range(n0, len(self.keys)):
+                a = max(int(f_now), self.oe[j])
+                if j:
+                    a = max(a, self.av[j - 1])
+                self.av[j] = a
+        return True
+
+    @torch.no_grad()
+    def encode_upto(self, f_done, h_at, f_start=None, settle=4):
+        """encode every word whose own frames are all decided (oe < f_done). A word's own end depends on the NEXT word's start (a
+        shared frame belongs to the next word, as in the training tables), so the newest word is encoded only once a later word
+        is known or `settle` frames have passed after its last frame (no word can still start inside it). Words encoded in this
+        call are readable from f_start (the first frame of this feed) at the earliest. h_at(frames) -> (n, D) audio trunk states."""
+        n_new = 0
+        while self.n_x + n_new < len(self.keys) and self.oe[self.n_x + n_new] < f_done:
+            k = self.n_x + n_new
+            if k + 1 >= len(self.keys) and f_done <= self.f1[k] + settle:
+                break
+            n_new += 1
+        if not n_new:
+            return
+        if f_start is not None:
+            for k in range(self.n_x, len(self.av)):
+                lo_ = f_start if k < self.n_x + n_new else self.av[k - 1]
+                self.av[k] = max(self.av[k], lo_, self.av[k - 1] if k else lo_)
+        J = list(range(self.n_x, self.n_x + n_new))
+        f0 = np.asarray([self.f0[j] for j in J]); oe = np.asarray([self.oe[j] for j in J])
+        tok, npc, flags = word_rows([self.txt[j] for j in J], f0, oe, None, self.P, prev_oe=self.oe[J[0] - 1] if J[0] > 0 else None)
+        hm, he = [], []
+        for a, b in zip(f0, oe):
+            H = h_at(list(range(int(a), int(b) + 1))).float()
+            hm.append(H.mean(0)); he.append(H[-1])
+        T = lambda a, dt=torch.long: torch.as_tensor(np.asarray(a), device=self.dev, dtype=dt)[None]
+        x = self.w.embed_in(T(tok), T(npc), T(flags), T([self.src[j] for j in J]), torch.stack(hm)[None], torch.stack(he)[None],
+                            torch.ones(1, n_new, dtype=torch.bool, device=self.dev))[0]
+        self.X = x if self.X is None or not len(self.X) else torch.cat([self.X, x], 0)
+        self.n_x += n_new
+        # outputs for words [n_u, n_x): run the encoder over the receptive field before them
+        n_u = 0 if self.U is None else len(self.U)
+        lo = max(0, n_u - RF)
+        xs = self.X[lo:self.n_x][None]
+        u = self.w.run_enc(xs, torch.ones(xs.shape[:2], dtype=torch.bool, device=self.dev), pos0=lo)[0][n_u - lo:]
+        self.U = u if self.U is None or not len(self.U) else torch.cat([self.U, u], 0)
+
+    @torch.no_grad()
+    def fuse(self, h, f_start):
+        """h (T, D) audio trunk states of frames f_start..; -> h' = h + c (T, D)."""
+        T = h.shape[0]; n_u = 0 if self.U is None else len(self.U)
+        if not n_u:
+            return h
+        f0 = np.asarray(self.f0[:n_u], np.int64); av = np.asarray(self.av[:n_u], np.int64)
+        idx = frame_window(f0, av, f_start, f_start + T)
+        if (idx < 0).all():
+            return h
+        lo = int(idx[idx >= 0].min()); hi = int(idx.max()) + 1
+        loc = np.where(idx >= 0, idx - lo, -1)
+        dv = self.dev
+        u = self.U[lo:hi][None]
+        F0 = torch.as_tensor(f0[lo:hi], device=dv)[None]; OE = torch.as_tensor(np.asarray(self.oe[lo:hi], np.int64), device=dv)[None]
+        c = self.w.frame_ctx(h[None], u, F0, OE, torch.as_tensor(loc, device=dv)[None], t0=f_start)[0]
+        return h + c.to(h.dtype)

arctag/_core/arc_text.py

@@ -0,0 +1,496 @@
+"""Arc Tagger = the Arc of Arc Transcriber (arc_audio 1.4.2 package, PK8a weights) with the audio path removed: the Arc's own word
+channel is the input. arc_audio is imported read-only from the PK8a trainer's package (pk_config.json "arc_pkg" =
+release_staging/arc_v142); nothing in it is edited.
+
+per word j (one trunk step per word, no audio frames):
+  e_j  = arc.word.embed_in(pieces, flags, src, hm=0, he=0, have=0)     the word channel's input with "no audio" (exactly the
+         Arc's no-audio path): frozen 32768 x 320 piece table tok_E (moonshine-streaming-tiny tokenizer.json, = Pieces), first
+         7 + last piece, flags (has digit, all digits, n pieces, duration, gap in 100 ms frames), word source row
+         + zero-init provider features (gap / duration buckets, final / interim, confidence, speaker change, asr_source)
+  u_j  = arc.word.run_enc(e)                                          2 causal local self-attention layers (32-word window)
+  x_j  = inp(in_norm(u_j))                                            NEW 256 -> 512 input projection replacing the audio one
+                                                                      (inp init = word.fo . word.fv, the value path by which the
+                                                                      Arc adds a word to its frame state)
+  h_j  = arc.trunk (5 x [FiLM -> Mamba-3 -> MLP], domain adapters)   chunked SSD in training, Trunk.step + carried state streaming
+read-outs: arc.head TagHead on h (definition keys, pooled utt / speaker / item with the Arc carry); span_beg (the B of BIO) =
+  arc.word.war_* B read-out on u with right context k (T0 k=0, T1/T2 k<=2); entity KB match = arc.ent.ptr_logits (NameEncoder
+  candidates, [new, generic, KB...]); latency tiers = delayed read-outs of the same causal pass (word i read at i+d with a delay
+  embedding, T1 d<=3, T2 d = to the utterance end <= 8).
+dropped: arc.diar (voice embedding needs audio), arc.bio (vocal events), trunk.inp/in_norm audio 640-d projection.
+"""
+import math, os, sys
+import numpy as np
+import torch, torch.nn as nn, torch.nn.functional as F
+
+ARC_PKG = os.path.dirname(os.path.abspath(__file__))          # [release] vendored arc_audio subset lives here
+if ARC_PKG not in sys.path:
+    sys.path.insert(0, ARC_PKG)
+from arc_audio.model import Trunk, TagHead, EntityHeads, zero_carry, encode_names, NEG    # noqa: E402
+from arc_audio.ontology import QUERIES as ARC_Q, POOLED                                  # noqa: E402
+from arc_audio.wordch import WordChannel, P_MAX, SRC_ROW                                 # noqa: E402
+
+HERE = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, HERE)
+from tagtable import FINE_DOMAINS, ASR_SOURCES                                           # noqa: E402
+
+TOK_PATH = os.path.join(ARC_PKG, "moonshine-streaming-tiny", "tokenizer.json")
+SCRATCH_10M = dict(d=256, d_w=256, d_e=256, vocab=8192, n_layers=6, adapters=16, d_state=64, head_dim=64, dropout=0.1)   # 9.79 M
+WORD_Q = ("frame", "span_beg", "span_in", "span_end", "ev_at", "ev_imm", "ev_rec")
+FEATS = {"gap": 10, "dur": 8, "fin": 3, "conf": 7, "chg": 3, "asr": 8}
+D_MAX = 8
+TIER_W = {"T0": 1.0, "T1": 0.5, "T2": 0.5}
+ARC_CFG = dict(d_in=640, d_model=512, n_layers=5, d_state=64, head_dim=64, dropout=0.2, in_drop=0.2, chunk=128, n_dom=16,
+               delays=[0, 1, 2, 3, 5, 8], adapters=64, long_mem=True)        # release config.json "model.trunk"
+
+
+BPE8K = os.path.join(os.environ.get("ARCTAG_HOME", os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), "tok", "bpe8k", "tokenizer.json")   # [release]   # own byte-level BPE (train_bpe.py), words as " w"
+
+
+class ArcPieces:
+    """tokenizer adapter for data.Windows / the streaming engines: word -> piece ids, no specials.
+    moonshine (the Arc package's tokenizer): the word itself (its normalizer prepends the space marker);
+    own byte-level BPE (tok/bpe8k): " " + word, exactly as train_bpe.py fed it."""
+    def __init__(self, path=TOK_PATH):
+        from tokenizers import Tokenizer
+        self.path = path
+        self.t = Tokenizer.from_file(path)
+        self.vocab_size = self.t.get_vocab_size()
+        self.space = "bpe8k" in path or "/tok/" in path
+        self.cls_token_id = 0; self.pad_token_id = 0; self.unk_token_id = 0; self.mask_token_id = 0
+
+    def word(self, w):
+        return (" " + w.strip()) if self.space else w.strip()
+
+    def __call__(self, texts, add_special_tokens=False):
+        enc = self.t.encode_batch([self.word(t) for t in texts], add_special_tokens=False)
+        return {"input_ids": [e.ids for e in enc]}
+
+    def encode_word(self, w):
+        return self.t.encode(self.word(w), add_special_tokens=False).ids
+
+
+class TableShim:
+    def __init__(self, T):
+        self.n = T.n; self.valid = T.valid
+        self.cols = {q: list(T.cols[q]) for q in ARC_Q}
+
+
+def frames(t):
+    return int(round(t * 10)) if t is not None else None
+
+
+class ArcText(nn.Module):
+    def __init__(self, table, n_in_war=41, n_units=83, d_in_new=256, train_tokE=True, scratch=None):
+        """scratch=None: the Arc's own dims (34.1 M, warm start from PK8a). scratch=dict(d, d_w, d_e, vocab, n_layers, adapters,
+        d_state, head_dim): the same block types, narrower, trained from scratch (audio-only / unused sub-modules removed)."""
+        super().__init__()
+        self.T = table; self.scratch = scratch
+        tc = dict(ARC_CFG); tc["d_in"] = d_in_new
+        if scratch:
+            tc.update(d_in=scratch["d_w"], d_model=scratch["d"], n_layers=scratch["n_layers"], adapters=scratch["adapters"],
+                      d_state=scratch.get("d_state", 64), head_dim=scratch.get("head_dim", 64), delays=[0], dropout=scratch.get("dropout", 0.1))
+        self.trunk = Trunk(**tc)
+        d = tc["d_model"]; self.d = d
+        nn.init.ones_(self.trunk.in_norm.weight); nn.init.zeros_(self.trunk.in_norm.bias)
+        self.head = TagHead(TableShim(table), n_text=table.E.shape[1], d_model=d, d_k=d)
+        with torch.no_grad():
+            self.head.E.copy_(torch.as_tensor(table.E))
+        if scratch:
+            self.ent = EntityHeads(d_model=d, d_key=d, n_units=n_units, span=False, span_new=False, sim=False)
+            for k in ("nwt", "unit", "clock", "approx"):                     # numeral value heads: no stage-1 target -> removed
+                delattr(self.ent, k)
+            self.word = WordChannel(d_model=d, d_k=d, n_in=n_in_war, k_slots=4, n_nwt=43, d_w=scratch["d_w"], n_tok=scratch["vocab"],
+                                    d_e=scratch["d_e"], d_c=64)
+            for k in ("fq", "fk", "fv", "fo", "R0", "R1", "QI", "bI", "bB", "nwt", "Wa", "hav"):   # audio-fusion / unused read-outs
+                delattr(self.word, k)
+            del self.word._buffers["tok_E"]
+            self.word.tok_E = nn.Parameter(torch.randn(scratch["vocab"], scratch["d_e"]) * 0.02)
+        else:
+            self.ent = EntityHeads(d_model=d, d_key=d, n_units=n_units, span=True, span_new=True, sim=True)
+            self.word = WordChannel(d_model=d, d_k=d, n_in=n_in_war, k_slots=4, n_nwt=43)
+            if train_tokE:                                                   # end-to-end: the piece table becomes a trainable fp32 parameter
+                del self.word._buffers["tok_E"]
+                self.word.tok_E = nn.Parameter(torch.zeros(32768, 320))
+        self.feat = nn.ModuleDict({k: nn.Embedding(n, self.word.ln.normalized_shape[0]) for k, n in FEATS.items()})
+        for m in self.feat.values():
+            nn.init.zeros_(m.weight)
+        self.delay_q = nn.Embedding(D_MAX + 1, d); nn.init.zeros_(self.delay_q.weight)
+        self.bB_all = nn.Parameter(torch.zeros(table.n))                      # span_beg bias per r8 tag (war bB on the r5 columns)
+        self.register_buffer("cols_span_beg", torch.tensor(table.cols["span_beg"], dtype=torch.long), persistent=False)
+        self.ver = nn.Linear(2 * d, 1)
+        self.lm = nn.Linear(d, self.word.tok_E.shape[1])
+
+    # ------------------------------------------------------------------ word input
+    def word_inputs(self, b):
+        """token-level batch (data.collate, ArcPieces ids) -> pieces (B,W,P_MAX) (first 7 + last), counts, word features"""
+        ids, tok_word, feats = b["ids"], b["tok_word"], b["feats"]; W = b["wvalid"].shape[1]
+        B, S = ids.shape; dev = ids.device
+        valid = tok_word >= 0; tw = tok_word.clamp(min=0)
+        pos = torch.arange(S, device=dev)[None].expand(B, S)
+        first = torch.full((B, W), S, dtype=torch.long, device=dev).scatter_reduce(1, tw, torch.where(valid, pos, torch.full_like(pos, S)), "amin")
+        rank = pos - torch.gather(first, 1, tw)
+        npc_all = torch.zeros(B, W, dtype=torch.long, device=dev).scatter_add(1, tw, valid.long())
+        nw = torch.gather(npc_all, 1, tw)
+        keep = valid & ((rank < P_MAX - 1) | (rank == nw - 1))
+        slot = torch.where(rank < P_MAX - 1, rank, torch.full_like(rank, P_MAX - 1))
+        pcs = torch.zeros(B, W, P_MAX, dtype=torch.long, device=dev)
+        bi = torch.arange(B, device=dev)[:, None].expand(B, S)
+        pcs[bi[keep], tw[keep], slot[keep]] = ids[keep]
+        npc = npc_all.clamp(max=P_MAX)
+        sc = lambda src: torch.zeros(B, W, dtype=torch.long, device=dev).scatter(1, torch.where(valid, tw, torch.zeros_like(tw)),
+                                                                                  torch.where(valid, src, torch.zeros_like(src)))
+        fw = {k: sc(feats[k]) for k in self.feat if k in feats}
+        spk = sc(feats["spk"]) if "spk" in feats else torch.zeros(B, W, dtype=torch.long, device=dev)
+        return pcs, npc, fw, spk
+
+    def embed(self, pcs, npc, fw, wfl, wsrc):
+        """arc.word.embed_in with no audio + provider features. wfl (B,W,5) Arc flags (npc filled here), wsrc (B,W) source row"""
+        fl = wfl.clone(); fl[..., 2] = npc.clamp(0, P_MAX)
+        if self.scratch:                                                     # = WordChannel.embed_in without the audio terms
+            w = self.word; Et = w.tok_E[pcs].float()
+            pm = (torch.arange(P_MAX, device=pcs.device)[None, None] < npc[..., None]).float()[..., None]
+            mean = ((Et + w.ppos) * pm).sum(2) / npc.clamp(min=1)[..., None].float()
+            last = torch.gather(Et, 2, (npc - 1).clamp(min=0)[..., None, None].expand(-1, -1, 1, Et.shape[-1]))[:, :, 0]
+            e = w.ln(w.W1(mean) + w.W2(last))
+            for i, m in enumerate(w.fl):
+                e = e + m(fl[..., i].clamp(0, m.num_embeddings - 1))
+            e = e + w.src(wsrc.clamp(0, w.src.num_embeddings - 1))
+        else:
+            z = torch.zeros(*pcs.shape[:2], self.d, device=pcs.device)
+            e = self.word.embed_in(pcs, npc, fl, wsrc, z, z, torch.zeros(pcs.shape[:2], dtype=torch.bool, device=pcs.device))
+        for k, m in self.feat.items():
+            if k in fw:
+                e = e + m(fw[k].clamp(0, m.num_embeddings - 1))
+        return e
+
+    def keys(self):
+        return self.head.keys()
+
+    def span_beg_logits(self, u, valid, k_i, dom, K):
+        """war B read-out (arc.word.war_*) on the word-channel output with right context k (per word, 0..2) = WordChannel.war_logits
+        with the text table's span_beg columns and a per-r8-tag bias"""
+        c = self.cols_span_beg; Kb = K[c]; w = self.word; W = u.shape[1]
+        out = None
+        for k in range(3):
+            us = F.pad(u[:, k:], (0, 0, 0, k)) if k else torch.zeros_like(u)
+            v = w.war_ln(u + w.war_r(us) + w.war_k.weight[k])
+            zB = w.QB(v).float() @ Kb.T + self.bB_all[c]
+            out = zB if out is None else torch.where((k_i == k)[..., None], zB, out)
+        return out.masked_fill(~self.head.valid[dom][:, c][:, None, :], NEG)
+
+    @staticmethod
+    def utt_last(uend, wvalid):
+        B, W = uend.shape; dev = uend.device
+        t = torch.arange(W, device=dev)[None].expand(B, W)
+        last_valid = (wvalid.long().sum(1) - 1).clamp(min=0)[:, None]
+        e = torch.where(uend | (t == last_valid), t, torch.full_like(t, W - 1))
+        return torch.flip(torch.cummin(torch.flip(e, [1]), 1).values, [1])
+
+    def tier_delays(self, tier, uend, wvalid):
+        B, W = uend.shape
+        t = torch.arange(W, device=uend.device)[None].expand(B, W)
+        to_end = self.utt_last(uend, wvalid) - t
+        if tier == "T0":
+            return torch.zeros_like(t)
+        return to_end.clamp(min=0, max=3 if tier == "T1" else D_MAX)
+
+    def delayed(self, h, d_i):
+        """read-out state of word i at delay d: h_{i+d} + delay_q(d) [+ stage 2 dmix(h_i) for d > 0: the delayed read-out sees word i
+        itself; stage 1 drifted toward word i+d (T1 person spans peaked +3 words off)]"""
+        W = h.shape[1]
+        j = (torch.arange(W, device=h.device)[None] + d_i).clamp(max=W - 1)
+        out = torch.gather(h, 1, j[..., None].expand(-1, -1, h.shape[-1])) + self.delay_q(d_i.clamp(0, D_MAX))
+        if hasattr(self, "dattn"):                                         # v2 A2: attention over h_i..h_{i+d} (zero-init output)
+            out = out + self.dattn(h, d_i)
+        elif hasattr(self, "dmix"):
+            out = out + self.dmix(h) * (d_i > 0)[..., None].to(h.dtype)
+        return out
+
+    def add_delay_mix(self):
+        """stage 2: W_d (d x d, no bias, zero-init -> stage-1 outputs unchanged; T0 never affected) adding word i's own state to its
+        delayed read-out"""
+        self.dmix = nn.Linear(self.d, self.d, bias=False); nn.init.zeros_(self.dmix.weight)
+        return self
+
+    def add_v2(self, gattn_at=(1, 3), delay_attn=True, span_head=True, utt_att=True, gattn_bias=False):
+        """AT10M-v2 (v2.py): global attention layers after trunk blocks gattn_at, attention delayed read-out, per-family span head,
+        attention-pooled utterance read-out. Every new output path is zero-initialised except the span head (new read-out)."""
+        import v2 as V2
+        if gattn_at:
+            self.gat = nn.ModuleDict({str(i): V2.GlobalAttn(self.d, biased=gattn_bias) for i in gattn_at})   # v2.1: ALiBi + key prior + gate
+        if delay_attn:
+            self.dattn = V2.DelayAttn(self.d, D_MAX)
+            if hasattr(self, "dmix"):
+                del self.dmix
+        if span_head:
+            self.span_head = V2.SpanHead(self.d, self.T)
+        if utt_att:
+            self.utt_att = nn.Linear(self.d, 1); nn.init.zeros_(self.utt_att.weight); nn.init.zeros_(self.utt_att.bias)
+        return self
+
+    def add_memread(self, d_ext=0):
+        """v2.2 (a): session mention memory read into the span read-outs (v2.MemRead; exact no-op at insertion)"""
+        import v2 as V2
+        self.memr = V2.MemRead(self.d, d_ext=d_ext, d_u=len(self.cols_span_beg))
+        return self
+
+    def utt_pool_att(self, h, ustart):
+        """v2 A6: utterance pool = softmax(utt_att(z))-weighted running mean of z = h + pool_in.utt(h) within the utterance (zero init
+        -> the plain running mean of pool_segment)"""
+        z = (h + self.head.pool_in["utt"](h)).float()
+        w = torch.exp(self.utt_att(z).clamp(-10, 10))                     # (B, W, 1)
+        cw = torch.cumsum(w, 1); cz = torch.cumsum(w * z, 1)
+        idx = torch.arange(h.shape[1], device=h.device)[None].expand_as(ustart)
+        st = torch.cummax(torch.where(ustart, idx, torch.zeros_like(idx)), 1).values   # start index of the current utterance
+        prev = (st - 1).clamp(min=0)
+        has = (st > 0)[..., None]
+        cw0 = torch.where(has, torch.gather(cw, 1, prev[..., None]), torch.zeros_like(cw))
+        cz0 = torch.where(has, torch.gather(cz, 1, prev[..., None].expand(-1, -1, z.shape[-1])), torch.zeros_like(cz))
+        return (cz - cz0) / (cw - cw0).clamp(min=1e-9)
+
+    def encode(self, b):
+        pcs, npc, fw, spk = self.word_inputs(b)
+        e = self.embed(pcs, npc, fw, b["wfl"], b["wsrc"])
+        u = self.word.run_enc(e, b["wvalid"])
+        dom = b["dom"]
+        if hasattr(self, "gat"):
+            import v2 as V2
+            h = V2.trunk_forward(self.trunk, self.gat, self.trunk_in(u), dom, torch.zeros_like(dom),
+                                 ckpt=getattr(self, "grad_ckpt", False) and self.training)
+            return pcs, spk, u, h
+        if getattr(self, "grad_ckpt", False) and self.training:
+            h = self.trunk_ckpt(self.trunk_in(u), dom, torch.zeros_like(dom))
+        else:
+            h, _ = self.trunk(self.trunk_in(u), dom, torch.zeros_like(dom))
+        return pcs, spk, u, h
+
+    def trunk_ckpt(self, x, dom, delay):
+        """= Trunk.forward (state None) with each [FiLM -> Mamba-3 -> MLP] block activation-checkpointed (training memory)"""
+        from torch.utils.checkpoint import checkpoint
+        from arc_audio.mamba3.ssd import mamba3_layer_ssd_state, _f
+        tr = self.trunk; h, e = tr._cond(x, dom, delay)
+        for i, (lay, mlp) in enumerate(zip(tr.layers, tr.mlps)):
+            def blk(h, e, i=i, lay=lay, mlp=mlp):
+                gb = tr.film[i](e)[:, None]; h = h * (1 + gb[..., :h.shape[-1]]) + gb[..., h.shape[-1]:]
+                s, _ = mamba3_layer_ssd_state(lay, _f(h), None, tr.chunk)
+                h = h + s
+                return h + mlp(h)
+            h = checkpoint(blk, h, e, use_reentrant=False)
+        return tr.out_norm(tr._adapt(h, dom))
+
+    def trunk_in(self, u):
+        return u          # Trunk applies in_norm + inp itself (d_in = 256)
+
+    def forward(self, b, tiers=("T0",), lm=False):
+        pcs, spk, u, h = self.encode(b)
+        dom = b["dom"]; W = h.shape[1]
+        spk_i = torch.where(spk > 0, spk - 1, torch.full_like(spk, -1))
+        st = b["ustart"].clone(); st[:, 0] = True
+        zin, _ = self.head.pooled(h, spk_i, st, zero_carry(h.shape[0], h.shape[-1], h.device))
+        if hasattr(self, "utt_att"):
+            zin["utt"] = self.utt_pool_att(h, st)
+        K = self.keys()
+        pooled = {q: self.head.logits_q(q, zin[q], dom, K) for q in POOLED}
+        out = {"h": h, "u": u, "tiers": {}}; mkv = None
+        for tier in tiers:
+            dl = self.tier_delays(tier, b["uend"], b["wvalid"])
+            hd = self.delayed(h, dl)
+            if hasattr(self, "memr"):                                      # v2.2 (a): earlier-mention read (keys shared across tiers)
+                if mkv is None:
+                    mkv = self.memr.keys(h, b.get("cng"))
+                if self.training and torch.is_grad_enabled():         # activation-checkpointed: GPU headroom is small
+                    from torch.utils.checkpoint import checkpoint
+                    y_, a_ = checkpoint(lambda hd_, h_, k_, v_, m_: self.memr(hd_, h_, kv=(k_, v_, m_)), hd, h, *mkv, use_reentrant=False)
+                else:
+                    y_, a_ = self.memr(hd, h, kv=mkv)
+                hd = hd + y_.to(hd.dtype)
+            lg = {q: self.head.logits_q(q, hd, dom, K) for q in WORD_Q if q != "span_beg"}
+            lg["span_beg"] = self.span_beg_logits(u, b["wvalid"], dl.clamp(max=2), dom, K)
+            if hasattr(self, "memr"):                                      # v2.2 (a): gated span_beg logit increment (read at word i)
+                lg["span_beg"] = lg["span_beg"] + self.memr.to_u(a_).float()
+            lg.update(pooled)
+            out["tiers"][tier] = lg
+            if hasattr(self, "kb_spot"):
+                out.setdefault("kb_spot", {})[tier] = self.kb_spot(hd).float()
+            if hasattr(self, "span_head"):
+                out.setdefault("span_head", {})[tier] = self.span_head(hd)
+            if tier == "T0" and hasattr(self, "pii_imm"):              # stage 2: "a PII span starts within 1..3 words" (T0 only)
+                out["pii_imm"] = self.pii_imm(hd).float()
+        out["logits"] = out["tiers"][tiers[0]]
+        if lm:
+            out["lm_h"] = self.lm(h)                                          # logits at the selected rows only (train.py)
+            out["lm"] = None if getattr(self, "lm_rows", False) else out["lm_h"] @ self.word.tok_E.float().t()
+            out["lm_target"] = torch.cat([pcs[:, 1:, 0], torch.zeros_like(pcs[:, :1, 0])], 1)
+            out["lm_valid"] = torch.cat([b["wvalid"][:, 1:], torch.zeros_like(b["wvalid"][:, :1])], 1)
+        return out
+
+    # ------------------------------------------------------------------ entities: verification + KB pointer
+    def span_rep(self, hw, spans):
+        cs = F.pad(torch.cumsum(hw, 0), (0, 0, 1, 0))
+        mean = (cs[spans[:, 1]] - cs[spans[:, 0]]) / (spans[:, 1] - spans[:, 0]).clamp(min=1)[:, None].to(hw.dtype)
+        return torch.cat([mean, hw[spans[:, 1] - 1]], -1)
+
+    # ------------------------------------------------------------------ long-list KB (stage 2, DESIGN.md section 13; +~136 k params)
+    def add_kb_long(self):
+        import kb as KBM
+        D = self.d; dn = self.ent.names.d
+        self.kb_len = nn.Embedding(7, D)                      # query span length 1..6 (+ longer) -> added to the span-end state
+        self.kb_type = nn.Embedding(len(KBM.TYPES), dn)        # entry type -> added to the NameEncoder entry vector
+        self.kb_spot = nn.Linear(D, 6)                         # "a KB entry of length L ends at this word" (L = 1..6), per tier
+        self.ver_e = nn.Sequential(nn.Linear(2 * D + 4 * dn, 128), nn.GELU(), nn.Linear(128, 1))   # entry-aware verification
+        for m in (self.kb_len, self.kb_type):
+            nn.init.zeros_(m.weight)
+        nn.init.zeros_(self.kb_spot.weight); nn.init.constant_(self.kb_spot.bias, -4.0)
+        return self
+
+    def add_kb_str(self, cng=True, vstr=True):
+        """v2.2 (b) char-n-gram mention key (hashed char 3-grams of the heard span -> added to the KB query before normalisation) and
+        (c) ver_e string features (span text vs entry text similarity -> added to ver_e's hidden pre-activation). Both zero-init:
+        outputs are bit-identical until trained; used only where the caller passes the span texts / features (else skipped)."""
+        import kb_strfeat as KS
+        dn = self.ent.names.d
+        if cng:
+            self.kb_cng = nn.EmbeddingBag(KS.NB, dn, mode="mean"); nn.init.zeros_(self.kb_cng.weight)
+        if vstr:
+            self.ver_s = nn.Linear(KS.NF, self.ver_e[0].out_features, bias=False); nn.init.zeros_(self.ver_s.weight)
+        return self
+
+    def add_redact(self, n=8):
+        """stage 2 real-time redaction (DESIGN.md section 14): the PII-imminent read-out, one logit per redact.PII type
+        (D x 8 + 8 = 2,056 params); zero weight, bias -4 (rare event) so the stage-1 outputs are unchanged at the start"""
+        self.pii_imm = nn.Linear(self.d, n)
+        nn.init.zeros_(self.pii_imm.weight); nn.init.constant_(self.pii_imm.bias, -4.0)
+        return self
+
+    def add_spk_link(self, flux=True):
+        """stage 2 text speaker linking (DESIGN.md 14.8, speaker_link.SpkLink) + input features "spksrc" (speaker ids: none / gold /
+        ASR-like) and, with flux, "turn" (Flux-like turn events). New embeddings are zero-initialised (stage-1 outputs unchanged)."""
+        import speaker_link as SL
+        n_roles = len(role_cols(self.T))
+        self.spk_link = SL.SpkLink(self.d, n_roles)
+        dw = self.word.ln.normalized_shape[0]
+        self.feat["spksrc"] = nn.Embedding(3, dw); nn.init.zeros_(self.feat["spksrc"].weight)
+        if flux:
+            self.feat["turn"] = nn.Embedding(4, dw); nn.init.zeros_(self.feat["turn"].weight)
+        return self
+
+    def kb_entry_vecs(self, texts, types):
+        """entry texts + type ids -> (N, dn) unit vectors (the model's NameEncoder + a type embedding); precomputable per list"""
+        dev = self.bB_all.device
+        ids = torch.as_tensor(encode_names(texts), device=dev)
+        e = self.ent.names(ids, torch.zeros(len(texts), dtype=torch.long, device=dev))
+        return F.normalize(e + self.kb_type(torch.as_tensor(types, device=dev)), dim=-1)
+
+    def kb_query(self, h, bi, a, b, qtext=None):
+        """span queries: h (B, W, D); bi / a / b (n,) -> rep (n, 2D) [mean, end], q (n, dn) unit; qtext: heard span texts (v2.2 b)"""
+        cs = F.pad(torch.cumsum(h, 1), (0, 0, 1, 0))
+        mean = (cs[bi, b] - cs[bi, a]) / (b - a).clamp(min=1)[:, None].to(h.dtype)
+        hend = h[bi, b - 1]; rep = torch.cat([mean, hend], -1)
+        q0 = self.ent.query(hend.float() + self.kb_len((b - a).clamp(1, 6)))
+        if qtext is not None and hasattr(self, "kb_cng"):
+            import kb_strfeat as KS
+            ids, off = KS.cng_batch(qtext)
+            q0 = q0 + self.kb_cng(torch.as_tensor(ids, device=h.device), torch.as_tensor(off, device=h.device)).to(q0.dtype)
+        q = F.normalize(q0, dim=-1)
+        return rep, q
+
+    def kb_scores(self, q, E):
+        """q (n, dn), E (n, K, dn) or (M, dn) -> candidate logits; null logit (n,) = the pointer's 'new' slot"""
+        sc = self.ent.log_scale.clamp(max=math.log(50.0)).exp()
+        s = torch.einsum("nd,nkd->nk", q, E) if E.dim() == 3 else q @ E.t()
+        null = (q @ F.normalize(self.ent.special, dim=-1)[0]) * sc + self.ent.special_b[0]
+        return s * sc, null
+
+    def kb_ver_e(self, rep, q, e, sf=None):
+        """entry-aware verification logit: the span heard as exactly this entry (rejects sound-alikes / wrong digits / qualifiers);
+        sf (..., NF): v2.2 (c) string features of (span text, entry text)"""
+        x = torch.cat([rep.float(), q, e, q * e, (q - e).abs()], -1)
+        if sf is None or not hasattr(self, "ver_s"):
+            return self.ver_e(x)[..., 0]
+        z = self.ver_e[0](x) + self.ver_s(sf.to(x.dtype))
+        return self.ver_e[2](self.ver_e[1](z))[..., 0]
+
+    def kb_embed(self, names, kinds=None):
+        ids = torch.as_tensor(encode_names(names), device=self.bB_all.device)
+        kk = torch.zeros(len(names), dtype=torch.long, device=ids.device) if kinds is None else kinds
+        return self.ent.names(ids, kk)
+
+    def kb_logits(self, h_end, kbe):
+        """h at the span's last word (n,D), KB embeddings (R,d) -> (n, 2+R) [new, generic, candidates] (arc ptr head)"""
+        n = h_end.shape[0]
+        cm = torch.ones(1, n, kbe.shape[0], dtype=torch.bool, device=h_end.device)
+        return self.ent.ptr_logits(h_end[None], torch.zeros(1, n, self.d, device=h_end.device), kbe[None], cm)[0]
+
+
+# ---------------------------------------------------------------------- warm start
+@torch.no_grad()
+def role_cols(T):
+    """positions (within the speaker read-out's columns) of the speaker.role.* tags"""
+    return [k for k, i in enumerate(T.cols["speaker"]) if T.ids[i].startswith("speaker.role.")]
+
+
+def warm_start(model, export_dir, verbose=True):
+    """PK8a export (arc_pk.safetensors, keys 'arc.*') -> ArcText. Same dims everywhere except the replaced input projection."""
+    from safetensors.torch import load_file
+    sd = load_file(os.path.join(export_dir, "arc_pk.safetensors"))
+    sd = {k[4:]: v for k, v in sd.items() if k.startswith("arc.")}
+    tgt = dict(model.named_parameters()); tgt.update({k: v for k, v in model.named_buffers() if k in ("head.E", "word.tok_E")})
+    rep = {"exact": [], "derived": [], "new": [], "dropped": [], "shape_mismatch": []}
+    for k, v in sd.items():
+        if k in ("trunk.inp.weight", "trunk.inp.bias", "trunk.in_norm.weight", "trunk.in_norm.bias") or k.startswith(("diar.", "bio.")):
+            rep["dropped"].append(k); continue
+        if k not in tgt:
+            rep["dropped"].append(k); continue
+        if tgt[k].shape != v.shape:
+            rep["shape_mismatch"].append((k, tuple(v.shape), tuple(tgt[k].shape))); continue
+        tgt[k].copy_(v.to(tgt[k].dtype)); rep["exact"].append(k)
+    # input projection = the word -> frame value path of the word channel (fo . fv)
+    fo_w, fo_b, fv_w, fv_b = (sd[x].float() for x in ("word.fo.weight", "word.fo.bias", "word.fv.weight", "word.fv.bias"))
+    tgt["trunk.inp.weight"].copy_(fo_w @ fv_w); tgt["trunk.inp.bias"].copy_(fo_w @ fv_b + fo_b)
+    rep["derived"] += ["trunk.inp.weight", "trunk.inp.bias"]
+    # span_beg bias per r8 tag from war bB (r5 span_in column order = the grown table's first n5 span_in columns)
+    import json
+    g = json.load(open(os.path.join(export_dir, "arc_grow.json")))
+    c5 = g["cols"]["span_in"][: sd["word.bB"].shape[0]]
+    for j, i in enumerate(c5):
+        model.bB_all[i] = sd["word.bB"][j]
+    rep["derived"].append("bB_all")
+    done = set(rep["exact"]) | set(rep["derived"])
+    rep["new"] = [n for n, _ in model.named_parameters() if n not in done]
+    P = dict(model.named_parameters())
+    cnt = lambda names: sum((P[n] if n in P else tgt[n]).numel() for n in names if n in P or n in tgt)
+    rep["summary"] = {k: dict(tensors=len(rep[k]), params=cnt(rep[k])) for k in ("exact", "derived", "new")}
+    rep["summary"]["dropped"] = dict(tensors=len(rep["dropped"]), params=sum(sd[k].numel() for k in rep["dropped"]))
+    rep["summary"]["export_total"] = sum(v.numel() for v in sd.values())
+    if verbose:
+        print("warm start", rep["summary"], "mismatch", rep["shape_mismatch"])
+    return rep
+
+
+def param_table(m):
+    out = {}
+    extra = [] if any(n == "word.tok_E" for n, _ in m.named_parameters()) else [("word.tok_E", m.word.tok_E)]
+    for n, p in list(m.named_parameters()) + extra:
+        g = n.split(".")[0]
+        if g == "trunk":
+            g = "trunk." + ("layers (Mamba-3)" if ".layers." in n else "mlps" if ".mlps." in n else "film" if ".film." in n else
+                            "adapt" if "adapt" in n else "inp/in_norm (new text input proj)" if (".inp." in n or "in_norm" in n) else "emb/norm")
+        elif g == "word":
+            g = "word." + ("tok_E (piece table)" if n == "word.tok_E" else "enc (2 local attn layers)" if ".enc." in n else
+                           "war/nwt (word read-outs)" if any(x in n for x in ("war", "QB", "QI", "bB", "bI", "nwt")) else
+                           "fusion (fq/fk/fv/fo/R0/R1)" if any(x in n for x in (".fq", ".fk", ".fv", ".fo", ".R0", ".R1")) else "embed_in")
+        out[g] = out.get(g, 0) + p.numel()
+    out["TOTAL"] = sum(v for v in out.values())
+    out["trainable"] = sum(p.numel() for p in m.parameters() if p.requires_grad)
+    return dict(sorted(out.items()))
+
+
+def rmsnorm_fp32():
+    """bf16 autocast: nn.RMSNorm (fp32 weights) on a bf16 input falls off the fused kernel (dtype-mismatch warning). Run it in fp32 on
+    the fused kernel and cast the result back to the input dtype."""
+    if getattr(nn.RMSNorm, "_fp32_patch", False):
+        return
+
+    def fwd(self, x):
+        with torch.autocast("cuda", enabled=False):
+            return F.rms_norm(x.float(), self.normalized_shape, self.weight, self.eps).to(x.dtype)
+    nn.RMSNorm.forward = fwd; nn.RMSNorm._fp32_patch = True

arctag/_core/engine.py

@@ -0,0 +1,888 @@
+"""Arc Tagger live demo: checkpoint manager (hot reload of the newest ckpt/AT10M/step<N>) + per-session streaming engine.
+
+ArcSession wraps code/fast_arc.FastArc (the parity-tested streaming engine; snapshot revision, T0 / T1 / T2 delayed read-outs)
+and turns its raw logits into UI-ready objects:
+  spans   per tier, decoded from span_beg / span_in / span_end (W tags: entities, numbers, PII, profanity, toxicity), with the
+          span verification head (model.ver) and, when a session KB is given, the KB pointer (model.kb_logits) for norm tags
+  events  ev_at / ev_imm / ev_rec above threshold, peak-picked (speaker change, turn yield, command complete, sports, legal, meeting)
+  utt     running pooled utt / speaker / item read-outs every word, final ones at the endpoint (intent, act, safety U tags, role)
+  frame   F / G states of the newest word (content mode / phase / game state / subject)
+Thresholds: step config.json["thresholds"][tag@q][fine domain] with a fallback chain (see Thresholds).
+CPU only; nothing here touches the GPU or writes into the checkpoint dir.
+"""
+import glob, json, math, os, re, sys, threading, time
+import numpy as np
+
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+import torch
+import torch.nn.functional as F
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+B = os.path.dirname(HERE)
+CODE = HERE   # [release] vendored modules live next to engine.py
+sys.path.insert(0, CODE)
+import arc_text as AT                                   # noqa: E402
+import stream_arc as SA                                 # noqa: E402
+import fast_arc as FA                                   # noqa: E402
+from tagtable import TextTable, FINE_DOMAINS, domain_index   # noqa: E402
+from losses import GroupIndex                           # noqa: E402
+from arc_text import POOLED                             # noqa: E402
+
+PROD_CKPT = os.path.join(B, "ckpt", "AT10M")
+CKPT_DIR = os.environ.get("ARCTAG_CKPT", PROD_CKPT)
+# prod stays on ckpt/AT10M (v2 pilots live in other dirs): another dir on prod needs the explicit ARCTAG_CKPT_SWITCH=1
+if os.environ.get("ARCTAG_ROLE", "prod") == "prod" and os.path.realpath(CKPT_DIR) != os.path.realpath(PROD_CKPT) \
+        and os.environ.get("ARCTAG_CKPT_SWITCH") != "1":
+    raise RuntimeError(f"prod ckpt dir must be {PROD_CKPT} (got {CKPT_DIR}); set ARCTAG_CKPT_SWITCH=1 to switch explicitly")
+# hot reload never changes the model architecture (v1 span_in <-> v2 span head) unless ARCTAG_ARCH_SWITCH=1 (then restart instead)
+ARCH_SWITCH = os.environ.get("ARCTAG_ARCH_SWITCH") == "1"
+# v2 checkpoints: display spans / redaction from the span head. auto (default): per family, the head where its dev exact F1 (T0)
+# is >= the span_in typed decode's - HEAD_TOL (eval.json / ARCTAG_SPAN_THR dev_exact_f1), span_in otherwise; 1 = every family; 0 = never
+SPAN_HEAD_MODE = os.environ.get("ARCTAG_SPAN_HEAD", "auto")
+SPAN_HEAD = SPAN_HEAD_MODE != "0"
+HEAD_TOL = 0.05      # the dev "fam:" decode is span_in at its oracle threshold, an upper bound of the demo's span_in path (pilot step2000
+                     # probe set: head better for number / pii / profanity / toxicity at dev deficits <= .03, worse for mentions at -.09)
+# ARCTAG_SPAN_THR=<json with thresholds_span> (staging tests: a checkpoint whose eval has not written thresholds_span yet)
+SPAN_THR_FILE = json.load(open(os.environ["ARCTAG_SPAN_THR"])) if os.environ.get("ARCTAG_SPAN_THR") else {}
+
+# UI domain -> fine domain (thresholds are keyed by fine domains) + fallbacks for missing threshold cells
+UI_DOMAINS = {"auto": ["unknown"], "talk": ["talk", "stream_chat", "news"], "meeting": ["meeting"],
+              "court": ["court_appellate", "court_trial"], "assistant": ["assistant", "assistant_noisy"],
+              "sports": ["nba", "nhl", "nfl", "soccer", "cricket"], "news": ["news", "talk"]}
+
+# fine engine domain within a vertical (FiLM / adapters / validity / thresholds are per fine domain; the UI only picks the vertical):
+# lexicon cues over the words so far, decided once the best domain has FINE_MIN hits and FINE_RATIO x the runner-up
+FINE_CUES = {
+    "sports": {
+        "nba": "basketball nba rebound rebounds dunk layup three-pointer threes jumper assist assists paint free-throw buzzer "
+               "lakers celtics warriors knicks bulls heat nuggets bucks suns mavericks clippers rockets spurs nets sixers raptors "
+               "dribble backcourt fourth-quarter timeout",
+        "nhl": "hockey nhl puck goalie power-play penalty-kill slapshot slap faceoff icing blue-line rink skate skates "
+               "bruins rangers penguins maple leafs canadiens blackhawks oilers flames capitals lightning stick",
+        "nfl": "nfl football touchdown quarterback yard yards field-goal end-zone interception sack sacks fumble punt kickoff "
+               "linebacker receiver tight-end scrimmage patriots cowboys packers chiefs eagles steelers 49ers broncos down-and",
+        "soccer": "soccer goalkeeper keeper offside midfield midfielder striker penalty-area corner free-kick pitch premier "
+                  "league champions-league fifa uefa header crossbar extra-time stoppage arsenal chelsea liverpool madrid "
+                  "barcelona united city tottenham",
+        "cricket": "cricket wicket wickets bowler bowling batsman batter innings over overs boundary six sixes lbw stumps "
+                   "crease spinner pace yorker bouncer test odi t20 ipl umpire maiden run-out century fifty",
+    },
+    "court": {
+        "court_appellate": "chief-justice justice justices petitioner petitioners respondent respondents argument amicus "
+                           "certiorari circuit appellate appeal brief briefs precedent statute statutory may-it-please "
+                           "constitutional rebuttal counsel",
+        "court_trial": "witness witnesses objection sustained overruled jury jurors exhibit cross-examination direct-examination "
+                       "redirect testimony testify testified sworn defendant plaintiff prosecution verdict sidebar approach "
+                       "bailiff deposition",
+    },
+}
+FINE_MIN, FINE_RATIO = 4, 2.0
+FINE_REDOMAIN = os.environ.get("ARCTAG_FINE_REDOMAIN", "1") != "0"
+
+
+def _fine_index():
+    ix = {}
+    for ui, d in FINE_CUES.items():
+        for fine, cues in d.items():
+            for c in cues.split():
+                ix.setdefault(ui, {})[tuple(c.replace("-", " ").split())] = fine
+    return ix
+
+
+FINE_IX = _fine_index()
+
+
+def fine_hint(ui, words, last=3000):
+    """the fine domain the words' lexicon cues point to within UI domain `ui` (sports / court), or None (undecided / other ui)"""
+    ix = FINE_IX.get(ui)
+    if not ix or not words:
+        return None
+    toks = [re.sub(r"[^a-z0-9]", "", w["w"].lower()) for w in words[-last:]]
+    cnt = {}
+    for i in range(len(toks)):
+        for L in (1, 2, 3):
+            f = ix.get(tuple(toks[i:i + L])) if i + L <= len(toks) else None
+            if f:
+                cnt[f] = cnt.get(f, 0) + 1
+    if not cnt:
+        return None
+    r = sorted(cnt.items(), key=lambda x: -x[1])
+    if r[0][1] >= FINE_MIN and (len(r) == 1 or r[0][1] >= FINE_RATIO * r[1][1]):
+        return r[0][0]
+    return None
+
+
+FAMILY = [("pii", ("safety.pii.",)), ("profanity", ("safety.profanity.", "safety.toxicity.")),
+          ("number", ("entity.number.", "entity.numeral")), ("entity", ("entity.mention.",))]
+SKIP_PREFIX = ("affect.",)                               # not trained in AT10M (DESIGN 12: masked)
+STATE_GROUPS = ["content.mode", "content.game_state", "content.legal_phase", "content.meeting_phase", "content.subject"]
+
+
+DECODE_FAM = [("number", ("entity.number.",)), ("entity", ("entity.mention.",)), ("pii", ("safety.pii.",)),
+              ("profanity", ("safety.profanity.",)), ("toxicity", ("safety.toxicity.",))]
+SPAN_EXCL = os.environ.get("ARCTAG_SPAN_EXCL", "1") != "0"   # type-exclusive span decode per family (ARCTAG_SPAN_EXCL=0: per-tag runs)
+
+
+UMBRELLA = ("entity.numeral", "safety.toxicity.toxic_any")    # co-labelled WITH a subtype in the gold (82 of 253 news number spans
+                                                             # overlap a numeral): never merged, never compete with the subtypes
+
+
+def decode_family(tag):
+    if tag in UMBRELLA:
+        return None
+    for fam, pre in DECODE_FAM:
+        if tag.startswith(pre):
+            return fam
+    return None
+
+
+def family_of(tag):
+    for fam, pre in FAMILY:
+        if tag.startswith(pre):
+            return fam
+    return "other"
+
+
+def list_steps(d=CKPT_DIR):
+    out = []
+    for p in glob.glob(os.path.join(d, "step*")):
+        m = re.fullmatch(r"step(\d+)", os.path.basename(p))
+        f = os.path.join(p, "model.safetensors")
+        if m and os.path.exists(f):
+            out.append((int(m.group(1)), p))
+    return sorted(out)
+
+
+PROD_THR_DIR = os.environ.get("ARCTAG_PROD_THR", os.path.join(HERE, "prod_thresholds"))
+DISPLAY_TIER = os.environ.get("ARCTAG_DISPLAY_TIER", "T0")
+AUTO_REDOMAIN = os.environ.get("ARCTAG_AUTO_REDOMAIN", "1") != "0"   # auto: re-condition the engine on the detected vertical's domain   # tier the demo decodes spans / events / states from by default
+SAFETY_DOMS = {"talk", "stream_chat", "short_clips", "esports"}   # fine domains of the safety verticals (contact centre, T&S, gaming)
+CLAMP_SPAN = 0.3                  # display floor of span_beg / span_in thresholds (T0-fitted values of rare tags can be tiny)
+CLAMP_TOX = 0.5                   # safety.toxicity.* span / utt floor outside the safety verticals (news, sports, meetings, legal, auto)
+
+
+SAME_VERT = [{"nba", "nhl", "nfl", "cricket", "soccer"}, {"court_appellate", "court_trial"}, {"assistant", "assistant_noisy"}]
+THR_FALLBACK = os.environ.get("ARCTAG_THR_FALLBACK", "pooled")     # pooled (default) | sibling (pre-19:05 behaviour)
+
+
+def lookup_thr(v, doms):
+    """one threshold cell {fine domain | 'unknown' | '_all': thr} for the session's fine domains (doms[0] = primary).
+    Tables with a pooled '_all' (the demo's T0 refit): the primary domain or a same-vertical sibling (sports leagues, the two
+    court domains, assistant / assistant_noisy), else the POOLED value (a domain with < 10 dev positives has no key; a sibling
+    from another vertical, e.g. news -> talk, is not a valid stand-in). Other tables: doms -> 'unknown' -> median."""
+    if not v:
+        return None
+    if THR_FALLBACK == "pooled" and "_all" in v:
+        d0 = doms[0] if doms else "unknown"
+        sib = next((g for g in SAME_VERT if d0 in g), {d0})
+        for d in [d0] + [d for d in doms if d in sib and d != d0]:
+            if d in v:
+                return float(v[d])
+        return float(v["_all"])
+    for d in list(doms) + ["unknown", "_all"]:
+        if d in v:
+            return float(v[d])
+    return float(np.median(list(v.values())))
+
+
+def clamp_thr(col, doms, x):
+    """display clamps on a looked-up threshold (doms = the fine domains of the session, first = primary)"""
+    tag, _, q = col.rpartition("@")
+    if q in ("span_in", "span_beg"):
+        x = max(x, CLAMP_SPAN)
+    if tag.startswith("safety.toxicity.") and (not doms or doms[0] not in SAFETY_DOMS):
+        x = max(x, CLAMP_TOX)
+    return x
+
+
+def load_prod_thresholds(step, d=None):
+    """demo/prod_thresholds/step<N>.json (T0-fitted on dev, refit_thr) for the newest N <= step -> (thresholds, N) or ({}, None)"""
+    d = d or PROD_THR_DIR
+    best = None
+    for p in glob.glob(os.path.join(d, "step*.json")):
+        m = re.fullmatch(r"step(\d+)\.json", os.path.basename(p))
+        if m and int(m.group(1)) <= step and (best is None or int(m.group(1)) > best[0]):
+            best = (int(m.group(1)), p)
+    if best is None:
+        return {}, None
+    try:
+        return json.load(open(best[1])).get("thresholds") or {}, best[0]
+    except Exception:
+        return {}, None
+
+
+def load_event_tiers(step, d=None):
+    """demo/prod_thresholds/events_step<N>.json (newest N <= step; tests ev_tier fit): {"tier": {tag: T0|T1}, "hide": [tags],
+    "thresholds": {tag@q: {domain|_all: thr}} fitted AT each tag's tier} -> (doc, N) or ({}, None)"""
+    d = d or PROD_THR_DIR
+    best = None
+    for p in glob.glob(os.path.join(d, "events_step*.json")):
+        m = re.fullmatch(r"events_step(\d+)\.json", os.path.basename(p))
+        if m and int(m.group(1)) <= step and (best is None or int(m.group(1)) > best[0]):
+            best = (int(m.group(1)), p)
+    if best is None:
+        return {}, None
+    try:
+        return json.load(open(best[1])), best[0]
+    except Exception:
+        return {}, None
+
+
+class Thresholds:
+    """thr(tag@q, ui_domain): the demo's T0-fitted table (prod_thresholds/step<=N>.json) -> the step's own per-domain value ->
+    the newest older step that has thresholds; within each table: the fine domains of the UI domain -> 'unknown' -> '_all'
+    (pooled) -> median over the domains it has; default 0.5. Display clamps (clamp_thr) on top."""
+    def __init__(self, own, older=None, older_step=None, default=0.5, pref=None, pref_step=None):
+        self.own = own or {}; self.older = older or {}; self.older_step = older_step; self.default = default
+        self.pref = pref or {}; self.pref_step = pref_step
+        self.ev = {}; self.ev_step = None                   # per-event display tier + tier-fitted event thresholds (load_event_tiers)
+
+    def set_events(self, doc, step):
+        self.ev = doc or {}; self.ev_step = step
+        if self.ev.get("thresholds"):                        # event thresholds fitted at the event's tier take precedence
+            self.pref = dict(self.pref, **self.ev["thresholds"])
+
+    def ev_tier(self, tag, default="T0"):
+        return (self.ev.get("tier") or {}).get(tag, default)
+
+    def ev_hidden(self, tag):
+        return tag in set(self.ev.get("hide") or ())
+
+    @staticmethod
+    def _look(th, col, doms):
+        return lookup_thr(th.get(col), doms)
+
+    def tables(self):
+        return [t for t in (self.pref, self.own, self.older) if t]
+
+    def __call__(self, col, ui_dom):
+        doms = list(ui_dom) if isinstance(ui_dom, (list, tuple)) else UI_DOMAINS.get(ui_dom, ["unknown"])
+        x = None
+        for th in self.tables():
+            x = self._look(th, col, doms)
+            if x is not None:
+                break
+        return clamp_thr(col, doms, self.default if x is None else x)
+
+    def source(self):
+        s = f"prod_thresholds/step{self.pref_step} (T0) + " if self.pref_step is not None else ""
+        if self.ev_step is not None:
+            s += f"events_step{self.ev_step} + "
+        if self.own:
+            return s + "own"
+        return s + (f"step{self.older_step}" if self.older else "default 0.5")
+
+
+class Loaded:
+    def __init__(self, step, path, model, table, thr, cfg):
+        self.step, self.path, self.model, self.table, self.thr, self.cfg = step, path, model, table, thr, cfg
+        self.GI = GroupIndex(table)
+        self.has_kb = hasattr(model, "ent") and hasattr(model.ent, "ptr_logits")
+        self.has_head = hasattr(model, "span_head")
+        self.arch = tuple(sorted(k for k in ("span_head", "gat", "dattn", "dmix", "kb_spot", "pii_imm", "spk_link") if hasattr(model, k)))
+        self.head_tags = {f: list(v[2]) for f, v in model.span_head.fm.items() if v[0]} if self.has_head else {}
+        self.head_use, self.head_why = self._head_families()
+        self.loaded_at = time.time()
+
+    def _head_families(self):
+        """families decoded from the span head (see SPAN_HEAD_MODE) + why"""
+        if not self.head_tags or SPAN_HEAD_MODE == "0":
+            return set(), "off"
+        if SPAN_HEAD_MODE == "1":
+            return set(self.head_tags), "all"
+        f1 = SPAN_THR_FILE.get("dev_exact_f1", {}).get("T0") if SPAN_THR_FILE.get("dev_exact_f1") else None
+        if f1 is None:
+            try:
+                sp = json.load(open(os.path.join(self.path, "eval.json")))["dev"]["span_f1"]["spans"]["T0"]
+                f1 = {k: (v.get("all") or {}).get("exact_f1") for k, v in sp.items()}
+            except Exception:
+                f1 = None
+        if not f1:
+            return set(self.head_tags), "all (no dev F1)"
+        use = {f for f in self.head_tags if f1.get("head:" + f) is not None and (f1.get("fam:" + f) is None or f1["head:" + f] >= f1["fam:" + f] - HEAD_TOL)}
+        return use, {f: (f1.get("head:" + f), f1.get("fam:" + f)) for f in self.head_tags}
+
+    def head_thr(self, tier, fam):
+        """v2 span-head threshold on 1 - P(O) for a family (config.json thresholds_span {tier: {head:<family>: th}}, dev-fitted
+        by the trainer's eval), tier -> T0 -> 0.5"""
+        ts = (self.cfg or {}).get("thresholds_span") or SPAN_THR_FILE.get("thresholds_span") or {}
+        for t in (tier, "T0"):
+            v = (ts.get(t) or {}).get("head:" + fam)
+            if v is not None:
+                return float(v)
+        return 0.5
+
+
+class ModelManager:
+    """loads the newest complete step; a watcher thread polls for newer ones (file size stable + loadable) and swaps them in."""
+    def __init__(self, ckpt_dir=CKPT_DIR, poll_s=20.0, log=print):
+        self.dir = ckpt_dir; self.poll = poll_s; self.log = log
+        self.table = TextTable(); self.table.load_E()
+        self.cur = None; self.lock = threading.Lock(); self.listeners = []; self._sizes = {}
+        steps = list_steps(ckpt_dir)
+        if not steps:
+            raise RuntimeError(f"no step*/model.safetensors in {ckpt_dir}")
+        for st, p in reversed(steps):                      # newest that loads
+            try:
+                self.cur = self._load(st, p); break
+            except Exception as e:                         # partially written -> try older
+                self.log(f"[ckpt] step{st} not loadable yet: {e!r}")
+        if self.cur is None:
+            raise RuntimeError("no loadable checkpoint")
+        self.log(f"[ckpt] loaded step{self.cur.step} (thresholds: {self.cur.thr.source()})"
+                 + (f" span head: {sorted(self.cur.head_use)} ({self.cur.head_why})" if self.cur.has_head else ""))
+
+    def _thresholds(self, step):
+        own = {}
+        cj = os.path.join(self.dir, f"step{step}", "config.json")
+        try:
+            own = json.load(open(cj)).get("thresholds") or {}
+        except Exception:
+            pass
+        older, ostep = {}, None
+        if not own:
+            for st, p in reversed(list_steps(self.dir)):
+                if st >= step:
+                    continue
+                try:
+                    th = json.load(open(os.path.join(p, "config.json"))).get("thresholds") or {}
+                except Exception:
+                    th = {}
+                if th:
+                    older, ostep = th, st; break
+        pref, pstep = load_prod_thresholds(step)
+        t = Thresholds(own, older, ostep, pref=pref, pref_step=pstep)
+        t.set_events(*load_event_tiers(step))
+        return t
+
+    def _load(self, step, path):
+        from safetensors.torch import load_file
+        cfg = {}
+        try:
+            cfg = json.load(open(os.path.join(path, "config.json")))
+        except Exception:
+            pass
+        sc = cfg.get("scratch") or AT.SCRATCH_10M
+        m = AT.ArcText(self.table, scratch=sc)
+        sd = load_file(os.path.join(path, "model.safetensors"))
+        SA.add_modules_for(m, sd)                           # every stage-2 / v2 module the keys imply (kb_long, redact, spk_link, v2)
+        missing, unexpected = m.load_state_dict(sd, strict=False)
+        if unexpected:
+            raise RuntimeError(f"step{step}: checkpoint keys without a module: {unexpected[:6]}")
+        bad = [k for k in missing if not k.startswith(("head.E",))]
+        if bad:
+            self.log(f"[ckpt] step{step}: {len(bad)} missing keys (e.g. {bad[:3]})")
+        m.eval()
+        return Loaded(step, path, m, self.table, self._thresholds(step), cfg)
+
+    def refresh_thresholds(self):
+        """the async eval writes thresholds into config.json after the step appears: pick them up"""
+        c = self.cur
+        t = self._thresholds(c.step)
+        if (bool(t.own) != bool(c.thr.own)) or (t.older_step != c.thr.older_step) or (t.pref_step != c.thr.pref_step) \
+                or (t.ev_step != c.thr.ev_step):
+            c.thr = t; return True
+        return False
+
+    def check_new(self):
+        steps = list_steps(self.dir)
+        changed = self.refresh_thresholds()
+        if not steps or steps[-1][0] <= self.cur.step:
+            return changed
+        st, p = steps[-1]
+        f = os.path.join(p, "model.safetensors")
+        sz = os.path.getsize(f); prev = self._sizes.get(st); self._sizes[st] = sz
+        if prev != sz or time.time() - os.path.getmtime(f) < 30:
+            return changed                                 # still being written: wait for a stable size
+        if self._sizes.get(("arch", st)):
+            return changed                                 # already rejected (architecture change): do not reload it every poll
+        try:
+            new = self._load(st, p)
+        except Exception as e:
+            self.log(f"[ckpt] step{st} load failed: {e!r}"); return changed
+        if new.arch != self.cur.arch and not ARCH_SWITCH:
+            if self._sizes.get(("arch", st)) is None:
+                self.log(f"[ckpt] step{st}: architecture {new.arch} != running {self.cur.arch}: NOT hot-reloaded (ARCTAG_ARCH_SWITCH=1 + restart)")
+            self._sizes[("arch", st)] = 1
+            return changed
+        with self.lock:
+            self.cur = new
+        self.log(f"[ckpt] hot-reloaded step{st} (thresholds: {new.thr.source()})")
+        for cb in list(self.listeners):
+            try:
+                cb(new)
+            except Exception as e:
+                self.log(f"[ckpt] listener error {e!r}")
+        return True
+
+    def watch(self, stop_evt, on_change=None):
+        while not stop_evt.wait(self.poll):
+            try:
+                if self.check_new() and on_change:
+                    on_change(self.cur)
+            except Exception as e:
+                self.log(f"[ckpt] watch error {e!r}")
+
+
+# ---------------------------------------------------------------------------------------------------------------- session
+class ArcSession:
+    """one live stream: provider-agnostic events (adapters.py schema) -> FastArc -> UI updates."""
+    SPAN_WIN = 96                     # words re-decoded per update (revisions never reach further back than FastArc.KEEP = 64)
+
+    def __init__(self, loaded, ui_domain="auto", asr="deepgram", kb=None, int8=False, floor=0.0):
+        self.ui_domain = ui_domain if ui_domain in UI_DOMAINS else "auto"
+        self.fine = UI_DOMAINS[self.ui_domain][0]; self.asr = asr; self.int8 = int8; self.floor = float(floor or 0.0)
+        self.kb_names = [k.strip() for k in (kb or []) if k and k.strip()][:256]
+        self.words = []              # provider word dicts (ours; mirror of engine.words)
+        self.utt_bounds = []         # word indices where an utterance ended (endpoint at k)
+        self.events = []; self._ev_keys = set(); self.alerts = []; self._al_keys = set()
+        self.utt_final = []          # finished utterances: {w0, w1, readouts}
+        self.lat = []                # (ms, n_words) per engine update
+        self._spksrc = None          # stage-2 / v2 input "spksrc" (0 no speaker ids, 1 gold, 2 ASR-like): set at the first words
+        self._bind(loaded)
+
+    # ---- model binding (also used by hot reload: rebuild + replay)
+    def _bind(self, loaded):
+        self.L = loaded; T = loaded.table
+        self.e = FA.FastArc(loaded.model, T, domain=self.fine, asr=self.asr)
+        if self.int8:
+            FA.quantize_fast(self.e)
+        self._set_xfeat()
+        self.store = {"T0": {}, "T1": {}, "T2": {}}   # (word) -> {q: probs np}
+        self.pooled_now = None
+        # column metadata per read-out
+        self.cn = {q: [n.split("@")[0] for n in T.col_names(q)] for q in ("frame", "span_beg", "span_in", "span_end", "ev_at", "ev_imm", "ev_rec", "utt", "speaker", "item")}
+        thr = loaded.thr; d = self.doms
+        self.th_raw = {q: np.array([thr(f"{t}@{q}", d) for t in names], np.float32) for q, names in self.cn.items()}
+        self.th = {q: np.maximum(v, self.floor) for q, v in self.th_raw.items()}      # optional display floor (UI), 0 = pure per-domain
+        valid = T.valid[domain_index(self.fine)]
+        self.valid = {q: np.array([bool(valid[T.idx[t]]) for t in names]) for q, names in self.cn.items()}
+        bi = {t: i for i, t in enumerate(self.cn["span_beg"])}
+        self.beg_of_in = np.array([bi.get(t, -1) for t in self.cn["span_in"]])
+        norm = {T.ids[i] for i in T.norm_tags}
+        self.is_norm = np.array([t in norm for t in self.cn["span_in"]])
+        self.kbe = None
+        if self.kb_names and loaded.has_kb:
+            with torch.inference_mode():
+                self.kbe = loaded.model.kb_embed(self.kb_names)
+
+    def _set_xfeat(self):
+        """session-constant stage-2 inputs: spksrc as the trainer's data.Windows sets it (2 = ASR-like ids, the live providers'
+        diarization; 0 = no ids) and turn = 0 (no Flux turn events). Without it FastArc embeds spksrc 0 for every word while the
+        model saw 2 on ASR sessions. ARCTAG_SPKSRC overrides (parity on gold-id sessions: 1)."""
+        if self._spksrc is not None and ("spksrc" in self.e.feat or "turn" in self.e.feat):
+            self.e.xfeat = {"spksrc": self._spksrc, "turn": 0}
+
+    def set_kb(self, names):
+        self.kb_names = [k.strip() for k in (names or []) if k and k.strip()][:256]
+        self.kbe = None
+        if self.kb_names and self.L.has_kb:
+            with torch.inference_mode():
+                self.kbe = self.L.model.kb_embed(self.kb_names)
+
+    def rebind(self, loaded):
+        """hot reload: new FastArc on the new weights, replay the session utterance by utterance (one chunk pass each)"""
+        ws = [dict(w) for w in self.words]; ub = list(self.utt_bounds)
+        self.utt_bounds = []
+        self._bind(loaded); self.words = ws
+        t = time.perf_counter(); lo = 0
+        for k in ub + [len(ws)]:
+            if k > lo:
+                self._absorb(self.e.update(lo, ws[lo:k]))
+            if k in ub:
+                self._absorb(self.e.endpoint()); self.utt_bounds.append(k)
+            lo = k
+        self.words = ws
+        return round(1000 * (time.perf_counter() - t), 1)
+
+    @property
+    def doms(self):
+        """threshold domains: the engine's fine domain first, then the UI domain's other codes"""
+        return [self.fine] + [d for d in UI_DOMAINS[self.ui_domain] if d != self.fine]
+
+    def set_domain(self, ui_domain, fine=None):
+        """re-condition the engine on another UI domain (auto mode, after the vertical detector locks): FastArc's FiLM /
+        adapters / validity are per fine domain and the model never sees 'unknown' for domain-labelled data, so 'auto' running on
+        'unknown' all session shifts every read-out (parity: mean |dp| frame .11, utt max .99). Rebuild + replay (= hot reload).
+        Returns the replay ms, or None when unchanged."""
+        ui = ui_domain if ui_domain in UI_DOMAINS else "auto"
+        fine = fine if fine in UI_DOMAINS[ui] else UI_DOMAINS[ui][0]
+        if ui == self.ui_domain and fine == self.fine:
+            return None
+        self.ui_domain = ui; self.fine = fine
+        return self.rebind(self.L)
+
+    def _update(self, k, ws):
+        """FastArc.update with a snapshot AT k. FastArc snapshots only at chunk starts, so a revision from k replays words[s0:k] +
+        ws as one chunk from the nearest older snapshot s0 and leaves none at k. Streams that keep revising from about the same
+        point (Flux re-sends the whole turn every Update; nova-3 the whole interim segment) then grow that chunk on every message
+        until s0 falls out of the KEEP window, and from then on every update is a full-session replay from word 0 (O(session)
+        per message). Running the unchanged prefix words[s0:k] as its own chunk first puts a snapshot at k: revisions stay
+        O(revised words). Same outputs (chunked == sequential), both read-out batches are returned."""
+        e = self.e
+        if k < e.n:
+            s0 = max([x for x in e.snap if x <= k], default=None)
+            if s0 is None or s0 < k:
+                d1 = dict(e.update(k, []))
+                d2 = e.update(k, ws)
+                d1.update(d2)
+                return d1
+        return e.update(k, ws)
+
+    # ---- event application (= adapters.TaggerSession.apply, with timing + UI materialisation)
+    def apply(self, evs):
+        """returns (lo, changed) where lo = lowest word index whose display may have changed, or None"""
+        lo = None; utt_done = []
+        for x in evs:
+            t = x["type"]
+            if t == "words":
+                k = x["from"]; ws = x["words"]; j = 0
+                while j < len(ws) and k + j < len(self.e.words) and ws[j] == self.e.words[k + j]:
+                    j += 1
+                if j == len(ws) and k + j == len(self.e.words):
+                    continue
+                k2 = k + j
+                if self._spksrc is None and "spksrc" in self.e.feat:
+                    ov = os.environ.get("ARCTAG_SPKSRC")
+                    self._spksrc = int(ov) if ov else (2 if any(w.get("spk") is not None for w in ws) else 0)
+                    self._set_xfeat()
+                t0 = time.perf_counter()
+                due = self._update(k2, ws[j:])
+                self.lat.append((1000 * (time.perf_counter() - t0), max(1, len(self.e.words) - min(k2, len(self.words)))))
+                self.words = list(self.e.words)
+                self._prune(); self._absorb(due); self._running()
+                lo = k2 if lo is None else min(lo, k2)
+            elif t == "finalize":
+                ws = self.e.words; up = min(x["upto"], len(ws))
+                first = next((i for i in range(up) if not ws[i].get("final", True)), None)
+                if first is not None:
+                    t0 = time.perf_counter()
+                    due = self._update(first, [dict(w, final=True) for w in ws[first:up]] + ws[up:])
+                    self.lat.append((1000 * (time.perf_counter() - t0), max(1, len(ws) - first)))
+                    self.words = list(self.e.words)
+                    self._prune(); self._absorb(due); self._running()
+                    lo = first if lo is None else min(lo, first)
+            elif t == "endpoint":
+                if not self.e.n or (self.utt_bounds and self.utt_bounds[-1] >= self.e.n):
+                    continue
+                due = self.e.endpoint()
+                self._absorb(due)
+                s = self.utt_bounds[-1] if self.utt_bounds else 0
+                self.utt_bounds.append(self.e.n)
+                ro = self.readouts(self.e.n - 1, final=True)
+                u = {"w0": s, "w1": self.e.n, "readouts": ro, "text": " ".join(w["w"] for w in self.words[s:self.e.n])}
+                self.utt_final.append(u); utt_done.append(u)
+                self._alerts_utt(u)
+                lo = s if lo is None else min(lo, s)
+        return lo, utt_done
+
+    def _prune(self):
+        for tier, st in self.store.items():
+            have = self.e.out[tier]
+            for i in [i for i in st if i not in have or i >= self.e.n]:
+                del st[i]
+
+    def _absorb(self, due):
+        GI = self.L.GI
+        for key, lg in due.items():
+            tier, i = key
+            if tier == "pooled":
+                continue
+            P = {}
+            for q, z in lg.items():
+                if q == "span_head":                       # v2 per-family (type, BIO) logits -> softmax probabilities
+                    P[q] = {f: (zt.softmax(-1).numpy(), zb.softmax(-1).numpy()) for f, (zt, zb) in z.items()}
+                    continue
+                P[q] = GI.probs(q, z[None])[0].numpy() if z.numel() else np.zeros(0, np.float32)
+            self.store[tier][i] = P
+            self._events(tier, i, P)
+
+    def _running(self):
+        if self.e.n:
+            self.pooled_now = self.e.n - 1
+
+    # ---- UI materialisation
+    def spans(self, tier, lo=None):
+        """decode W-tag spans of `tier` over the last SPAN_WIN words: runs of span_in >= thr, split where span_beg >= thr
+        (a new B inside a run). Returns (a0, spans): every span with end > a0 (the client replaces its spans with end > a0)."""
+        st = self.store[tier]; n = self.e.n
+        a0 = max(0, n - self.SPAN_WIN)
+        idx = [i for i in range(max(0, a0 - 16), n) if i in st]
+        if not idx:
+            return a0, []
+        ON = np.stack([st[i]["span_in"] for i in idx]); PIN = ON
+        ON = (ON >= self.th["span_in"]) & self.valid["span_in"]
+        bo = self.beg_of_in; hb = bo >= 0
+        PB = np.stack([st[i]["span_beg"] for i in idx])
+        BEG = np.zeros_like(ON)
+        BEG[:, hb] = PB[:, bo[hb]] >= self.th["span_beg"][bo[hb]]
+        us = np.array([bool(self.e.ustart[i]) if i < len(self.e.ustart) else False for i in idx])
+        BEG[us] = True                                                       # spans never cross an utterance start
+        head = self._head_cols() if (self.L.has_head and SPAN_HEAD) else {}
+        if head:
+            ON[:, [c for cs in head.values() for c in cs]] = False             # head families: decoded from the span head below
+        res = []
+        for c in np.nonzero(ON.any(0))[0]:
+            a = None
+            for r in range(len(idx) + 1):
+                on = r < len(idx) and ON[r, c]
+                gap = r < len(idx) and r > 0 and idx[r] != idx[r - 1] + 1
+                if a is not None and (not on or gap or BEG[r, c]):
+                    w0, w1 = idx[a], idx[r - 1] + 1
+                    if w1 > a0:
+                        res.append((c, w0, w1, float(PIN[a:r, c].max())))
+                    a = None
+                if on and a is None:
+                    a = r
+        alt = {}
+        if SPAN_EXCL:
+            res, alt = self._exclusive(res, idx, PIN)
+        hth = {}
+        if head:
+            hres, alt2, hth = self._head_spans(tier, idx, us, a0, head)
+            res = res + hres; alt.update(alt2)
+        out = []
+        th_in = self.th["span_in"]
+        for c, a, b, p in res:
+            tag = self.cn["span_in"][c]
+            s = {"tag": tag, "fam": family_of(tag), "a": int(a), "b": int(b), "p": round(p, 3),
+                 "thr": round(float(hth.get(c, th_in[c])), 3), "text": " ".join(w["w"] for w in self.words[a:b])}
+            if c in hth:
+                s["src"] = "head"
+            if (c, a, b) in alt:
+                s["alt"] = alt[(c, a, b)]
+            if self.is_norm[c]:
+                s.update(self._verify(a, b))
+            out.append(s)
+        out.sort(key=lambda s: (s["a"], -s["b"]))
+        return a0, out
+
+    def _head_cols(self):
+        """{family: span_in columns of its VALID members} for the v2 span head (cached per binding)"""
+        hc = getattr(self, "_hc", None)
+        if hc is None or hc[0] is not self.L:
+            ci = {t: c for c, t in enumerate(self.cn["span_in"])}
+            m = {f: [ci[t] for t in tags if t in ci and self.valid["span_in"][ci[t]]] for f, tags in self.L.head_tags.items()
+                 if f in self.L.head_use}
+            hc = self._hc = (self.L, {f: cs for f, cs in m.items() if cs})
+        return hc[1]
+
+    def head_thr(self, tier, fam):
+        f_ = self.L.head_tags.get(fam) or [fam + ".x"]
+        return clamp_thr(f"{f_[0]}@span_in", self.doms, max(self.L.head_thr(tier, fam), self.floor))
+
+    def _head_spans(self, tier, idx, us, a0, head):
+        """v2 span head decode per family over the window: in-span = 1 - P(O) >= th (family threshold). T0 greedy (left to
+        right: a span starts at B > I, after a non-span word, a gap or an utterance start); T1 / T2 (read-outs already see the
+        delay window) Viterbi over O / B / I with O -> I forbidden, in-states biased by log((1 - th) / th) so a lone word follows
+        the same threshold. One type per span: argmax of the summed log type probabilities over the family's VALID members."""
+        st = self.store[tier]; ci = {t: c for c, t in enumerate(self.cn["span_in"])}
+        res, alt, hth = [], {}, {}
+        brk = np.array([bool(us[r]) or (r > 0 and idx[r] != idx[r - 1] + 1) for r in range(len(idx))])
+        for f, cols in head.items():
+            rows = [r for r, i in enumerate(idx) if "span_head" in st[i] and f in st[i]["span_head"]]
+            if len(rows) != len(idx):
+                continue
+            Pt = np.stack([st[i]["span_head"][f][0] for i in idx]); Pb = np.stack([st[i]["span_head"][f][1] for i in idx])
+            tags = self.L.head_tags[f]; mem = [k for k, t in enumerate(tags) if t in ci and ci[t] in cols]
+            th = self.head_thr(tier, f); pin = 1.0 - Pb[:, 0]
+            if tier == "T0":
+                lab = np.where(pin < th, 0, np.where(brk | (Pb[:, 1] > Pb[:, 2]), 1, 2))
+                for r in range(1, len(idx)):
+                    if lab[r] == 2 and lab[r - 1] == 0:
+                        lab[r] = 1
+                if len(idx):
+                    lab[0] = 1 if lab[0] else 0
+            else:
+                lab = self._viterbi(Pb, th, brk)
+            lt = np.log(np.clip(Pt[:, mem], 1e-9, 1.0))
+            r = 0; W = len(idx)
+            while r < W:
+                if lab[r] == 0:
+                    r += 1; continue
+                e = r + 1
+                while e < W and lab[e] == 2:
+                    e += 1
+                w0, w1 = idx[r], idx[e - 1] + 1
+                if w1 > a0:
+                    sc = lt[r:e].sum(0); o = np.argsort(-sc); c = ci[tags[mem[int(o[0])]]]
+                    res.append((c, w0, w1, float(pin[r:e].max()))); hth[c] = th
+                    if len(o) > 1:
+                        pm = Pt[r:e][:, mem].mean(0)
+                        alt[(c, w0, w1)] = {"tag": tags[mem[int(o[1])]], "p": round(float(pm[o[1]]), 3)}
+                r = e
+        return res, alt, hth
+
+    @staticmethod
+    def _viterbi(Pb, th, brk):
+        """BIO labels (0 O, 1 B, 2 I) maximising sum of log scores with O -> I and break -> I forbidden: O = log P(O), B / I =
+        log P(in) + log((1 - th) / th) + log(P(B or I) / max(P(B), P(I))) (a lone word is in-span exactly when P(in) >= th; B / I
+        only costs where the BIO constraint overrides the per-word choice)"""
+        W = len(Pb)
+        if not W:
+            return np.zeros(0, int)
+        pin = np.clip(1.0 - Pb[:, 0], 1e-9, 1.0); bias = math.log((1.0 - th) / max(th, 1e-6))
+        cb = np.clip(Pb[:, 1:], 1e-9, 1.0); cb = np.log(cb / cb.max(1, keepdims=True))      # B vs I, the better one costs 0
+        L = np.empty((W, 3)); L[:, 0] = np.log(np.clip(Pb[:, 0], 1e-9, 1.0)); L[:, 1:] = (np.log(pin) + bias)[:, None] + cb
+        NEG = -1e9; dp = np.full((W, 3), NEG); bp = np.zeros((W, 3), int)
+        dp[0, 0], dp[0, 1] = L[0, 0], L[0, 1]
+        for r in range(1, W):
+            p = dp[r - 1]
+            for s_ in (0, 1):
+                k = int(p.argmax()); dp[r, s_] = p[k] + L[r, s_]; bp[r, s_] = k
+            if not brk[r]:
+                k = 1 if p[1] >= p[2] else 2; dp[r, 2] = p[k] + L[r, 2]; bp[r, 2] = k
+        lab = np.zeros(W, int); lab[-1] = int(dp[-1].argmax())
+        for r in range(W - 1, 0, -1):
+            lab[r - 1] = bp[r, lab[r]]
+        return lab
+
+    def _exclusive(self, res, idx, PIN):
+        """type-exclusive decode: overlapping spans of one family (number / entity / pii / profanity / toxicity) merge into ONE span
+        over their union, typed by the argmax over the family's valid columns of the mean span_in over the span's words; the
+        runner-up type is kept ("alt"). Spans of different families may still overlap (pii.private_name + entity.mention.person)."""
+        fam_cols = getattr(self, "_fam_cols", None)
+        if fam_cols is None:
+            fam_cols = self._fam_cols = {}
+            for c, t in enumerate(self.cn["span_in"]):
+                f = decode_family(t)
+                if f and self.valid["span_in"][c]:
+                    fam_cols.setdefault(f, []).append(c)
+        row = {w: r for r, w in enumerate(idx)}
+        by = {}; keep = []
+        for x in res:
+            f = decode_family(self.cn["span_in"][x[0]])
+            (by.setdefault(f, []) if f else keep).append(x)
+        out = list(keep); alt = {}
+        for f, xs in by.items():
+            xs.sort(key=lambda x: (x[1], x[2])); groups = []
+            for x in xs:
+                if groups and x[1] < groups[-1][1]:
+                    groups[-1][1] = max(groups[-1][1], x[2]); groups[-1][2].append(x)
+                else:
+                    groups.append([x[1], x[2], [x]])
+            cols = fam_cols.get(f, [])
+            for a, b, mem in groups:
+                if len(mem) == 1 or not cols:
+                    out.append(mem[0]); continue
+                rr = [row[w] for w in range(a, b) if w in row]
+                M = PIN[rr][:, cols].mean(0)
+                o = np.argsort(-M); c = cols[int(o[0])]
+                out.append((c, a, b, float(PIN[rr, c].max())))
+                if len(o) > 1:
+                    c2 = cols[int(o[1])]
+                    alt[(c, a, b)] = {"tag": self.cn["span_in"][c2], "p": round(float(M[o[1]]), 3)}
+        return out, alt
+
+    @torch.inference_mode()
+    def _verify(self, a, b):
+        """span verification (+ KB pointer), cached per span: the key holds the word dicts of [a, b) by identity (FastArc
+        re-creates them whenever it re-runs those words), so a revision or a re-bind invalidates the entry."""
+        e = self.e; ws = e.words[a:b]
+        c = getattr(self, "_vcache", None)
+        if c is None or c[0] is not e or c[1] is not self.kbe:
+            c = self._vcache = (e, self.kbe, {})
+        hit = c[2].get((a, b))
+        if hit is not None and len(hit[0]) == len(ws) and all(x is y for x, y in zip(hit[0], ws)):
+            return dict(hit[1])
+        out = self._verify_raw(a, b)
+        if len(c[2]) > 4000:
+            c[2].clear()
+        c[2][(a, b)] = (tuple(ws), out)
+        return dict(out)
+
+    def _verify_raw(self, a, b):
+        m = self.L.model; H = self.e.Hbuf
+        sp = torch.tensor([[0, b - a]])                 # local window: same mean / end state as over H[:n], O(span) not O(session)
+        rep = m.span_rep(H[a:b], sp)
+        out = {"ver": round(float(torch.sigmoid(m.ver(rep)[0, 0])), 3)}
+        if self.kbe is not None:
+            z = m.kb_logits(H[b - 1][None], self.kbe)[0].float()
+            p = torch.softmax(z, -1)
+            j = int(p[2:].argmax()) if p.numel() > 2 else -1
+            out["kb"] = {"p_new": round(float(p[0]), 3), "p_generic": round(float(p[1]), 3)}
+            if j >= 0:
+                out["kb"].update(entry=self.kb_names[j], p=round(float(p[2 + j]), 3), match=bool(p[2 + j] > max(p[0], p[1])))
+        return out
+
+    def _events(self, tier, i, P):
+        if tier == "T2":
+            return
+        w = self.words[i] if i < len(self.words) else {}
+        for q in ("ev_at", "ev_imm", "ev_rec"):
+            p = P.get(q)
+            if p is None or not len(p):
+                continue
+            on = np.nonzero((p >= self.th[q]) & self.valid[q])[0]
+            for c in on:
+                tag = self.cn[q][c]
+                if any((tier, tag, j) in self._ev_keys for j in range(i - 3, i + 1)):
+                    continue
+                self._ev_keys.add((tier, tag, i))
+                self.events.append({"tier": tier, "tag": tag, "q": q, "i": int(i), "p": round(float(p[c]), 3),
+                                    "thr": round(float(self.th[q][c]), 3), "t": w.get("t1"), "wall": time.time(),
+                                    "word": w.get("w", "")})
+
+    def _alerts_utt(self, u):
+        for r in u["readouts"].get("safety", []):
+            k = (r["tag"], u["w0"])
+            if r["on"] and k not in self._al_keys:
+                self._al_keys.add(k)
+                self.alerts.append({"tag": r["tag"], "p": r["p"], "thr": r["thr"], "w0": u["w0"], "w1": u["w1"],
+                                    "text": u["text"][-160:], "wall": time.time(), "kind": "utterance"})
+
+    def span_alerts(self, spans, tier):
+        for s in spans:
+            if s["fam"] in ("pii", "profanity"):
+                k = (tier, s["tag"], s["a"])
+                if k not in self._al_keys:
+                    self._al_keys.add(k)
+                    self.alerts.append({"tag": s["tag"], "p": s["p"], "thr": s["thr"], "w0": s["a"], "w1": s["b"], "text": s["text"],
+                                        "wall": time.time(), "kind": f"span {tier}"})
+
+    def _top(self, q, P, prefix, k=3):
+        names = self.cn[q]
+        idx = [c for c, t in enumerate(names) if t.startswith(prefix) and self.valid[q][c]]
+        idx.sort(key=lambda c: -P[c])
+        return [{"tag": names[c], "p": round(float(P[c]), 3), "thr": round(float(self.th[q][c]), 3), "on": bool(P[c] >= self.th[q][c])} for c in idx[:k]]
+
+    @torch.inference_mode()
+    def readouts(self, j=None, final=False, tier="T0"):
+        """utterance read-outs (pooled utt / speaker / item at word j) + frame states of the newest word of `tier`"""
+        if j is None:
+            j = self.pooled_now
+        if j is None or j not in self.e.pooled:
+            return {}
+        GI = self.L.GI
+        pl = self.e.pooled_logits(j)
+        U = GI.probs("utt", pl["utt"][None])[0].numpy(); S = GI.probs("speaker", pl["speaker"][None])[0].numpy()
+        I = GI.probs("item", pl["item"][None])[0].numpy()
+        sp = self.e.spk[j] if j < len(self.e.spk) else 0
+        ro = {"j": int(j), "final": final,
+              "intent": self._top("utt", U, "intent.", 3) + self._top("utt", U, "discourse.intent.", 3),
+              "act": self._top("utt", U, "discourse.act.", 3) + self._top("utt", U, "discourse.answer", 1) + self._top("utt", U, "discourse.backchannel", 1)
+                     + self._top("utt", U, "discourse.assistant.", 1),
+              "role": self._top("speaker", S, "speaker.role.", 3) if sp > 0 else [],
+              "role_note": None if sp > 0 else "speaker role needs diarization (no speaker id on words)",
+              "safety": [r for r in self._top("utt", U, "safety.", 20) if r["p"] >= 0.5 * r["thr"]][:6],
+              "item": self._top("item", I, "item.", 3)}
+        ro["intent"].sort(key=lambda r: -r["p"]); ro["intent"] = ro["intent"][:4]
+        fr = {}
+        jt = max([i for i in self.store[tier] if i <= j], default=None)
+        if jt is not None and "frame" in self.store[tier][jt]:
+            Fp = self.store[tier][jt]["frame"]
+            for g in STATE_GROUPS:
+                top = self._top("frame", Fp, g + ".", 3)
+                if top:
+                    fr[g] = top
+        ro["state"] = fr
+        return ro
+
+    def stats(self):
+        L = self.lat[-500:]
+        if not L:
+            return {}
+        per_upd = np.array([x[0] for x in L]); per_w = np.array([x[0] / x[1] for x in L])
+        return {"upd_p50": round(float(np.percentile(per_upd, 50)), 2), "upd_p95": round(float(np.percentile(per_upd, 95)), 2),
+                "word_p50": round(float(np.percentile(per_w, 50)), 2), "word_p95": round(float(np.percentile(per_w, 95)), 2),
+                "n_updates": len(self.lat)}

arctag/_core/fast_arc.py

@@ -0,0 +1,575 @@
+"""Fast streaming engine for arc_text.ArcText (same maths as stream_arc.ArcStream / the batch forward; parity-tested).
+
+What makes it fast (vs stream_arc.py, which calls the arc_audio modules per word):
+  (a) Mamba-3 in true recurrent mode for any T >= 1: ALL projections of a layer (in, x, B, C, theta, dt, lambda) are ONE fused
+      matmul over the T new words; the state recurrence is a T-step loop on (H, N, P) tensors; out_proj / MLP batched over T.
+      FiLM (domain + delay 0 constants) and the trunk input bias are folded per session; Trunk.step recomputed them every word.
+  (b) a revision of k words = ONE pass over the k-word chunk from the restored state (not k single steps).
+  (c) local attention: per layer a preallocated ring buffer of the window's K / V (32 words); only the new words' q / k / v.
+  (d) read-outs: Q_q and the tag keys are folded into one (d x C) matrix per read-out group (only the text tag columns), so the
+      T0 read of word t and the T1 read of word t-3 (both read h_t, differ by a delay-embedding constant) cost one matvec;
+      utt / speaker / item and T2 are computed only at an endpoint (batched); the pooled carries are updated every word.
+  (e) preallocated buffers, torch.inference_mode, no module-call overhead.
+Snapshots for revision: per word the (state, carry, ring, slots) references of the last KEEP words.
+"""
+import argparse, bisect, json, math, os, sys, time
+import numpy as np, torch, torch.nn.functional as F
+
+HERE = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, HERE)
+import arc_text as AT                                                       # noqa: E402
+from arc_text import WORD_Q, POOLED, D_MAX, P_MAX, NEG                      # noqa: E402
+from arc_audio.mamba3.layer import rotate_pairs                             # noqa: E402
+from tagtable import domain_index, asr_index                                # noqa: E402
+from losses import GroupIndex                                               # noqa: E402
+import stream_arc as SA                                                     # noqa: E402
+
+KEEP = 64
+WQ = [q for q in WORD_Q if q != "span_beg"]
+
+
+class FastArc:
+    @torch.inference_mode()
+    def __init__(self, model, table, domain=None, asr=None, max_words=20000):
+        m = self.m = model.eval(); self.T = table; self.GI = GroupIndex(table)
+        self.asr_name = asr or "unknown"; self.asr = asr_index(asr); di = domain_index(domain)
+        self.dom = torch.tensor([di])
+        tr = m.trunk; d = m.d; self.d = d
+        self.tok = AT.ArcPieces(AT.BPE8K if getattr(model, "scratch", None) else AT.TOK_PATH); self.cache = {}
+        f32 = lambda t: t.detach().float().contiguous()
+        # ---- word input constants (no audio: hm = he = 0, have = 0)
+        w = m.word
+        self.tokE = w.tok_E; self.ppos = f32(w.ppos); self.W1 = f32(w.W1.weight).t().contiguous(); self.b1 = f32(w.W1.bias)
+        self.W2 = f32(w.W2.weight).t().contiguous(); self.ln = (f32(w.ln.weight), f32(w.ln.bias), w.ln.eps)
+        self.fl = [f32(e.weight) for e in w.fl]
+        src = 1 if self.asr_name == "arc_transcriber" else 0
+        self.c_in = f32(w.src.weight[src]) + ((f32(w.Wa.bias) + f32(w.hav.weight[0])) if hasattr(w, "Wa") else 0)
+        self.feat = {k: f32(e.weight) for k, e in m.feat.items()}
+        # ---- local attention layers
+        self.enc = []
+        for lay in w.enc:
+            self.enc.append(dict(n1=(f32(lay.n1.weight), f32(lay.n1.bias), lay.n1.eps), Wqkv=f32(lay.qkv.weight).t().contiguous(), bqkv=f32(lay.qkv.bias),
+                                 Wo=f32(lay.o.weight).t().contiguous(), bo=f32(lay.o.bias), n2=(f32(lay.n2.weight), f32(lay.n2.bias), lay.n2.eps),
+                                 F1=f32(lay.ff[0].weight).t().contiguous(), f1=f32(lay.ff[0].bias), F2=f32(lay.ff[3].weight).t().contiguous(),
+                                 f2=f32(lay.ff[3].bias), h=lay.h, win=lay.win))
+        self.enc_norm = (f32(w.enc_norm.weight), f32(w.enc_norm.bias), w.enc_norm.eps)
+        dw = w.ln.normalized_shape[0]; hE = self.enc[0]["h"]; dh = dw // hE; self.dh = dh; self.hE = hE; self.dw = dw
+        self.win = self.enc[0]["win"]
+        inv = 1.0 / (10000 ** (torch.arange(0, dh, 2, dtype=torch.float32) / dh)); self.rope_inv = inv
+        ang = torch.arange(max_words, dtype=torch.float32)[:, None] * inv          # RoPE table (positions x dh/2)
+        self.rope_c, self.rope_s = torch.cos(ang)[:, None, :], torch.sin(ang)[:, None, :]
+        # ---- trunk: input projection, FiLM / domain constants folded
+        di_ = tr.delay_lut[0]
+        e_ = f32(tr.dom_emb.weight[di]) + f32(tr.delay_film.weight[di_])
+        self.in_norm = (f32(tr.in_norm.weight), f32(tr.in_norm.bias), tr.in_norm.eps)
+        self.Winp = f32(tr.inp.weight).t().contiguous(); self.binp = f32(tr.inp.bias) + f32(tr.dom_emb.weight[di]) + f32(tr.delay_emb.weight[di_])
+        self.layers = []
+        for i, (lay, mlp) in enumerate(zip(tr.layers, tr.mlps)):
+            gb = f32(tr.film[i](e_)); g, bb = 1 + gb[:d], gb[d:]
+            H, N, P = lay.H, lay.N, lay.P
+            Ws = [lay.in_proj, lay.x_proj, lay.B_proj, lay.C_proj, lay.theta_proj, lay.dt_proj, lay.lam_proj]
+            Wf = torch.cat([f32(l.weight) for l in Ws], 0).t().contiguous(); bf = torch.cat([f32(l.bias) for l in Ws])
+            sizes = [l.out_features for l in Ws]
+            self.layers.append(dict(g=g, b=bb, nw=f32(lay.norm.weight), neps=lay.norm.eps, Wf=Wf, bf=bf, sizes=sizes, H=H, N=N, P=P,
+                                    A=-torch.exp(f32(lay.A_log)), Bn=(f32(lay.B_norm.weight), lay.B_norm.eps), Cn=(f32(lay.C_norm.weight), lay.C_norm.eps),
+                                    Bb=f32(lay.B_bias)[..., 0], Cb=f32(lay.C_bias)[..., 0],
+                                    BCw=torch.stack([f32(lay.B_norm.weight), f32(lay.C_norm.weight)])[:, None, :],
+                                    BCb=torch.stack([f32(lay.B_bias)[..., 0], f32(lay.C_bias)[..., 0]]), Wout=f32(lay.out_proj.weight).t().contiguous(), bout=f32(lay.out_proj.bias),
+                                    mn=(f32(mlp[0].weight), mlp[0].eps), M1=f32(mlp[1].weight).t().contiguous(), m1=f32(mlp[1].bias),
+                                    M2=f32(mlp[4].weight).t().contiguous(), m2=f32(mlp[4].bias)))
+        self.adD = f32(tr.adapt_down[di]); self.adU = f32(tr.adapt_up[di]); self.on = (f32(tr.out_norm.weight), tr.out_norm.eps)
+        # ---- AT10M-v2: global attention layers (KV memory + chunk summaries, v2.GlobalAttn), delayed attention, span head, utt pool
+        import v2 as V2; self.V2 = V2
+        self.gat = {}
+        for k_, G in (getattr(m, "gat", None) or {}).items():
+            self.gat[int(k_)] = dict(nw=f32(G.norm.weight), neps=G.norm.eps, Wqkv=f32(G.qkv.weight).t().contiguous(), qn=(f32(G.qn.weight), G.qn.eps),
+                                     kn=(f32(G.kn.weight), G.kn.eps), Wo=f32(G.o.weight).t().contiguous(), h=G.h, dh=G.dh, mem=G.mem, ch=G.chunk, slots=G.slots,
+                                     K=torch.zeros(max_words, G.h, G.dh), Kr=torch.zeros(max_words, G.h, G.dh), V=torch.zeros(max_words, G.h, G.dh))
+            if getattr(G, "biased", False):                     # v2.1: ALiBi slopes, span-bearing key prior, tanh output gate
+                self.gat[int(k_)].update(biased=True, slope=f32(G.lslope).exp(), kb=f32(G.kb), kmW=f32(G.kmark.weight)[0], kmb=float(G.kmark.bias),
+                                         gate=float(torch.tanh(G.gate.detach().float())), M=torch.zeros(max_words))
+                self.gat[int(k_)]["Wo"] = self.gat[int(k_)]["Wo"] * self.gat[int(k_)]["gate"]
+        if self.gat:
+            g0 = next(iter(self.gat.values()))
+            self.grc, self.grs = V2.rope_tab(torch.arange(max_words), g0["dh"])
+        self.da = None
+        if hasattr(m, "dattn"):
+            A = m.dattn
+            self.da = dict(Wq=f32(A.q.weight).t().contiguous(), bq=f32(A.q.bias), Wk=f32(A.k.weight).t().contiguous(), bk=f32(A.k.bias),
+                           Wv=f32(A.v.weight).t().contiguous(), bv=f32(A.v.bias), qd=f32(A.qd.weight), ko=f32(A.koff.weight), Wo=f32(A.o.weight).t().contiguous(), dk=A.dk)
+        self.sh = None
+        if hasattr(m, "span_head"):
+            S_ = m.span_head; ws_, bs_, self.sh_sl = [], [], {}; o_ = 0
+            for f in S_.fm:
+                for kind, mod in (("t", S_.typ[f.replace(".", "_")]), ("b", S_.bio[f.replace(".", "_")])):
+                    ws_.append(f32(mod.weight).t()); bs_.append(f32(mod.bias)); self.sh_sl[(f, kind)] = (o_, o_ + mod.out_features); o_ += mod.out_features
+            self.sh = (torch.cat(ws_, 1).contiguous(), torch.cat(bs_)); self.sh_dq = f32(m.delay_q.weight) @ self.sh[0]
+            self.Zs = torch.zeros(max_words, o_)
+        self.ua = (f32(m.utt_att.weight)[0], float(m.utt_att.bias)) if hasattr(m, "utt_att") else None
+        # ---- read-outs: fold Q_q with the keys of its (text) columns; one matrix for the word-level read-outs
+        h_ = m.head; K = f32(m.keys()); valid = h_.valid[di]
+        cols, Ms, bs, self.wslices = [], [], [], {}
+        o = 0
+        for q in WQ:
+            c = h_.cols(q); Kc = K[c]
+            Ms.append(f32(h_.Q[q].weight).t() @ Kc.t()); bs.append(f32(h_.Q[q].bias) @ Kc.t() + f32(h_.b + h_.bq[q])[c])
+            msk = ~valid[c]; bs[-1] = bs[-1].masked_fill(msk, NEG)
+            self.wslices[q] = (o, o + len(c)); o += len(c)
+        self.Mw = torch.cat(Ms, 1).contiguous(); self.bw = torch.cat(bs)
+        self.dq = f32(m.delay_q.weight) @ self.Mw                             # delay-embedding constant per d (D_MAX+1, Cw)
+        self.Md = (f32(m.dmix.weight).t() @ self.Mw).contiguous() if hasattr(m, "dmix") else None   # stage 2: word i's own state (d > 0)
+        self.dq[:, self.bw <= NEG / 2] = 0
+        cb = m.cols_span_beg; Kb = K[cb]
+        self.war = dict(r=(f32(w.war_r.weight).t().contiguous(), f32(w.war_r.bias)), k=f32(w.war_k.weight), ln=(f32(w.war_ln.weight), f32(w.war_ln.bias), w.war_ln.eps),
+                        M=(f32(w.QB.weight).t() @ Kb.t()).contiguous(), b=(f32(w.QB.bias) @ Kb.t() + f32(m.bB_all)[cb]).masked_fill(~valid[cb], NEG))
+        self.pool = {}
+        P1 = torch.cat([f32(h_.pool_in[q][0].weight) for q in POOLED], 0).t().contiguous(); p1 = torch.cat([f32(h_.pool_in[q][0].bias) for q in POOLED])
+        self.P1, self.p1 = P1, p1
+        self.P2 = [(f32(h_.pool_in[q][2].weight).t().contiguous(), f32(h_.pool_in[q][2].bias)) for q in POOLED]
+        for q in POOLED:
+            c = h_.cols(q); Kc = K[c]
+            self.pool[q] = ((f32(h_.Q[q].weight).t() @ Kc.t()).contiguous(), (f32(h_.Q[q].bias) @ Kc.t() + f32(h_.b + h_.bq[q])[c]).masked_fill(~valid[c], NEG))
+        # ---- buffers
+        self.Hbuf = torch.zeros(max_words, d); self.Ubuf = torch.zeros(max_words, dw); self.Zw = torch.zeros(max_words, self.Mw.shape[1])
+        self.Zd = torch.zeros(max_words, self.Mw.shape[1]) if self.Md is not None else None
+        self.reset()
+
+    def reset(self):
+        self.mr = getattr(self.m, "memr", None)                               # v2.2 (a) MemRead: key cache over the words so far
+        if self.mr is not None:
+            dk = self.mr.dk; mw = self.Hbuf.shape[0]
+            self.Kc = torch.zeros(mw, dk); self.Vc = torch.zeros(mw, dk); self.Bc = torch.zeros(mw); self._mk = 0; self._mrc = {}
+            self.dqv = self.m.delay_q.weight.detach().float(); self.mr_g = math.tanh(float(self.mr.gate)); self.mr_gu = math.tanh(float(self.mr.gate_u))
+        self.n = 0; self.words = []; self.state = None; self.ustart = []
+        L = len(self.enc)
+        self.ring = [(torch.zeros(self.win, self.hE, self.dh), torch.zeros(self.win, self.hE, self.dh)) for _ in range(L)]   # K, V by word % win
+        self.ring_ver = 0
+        self.carry = dict(utt=(torch.zeros(self.d, dtype=torch.float64), 0.0), item=(torch.zeros(self.d, dtype=torch.float64), 0.0),
+                          speaker=(torch.zeros(10, self.d, dtype=torch.float64), torch.zeros(10, dtype=torch.float64)))
+        self.slots = {}; self.prev = None; self.next_ustart = True; self.spk = []
+        self.snap = {}; self.out = {"T0": {}, "T1": {}, "T2": {}}; self.pooled = {}
+        self.utt0 = 0
+        self.gsum = {li: {} for li in self.gat}                     # chunk summaries per global layer: chunk -> (k roped, v)
+
+    # ------------------------------------------------------------------ primitives
+    @staticmethod
+    def _mm(x, W):
+        return W.mm(x) if isinstance(W, QLin) else x @ W
+
+    @staticmethod
+    def _ln(x, p):
+        return F.layer_norm(x, (x.shape[-1],), p[0], p[1], p[2])
+
+    @staticmethod
+    def _rms(x, w, eps):
+        return F.rms_norm(x, (x.shape[-1],), w, eps)
+
+    def _rope(self, x, pos):
+        c, s = self.rope_c[pos], self.rope_s[pos]; x1, x2 = x[..., 0::2], x[..., 1::2]
+        return torch.stack([x1 * c - x2 * s, x1 * s + x2 * c], -1).flatten(-2)
+
+    def _embed(self, ws, fes, fls):
+        T = len(ws); Pm = torch.zeros(T, P_MAX, dtype=torch.long); npc = torch.zeros(T, dtype=torch.long)
+        for i, p in enumerate(ws):
+            Pm[i, :len(p)] = torch.tensor(p); npc[i] = len(p)
+        Et = self.tokE[Pm].float()
+        pm = (torch.arange(P_MAX)[None] < npc[:, None]).float()[..., None]
+        mean = ((Et + self.ppos) * pm).sum(1) / npc.clamp(min=1)[:, None].float()
+        last = Et[torch.arange(T), (npc - 1).clamp(min=0)]
+        e = self._ln(mean @ self.W1 + self.b1 + last @ self.W2, self.ln) + self.c_in
+        fl = torch.tensor(fls); fl[:, 2] = npc.clamp(0, P_MAX)
+        for i, tab in enumerate(self.fl):
+            e = e + tab[fl[:, i].clamp(0, tab.shape[0] - 1)]
+        for k, tab in self.feat.items():
+            xf = getattr(self, "xfeat", None) or {}            # stage-2 session-constant inputs (spksrc; turn = 0 without Flux events)
+            e = e + tab[torch.tensor([f[k] if k in f else xf.get(k, 0) for f in fes]).clamp(0, tab.shape[0] - 1)]
+        return e
+
+    def _enc(self, x, j0):
+        """local causal attention for words j0..j0+T-1 with ring-buffer K / V of the previous win-1 words"""
+        T = x.shape[0]; hE, dh, win = self.hE, self.dh, self.win
+        pos = torch.arange(j0, j0 + T)
+        for li, L in enumerate(self.enc):
+            q, k, v = (self._mm(self._ln(x, L["n1"]), L["Wqkv"]) + L["bqkv"]).view(T, 3, hE, dh).unbind(1)
+            q = self._rope(q, pos); k = self._rope(k, pos)
+            Kr, Vr = self.ring[li]
+            npast = min(j0, win - 1)
+            pidx = torch.arange(j0 - npast, j0) % win
+            Kall = torch.cat([Kr[pidx], k], 0); Vall = torch.cat([Vr[pidx], v], 0)          # (npast+T, h, dh)
+            kp = torch.arange(j0 - npast, j0 + T)
+            msk = (kp[None] <= pos[:, None]) & (kp[None] > pos[:, None] - win)
+            a = F.scaled_dot_product_attention(q.transpose(0, 1)[None], Kall.transpose(0, 1)[None], Vall.transpose(0, 1)[None], attn_mask=msk[None, None])
+            a = a[0].transpose(0, 1).reshape(T, hE * dh)
+            # write the new K / V into a COPY of the ring (snapshots keep the old one)
+            Kn, Vn = Kr.clone(), Vr.clone(); Kn[pos % win] = k; Vn[pos % win] = v
+            self.ring[li] = (Kn, Vn)
+            x = x + self._mm(a, L["Wo"]) + L["bo"]
+            x = x + self._mm(F.gelu(self._mm(self._ln(x, L["n2"]), L["F1"]) + L["f1"]), L["F2"]) + L["f2"]
+        return self._ln(x, self.enc_norm)
+
+    def _trunk(self, u):
+        T = u.shape[0]
+        h = self._mm(self._ln(u, self.in_norm), self.Winp) + self.binp
+        state = self.state; new = []
+        for i, L in enumerate(self.layers):
+            h = h * L["g"] + L["b"]
+            nu = self._rms(h, L["nw"], L["neps"])
+            pr = self._mm(nu, L["Wf"]) + L["bf"]
+            zi, xx, Bp, Cp, th, dtp, lmp = torch.split(pr, L["sizes"], -1)
+            H, N, P = L["H"], L["N"], L["P"]
+            dt = F.softplus(dtp); alpha = torch.exp(dt * L["A"]); lam = torch.sigmoid(lmp)
+            gamma = lam * dt; beta = (1 - lam) * dt * alpha
+            BC = F.rms_norm(torch.stack([Bp.view(T, H, N), Cp.view(T, H, N)], 1), (N,), None, L["Bn"][1]) * L["BCw"] + L["BCb"]
+            Bm, Cm = BC[:, 0], BC[:, 1]
+            xv = xx.view(T, H, P); theta = th.view(T, H, N // 2) * dt[..., None]
+            bx = Bm[..., :, None] * xv[..., None, :]                          # (T,H,N,P)
+            st = None if state is None else state[i]
+            cth, sth = torch.cos(theta)[..., None], torch.sin(theta)[..., None]   # (T,H,N/2,1)
+            ys = []
+            for t in range(T):
+                if st is None:
+                    hs = gamma[t, :, None, None] * bx[t]
+                else:
+                    pre = alpha[t, :, None, None] * st[0] + beta[t, :, None, None] * st[1]      # (H,N,P); rotate the N pairs
+                    c_, s_ = cth[t], sth[t]
+                    a_, b_ = pre[:, 0::2], pre[:, 1::2]
+                    hs = torch.stack([c_ * a_ - s_ * b_, s_ * a_ + c_ * b_], 2).reshape(H, N, P) + gamma[t, :, None, None] * bx[t]
+                st = (hs, bx[t])
+                ys.append((Cm[t][:, None, :] @ hs).reshape(-1))
+            new.append(st)
+            y = torch.stack(ys)
+            h = h + self._mm(y * F.silu(zi), L["Wout"]) + L["bout"]
+            h = h + self._mm(F.gelu(self._mm(self._rms(h, *L["mn"]), L["M1"]) + L["m1"]), L["M2"]) + L["m2"]
+            if i in self.gat:
+                h = h + self._gattn(i, h)
+        self.state = new
+        h = h + F.gelu(h @ self.adD) @ self.adU
+        return self._rms(h, *self.on)
+
+    def _gattn(self, li, h, qb=256):
+        """v2.GlobalAttn for words j0..j0+T-1 (j0 = self.n): q / k / v of the new words into the position-indexed buffers, then each
+        query attends over its direct span (last mem words) + the newest `slots` complete chunk summaries older than that"""
+        G = self.gat[li]; T = h.shape[0]; j0 = self.n_run; H, dh, mem, ch = G["h"], G["dh"], G["mem"], G["ch"]
+        hn = self._rms(h, G["nw"], G["neps"])
+        q, k, v = self._mm(hn, G["Wqkv"]).view(T, 3, H, dh).unbind(1)
+        bi = G.get("biased", False)
+        q = self._rms(q, *G["qn"]); k = self._rms(k, *G["kn"])
+        pos = torch.arange(j0, j0 + T)
+        c, s = self.grc[pos], self.grs[pos]
+        qr = self.V2.rope(q, c, s); G["K"][pos] = k; G["Kr"][pos] = self.V2.rope(k, c, s); G["V"][pos] = v
+        if bi:
+            G["M"][pos] = torch.sigmoid(hn @ G["kmW"] + G["kmb"])
+        cache = self.gsum[li]
+        for cc in [cc for cc in cache if (cc + 1) * ch > j0]:            # chunks touched by (re)written words
+            del cache[cc]
+        outs = []
+        for a in range(0, T, qb):
+            b = min(T, a + qb); i0, i1 = j0 + a, j0 + b                      # queries i0..i1-1
+            lo = max(0, i0 - mem + 1)
+            dm, sm = self.V2.summary_mask(i1 - i0, i1, i0, mem, ch, G["slots"])
+            dm = dm[:, lo:]
+            cs = sm.any(0).nonzero().flatten().tolist()
+            Ks, Vs = [G["Kr"][lo:i1]], [G["V"][lo:i1]]
+            if cs:
+                for cc in cs:
+                    if cc not in cache:
+                        kc = G["K"][cc * ch:(cc + 1) * ch].mean(0, keepdim=True); vc = G["V"][cc * ch:(cc + 1) * ch].mean(0, keepdim=True)
+                        c1, s1 = self.V2.rope_tab(torch.tensor([cc * ch + ch // 2]), dh)
+                        cache[cc] = (self.V2.rope(kc, c1, s1)[0], vc[0])
+                Ks.append(torch.stack([cache[cc][0] for cc in cs])); Vs.append(torch.stack([cache[cc][1] for cc in cs]))
+                dm = torch.cat([dm, sm[:, cs]], 1)
+            Kall = torch.cat(Ks, 0); Vall = torch.cat(Vs, 0)
+            msk = dm[None, None]
+            if bi:                                                           # -slope_h (i - j) + kb_h m_j (j = chunk centre / mean m)
+                jp = torch.cat([torch.arange(lo, i1)] + ([torch.tensor([cc * ch + ch // 2 for cc in cs])] if cs else [])).float()
+                mj = torch.cat([G["M"][lo:i1]] + ([torch.stack([G["M"][cc * ch:(cc + 1) * ch].mean() for cc in cs])] if cs else []))
+                ii = torch.arange(i0, i1).float()
+                bias = -G["slope"][:, None, None] * (ii[None, :, None] - jp[None, None, :]) + G["kb"][:, None, None] * mj[None, None, :]
+                msk = torch.where(dm[None], bias, torch.full_like(bias, float("-inf")))[None]
+            o = F.scaled_dot_product_attention(qr[a:b].transpose(0, 1)[None], Kall.transpose(0, 1)[None], Vall.transpose(0, 1)[None],
+                                               attn_mask=msk)
+            outs.append(o[0].transpose(0, 1).reshape(b - a, H * dh))
+        return self._mm(torch.cat(outs, 0), G["Wo"])
+
+    # ------------------------------------------------------------------ public API
+    def _snapshot(self):
+        return dict(state=self.state, carry=self.carry, ring=list(self.ring), slots=dict(self.slots), prev=self.prev, nu=self.next_ustart)
+
+    def _restore(self, k):
+        if k >= self.n:
+            return True
+        sn = self.snap.get(k)
+        if sn is None:
+            return False
+        self.state, self.carry, self.ring, self.slots, self.prev, self.next_ustart = sn["state"], sn["carry"], list(sn["ring"]), dict(sn["slots"]), sn["prev"], sn["nu"]
+        self.n = k; del self.words[k:], self.ustart[k:], self.spk[k:]
+        if self.mr is not None:
+            self._mk = min(self._mk, k); self._mrc = {c: v for c, v in self._mrc.items() if max(c) < k}
+        for t in self.out.values():                           # outputs that READ a word >= k (out[tier][i] = read word index)
+            for j in [j for j, jr in t.items() if jr >= k]:
+                del t[j]
+        for j in [j for j in self.pooled if j >= k]:
+            del self.pooled[j]
+        for j in [j for j in self.snap if j > k]:
+            del self.snap[j]
+        return True
+
+    def pieces(self, w):
+        x = self.cache.get(w)
+        if x is None:
+            ids = self.tok.encode_word(w) or [0]
+            x = self.cache[w] = ids[:P_MAX - 1] + ids[-1:] if len(ids) > P_MAX else ids
+        return x
+
+    @torch.inference_mode()
+    def update(self, k, words):
+        """words[k:] := words. Restores the nearest snapshot s <= k and runs words[s:k] + words as ONE chunk."""
+        k = min(k, self.n)
+        s0 = max([x for x in self.snap if x <= k], default=None)
+        if k < self.n and s0 is None:
+            ws = self.words[:k]; self.reset(); s0 = 0; rep_ = ws
+        else:
+            s0 = k if k >= self.n else s0
+            rep_ = self.words[s0:k]
+            self._restore(s0)
+        self._run(rep_ + [dict(w) for w in words])
+        return self.due
+
+    def _run(self, words):
+        """ONE pass over the new words (append: 1 word; revision: the revised suffix as one chunk)"""
+        if not words:
+            return
+        j0 = self.n; T = len(words)
+        self.snap[j0] = self._snapshot()                                      # state before the chunk (words inside: replay)
+        fes, fls, pcs, spks, uss = [], [], [], [], []
+        for i, w in enumerate(words):
+            fe, fl, spk = SA.ArcStream.feats(self, w)
+            fes.append(fe); fls.append(fl); pcs.append(self.pieces(w["w"])); spks.append(spk)
+            uss.append(self.next_ustart or (j0 + i) == 0); self.next_ustart = False
+        e = self._embed(pcs, fes, fls)
+        self.n_run = j0
+        u = self._enc(e, j0)
+        h = self._trunk(u)
+        self._after(j0, words, u, h, uss, spks)
+
+    def _after(self, j0, words, u, h, uss, spks):
+        T = len(words)
+        self.Hbuf[j0:j0 + T] = h; self.Ubuf[j0:j0 + T] = u
+        if self.mr is not None:                                              # keys of rewritten words are recomputed lazily
+            self._mk = min(self._mk, j0); self._mrc = {c: v for c, v in self._mrc.items() if max(c) < j0}
+        self.Zw[j0:j0 + T] = h @ self.Mw + self.bw
+        if self.Zd is not None:
+            self.Zd[j0:j0 + T] = h @ self.Md
+        if self.sh is not None:
+            self.Zs[j0:j0 + T] = h @ self.sh[0] + self.sh[1]
+        z1 = F.gelu(self._mm(h, self.P1) + self.p1)
+        zs = [h + self._mm(z1[:, i * self.d:(i + 1) * self.d], W) + b for i, (W, b) in enumerate(self.P2)]
+        for i in range(T):
+            j = j0 + i
+            us = uss[i]; sp = spks[i]
+            cu, nu_ = self.carry["utt"]; ci, ni = self.carry["item"]; cs, ns = self.carry["speaker"]
+            w_ = 1.0 if self.ua is None else float(torch.exp((zs[0][i] @ self.ua[0] + self.ua[1]).clamp(-10, 10)))
+            cu = w_ * zs[0][i].double() + (0 if us else cu); nu_ = w_ + (0 if us else nu_)
+            ci = ci + zs[2][i].double(); ni = ni + 1
+            if sp > 0:
+                cs = cs.clone(); ns = ns.clone(); cs[sp - 1] += zs[1][i].double(); ns[sp - 1] += 1
+            self.carry = dict(utt=(cu, nu_), item=(ci, ni), speaker=(cs, ns))
+            self.pooled[j] = dict(utt=(cu / nu_).float(), item=(ci / ni).float(),
+                                  speaker=((cs[sp - 1] / ns[sp - 1]).float() if sp > 0 else torch.zeros(self.d)))
+            self.words.append(words[i]); self.ustart.append(us); self.spk.append(sp)
+        self.n = j0 + T
+        for x in [x for x in self.snap if x < self.n - KEEP]:
+            del self.snap[x]
+        # T0 for the new words; T1 for words whose read word is now available
+        new1 = [i for i in range(max(0, j0 - D_MAX), self.n) if i not in self.out["T1"]]
+        for j in range(j0, self.n):
+            self.out["T0"][j] = j
+        self._delayed(final=False)
+        self.due = {("T0", j): self.word_logits(j, "T0") for j in range(j0, self.n)}           # read-outs due now
+        self.due.update({("T1", i): self.word_logits(i, "T1") for i in new1 if i in self.out["T1"]})
+
+    def _delayed(self, final, lo=None):
+        n = self.n
+        lo = max(0, n - 2 * D_MAX - 2) if lo is None else lo
+        starts = [i for i in range(max(0, min(lo, n - 3 * D_MAX)), n) if self.ustart[i]]
+        for tier, cap in (("T1", 3), ("T2", D_MAX)):
+            if tier == "T2" and not final:
+                continue                                                         # T2: computed at the endpoint only
+            for i in range(lo, n):
+                if i in self.out[tier]:
+                    continue
+                s_next = next((s for s in starts if s > i), None)
+                end = (s_next - 1) if s_next is not None else (n - 1 if final else None)
+                jr = min(i + cap, end) if end is not None else (i + cap if i + cap <= n - 1 else None)
+                if jr is None:
+                    continue
+                self.out[tier][i] = jr
+
+    @torch.inference_mode()
+    def endpoint(self):
+        """utterance end: T1 / T2 of the closing utterance + its pooled read-outs (utt / speaker / item at the last word)"""
+        self.next_ustart = True
+        if not self.n:
+            return {}
+        s = max(i for i in range(self.n) if self.ustart[i])
+        n1 = [i for i in range(s, self.n) if i not in self.out["T1"]]; n2 = [i for i in range(s, self.n) if i not in self.out["T2"]]
+        self._delayed(final=True, lo=s)
+        due = {("T1", i): self.word_logits(i, "T1") for i in n1}
+        due.update({("T2", i): self.word_logits(i, "T2") for i in n2})
+        due[("pooled", self.n - 1)] = self.pooled_logits(self.n - 1)
+        return due
+
+    # ------------------------------------------------------------------ materialise logits (lazy; what an API call returns)
+    def word_logits(self, i, tier):
+        jr = self.out[tier][i]; d = jr - i
+        z = self.Zw[jr] + self.dq[min(d, D_MAX)]
+        if d > 0 and self.Zd is not None:
+            z = z + self.Zd[i]
+        ctx = None
+        if d > 0 and self.da is not None:                                  # v2 A2: attention of word i over h_i..h_{i+d}
+            A = self.da; Hs = self.Hbuf[i:jr + 1]
+            kk = Hs @ A["Wk"] + A["bk"] + A["ko"][:d + 1]; vv = Hs @ A["Wv"] + A["bv"]
+            qv = self.Hbuf[i] @ A["Wq"] + A["bq"] + A["qd"][min(d, D_MAX)]
+            ctx = ((kk @ qv) / math.sqrt(A["dk"])).softmax(0) @ vv @ A["Wo"]   # (d,) state increment
+            z = z + ctx @ self.Mw
+        ymr = amr = None
+        if self.mr is not None:                                              # v2.2 (a): read over words < i, query = hd_i
+            ck_ = (i, jr)                                                    # cached per (word, read position); dropped on rewrites
+            if ck_ not in self._mrc:
+                self._mrc[ck_] = self._memread(i, self.Hbuf[jr] + self.dqv[min(d, D_MAX)] + (ctx if ctx is not None else 0))
+            ymr, amr = self._mrc[ck_]
+            z = z + ymr @ self.Mw
+        out = {q: z[a:b] for q, (a, b) in self.wslices.items()}
+        if self.sh is not None:
+            zs_ = self.Zs[jr] + self.sh_dq[min(d, D_MAX)] + (ctx @ self.sh[0] if ctx is not None else 0) + (ymr @ self.sh[0] if ymr is not None else 0)
+            out["span_head"] = {f: (zs_[self.sh_sl[(f, "t")][0]:self.sh_sl[(f, "t")][1]], zs_[self.sh_sl[(f, "b")][0]:self.sh_sl[(f, "b")][1]])
+                                for f in self.V2.SPAN_FAMS}
+        k = min(d, 2); W = self.war
+        us = (self.Ubuf[i + k] @ W["r"][0] + W["r"][1]) if k else W["r"][1]
+        v = self._ln(self.Ubuf[i] + us + W["k"][k], W["ln"])
+        out["span_beg"] = v @ W["M"] + W["b"]
+        if amr is not None:
+            out["span_beg"] = out["span_beg"] + amr
+        return out
+
+    @torch.inference_mode()
+    def _memread(self, i, hd):
+        R = self.mr; n = self.n
+        if self._mk < n:                                                     # extend / refresh the key cache
+            hn = R.norm(self.Hbuf[self._mk:n]); self.Kc[self._mk:n] = R.kn(R.k(hn)); self.Vc[self._mk:n] = R.v(hn)
+            self.Bc[self._mk:n] = float(R.kb) * F.logsigmoid(R.kmark(hn)[:, 0]); self._mk = n
+        q = R.qn(R.q(R.norm(hd)))
+        sc = torch.cat([torch.zeros(1), self.Kc[:i] @ q / math.sqrt(R.dk) + self.Bc[:i]])
+        p = sc.softmax(0)[1:]                                                # slot 0 = null key (zero value)
+        a = p @ self.Vc[:i]
+        return self.mr_g * R.o(a), self.mr_gu * R.ou(a)       # state increment, span_beg logit increment
+
+    def pooled_logits(self, j):
+        return {q: self.pooled[j][q] @ self.pool[q][0] + self.pool[q][1] for q in POOLED}
+
+
+# ---------------------------------------------------------------------- parity + bench
+def parity(sample, ckpt=None, init=None, n_words=150, scratch=None):
+    from data import Windows, Cfg, collate
+    m, T = SA.load(ckpt, init, scratch)
+    s = json.loads(open(sample).readline())
+    cfg = Cfg(); cfg.mlm = False; cfg.max_sub = 32
+    ex = Windows(T, AT.ArcPieces(AT.BPE8K if m.scratch else AT.TOK_PATH), cfg, train=False).example(s, 0, n_words, tier="T0"); b = collate([ex])
+    with torch.inference_mode():
+        ref = m(b, tiers=("T0", "T1", "T2"))["tiers"]
+    st = FastArc(m, T, domain=s.get("domain"), asr=s.get("asr"))
+    ust = set(np.nonzero(ex["ustart"])[0].tolist())
+    rv = next(i for i in range(n_words - 20, 10, -1) if not any(j in ust for j in range(i - 4, i + 1)))
+    for i, w in enumerate(s["words"][:n_words]):
+        if i in ust and i > 0:
+            st.endpoint()
+        st.update(i, [w])
+        if i == rv:
+            st.update(i - 4, [{"w": "xx"}, {"w": "yy"}, {"w": "zz"}])
+            st.update(i - 4, s["words"][i - 4: i + 1])
+    st.endpoint()
+    res = {}
+    for tier in ("T0", "T1", "T2"):
+        for q in WORD_Q:
+            R = ref[tier][q][0]
+            if R.shape[-1] == 0:
+                continue
+            got = torch.stack([st.word_logits(i, tier)[q] for i in range(n_words)])
+            fin = R > NEG / 2
+            if fin.any():
+                res[f"{tier}:{q}"] = round(float((R - got).abs()[fin].max()), 6)
+    for q in POOLED:
+        R = ref["T0"][q][0]; got = torch.stack([st.pooled_logits(i)[q] for i in range(n_words)]); fin = R > NEG / 2
+        if fin.any():
+            res[q] = round(float((R - got).abs()[fin].max()), 6)
+    return res, max(res.values())
+
+
+def bench(sample, prefix=1000, n=40, threads=1, int8=False, init=None, engine="fast"):
+    torch.set_num_threads(threads)
+    m, T = SA.load(None, init)
+    if int8:
+        import torch.ao.quantization as Q
+        if engine == "fast":
+            m = m  # fast engine: weights are pre-folded fp32 matrices; int8 = quantised matmuls (qint8 per-tensor) below
+        else:
+            for name in ("trunk", "word"):
+                setattr(m, name, Q.quantize_dynamic(getattr(m, name), {torch.nn.Linear}, dtype=torch.qint8))
+    words = []
+    for line in open(sample):
+        words += json.loads(line)["words"]
+        if len(words) >= prefix + 2 * n:
+            break
+    st = (FastArc if engine == "fast" else SA.ArcStream)(m, T, domain="court_appellate", asr="deepgram")
+    if int8 and engine == "fast":
+        quantize_fast(st)
+    t = time.time()
+    for i in range(0, prefix, 1):
+        st.update(i, [words[i]])
+    tp = time.time() - t
+    ta, tr = [], []
+    for i in range(prefix, prefix + n):
+        t = time.perf_counter(); st.update(i, [words[i]]); ta.append(time.perf_counter() - t)
+    for _ in range(n // 2):
+        k = len(st.words) - 3
+        t = time.perf_counter(); st.update(k, [dict(w, final=True) for w in st.words[k:]]); tr.append(time.perf_counter() - t)
+    q = lambda x: round(1000 * float(np.median(x)), 2)
+    return dict(engine=engine, threads=threads, int8=int8, append_p50_ms=q(ta), append_p90_ms=round(1000 * float(np.percentile(ta, 90)), 2),
+                revise3_p50_ms=q(tr), revise3_p90_ms=round(1000 * float(np.percentile(tr, 90)), 2), prefix_s=round(tp, 1))
+
+
+class QLin:
+    """int8 weight (per output channel) x fp32 activation via torch.ops.quantized dynamic linear"""
+    def __init__(self, W_t, b=None):
+        W = W_t.t().contiguous()
+        self.p = torch.ops.quantized.linear_prepack(torch.quantize_per_channel(W, W.abs().amax(1).clamp(min=1e-8) / 127,
+                                                                               torch.zeros(W.shape[0], dtype=torch.long), 0, torch.qint8), b)
+
+    def mm(self, x):
+        return torch.ops.quantized.linear_dynamic(x.contiguous(), self.p)
+
+
+def quantize_fast(st):
+    """int8 for the big matrices of the fast engine (trunk fused / out / MLP, input proj, encoder); read-out matrices stay fp32"""
+    for L in st.layers:
+        for k in ("Wf", "Wout", "M1", "M2"):
+            L[k] = QLin(L[k])
+    st.Winp = QLin(st.Winp)
+    for L in st.enc:
+        for k in ("Wqkv", "Wo", "F1", "F2"):
+            L[k] = QLin(L[k])
+    st.P1 = QLin(st.P1); st.P2 = [(QLin(W), b) for W, b in st.P2]
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--sample", default=os.path.join(os.path.dirname(HERE), "cache/sample/real.jsonl"))
+    ap.add_argument("--ckpt"); ap.add_argument("--init", default="/media/karan/Disk1/sidescreen-audiodec/v14/edge/ckpt/PK8a/step24000")
+    ap.add_argument("--parity", action="store_true"); ap.add_argument("--bench", action="store_true"); ap.add_argument("--int8", action="store_true")
+    ap.add_argument("--threads", type=int, default=1); ap.add_argument("--engine", default="fast"); ap.add_argument("--prefix", type=int, default=1000)
+    a = ap.parse_args()
+    torch.set_num_threads(a.threads)
+    if a.parity:
+        r, mx = parity(a.sample, a.ckpt, a.init)
+        print("fast parity max|dlogit|", mx, json.dumps(r))
+    if a.bench:
+        print("bench", json.dumps(bench(a.sample, a.prefix, threads=a.threads, int8=a.int8, init=a.init, engine=a.engine)))

arctag/_core/kb.py

@@ -0,0 +1,290 @@
+"""Arc Tagger long-list KB ("live RAG"): typed entries, lexical + phonetic index, distractors, per-session entry extraction.
+The trained part (span query -> pointer over a K-shortlist + null, entry-aware verification, KB spotting) is in arc_text.py;
+this module is pure CPU / numpy (data workers, eval, inference-side index). DESIGN.md section 13.
+
+Entry: {"id", "type" (TYPES), "text", "aliases": [...], "value": optional canonical value, "fields": optional dict}
+"""
+import math, re, zlib
+import numpy as np
+
+TYPES = ("name", "number", "keyphrase", "term", "product", "legal_ref", "slot_value", "fact")
+TYPE_ID = {t: i for i, t in enumerate(TYPES)}
+NAME_TAGS = {"person", "team", "organization", "place", "event_name", "game_entity", "game_map"}
+STOP = set("""a an the and or but of to in on at for from by with as is are was were be been being it its this that these those i you he she
+we they me him her us them my your his our their not no so if then than there here what which who whom whose when where why how all any
+each some such very just also do does did done have has had having will would can could should may might must shall up down out over
+into onto about after before again more most other own same too only s t don now ok okay yeah yes uh um oh like well right know think
+going get got go say said see let one two three""".split())
+WEB_SLOT_SRC = ("web_sgd", "web_mwoz", "web_abcd", "web_tm1", "web_tm2", "web_tm3")
+
+
+# ---------------------------------------------------------------------------------------------------------- normalisation / keys
+def norm(t):
+    return " ".join(re.findall(r"[a-z0-9]+", str(t).lower()))
+
+
+_PH = [(r"ph", "f"), (r"gh(?=[^aeiou]|$)", ""), (r"kn|gn|pn|wr|ps", lambda m: m.group(0)[1]), (r"ck|q", "k"), (r"x", "ks"),
+       (r"c(?=[iey])", "s"), (r"c", "k"), (r"dg(?=[iey])", "j"), (r"g(?=[iey])", "j"), (r"sch", "sk"), (r"sh|ch|tch", "x"),
+       (r"th", "0"), (r"tion|sion", "xn"), (r"z", "s"), (r"v", "f"), (r"w(?=[^aeiou]|$)", ""), (r"h(?=[^aeiou]|$)", ""),
+       (r"b|p", "p"), (r"d|t", "t"), (r"m", "m"), (r"(.)\1+", r"\1")]
+_DIG = {"0": "zero", "1": "one", "2": "two", "3": "three", "4": "four", "5": "five", "6": "six", "7": "seven", "8": "eight", "9": "nine"}
+
+
+_PHC = [(re.compile(a), b) for a, b in _PH]
+_NONA = re.compile(r"[^a-z]"); _VOW = re.compile(r"[aeiouy]")
+
+
+@__import__("functools").lru_cache(maxsize=1 << 18)
+def phon_word(w):
+    """compact metaphone-style key of one word: first letter kept, vowels dropped after it, consonant classes merged"""
+    w = _NONA.sub("", w.lower())
+    if not w:
+        return ""
+    k = w
+    for a, b in _PHC:
+        k = a.sub(b, k)
+    return (k[0] + _VOW.sub("", k[1:]))[:6] if k else w[:1]
+
+
+def phon(t):
+    """phonetic key of a phrase (digits spelled out, so '2' ~ 'two')"""
+    ws = []
+    for w in re.findall(r"[a-z]+|[0-9]", str(t).lower()):
+        ws.append(phon_word(_DIG.get(w, w)))
+    return " ".join(x for x in ws if x)
+
+
+def ngrams(t, n=3):
+    s = f" {norm(t)} "
+    return [s[i:i + n] for i in range(max(1, len(s) - n + 1))]
+
+
+H = 1 << 18
+
+
+def _h(g):
+    return zlib.crc32(g.encode()) & (H - 1)
+
+
+@__import__("functools").lru_cache(maxsize=1 << 19)
+def grams_of(t):
+    """hashed char-3gram ids of a text (cached: pool entries are re-indexed in every sampled list)"""
+    return np.unique(np.fromiter((_h(x) for x in ngrams(t)), np.int64))
+
+
+@__import__("functools").lru_cache(maxsize=1 << 19)
+def phon_c(t):
+    return phon(t)
+
+
+# ---------------------------------------------------------------------------------------------------------- lexical + phonetic index
+class LexIndex:
+    """char-3gram TF-IDF (hashed, sparse) over entry texts + aliases, plus phonetic-key buckets (exact phrase key and per-word keys).
+    query(texts, K) -> (n, K) entry indices (-1 pad) ranked by max(cosine, phonetic bonus)."""
+    def __init__(self, entries):
+        import scipy.sparse as sp
+        self.entries = entries; rows, cols, vals, owner = [], [], [], []
+        surf = []
+        for i, e in enumerate(entries):
+            for t in [e["text"]] + list(e.get("aliases") or []):
+                surf.append(t); owner.append(i)
+        self.owner = np.array(owner, np.int64)
+        grams = [grams_of(t) for t in surf]
+        df = np.bincount(np.concatenate(grams), minlength=H).astype(np.float32) if grams else np.zeros(H, np.float32)
+        self.idf = np.log((1 + len(surf)) / (1 + df)).astype(np.float32) + 1
+        lens = np.array([len(g) for g in grams], np.int64); cols = np.concatenate(grams) if grams else np.zeros(0, np.int64)
+        rows = np.repeat(np.arange(len(surf)), lens); w = self.idf[cols]
+        nrm = np.sqrt(np.bincount(rows, w ** 2, minlength=len(surf))); w = w / np.maximum(nrm[rows], 1e-6)
+        self.M = sp.csr_matrix((w, (rows, cols)), shape=(len(surf), H), dtype=np.float32)
+        self.MT = self.M.T.tocsr()
+        self.pk = {}
+        for r, t in enumerate(surf):
+            self.pk.setdefault(phon_c(t), set()).add(int(owner[r]))
+
+    def query(self, texts, K=64):
+        import scipy.sparse as sp
+        rows, cols, vals = [], [], []
+        for r, t in enumerate(texts):
+            g = grams_of(t); w = self.idf[g]; w = w / max(1e-6, float(np.sqrt((w ** 2).sum())))
+            rows += [r] * len(g); cols += g.tolist(); vals += w.tolist()
+        Q = sp.csr_matrix((vals, (rows, cols)), shape=(len(texts), H), dtype=np.float32)
+        S = (Q @ self.MT).toarray() if len(texts) * self.M.shape[0] <= 4e7 else None
+        out = np.full((len(texts), K), -1, np.int64); sc = np.zeros((len(texts), K), np.float32)
+        for r, t in enumerate(texts):
+            s = S[r] if S is not None else (Q[r] @ self.MT).toarray()[0]
+            best = {}
+            nz = np.nonzero(s > 0.05)[0]
+            if len(nz) > 4 * K:
+                nz = nz[np.argpartition(-s[nz], 4 * K)[:4 * K]]
+            for j in nz:
+                o = int(self.owner[j]); best[o] = max(best.get(o, 0.0), float(s[j]))
+            for o in self.pk.get(phon_c(t), ()):                              # sound-alike: phonetic key equal
+                best[o] = max(best.get(o, 0.0), 0.5 + 0.5 * best.get(o, 0.0))
+            top = sorted(best.items(), key=lambda x: -x[1])[:K]
+            for k, (o, v) in enumerate(top):
+                out[r, k] = o; sc[r, k] = v
+        return out, sc
+
+
+# ---------------------------------------------------------------------------------------------------------- distractors
+def asr_corrupt(t, rng):
+    """an ASR-style variant: char substitution / deletion / merge-split / homophone-ish vowel change / word drop"""
+    ws = t.split()
+    if not ws:
+        return t
+    r = rng.random(); i = rng.randrange(len(ws)); w = ws[i]
+    if r < 0.25 and len(w) > 3:
+        j = rng.randrange(1, len(w)); w = w[:j] + rng.choice("aeioubdgkmnprst") + w[j + 1:]
+    elif r < 0.45 and len(w) > 3:
+        j = rng.randrange(1, len(w)); w = w[:j] + w[j + 1:]
+    elif r < 0.6 and len(ws) > 1:
+        ws = ws[:i] + ws[i + 1:]; return " ".join(ws)
+    elif r < 0.75 and len(w) > 5:
+        j = rng.randrange(2, len(w) - 2); w = w[:j] + " " + w[j:]
+    else:
+        w = re.sub(r"[aeiou]", lambda m: rng.choice("aeiou"), w, count=1)
+    ws[i] = w
+    return " ".join(ws)
+
+
+def near_number(t, rng):
+    """near-miss number: one digit changed, two digits swapped, or the scale changed"""
+    ds = [m.start() for m in re.finditer(r"\d", t)]
+    if not ds:
+        return asr_corrupt(t, rng)
+    r = rng.random(); s = list(t)
+    if r < 0.5:
+        j = rng.choice(ds); s[j] = rng.choice([d for d in "0123456789" if d != s[j]])
+    elif r < 0.8 and len(ds) > 1:
+        j = rng.randrange(len(ds) - 1); a, b = ds[j], ds[j + 1]; s[a], s[b] = s[b], s[a]
+    else:
+        return t + "0" if rng.random() < 0.5 else (t[:-1] if len(t) > 1 else t + "1")
+    return "".join(s)
+
+
+def wrong_qualifier(t, rng):
+    """'close but wrong' variant that keeps the head: qualifier swapped (Inc / Corp, north / south, first / second, v. reversed)"""
+    swaps = [("inc", "corp"), ("north", "south"), ("east", "west"), ("first", "second"), ("senior", "junior"), ("new", "old"),
+             ("upper", "lower"), ("international", "national"), ("v.", "vs."), ("i", "ii")]
+    low = t.lower()
+    for a, b in swaps:
+        if re.search(rf"\b{re.escape(a)}\b", low):
+            return re.sub(rf"\b{re.escape(a)}\b", b, low, count=1)
+        if re.search(rf"\b{re.escape(b)}\b", low):
+            return re.sub(rf"\b{re.escape(b)}\b", a, low, count=1)
+    return None
+
+
+def distractors(entry, pool_by_type, pool_by_phon, rng, n=8):
+    """hard distractors for one entry: phonetic neighbours (pool), same-type entries (pool), ASR-corrupted / near-miss / wrong-qualifier
+    variants of its own text. Returns entries tagged with 'dkind'."""
+    out = []
+    t = entry["text"]; ty = entry["type"]
+    for o in list(pool_by_phon.get(phon(t), ()))[:3]:
+        if norm(o["text"]) != norm(t):
+            out.append(dict(o, dkind="phonetic"))
+    same = pool_by_type.get(ty) or []
+    for _ in range(min(3, len(same))):
+        o = same[rng.randrange(len(same))]
+        if norm(o["text"]) != norm(t):
+            out.append(dict(o, dkind="same_type"))
+    v = near_number(t, rng) if ty == "number" else asr_corrupt(t, rng)
+    if norm(v) and norm(v) != norm(t):
+        out.append({"text": v, "type": ty, "aliases": [], "dkind": "near_miss" if ty == "number" else "asr_variant"})
+    q = wrong_qualifier(t, rng)
+    if q and norm(q) != norm(t):
+        out.append({"text": q, "type": ty, "aliases": [], "dkind": "wrong_qualifier"})
+    return out[:n]
+
+
+def list_size(rng, lo=16, hi=10000):
+    return int(round(math.exp(rng.uniform(math.log(lo), math.log(hi)))))
+
+
+# ---------------------------------------------------------------------------------------------------------- per-session extraction
+def _type_of(tag, src_file, ex):
+    if tag.startswith("entity.number.") or tag == "entity.numeral":
+        return "number"
+    if not tag.startswith("entity.mention."):
+        return None
+    k = tag.split(".")[2]
+    if any(src_file.startswith(w) for w in WEB_SLOT_SRC) and (ex or {}).get("src") == "dataset":
+        return "slot_value"
+    return {"work_product": "product", "legal_ref": "legal_ref", "term": "term"}.get(k, "name" if k in NAME_TAGS else None)
+
+
+LEGAL_RE = re.compile(r"\b(\d+\s+U\.?\s?S\.?\s?C\.?|§+\s*\d+|[A-Z][a-z]+\s+v\.?\s+[A-Z][a-z]+|Section\s+\d+[a-z]?|Rule\s+\d+)", re.I)
+
+
+def keyphrases(words, ent_mask, idf, k=24, nmax=4):
+    """salient non-entity n-grams (1-4 words, no stop word at either end, >= 1 letter word of >= 3 chars), scored by
+    session tf x idf (idf over a sample of sessions; unseen = high); multi-word preferred; non-overlapping with entity spans"""
+    low = [norm(w) for w in words]; cand = {}
+    for n in range(1, nmax + 1):
+        for i in range(len(low) - n + 1):
+            g = low[i:i + n]
+            if any(ent_mask[i:i + n]) or not all(g) or g[0] in STOP or g[-1] in STOP or not any(len(x) >= 3 and x.isalpha() for x in g):
+                continue
+            if n == 1 and (len(g[0]) < 5 or not g[0].isalpha()):
+                continue
+            key = " ".join(g); c = cand.setdefault(key, [0, i, n]); c[0] += 1
+    scored = []
+    for key, (tf, i, n) in cand.items():
+        if n == 1 and tf < 2:
+            continue
+        s = (1 + math.log(tf)) * idf.get(key, 8.0) * (1.0 + 0.5 * (n - 1))
+        scored.append((s, key, i, n))
+    scored.sort(reverse=True); out, used = [], np.zeros(len(low), bool)
+    for s, key, i, n in scored:
+        if used[i:i + n].any():
+            continue
+        used[i:i + n] = True; out.append((key, i, i + n))
+        if len(out) >= k:
+            break
+    return out
+
+
+def session_entries(s, src_file, idf, k_key=24):
+    """true KB entries of a session from its own text + labels: [{"text", "type", "occ": [[a, b], ...], "value"?, "tag"?}]"""
+    W = [w["w"] for w in s.get("words") or []]; ent = {}; mask = np.zeros(len(W), bool)
+    for sp in (s.get("labels") or {}).get("span", []):
+        tag, a, b = sp[0], int(sp[1]), int(sp[2]); ex = sp[3] if len(sp) > 3 and isinstance(sp[3], dict) else {}
+        ty = _type_of(tag, src_file, ex)
+        if ty is None or b <= a:
+            continue
+        txt = " ".join(W[a:b]).strip(" .,;:!?\"'")
+        if len(norm(txt)) < 1 or (ty != "number" and len(norm(txt)) < 2):
+            continue
+        key = (ty, norm(txt)); e = ent.setdefault(key, {"text": txt, "type": ty, "occ": [], "tag": tag})
+        e["occ"].append([a, b]); mask[a:b] = True
+        if ex.get("canon") and ty == "number":
+            e["value"] = ex["canon"]
+    for e0 in (s.get("kb") or []):                                         # the session's own KB list (names), as aliases-free entries
+        t = e0[0] if isinstance(e0, (list, tuple)) else e0
+        key = ("name", norm(t))
+        if norm(t) and key not in ent:
+            ent[key] = {"text": str(t), "type": "name", "occ": [], "tag": "kb"}
+    txt = " ".join(W)
+    for m in LEGAL_RE.finditer(txt):
+        key = ("legal_ref", norm(m.group(0)))
+        if key not in ent and norm(m.group(0)):
+            ent[key] = {"text": m.group(0), "type": "legal_ref", "occ": [], "tag": "regex"}
+    for key, a, b in keyphrases(W, mask, idf, k=k_key):
+        ent.setdefault(("keyphrase", key), {"text": " ".join(W[a:b]).strip(" .,;:!?\"'"), "type": "keyphrase", "occ": [[a, b]], "tag": "salience"})
+    locate_all(list(ent.values()), W)
+    return list(ent.values())
+
+
+def locate_all(E, W, nmax=8):
+    """fill e["occ"] (if empty) with every occurrence (<= 20) of the entry's normalised text in the words: one n-gram -> positions map"""
+    need = [e for e in E if not e["occ"]]
+    if not need:
+        return
+    low = [norm(w) for w in W]; lens = {len(norm(e["text"]).split()) for e in need}; pos = {}
+    for n in lens:
+        if 1 <= n <= nmax:
+            for i in range(len(low) - n + 1):
+                pos.setdefault(" ".join(low[i:i + n]), []).append(i)
+    for e in need:
+        nt = norm(e["text"]); n = len(nt.split())
+        e["occ"] = [[i, i + n] for i in pos.get(nt, [])[:20]]

arctag/_core/kb_strfeat.py

@@ -0,0 +1,37 @@
+"""v2.2 (b) / (c) string features for the long-list KB (zero-init add-ons, arc_text.add_kb_str):
+  (b) char-n-gram mention key: hashed char 3-grams of the heard span text -> EmbeddingBag added to the KB query before normalisation;
+  (c) ver_e string features: explicit (span text, entry text) similarity features into the entry-aware verifier's hidden layer."""
+import zlib, difflib, functools
+import kb as KBM
+NB = 4096          # hash buckets
+NF = 8             # string features
+
+
+def grams(t):
+    t = f"#{KBM.norm(t)}#"
+    return [zlib.crc32(t[i:i + 3].encode()) % NB for i in range(max(1, len(t) - 2))]
+
+
+def cng_batch(texts):
+    """-> (ids list, offsets list) for nn.EmbeddingBag"""
+    ids, off = [], []
+    for t in texts:
+        off.append(len(ids)); ids += grams(t)
+    return ids, off
+
+
+def _g3(t):
+    t = f"#{t}#"; return {t[i:i + 3] for i in range(max(1, len(t) - 2))}
+
+
+@functools.lru_cache(maxsize=200000)
+def sfeat(q, e):
+    """(span text, entry text) -> NF floats in [0, 1]"""
+    qn, en = KBM.norm(q), KBM.norm(e)
+    if not qn or not en:
+        return [0.0] * NF
+    a, b = _g3(qn), _g3(en); qt, et = set(qn.split()), set(en.split())
+    qd, ed = "".join(c for c in qn if c.isdigit()), "".join(c for c in en if c.isdigit())
+    return [len(a & b) / max(1, len(a | b)), len(qt & et) / max(1, len(qt | et)), difflib.SequenceMatcher(None, qn, en).ratio(),
+            min(len(qn), len(en)) / max(len(qn), len(en)), float(qn == en), float(qn[:3] == en[:3]),
+            float(qd == ed) if (qd or ed) else 0.5, float(KBM.phon(qn) == KBM.phon(en))]

arctag/_core/losses.py

@@ -0,0 +1,254 @@
+"""Arc Tagger losses: masked BCE (independent tags), soft-CE with a null class over exclusive-group members (members whose target
+is NaN = not labelled-for drop out of the softmax, as pk_safety.readout_loss_null(excl=...)), KD on teacher probabilities with the
+same form, span verification / KB-match, MLM."""
+import torch, torch.nn.functional as F
+
+from tagtable import QUERIES
+
+
+class GroupIndex:
+    """per read-out: column positions of independent tags and of each exclusive group"""
+    def __init__(self, T, w_weak=0.3):
+        self.indep, self.groups, self.colw = {}, {}, {}
+        for q in QUERIES:
+            cols = T.cols[q]
+            self.indep[q] = torch.tensor([k for k, i in enumerate(cols) if i not in T.group_of], dtype=torch.long)
+            g = {}
+            for k, i in enumerate(cols):
+                if i in T.group_of:
+                    g.setdefault(T.group_of[i], []).append(k)
+            self.groups[q] = {n: torch.tensor(v, dtype=torch.long) for n, v in g.items()}
+            self.colw[q] = torch.tensor([w_weak if i in T.weak else 1.0 for i in cols])
+        self.posw = None              # stage 2: {q: (C,) false-negative weight} used at T0 only (set_fn_weights)
+        self.ls = {}                  # stage 2: {(q, group): label smoothing} for exclusive-group soft-CE (e.g. discourse.intent .1)
+        self.clsw = {}                # m4 (11-06): {(q, group): (M,) per-member positive-row multiplier} (e.g. ritual_procedural 2.0);
+                                      # normaliser unchanged (like posw). Empty = identical to before.
+        self.typw = {}                # m5/m6 (11-06): {q: (C,) positive-target multiplier, ALL tiers, independent columns} (--type_pos_w,
+                                      # e.g. entity.mention.organization 1.5); normaliser unchanged. Empty = identical to before.
+        self._dev = {}
+
+    def set_fn_weights(self, T, weights, queries=("span_beg", "span_in", "span_end")):
+        """asymmetric T0 loss: positive-target multiplier per tag (redact.FN_W_T0) on the span read-outs; 1 elsewhere"""
+        self.posw = {q: torch.tensor([float(weights.get(T.ids[i], 1.0)) if q in queries else 1.0 for i in T.cols[q]]) for q in QUERIES}
+        return self
+
+    def set_type_pos_w(self, T, spec, queries=("span_beg", "span_in", "span_end")):
+        """spec 'tag=w,...' -> positive-target multiplier on the span read-outs of independent (non-exclusive) columns, every tier"""
+        want = {}
+        for x in (spec or "").split(","):
+            if x.strip():
+                k, v = x.rsplit("=", 1); want[k.strip()] = float(v)
+        hit = set()
+        for q in queries:
+            ind = set(self.indep[q].tolist()); v = []
+            for k, i in enumerate(T.cols[q]):
+                x = want.get(T.ids[i], 1.0) if k in ind else 1.0; v.append(x)
+                if x != 1.0: hit.add(T.ids[i])
+            if any(x != 1.0 for x in v):
+                self.typw[q] = torch.tensor(v)
+        miss = set(want) - hit
+        if miss:
+            raise ValueError(f"--type_pos_w: not independent span columns: {sorted(miss)}")
+        return self
+
+    def set_class_w(self, T, spec):
+        """spec 'tag=w,tag=w' -> per-member multipliers of exclusive groups (tags outside exclusive groups are rejected)"""
+        want = {}
+        for x in (spec or "").split(","):
+            if x.strip():
+                k, v = x.rsplit("=", 1); want[k.strip()] = float(v)
+        hit = set()
+        for q in QUERIES:
+            cols = T.cols[q]
+            for n, m in self.groups[q].items():
+                v = [want.get(T.ids[cols[k]], 1.0) for k in m.tolist()]
+                if any(x != 1.0 for x in v):
+                    self.clsw[(q, n)] = torch.tensor(v); hit.update(T.ids[cols[k]] for k in m.tolist() if T.ids[cols[k]] in want)
+        miss = set(want) - hit
+        if miss:
+            raise ValueError(f"--class_w: not members of an exclusive group: {sorted(miss)}")
+        return self
+
+    def on(self, dev):
+        """index tensors resident on dev (cached: no per-call host->device copies)"""
+        dev = str(dev)
+        if dev not in self._dev:
+            self._dev[dev] = ({q: v.to(dev) for q, v in self.indep.items()},
+                              {q: {n: m.to(dev) for n, m in g.items()} for q, g in self.groups.items()},
+                              {q: v.to(dev) for q, v in self.colw.items()},
+                              {q: {n: float(self.colw[q][m][0]) for n, m in g.items()} for q, g in self.groups.items()})
+        return self._dev[dev]
+
+    def probs(self, q, z):
+        """logits (..., C) -> probabilities with the same layout (sigmoid / group softmax with null logit 0)"""
+        p = torch.sigmoid(z.float())
+        for n, m in self.groups[q].items():
+            zz = torch.cat([z[..., m].float(), torch.zeros_like(z[..., :1]).float()], -1)
+            p[..., m] = torch.softmax(zz, -1)[..., :-1]
+        return p
+
+
+def readout_loss(GI, q, z, tg, wq, wvalid, posw=None, cw=None):
+    """z, tg (B,W,C); wq (B,W) word weight; posw (C,) optional false-negative weight (positive-target multiplier, stage-2 T0);
+    returns (sum loss, sum weight). The normaliser stays the unweighted labelled weight, so posw raises the FN cost only."""
+    dev = z.device; z = z.float()
+    fin = torch.isfinite(tg); t = torch.nan_to_num(tg)
+    w_word = wq * wvalid
+    tot = z.new_zeros(()); den = z.new_zeros(())
+    ind = GI.indep[q].to(dev)
+    if len(ind):
+        zi, ti, fi = z[..., ind], t[..., ind], fin[..., ind]
+        w = fi * w_word[..., None] * GI.colw[q].to(dev)[ind]
+        if cw is not None:                                          # v2 per-cell weights (flagged gap tags, toxicity negatives)
+            w = w * cw.to(dev)[..., ind]
+        lw = w if posw is None else w * (1 + (posw.to(dev)[ind] - 1) * ti)
+        if q in GI.typw:
+            lw = lw * (1 + (GI.typw[q].to(dev)[ind] - 1) * ti)
+        tot = tot + (F.binary_cross_entropy_with_logits(zi, ti, reduction="none") * lw).sum(); den = den + w.sum()
+    for n, m in GI.groups[q].items():
+        m = m.to(dev); fm = fin[..., m]
+        row = fm.any(-1) & (w_word > 0)
+        if not row.any():
+            continue
+        zm = z[..., m].masked_fill(~fm, -1e4)
+        zz = torch.cat([zm, torch.zeros_like(zm[..., :1])], -1)[row]
+        tm = t[..., m][row]
+        s = tm.sum(-1, keepdim=True)
+        tt = torch.cat([tm, (1 - s).clamp(min=0)], -1)
+        tt = tt / tt.sum(-1, keepdim=True).clamp(min=1e-6)
+        eps_ = GI.ls.get((q, n), 0.0)
+        if eps_:                                                    # smoothing over the labelled members + null
+            av = torch.cat([fm[row].float(), torch.ones_like(tt[:, :1])], -1)
+            tt = (1 - eps_) * tt + eps_ * av / av.sum(-1, keepdim=True)
+        cw_g = GI.colw[q].to(dev)[m][0]
+        ww = w_word[row] * cw_g
+        if cw is not None:
+            ww = ww * cw.to(dev)[..., m].min(-1).values[row]
+        lw = ww if posw is None else ww * (1 + ((posw.to(dev)[m] - 1) * tm).sum(-1))
+        if (q, n) in GI.clsw:
+            lw = lw * (1 + ((GI.clsw[(q, n)].to(dev) - 1) * tm).sum(-1))
+        tot = tot + (-(tt * torch.log_softmax(zz, -1)).sum(-1) * lw).sum(); den = den + ww.sum()
+    return tot, den
+
+
+LAZY = False      # train.py --lazy 1: no host syncs (no .any()/float()/boolean indexing); per-readout losses stay tensors
+
+
+def readout_loss_lazy(GI, q, z, tg, wq, wvalid, posw=None, cw=None):
+    """= readout_loss with masked weights instead of row selection (same value, no device->host sync)"""
+    ind_, grp_, colw_, cw_ = GI.on(z.device); z = z.float()
+    fin = torch.isfinite(tg); t = torch.nan_to_num(tg)
+    w_word = wq * wvalid
+    tot = z.new_zeros(()); den = z.new_zeros(())
+    ind = ind_[q]
+    if len(ind):
+        zi, ti, fi = z[..., ind], t[..., ind], fin[..., ind]
+        w = fi * w_word[..., None] * colw_[q][ind]
+        if cw is not None:
+            w = w * cw[..., ind]
+        lw = w if posw is None else w * (1 + (posw.to(z.device)[ind] - 1) * ti)
+        if q in GI.typw:
+            lw = lw * (1 + (GI.typw[q].to(z.device)[ind] - 1) * ti)
+        tot = tot + (F.binary_cross_entropy_with_logits(zi, ti, reduction="none") * lw).sum(); den = den + w.sum()
+    for n, m in grp_[q].items():
+        fm = fin[..., m]
+        row = (fm.any(-1) & (w_word > 0)).float()
+        zm = z[..., m].masked_fill(~fm, -1e4)
+        zz = torch.cat([zm, torch.zeros_like(zm[..., :1])], -1)
+        tm = t[..., m]
+        s = tm.sum(-1, keepdim=True)
+        tt = torch.cat([tm, (1 - s).clamp(min=0)], -1)
+        tt = tt / tt.sum(-1, keepdim=True).clamp(min=1e-6)
+        eps_ = GI.ls.get((q, n), 0.0)
+        if eps_:
+            av = torch.cat([fm.float(), torch.ones_like(tt[..., :1])], -1)
+            tt = (1 - eps_) * tt + eps_ * av / av.sum(-1, keepdim=True).clamp(min=1.0)
+        ww = w_word * row * cw_[q][n]
+        if cw is not None:
+            ww = ww * cw[..., m].min(-1).values
+        lw = ww if posw is None else ww * (1 + ((posw.to(z.device)[m] - 1) * tm).sum(-1))
+        if (q, n) in GI.clsw:
+            lw = lw * (1 + ((GI.clsw[(q, n)].to(z.device) - 1) * tm).sum(-1))
+        tot = tot + (-(tt * torch.log_softmax(zz, -1)).sum(-1) * lw).sum(); den = den + ww.sum()
+    return tot, den
+
+
+def tag_loss(GI, logits, b, key="tg", fn_weighted=False):
+    """mean over all labelled (word, column) units of every read-out (Arc: one pooled tag loss); fn_weighted: GI.posw (stage-2 T0)"""
+    z0 = next(iter(logits.values()))
+    tot, den, per = z0.new_zeros((), dtype=torch.float32), z0.new_zeros((), dtype=torch.float32), {}
+    for q, z in logits.items():
+        pw = GI.posw[q] if (fn_weighted and GI.posw is not None) else None
+        if LAZY:
+            l, d = readout_loss_lazy(GI, q, z, b[key][q], b["wq"][q], b["wvalid"].float(), pw, (b.get("tgw") or {}).get(q) if key == "tg" else None)
+            per[q] = (l / d.clamp(min=1e-6)).detach(); tot = tot + l; den = den + d
+            continue
+        l, d = readout_loss(GI, q, z, b[key][q].to(z.device), b["wq"][q].to(z.device), b["wvalid"].to(z.device).float(), pw,
+                            (b.get("tgw") or {}).get(q) if key == "tg" else None)
+        if d > 0:
+            per[q] = float(l / d)
+        tot = tot + l; den = den + d
+    return tot / den.clamp(min=1.0), per
+
+
+def norm_loss(model, out, b):
+    """span verification BCE + KB-match CE over the session KB (+ null)"""
+    dev = out["h"].device; lv, nv, lk, nk = 0.0, 0, 0.0, 0
+    for k, spans in enumerate(b["spans"]):
+        if not spans:
+            continue
+        sp = torch.tensor([[x[0], x[1]] for x in spans], device=dev)
+        rep = model.norm.span_rep(out["h"][k], sp)
+        kbe = None
+        kbt = torch.tensor([x[3] for x in spans], device=dev)
+        vw = torch.tensor([x[4] if len(x) > 4 else 1.0 for x in spans], device=dev)
+        if (kbt != -100).any() and b["kb_ids"][k]:
+            kbe = model.encode_kb([torch.tensor(x, device=dev) for x in b["kb_ids"][k]])
+        v, m = model.norm(rep, kbe)
+        vt = torch.tensor([x[2] for x in spans], device=dev, dtype=torch.float32)
+        lv = lv + (F.binary_cross_entropy_with_logits(v.float(), vt, reduction="none") * vw).sum(); nv += float(vw.sum())
+        if m is not None:
+            sel = kbt != -100
+            tgt = torch.where((kbt >= 0) & (kbt < m.shape[1] - 1), kbt, torch.full_like(kbt, m.shape[1] - 1))
+            lk = lk + F.cross_entropy(m.float()[sel], tgt[sel], reduction="sum"); nk += int(sel.sum())
+    return (lv / max(nv, 1)), (lk / max(nk, 1) if nk else 0.0), nv, nk
+
+
+def mlm_loss(out, b):
+    t = b["mlm"].to(out["mlm"].device)
+    if not (t != -100).any():
+        return 0.0
+    return F.cross_entropy(out["mlm"].float().reshape(-1, out["mlm"].shape[-1]), t.reshape(-1), ignore_index=-100)
+
+
+SF1_Q = ("ev_at", "ev_imm", "ev_rec", "span_beg", "span_in", "span_end")
+
+
+def soft_f1_loss(GI, logits, b, queries=SF1_Q, min_pos=1.0, eps=0.1):
+    """stage 2 (--w_sf1): differentiable soft-F1 per tag column over the whole batch, 1 - F1 averaged over the columns with >= min_pos
+    labelled positives in the batch (no host sync). p = GI.probs (sigmoid / group softmax), masked where the target is NaN or the
+    word is padding / has zero read-out weight. soft TP = sum p*y, FP = sum p*(1-y), FN = sum (1-p)*y; eps smooths empty columns.
+    Complements BCE for rare, timing-sensitive events (score, substitution, speaker.change), where AUC is high but F1 is low."""
+    tot = None; n = None
+    for q in queries:
+        if q not in logits:
+            continue
+        z = logits[q]; tg = b["tg"][q].to(z.device)
+        if z.shape[-1] == 0:
+            continue
+        m = (torch.isfinite(tg) & ((b["wq"][q].to(z.device) * b["wvalid"].to(z.device).float()) > 0)[..., None]).float()
+        y = torch.nan_to_num(tg) * m
+        p = torch.sigmoid(z.float())                                  # = GI.probs without in-place writes (autograd-safe)
+        for _, gm in GI.groups[q].items():
+            gm = gm.to(z.device)
+            zz = torch.cat([z[..., gm].float(), torch.zeros_like(z[..., :1]).float()], -1)
+            p = p.index_copy(-1, gm, torch.softmax(zz, -1)[..., :-1])
+        p = p * m
+        tp = (p * y).sum((0, 1)); fp = (p * (1 - y) * m).sum((0, 1)); fn = ((1 - p) * y).sum((0, 1))
+        f1 = 2 * tp / (2 * tp + fp + fn + eps)
+        use = (y.sum((0, 1)) >= min_pos).float()
+        l = ((1 - f1) * use).sum(); c = use.sum()
+        tot = l if tot is None else tot + l; n = c if n is None else n + c
+    if tot is None:
+        return next(iter(logits.values())).sum() * 0.0
+    return tot / n.clamp(min=1.0)

arctag/_core/speaker_link.py

@@ -0,0 +1,288 @@
+"""Stage-2 text speaker linking (DESIGN.md 14.8): at every turn start, a pointer over the window's speaker memory (one slot per
+speaker cluster seen so far, <= data.LINK_K, = the mean of that cluster's turn embeddings) plus a NEW-speaker slot answers "same
+speaker as cluster c". Fusion with ASR diarization: the input speaker ids (gold, simulated-ASR or provider) enter the trunk through the
+existing spk / chg features and the pointer through two learned scalars (same-ASR-id evidence per cluster, unseen-ASR-id evidence for
+NEW). Per-cluster roles: a role read-out on the cluster memory (incl. the current turn), trained to the cluster's majority gold role.
+
+Training is teacher-forced (memory built from the gold clusters of earlier turns) and fully vectorised: with Z the turn embeddings and
+q the turn queries, score(j, k) = sum_{i<j, C_i=k} q_j . z_i / n_jk. Inference / eval (free-running) assigns turns one by one and
+builds the memory from its own decisions (`assign`).
+
+Losses: pointer CE (T0 / T1 / T2 queries = turn start, start+2, start+7 clamped to the turn end), supervised-contrastive InfoNCE on
+the turn embeddings (same speaker = positive, every other turn of the batch = negative), cluster-role soft-CE.
+"""
+import numpy as np
+import torch, torch.nn as nn, torch.nn.functional as F
+
+TIER_OFF = {"T0": 0, "T1": 2, "T2": 7}
+
+
+class SpkLink(nn.Module):
+    def __init__(self, d, n_roles, dz=128):
+        super().__init__()
+        self.att = nn.Linear(d, 1); self.zp = nn.Linear(d, dz); self.qp = nn.Linear(d, dz)
+        self.tier = nn.Embedding(3, dz); nn.init.zeros_(self.tier.weight)
+        self.new_v = nn.Parameter(torch.randn(dz) * 0.02); self.b_new = nn.Parameter(torch.zeros(()))
+        self.scale = nn.Parameter(torch.tensor(np.log(10.0), dtype=torch.float32))
+        self.w_same = nn.Parameter(torch.zeros(())); self.w_newid = nn.Parameter(torch.zeros(()))
+        self.role = nn.Linear(dz, max(1, n_roles)); self.n_roles = n_roles
+
+    # -------------------------------------------------------------- turn embeddings and queries (whole batch at once)
+    def turn_emb(self, h, turns):
+        """h (B,W,d); turns [(b, s, e)] -> Z (T, dz) unit vectors: attention pool of h over each turn's words"""
+        dev = h.device
+        bi = torch.cat([torch.full((e - s,), b, dtype=torch.long) for b, s, e in turns]).to(dev)
+        wi = torch.cat([torch.arange(s, e) for b, s, e in turns]).to(dev)
+        ti = torch.cat([torch.full((e - s,), t, dtype=torch.long) for t, (b, s, e) in enumerate(turns)]).to(dev)
+        x = h[bi, wi].float(); a = self.att(x)[:, 0]
+        T = len(turns)
+        mx = torch.full((T,), -1e30, device=dev).scatter_reduce(0, ti, a, "amax")
+        ea = torch.exp(a - mx[ti]); den = torch.zeros(T, device=dev).index_add(0, ti, ea)
+        pooled = torch.zeros(T, x.shape[1], device=dev).index_add(0, ti, x * (ea / den[ti])[:, None])
+        return F.normalize(self.zp(pooled), dim=-1)
+
+    def queries(self, h, turns, tier):
+        k = list(TIER_OFF).index(tier); off = TIER_OFF[tier]
+        bi = torch.tensor([b for b, s, e in turns], device=h.device)
+        pi = torch.tensor([min(s + off, e - 1) for b, s, e in turns], device=h.device)
+        return F.normalize(self.qp(h[bi, pi].float()) + self.tier.weight[k], dim=-1)
+
+    # -------------------------------------------------------------- pointer logits for one window (teacher forced)
+    def logits(self, q, Z, C, A):
+        """q, Z (T,dz); C (T,) memory clusters (-1 = not in memory); A (T,) input ASR-id slot at each turn start (0 = none)
+        -> (T, 1+K) logits [NEW, cluster 0..K-1], unseen clusters at -inf; n (T,K) memory counts"""
+        T = q.shape[0]; K = int(C.max().item()) + 1 if (C >= 0).any() else 1
+        O = F.one_hot(C.clamp(min=0), K).float() * (C >= 0).float()[:, None]                  # (T,K)
+        low = torch.tril(torch.ones(T, T, device=q.device), -1)                               # i < j
+        n = low @ O                                                                            # (T,K) earlier turns per cluster
+        S = (q @ Z.t()) * low
+        sc = (S @ O) / n.clamp(min=1)
+        same = ((A[:, None] == A[None, :]) & (A[:, None] > 0)).float() * low                  # earlier turns with the same ASR id
+        fs = (same @ O) / n.clamp(min=1)
+        newid = ((A > 0) & (same.sum(1) == 0)).float()
+        lc = self.scale.exp() * sc + self.w_same * fs
+        lc = lc.masked_fill(n == 0, -1e4)
+        ln = self.scale.exp() * (q @ self.new_v) + self.b_new + self.w_newid * newid
+        return torch.cat([ln[:, None], lc], 1), n
+
+    def role_logits(self, Z, C):
+        """cluster memory INCLUDING the current turn -> role logits (T, n_roles)"""
+        T = Z.shape[0]
+        same = ((C[:, None] == C[None, :]) & (C[:, None] >= 0)).float() * torch.tril(torch.ones(T, T, device=Z.device))
+        m = (same @ Z) / same.sum(1, keepdim=True).clamp(min=1)
+        return self.role(m)
+
+    # -------------------------------------------------------------- free-running assignment (eval / inference)
+    @torch.no_grad()
+    def assign(self, q, Z, A, K_max=16):
+        """turn-by-turn decisions, memory from the model's own clusters -> predicted cluster per turn (new ids in order)"""
+        T = q.shape[0]; pred = torch.full((T,), -1, dtype=torch.long, device=q.device); nk = 0
+        for j in range(T):
+            if j == 0:
+                pred[0] = 0; nk = 1; continue
+            lg, _ = self.logits(q[: j + 1], Z[: j + 1], torch.cat([pred[:j], pred.new_full((1,), -1)]), A[: j + 1])
+            lg = lg[j]
+            if nk >= K_max:
+                lg = lg.clone(); lg[0] = -1e9
+            c = int(lg.argmax())
+            if c == 0:
+                pred[j] = nk; nk += 1
+            else:
+                pred[j] = c - 1
+        return pred
+
+
+def batch_turns(b):
+    """collated batch -> flat turn list [(window, s, e)], gold clusters C, role targets R, window index per turn, ASR slot A"""
+    turns, C, R, wi, A = [], [], [], [], []
+    spk = b["spk"].cpu() if torch.is_tensor(b["spk"]) else torch.as_tensor(b["spk"])
+    for k, L in enumerate(b.get("link") or []):
+        if not L:
+            continue
+        roles = L.get("role") or [-1] * len(L["turns"])
+        for (s, e, c), r in zip(L["turns"], roles):
+            turns.append((k, s, e)); C.append(c); R.append(r); wi.append(k); A.append(int(spk[k, s]))
+    return turns, torch.tensor(C, dtype=torch.long), torch.tensor(R, dtype=torch.long), torch.tensor(wi, dtype=torch.long), \
+        torch.tensor(A, dtype=torch.long)
+
+
+def link_loss(model, h, b, tiers=("T0", "T1", "T2"), tier_w=None, tau=0.1, max_anchor=1024):
+    """-> (pointer CE, InfoNCE, role CE, stats); zero tensors when the batch has no linkable windows"""
+    m = model.spk_link; z0 = h.sum() * 0.0
+    turns, C, R, wi, A = batch_turns(b)
+    if len(turns) < 2:
+        return z0, z0, z0, {}
+    dev = h.device; C, R, wi, A = C.to(dev), R.to(dev), wi.to(dev), A.to(dev)
+    Z = m.turn_emb(h, turns)
+    Q = {t: m.queries(h, turns, t) for t in tiers}
+    tw = tier_w or {"T0": 1.0, "T1": 0.5, "T2": 0.5}
+    lp = z0; npt = 0; acc = {}
+    lr_ = z0; nr = 0
+    for k in torch.unique(wi).tolist():
+        sel = (wi == k).nonzero()[:, 0]
+        if len(sel) < 2:
+            continue
+        Ck, Ak, Zk = C[sel], A[sel], Z[sel]
+        for t in tiers:
+            lg, n = m.logits(Q[t][sel], Zk, Ck, Ak)
+            seen = n.gather(1, Ck.clamp(min=0)[:, None])[:, 0] > 0
+            y = torch.where(seen, Ck + 1, torch.zeros_like(Ck))
+            ok = Ck >= 0
+            lp = lp + tw[t] * F.cross_entropy(lg[ok], y[ok], reduction="sum")
+            if t == "T1":
+                acc.setdefault("n", 0); acc["n"] += int(ok.sum()); acc.setdefault("c", 0)
+                acc["c"] += int((lg[ok].argmax(1) == y[ok]).sum())
+        npt += int((Ck >= 0).sum())
+        Rk = R[sel]
+        if m.n_roles and (Rk >= 0).any():                         # cluster-majority role (whole window) per turn
+            rl = m.role_logits(Zk, Ck)
+            oh = F.one_hot(Rk.clamp(min=0), m.n_roles).float() * (Rk >= 0).float()[:, None]
+            same = ((Ck[:, None] == Ck[None, :]) & (Ck[:, None] >= 0)).float()
+            dist = same @ oh; has = dist.sum(1) > 0
+            if has.any():
+                dist = dist[has] / dist[has].sum(1, keepdim=True)
+                lr_ = lr_ + (-(dist * F.log_softmax(rl[has].float(), -1)).sum(-1)).sum(); nr += int(has.sum())
+    lp = lp / max(1, npt) / sum(tw[t] for t in tiers)
+    lr_ = lr_ / max(1, nr)
+    # supervised-contrastive InfoNCE over the batch's turns
+    ok = (C >= 0).nonzero()[:, 0]
+    if len(ok) > max_anchor:
+        ok = ok[torch.randperm(len(ok), device=dev)[:max_anchor]]
+    lab = wi[ok] * 1000 + C[ok]; Zs = Z[ok]
+    sim = (Zs @ Zs.t()) / tau
+    eye = torch.eye(len(ok), device=dev, dtype=torch.bool)
+    sim = sim.masked_fill(eye, -1e4)
+    pos = (lab[:, None] == lab[None, :]) & ~eye
+    hp = pos.any(1)
+    if hp.any():
+        lsm = F.log_softmax(sim, 1)
+        ln = -((lsm * pos.float()).sum(1)[hp] / pos.float().sum(1)[hp]).mean()
+    else:
+        ln = z0
+    st = {"link_acc_T1": acc.get("c", 0) / max(1, acc.get("n", 0)), "link_turns": npt}
+    return lp, ln, lr_, st
+
+
+# ------------------------------------------------------------------ eval
+def _heur(C):
+    """teacher-forced heuristics per turn: two-back (2-party alternation) and most-frequent-other-than-previous"""
+    T = len(C); out = {"two_back": [], "most_freq": []}
+    for j in range(T):
+        seen = set(C[:j])
+        y = C[j] if C[j] in seen else -1
+        p2 = C[j - 2] if j >= 2 else -1
+        cnt = {}
+        for c in C[:j]:
+            cnt[c] = cnt.get(c, 0) + 1
+        prev = C[j - 1] if j else None
+        cand = sorted(((v, c) for c, v in cnt.items() if c != prev), reverse=True)
+        pm = cand[0][1] if cand else -1
+        out["two_back"].append(p2 == y); out["most_freq"].append(pm == y)
+    return out
+
+
+def _asr_only(C, A):
+    """teacher-forced 'ASR ids alone': a seen ASR id -> the gold cluster most associated with it so far; unseen / none -> NEW"""
+    hist = {}; out = []
+    for j, (c, a) in enumerate(zip(C, A)):
+        y = c if c in set(C[:j]) else -1
+        if a > 0 and a in hist:
+            p = max(hist[a].items(), key=lambda x: x[1])[0]
+        else:
+            p = -1
+        out.append(p == y)
+        if a > 0:
+            hist.setdefault(a, {}); hist[a][c] = hist[a].get(c, 0) + 1
+    return out
+
+
+def der_like(gold, pred, w):
+    """word-weighted speaker error after the optimal one-to-one cluster mapping (Hungarian) + purity / coverage"""
+    from scipy.optimize import linear_sum_assignment
+    g = {c: i for i, c in enumerate(sorted(set(gold)))}; p = {c: i for i, c in enumerate(sorted(set(pred)))}
+    M = np.zeros((len(g), len(p)))
+    for a, b_, x in zip(gold, pred, w):
+        M[g[a], p[b_]] += x
+    r, c = linear_sum_assignment(-M)
+    tot = M.sum()
+    return dict(err=1 - M[r, c].sum() / max(tot, 1), purity=M.max(0).sum() / max(tot, 1), coverage=M.max(1).sum() / max(tot, 1))
+
+
+@torch.no_grad()
+def eval_windows(model, h, b, exs, tiers=("T0", "T1", "T2")):
+    """one record per linkable window of a collated eval batch: teacher-forced accuracy per tier, heuristics, ASR-id-alone,
+    free-running DER-like (T1 queries), role consistency"""
+    m = model.spk_link; recs = []
+    turns, C, R, wi, A = batch_turns(b)
+    if not turns:
+        return recs
+    dev = h.device
+    Z = m.turn_emb(h, turns); Q = {t: m.queries(h, turns, t) for t in tiers}
+    C, R, wi, A = C.to(dev), R.to(dev), wi.to(dev), A.to(dev)
+    for k in torch.unique(wi).tolist():
+        sel = (wi == k).nonzero()[:, 0]
+        if len(sel) < 2:
+            continue
+        Ck, Ak, Zk = C[sel], A[sel], Z[sel]
+        ok = (Ck >= 0).cpu().numpy()
+        rec = dict(dom=exs[k]["dom"], asr=exs[k]["asr"], sid=exs[k]["sid"], turns=int(ok.sum()))
+        for t in tiers:
+            lg, n = m.logits(Q[t][sel], Zk, Ck, Ak)
+            seen = n.gather(1, Ck.clamp(min=0)[:, None])[:, 0] > 0
+            y = torch.where(seen, Ck + 1, torch.zeros_like(Ck))
+            rec[f"acc_{t}"] = int(((lg.argmax(1) == y).cpu().numpy() & ok).sum())
+        Cl = Ck.cpu().tolist(); Al = Ak.cpu().tolist()
+        hz = _heur(Cl); rec["heur"] = {kk: int(np.sum(np.array(v) & ok)) for kk, v in hz.items()}
+        rec["asr_id"] = int(np.sum(np.array(_asr_only(Cl, Al)) & ok)) if any(a > 0 for a in Al) else None
+        pred = m.assign(Q["T1"][sel] if "T1" in Q else Q[tiers[0]][sel], Zk, Ak).cpu().tolist()
+        wts = [e - s for (_, s, e) in [turns[i] for i in sel.tolist()]]
+        gg = [(c, p_, w_) for c, p_, w_ in zip(Cl, pred, wts) if c >= 0]
+        rec["der"] = der_like([x[0] for x in gg], [x[1] for x in gg], [x[2] for x in gg]) if gg else None
+        if m.n_roles:
+            rp = m.role_logits(Zk, Ck).argmax(1).cpu().tolist(); Rl = R[sel].cpu().tolist()
+            cons = []; racc = []
+            for c in set(x for x in Cl if x >= 0):
+                idx = [i for i, x in enumerate(Cl) if x == c]
+                rr = [rp[i] for i in idx]; maj = max(set(rr), key=rr.count)
+                cons += [x == maj for x in rr]
+                gl = [Rl[i] for i in idx if Rl[i] >= 0]
+                if gl:
+                    gm = max(set(gl), key=gl.count); racc += [rp[i] == gm for i in idx]
+            rec["role_cons"] = (int(sum(cons)), len(cons)); rec["role_acc"] = (int(sum(racc)), len(racc))
+        recs.append(rec)
+    return recs
+
+
+def report(recs, dom_name=lambda d: str(d), tier="T1"):
+    """-> {"all" | "dom:<d>" | "asr:<a>": {turns, acc_T0/T1/T2, acc_text (= acc_<tier>), acc_heur (best heuristic), acc_asr_id,
+    der / purity / coverage (word-weighted), role_consistency, role_acc}}"""
+    groups = {}
+    for r in recs:
+        for g in ("all", f"dom:{dom_name(r['dom'])}", f"asr:{r['asr']}"):
+            groups.setdefault(g, []).append(r)
+    out = {}
+    for g, rs in groups.items():
+        n = sum(r["turns"] for r in rs)
+        if not n:
+            continue
+        x = {"turns": n, "windows": len(rs)}
+        for t in ("T0", "T1", "T2"):
+            if f"acc_{t}" in rs[0]:
+                x[f"acc_{t}"] = round(sum(r[f"acc_{t}"] for r in rs) / n, 4)
+        x["acc_text"] = x.get(f"acc_{tier}")
+        hk = rs[0]["heur"].keys()
+        x["acc_heur_each"] = {k: round(sum(r["heur"][k] for r in rs) / n, 4) for k in hk}
+        x["acc_heur"] = max(x["acc_heur_each"].values())
+        ra = [r for r in rs if r.get("asr_id") is not None]
+        x["acc_asr_id"] = round(sum(r["asr_id"] for r in ra) / max(1, sum(r["turns"] for r in ra)), 4) if ra else None
+        dd = [r for r in rs if r.get("der")]
+        if dd:
+            w = [r["turns"] for r in dd]
+            for k in ("err", "purity", "coverage"):
+                x["der_" + k] = round(float(np.average([r["der"][k] for r in dd], weights=w)), 4)
+        rc = [r["role_cons"] for r in rs if r.get("role_cons") and r["role_cons"][1]]
+        x["role_consistency"] = round(sum(a for a, _ in rc) / sum(b_ for _, b_ in rc), 4) if rc else None
+        ra_ = [r["role_acc"] for r in rs if r.get("role_acc") and r["role_acc"][1]]
+        x["role_acc"] = round(sum(a for a, _ in ra_) / sum(b_ for _, b_ in ra_), 4) if ra_ else None
+        out[g] = x
+    return out