declarative-attention 0.0.4__py3-none-any.whl

declarative_attention/__init__.py

@@ -0,0 +1,7 @@
+from declarative_attention.declarative_attention import (
+    DeclarativeStateMachine
+)
+
+__all__ = [
+    'DeclarativeStateMachine'
+]

declarative_attention/declarative_attention.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+from html.parser import HTMLParser
+from typing import Callable, Sequence
+
+import torch
+from torch import Tensor
+from statemachine import StateMachine, State
+from torch_einops_utils import maybe
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def default_decode_fn(token: int) -> str:
+    return chr(token) if 0 <= token < 128 else ""
+
+# types
+
+ChunkId = int
+StartPos = int
+EndPos = int
+Span = tuple[StartPos, EndPos]
+
+# chunk spans can be:
+# 1. dict mapping 1-indexed chunk id to token span:
+#    {1: (16, 32), 2: (32, 48)}
+# 2. sequence of spans (auto-indexed from 1):
+#    [(16, 32), (32, 48)]
+
+ChunkSpans = (
+    dict[ChunkId, Span] |
+    Sequence[Span]
+)
+
+TokenizerDecode = Callable[[int], str]
+
+# general streaming tag parser
+
+class StreamingTagParser(HTMLParser):
+    def __init__(self, on_tag: Callable[[bool, str, dict[str, str]], None]):
+        super().__init__()
+        self.on_tag = on_tag
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]):
+        self.on_tag(False, tag, dict(attrs))
+
+    def handle_endtag(self, tag: str):
+        self.on_tag(True, tag, {})
+
+# declarative attention state machine (Ho et al., 2026)
+
+class DeclarativeStateMachine(StateMachine):
+    """
+    Declarative Attention State Machine
+
+    - <global>              : attends to all context chunks (default)
+    - <focus chunks="1,2">  : attends only to specified chunk(s)
+    - <local>               : attends to 0 context chunks
+    """
+
+    global_mode = State(initial = True)
+    focus_mode  = State()
+    local_mode  = State()
+
+    focus  = focus_mode.from_(global_mode, local_mode)
+    local  = local_mode.from_(global_mode, focus_mode)
+    revert = global_mode.from_(focus_mode, local_mode)
+
+    def __init__(
+        self,
+        chunk_spans: ChunkSpans,
+        tokenizer_decode: TokenizerDecode | None = None
+    ):
+        if isinstance(chunk_spans, (list, tuple)):
+            chunk_spans = {i + 1: span for i, span in enumerate(chunk_spans)}
+
+        self.chunk_spans = chunk_spans
+        self.active_chunks = set(chunk_spans.keys())
+        self.parser = StreamingTagParser(self._handle_tag)
+        self.tokenizer_decode = default(tokenizer_decode, default_decode_fn)
+        self.handlers: dict[str, Callable] = dict()
+
+        super().__init__()
+
+    # register custom tag handlers for arbitrary researcher logic
+
+    def on(self, tag: str, fn: Callable | None = None):
+        def decorator(handler: Callable):
+            self.handlers[tag.lower()] = handler
+            return handler
+
+        return maybe(decorator, default = decorator)(fn)
+
+    # dispatch streaming tags to transitions or custom handlers
+
+    def _handle_tag(self, is_closing: bool, tag: str, attrs: dict[str, str]):
+        tag = tag.lower()
+
+        if is_closing:
+            if hasattr(self, f"revert_{tag}"):
+                self.send(f"revert_{tag}")
+            elif not self.is_global:
+                try: self.revert()
+                except Exception: pass
+            return
+
+        # opening tag: try transition event first, then custom handler
+
+        if tag in [e.id for e in self.events]:
+            try:
+                self.send(tag, **attrs)
+                return
+            except Exception:
+                pass
+
+        if tag in self.handlers:
+            self.handlers[tag](self, **attrs)
+
+    # transition callbacks
+
+    @focus.on
+    def _on_focus(
+        self,
+        chunks: str | Sequence[int] | None = None,
+        chunk: str | int | None = None,
+        magic_chunks: str | None = None
+    ):
+        target = default(chunks, default(chunk, magic_chunks))
+
+        if isinstance(target, str):
+            self.active_chunks = {int(c) for c in target.split(',') if c.strip().isdigit()}
+        elif isinstance(target, int):
+            self.active_chunks = {target}
+        elif exists(target):
+            self.active_chunks = set(target)
+        else:
+            self.active_chunks = set()
+
+    @local.on
+    def _on_local(self):
+        self.active_chunks = set()
+
+    @revert.on
+    def _on_revert(self):
+        self.active_chunks = set(self.chunk_spans.keys())
+
+    # properties
+
+    @property
+    def is_global(self):
+        return self.global_mode.is_active
+
+    @property
+    def is_focus(self):
+        return self.focus_mode.is_active
+
+    @property
+    def is_local(self):
+        return self.local_mode.is_active
+
+    @property
+    def mode(self):
+        return "global" if self.is_global else ("focus" if self.is_focus else "local")
+
+    @property
+    def is_streaming_tag(self) -> bool:
+        """True if the parser is currently mid-tag (e.g. between '<' and '>')"""
+        return "<" in self.parser.rawdata
+
+    # token streaming (accepts int, torch LongTensor, or str)
+
+    def step(self, token: int | Tensor | str):
+        if isinstance(token, Tensor):
+            token = token.item()
+
+        text = self.tokenizer_decode(token) if isinstance(token, int) else str(token)
+        self.parser.feed(text)
+
+    # attention mask generation
+
+    def get_mask(self, total_len: int, device = None) -> Tensor:
+        mask = torch.ones(total_len, dtype = torch.bool, device = device)
+
+        if self.is_global:
+            return mask
+
+        for chunk_id, (start, end) in self.chunk_spans.items():
+            if chunk_id not in self.active_chunks:
+                mask[start:end] = False
+
+        return mask
+
+    def __call__(self, total_len: int, device = None) -> Tensor:
+        return self.get_mask(total_len, device = device)

declarative_attention-0.0.4.dist-info/METADATA

@@ -0,0 +1,184 @@
+Metadata-Version: 2.5
+Name: declarative-attention
+Version: 0.0.4
+Summary: Declarative Attention - Pytorch
+Project-URL: Homepage, https://pypi.org/project/declarative-attention/
+Project-URL: Repository, https://github.com/lucidrains/declarative-attention
+Author-email: Phil Wang <lucidrains@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2026 Phil Wang
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: artificial intelligence,attention mechanisms,deep learning
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: einops>=0.8.1
+Requires-Dist: python-statemachine>=2.5
+Requires-Dist: torch-einops-utils
+Requires-Dist: torch>=2.5
+Provides-Extra: examples
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Description-Content-Type: text/markdown
+
+<img src="./fig1.png" width="400px"></img>
+
+## Declarative Attention - (wip)
+
+Implementation of the procedure in [Language Models Can Control Their Own Attention](https://arxiv.org/abs/2609.02737), from Namgyu Ho et al. of KAIST AI
+
+## Install
+
+```bash
+$ pip install declarative-attention
+```
+
+## Usage
+
+```python
+import torch
+from declarative_attention import DeclarativeStateMachine
+
+# Define context chunk spans (1-indexed)
+chunk_spans = [
+    (16, 1024),   # Chunk 1
+    (1024, 2048), # Chunk 2
+    (2048, 3072), # Chunk 3
+]
+
+state_machine = DeclarativeStateMachine(
+    chunk_spans = chunk_spans,
+    tokenizer_decode = tokenizer.decode  # or lambda t: tokenizer.decode([t])
+)
+
+# Decode loop: feed generated tokens (int, torch.long scalar, or str)
+state_machine.step('<focus chunks="1,2">')
+
+assert state_machine.is_focus
+assert state_machine.active_chunks == {1, 2}
+
+# Obtain 1D boolean attention mask for current decode step
+mask = state_machine.get_mask(total_len = 4096)  # True for kept keys, False for masked out
+
+# Model finishes chunk reasoning and reverts back to global
+state_machine.step('</focus>')
+assert state_machine.is_global
+```
+
+### Custom Attention Patterns
+
+Researchers can easily declare custom attention patterns and tags.
+
+#### 1. Logarithmically Spaced Attention (Recent to Past)
+
+Sample chunks with exponentially decaying density into the past (e.g. current, $t-1, t-2, t-4, t-8\dots$):
+
+```python
+state_machine = DeclarativeStateMachine(chunk_spans)
+
+@state_machine.on('log_sparse')
+def handle_log_sparse(machine, base = 2):
+    base = int(base)
+    total = len(machine.chunk_spans)
+
+    # Powers of base distance from current chunk: 0, 1, 2, 4, 8...
+    offsets = [0] + [base ** i for i in range(10)]
+    machine.active_chunks = {
+        total - d for d in offsets if (total - d) in machine.chunk_spans
+    }
+
+# When model emits <log_sparse> or <log_sparse base="3">
+state_machine.step('<log_sparse base="2">')
+
+# Reverts back to global when closed
+state_machine.step('</log_sparse>')
+```
+
+#### 2. Sliding Window (Last $K$ Chunks)
+
+Select the most recent $K$ context chunks via a custom `State`:
+
+```python
+from statemachine import State
+from declarative_attention import DeclarativeStateMachine
+
+class CustomAttentionMachine(DeclarativeStateMachine):
+    window_mode = State()
+
+    window = window_mode.from_(DeclarativeStateMachine.global_mode)
+    revert = DeclarativeStateMachine.global_mode.from_(
+        DeclarativeStateMachine.focus_mode,
+        DeclarativeStateMachine.local_mode,
+        window_mode
+    )
+
+    @window.on
+    def on_window(self, size = 2):
+        total = len(self.chunk_spans)
+        self.active_chunks = {total - i for i in range(int(size)) if total - i > 0}
+
+state_machine = CustomAttentionMachine(chunk_spans)
+
+# Model calls <window size="2">
+state_machine.step('<window size="2">')
+assert state_machine.is_window_mode
+```
+
+#### 3. Dynamic Temperature & In-Context Verification
+
+Models can dynamically adjust sampling temperature during decode—brainstorming at high temperature, then re-attending to those thoughts at zero temperature for critical verification:
+
+```python
+state_machine = DeclarativeStateMachine(chunk_spans)
+state_machine.temperature = 0.7  # default
+
+@state_machine.on('explore')
+def on_explore(machine, temp = 1.2):
+    machine.temperature = float(temp)
+
+@state_machine.on('verify')
+def on_verify(machine, temp = 0.0):
+    machine.temperature = float(temp)
+
+# High-temperature exploration
+state_machine.step('<explore temp="1.2">')
+
+# Switch to zero-temperature verification on Chunk 1
+state_machine.step('</explore><focus chunks="1"><verify temp="0.0">')
+```
+
+## Citations
+
+```bibtex
+@misc{ho2026languagemodelscontrolattention,
+    title    = {Language Models Can Control Their Own Attention}, 
+    author   = {Namgyu Ho and Huzama Ahmad and Woosung Koh and Se-Young Yun and Tal Schuster and Cicero Nogueira dos Santos},
+    year     = {2026},
+    eprint   = {2609.02737},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CL},
+    url      = {https://arxiv.org/abs/2609.02737}, 
+}
+```

declarative_attention-0.0.4.dist-info/RECORD

@@ -0,0 +1,6 @@
+declarative_attention/__init__.py,sha256=D4GVZyXHfONJrkYHbuj7Le2djgaAuuRdyhBjK8zdXkk,133
+declarative_attention/declarative_attention.py,sha256=qHjqa9hcUVHJ0nPGVUA1JRGPRZA9sMLyN8_7Ulikdl4,5592
+declarative_attention-0.0.4.dist-info/METADATA,sha256=5sL_KygZu4nugwMG51MNFbo10RESUzK5E28rnPbIrxM,6243
+declarative_attention-0.0.4.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
+declarative_attention-0.0.4.dist-info/licenses/LICENSE,sha256=e6AOF7Z8EFdK3IdcL0x0fLw4cY7Q0d0kNR0o0TmBewM,1066
+declarative_attention-0.0.4.dist-info/RECORD,,

declarative_attention-0.0.4.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.0
+Root-Is-Purelib: true
+Tag: py3-none-any

declarative_attention-0.0.4.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Phil Wang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.