declarative-attention 0.0.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,220 @@
1
+ *.sh
2
+
3
+ # Byte-compiled / optimized / DLL files
4
+ __pycache__/
5
+ *.py[codz]
6
+ *$py.class
7
+
8
+ # C extensions
9
+ *.so
10
+
11
+ # Distribution / packaging
12
+ .Python
13
+ build/
14
+ develop-eggs/
15
+ dist/
16
+ downloads/
17
+ eggs/
18
+ .eggs/
19
+ lib/
20
+ lib64/
21
+ parts/
22
+ sdist/
23
+ var/
24
+ wheels/
25
+ share/python-wheels/
26
+ *.egg-info/
27
+ .installed.cfg
28
+ *.egg
29
+ MANIFEST
30
+
31
+ # PyInstaller
32
+ # Usually these files are written by a python script from a template
33
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
34
+ *.manifest
35
+ *.spec
36
+
37
+ # Installer logs
38
+ pip-log.txt
39
+ pip-delete-this-directory.txt
40
+
41
+ # Unit test / coverage reports
42
+ htmlcov/
43
+ .tox/
44
+ .nox/
45
+ .coverage
46
+ .coverage.*
47
+ .cache
48
+ nosetests.xml
49
+ coverage.xml
50
+ *.cover
51
+ *.py.cover
52
+ .hypothesis/
53
+ .pytest_cache/
54
+ cover/
55
+
56
+ # Translations
57
+ *.mo
58
+ *.pot
59
+
60
+ # Django stuff:
61
+ *.log
62
+ local_settings.py
63
+ db.sqlite3
64
+ db.sqlite3-journal
65
+
66
+ # Flask stuff:
67
+ instance/
68
+ .webassets-cache
69
+
70
+ # Scrapy stuff:
71
+ .scrapy
72
+
73
+ # Sphinx documentation
74
+ docs/_build/
75
+
76
+ # PyBuilder
77
+ .pybuilder/
78
+ target/
79
+
80
+ # Jupyter Notebook
81
+ .ipynb_checkpoints
82
+
83
+ # IPython
84
+ profile_default/
85
+ ipython_config.py
86
+
87
+ # pyenv
88
+ # For a library or package, you might want to ignore these files since the code is
89
+ # intended to run in multiple environments; otherwise, check them in:
90
+ # .python-version
91
+
92
+ # pipenv
93
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
94
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
95
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
96
+ # install all needed dependencies.
97
+ # Pipfile.lock
98
+
99
+ # UV
100
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
101
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
102
+ # commonly ignored for libraries.
103
+ # uv.lock
104
+
105
+ # poetry
106
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
107
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
108
+ # commonly ignored for libraries.
109
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
110
+ # poetry.lock
111
+ # poetry.toml
112
+
113
+ # pdm
114
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
115
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
116
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
117
+ # pdm.lock
118
+ # pdm.toml
119
+ .pdm-python
120
+ .pdm-build/
121
+
122
+ # pixi
123
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
124
+ # pixi.lock
125
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
126
+ # in the .venv directory. It is recommended not to include this directory in version control.
127
+ .pixi
128
+
129
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
130
+ __pypackages__/
131
+
132
+ # Celery stuff
133
+ celerybeat-schedule
134
+ celerybeat.pid
135
+
136
+ # Redis
137
+ *.rdb
138
+ *.aof
139
+ *.pid
140
+
141
+ # RabbitMQ
142
+ mnesia/
143
+ rabbitmq/
144
+ rabbitmq-data/
145
+
146
+ # ActiveMQ
147
+ activemq-data/
148
+
149
+ # SageMath parsed files
150
+ *.sage.py
151
+
152
+ # Environments
153
+ .env
154
+ .envrc
155
+ .venv
156
+ env/
157
+ venv/
158
+ ENV/
159
+ env.bak/
160
+ venv.bak/
161
+
162
+ # Spyder project settings
163
+ .spyderproject
164
+ .spyproject
165
+
166
+ # Rope project settings
167
+ .ropeproject
168
+
169
+ # mkdocs documentation
170
+ /site
171
+
172
+ # mypy
173
+ .mypy_cache/
174
+ .dmypy.json
175
+ dmypy.json
176
+
177
+ # Pyre type checker
178
+ .pyre/
179
+
180
+ # pytype static type analyzer
181
+ .pytype/
182
+
183
+ # Cython debug symbols
184
+ cython_debug/
185
+
186
+ # PyCharm
187
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
188
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
189
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
190
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
191
+ # .idea/
192
+
193
+ # Abstra
194
+ # Abstra is an AI-powered process automation framework.
195
+ # Ignore directories containing user credentials, local state, and settings.
196
+ # Learn more at https://abstra.io/docs
197
+ .abstra/
198
+
199
+ # Visual Studio Code
200
+ # Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
201
+ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
202
+ # and can be added to the global gitignore or merged into this file. However, if you prefer,
203
+ # you could uncomment the following to ignore the entire vscode folder
204
+ # .vscode/
205
+ # Temporary file for partial code execution
206
+ tempCodeRunnerFile.py
207
+
208
+ # Ruff stuff:
209
+ .ruff_cache/
210
+
211
+ # PyPI configuration file
212
+ .pypirc
213
+
214
+ # Marimo
215
+ marimo/_static/
216
+ marimo/_lsp/
217
+ __marimo__/
218
+
219
+ # Streamlit
220
+ .streamlit/secrets.toml
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Phil Wang
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,184 @@
1
+ Metadata-Version: 2.5
2
+ Name: declarative-attention
3
+ Version: 0.0.4
4
+ Summary: Declarative Attention - Pytorch
5
+ Project-URL: Homepage, https://pypi.org/project/declarative-attention/
6
+ Project-URL: Repository, https://github.com/lucidrains/declarative-attention
7
+ Author-email: Phil Wang <lucidrains@gmail.com>
8
+ License: MIT License
9
+
10
+ Copyright (c) 2026 Phil Wang
11
+
12
+ Permission is hereby granted, free of charge, to any person obtaining a copy
13
+ of this software and associated documentation files (the "Software"), to deal
14
+ in the Software without restriction, including without limitation the rights
15
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
16
+ copies of the Software, and to permit persons to whom the Software is
17
+ furnished to do so, subject to the following conditions:
18
+
19
+ The above copyright notice and this permission notice shall be included in all
20
+ copies or substantial portions of the Software.
21
+
22
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
23
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
24
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
25
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
26
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
27
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
28
+ SOFTWARE.
29
+ License-File: LICENSE
30
+ Keywords: artificial intelligence,attention mechanisms,deep learning
31
+ Classifier: Development Status :: 4 - Beta
32
+ Classifier: Intended Audience :: Developers
33
+ Classifier: License :: OSI Approved :: MIT License
34
+ Classifier: Programming Language :: Python :: 3.10
35
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
36
+ Requires-Python: >=3.10
37
+ Requires-Dist: einops>=0.8.1
38
+ Requires-Dist: python-statemachine>=2.5
39
+ Requires-Dist: torch-einops-utils
40
+ Requires-Dist: torch>=2.5
41
+ Provides-Extra: examples
42
+ Provides-Extra: test
43
+ Requires-Dist: pytest; extra == 'test'
44
+ Description-Content-Type: text/markdown
45
+
46
+ <img src="./fig1.png" width="400px"></img>
47
+
48
+ ## Declarative Attention - (wip)
49
+
50
+ 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
51
+
52
+ ## Install
53
+
54
+ ```bash
55
+ $ pip install declarative-attention
56
+ ```
57
+
58
+ ## Usage
59
+
60
+ ```python
61
+ import torch
62
+ from declarative_attention import DeclarativeStateMachine
63
+
64
+ # Define context chunk spans (1-indexed)
65
+ chunk_spans = [
66
+ (16, 1024), # Chunk 1
67
+ (1024, 2048), # Chunk 2
68
+ (2048, 3072), # Chunk 3
69
+ ]
70
+
71
+ state_machine = DeclarativeStateMachine(
72
+ chunk_spans = chunk_spans,
73
+ tokenizer_decode = tokenizer.decode # or lambda t: tokenizer.decode([t])
74
+ )
75
+
76
+ # Decode loop: feed generated tokens (int, torch.long scalar, or str)
77
+ state_machine.step('<focus chunks="1,2">')
78
+
79
+ assert state_machine.is_focus
80
+ assert state_machine.active_chunks == {1, 2}
81
+
82
+ # Obtain 1D boolean attention mask for current decode step
83
+ mask = state_machine.get_mask(total_len = 4096) # True for kept keys, False for masked out
84
+
85
+ # Model finishes chunk reasoning and reverts back to global
86
+ state_machine.step('</focus>')
87
+ assert state_machine.is_global
88
+ ```
89
+
90
+ ### Custom Attention Patterns
91
+
92
+ Researchers can easily declare custom attention patterns and tags.
93
+
94
+ #### 1. Logarithmically Spaced Attention (Recent to Past)
95
+
96
+ Sample chunks with exponentially decaying density into the past (e.g. current, $t-1, t-2, t-4, t-8\dots$):
97
+
98
+ ```python
99
+ state_machine = DeclarativeStateMachine(chunk_spans)
100
+
101
+ @state_machine.on('log_sparse')
102
+ def handle_log_sparse(machine, base = 2):
103
+ base = int(base)
104
+ total = len(machine.chunk_spans)
105
+
106
+ # Powers of base distance from current chunk: 0, 1, 2, 4, 8...
107
+ offsets = [0] + [base ** i for i in range(10)]
108
+ machine.active_chunks = {
109
+ total - d for d in offsets if (total - d) in machine.chunk_spans
110
+ }
111
+
112
+ # When model emits <log_sparse> or <log_sparse base="3">
113
+ state_machine.step('<log_sparse base="2">')
114
+
115
+ # Reverts back to global when closed
116
+ state_machine.step('</log_sparse>')
117
+ ```
118
+
119
+ #### 2. Sliding Window (Last $K$ Chunks)
120
+
121
+ Select the most recent $K$ context chunks via a custom `State`:
122
+
123
+ ```python
124
+ from statemachine import State
125
+ from declarative_attention import DeclarativeStateMachine
126
+
127
+ class CustomAttentionMachine(DeclarativeStateMachine):
128
+ window_mode = State()
129
+
130
+ window = window_mode.from_(DeclarativeStateMachine.global_mode)
131
+ revert = DeclarativeStateMachine.global_mode.from_(
132
+ DeclarativeStateMachine.focus_mode,
133
+ DeclarativeStateMachine.local_mode,
134
+ window_mode
135
+ )
136
+
137
+ @window.on
138
+ def on_window(self, size = 2):
139
+ total = len(self.chunk_spans)
140
+ self.active_chunks = {total - i for i in range(int(size)) if total - i > 0}
141
+
142
+ state_machine = CustomAttentionMachine(chunk_spans)
143
+
144
+ # Model calls <window size="2">
145
+ state_machine.step('<window size="2">')
146
+ assert state_machine.is_window_mode
147
+ ```
148
+
149
+ #### 3. Dynamic Temperature & In-Context Verification
150
+
151
+ Models can dynamically adjust sampling temperature during decode—brainstorming at high temperature, then re-attending to those thoughts at zero temperature for critical verification:
152
+
153
+ ```python
154
+ state_machine = DeclarativeStateMachine(chunk_spans)
155
+ state_machine.temperature = 0.7 # default
156
+
157
+ @state_machine.on('explore')
158
+ def on_explore(machine, temp = 1.2):
159
+ machine.temperature = float(temp)
160
+
161
+ @state_machine.on('verify')
162
+ def on_verify(machine, temp = 0.0):
163
+ machine.temperature = float(temp)
164
+
165
+ # High-temperature exploration
166
+ state_machine.step('<explore temp="1.2">')
167
+
168
+ # Switch to zero-temperature verification on Chunk 1
169
+ state_machine.step('</explore><focus chunks="1"><verify temp="0.0">')
170
+ ```
171
+
172
+ ## Citations
173
+
174
+ ```bibtex
175
+ @misc{ho2026languagemodelscontrolattention,
176
+ title = {Language Models Can Control Their Own Attention},
177
+ author = {Namgyu Ho and Huzama Ahmad and Woosung Koh and Se-Young Yun and Tal Schuster and Cicero Nogueira dos Santos},
178
+ year = {2026},
179
+ eprint = {2609.02737},
180
+ archivePrefix = {arXiv},
181
+ primaryClass = {cs.CL},
182
+ url = {https://arxiv.org/abs/2609.02737},
183
+ }
184
+ ```
@@ -0,0 +1,139 @@
1
+ <img src="./fig1.png" width="400px"></img>
2
+
3
+ ## Declarative Attention - (wip)
4
+
5
+ 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
6
+
7
+ ## Install
8
+
9
+ ```bash
10
+ $ pip install declarative-attention
11
+ ```
12
+
13
+ ## Usage
14
+
15
+ ```python
16
+ import torch
17
+ from declarative_attention import DeclarativeStateMachine
18
+
19
+ # Define context chunk spans (1-indexed)
20
+ chunk_spans = [
21
+ (16, 1024), # Chunk 1
22
+ (1024, 2048), # Chunk 2
23
+ (2048, 3072), # Chunk 3
24
+ ]
25
+
26
+ state_machine = DeclarativeStateMachine(
27
+ chunk_spans = chunk_spans,
28
+ tokenizer_decode = tokenizer.decode # or lambda t: tokenizer.decode([t])
29
+ )
30
+
31
+ # Decode loop: feed generated tokens (int, torch.long scalar, or str)
32
+ state_machine.step('<focus chunks="1,2">')
33
+
34
+ assert state_machine.is_focus
35
+ assert state_machine.active_chunks == {1, 2}
36
+
37
+ # Obtain 1D boolean attention mask for current decode step
38
+ mask = state_machine.get_mask(total_len = 4096) # True for kept keys, False for masked out
39
+
40
+ # Model finishes chunk reasoning and reverts back to global
41
+ state_machine.step('</focus>')
42
+ assert state_machine.is_global
43
+ ```
44
+
45
+ ### Custom Attention Patterns
46
+
47
+ Researchers can easily declare custom attention patterns and tags.
48
+
49
+ #### 1. Logarithmically Spaced Attention (Recent to Past)
50
+
51
+ Sample chunks with exponentially decaying density into the past (e.g. current, $t-1, t-2, t-4, t-8\dots$):
52
+
53
+ ```python
54
+ state_machine = DeclarativeStateMachine(chunk_spans)
55
+
56
+ @state_machine.on('log_sparse')
57
+ def handle_log_sparse(machine, base = 2):
58
+ base = int(base)
59
+ total = len(machine.chunk_spans)
60
+
61
+ # Powers of base distance from current chunk: 0, 1, 2, 4, 8...
62
+ offsets = [0] + [base ** i for i in range(10)]
63
+ machine.active_chunks = {
64
+ total - d for d in offsets if (total - d) in machine.chunk_spans
65
+ }
66
+
67
+ # When model emits <log_sparse> or <log_sparse base="3">
68
+ state_machine.step('<log_sparse base="2">')
69
+
70
+ # Reverts back to global when closed
71
+ state_machine.step('</log_sparse>')
72
+ ```
73
+
74
+ #### 2. Sliding Window (Last $K$ Chunks)
75
+
76
+ Select the most recent $K$ context chunks via a custom `State`:
77
+
78
+ ```python
79
+ from statemachine import State
80
+ from declarative_attention import DeclarativeStateMachine
81
+
82
+ class CustomAttentionMachine(DeclarativeStateMachine):
83
+ window_mode = State()
84
+
85
+ window = window_mode.from_(DeclarativeStateMachine.global_mode)
86
+ revert = DeclarativeStateMachine.global_mode.from_(
87
+ DeclarativeStateMachine.focus_mode,
88
+ DeclarativeStateMachine.local_mode,
89
+ window_mode
90
+ )
91
+
92
+ @window.on
93
+ def on_window(self, size = 2):
94
+ total = len(self.chunk_spans)
95
+ self.active_chunks = {total - i for i in range(int(size)) if total - i > 0}
96
+
97
+ state_machine = CustomAttentionMachine(chunk_spans)
98
+
99
+ # Model calls <window size="2">
100
+ state_machine.step('<window size="2">')
101
+ assert state_machine.is_window_mode
102
+ ```
103
+
104
+ #### 3. Dynamic Temperature & In-Context Verification
105
+
106
+ Models can dynamically adjust sampling temperature during decode—brainstorming at high temperature, then re-attending to those thoughts at zero temperature for critical verification:
107
+
108
+ ```python
109
+ state_machine = DeclarativeStateMachine(chunk_spans)
110
+ state_machine.temperature = 0.7 # default
111
+
112
+ @state_machine.on('explore')
113
+ def on_explore(machine, temp = 1.2):
114
+ machine.temperature = float(temp)
115
+
116
+ @state_machine.on('verify')
117
+ def on_verify(machine, temp = 0.0):
118
+ machine.temperature = float(temp)
119
+
120
+ # High-temperature exploration
121
+ state_machine.step('<explore temp="1.2">')
122
+
123
+ # Switch to zero-temperature verification on Chunk 1
124
+ state_machine.step('</explore><focus chunks="1"><verify temp="0.0">')
125
+ ```
126
+
127
+ ## Citations
128
+
129
+ ```bibtex
130
+ @misc{ho2026languagemodelscontrolattention,
131
+ title = {Language Models Can Control Their Own Attention},
132
+ author = {Namgyu Ho and Huzama Ahmad and Woosung Koh and Se-Young Yun and Tal Schuster and Cicero Nogueira dos Santos},
133
+ year = {2026},
134
+ eprint = {2609.02737},
135
+ archivePrefix = {arXiv},
136
+ primaryClass = {cs.CL},
137
+ url = {https://arxiv.org/abs/2609.02737},
138
+ }
139
+ ```
@@ -0,0 +1,7 @@
1
+ from declarative_attention.declarative_attention import (
2
+ DeclarativeStateMachine
3
+ )
4
+
5
+ __all__ = [
6
+ 'DeclarativeStateMachine'
7
+ ]
@@ -0,0 +1,198 @@
1
+ from __future__ import annotations
2
+ from html.parser import HTMLParser
3
+ from typing import Callable, Sequence
4
+
5
+ import torch
6
+ from torch import Tensor
7
+ from statemachine import StateMachine, State
8
+ from torch_einops_utils import maybe
9
+
10
+ # helpers
11
+
12
+ def exists(v):
13
+ return v is not None
14
+
15
+ def default(v, d):
16
+ return v if exists(v) else d
17
+
18
+ def default_decode_fn(token: int) -> str:
19
+ return chr(token) if 0 <= token < 128 else ""
20
+
21
+ # types
22
+
23
+ ChunkId = int
24
+ StartPos = int
25
+ EndPos = int
26
+ Span = tuple[StartPos, EndPos]
27
+
28
+ # chunk spans can be:
29
+ # 1. dict mapping 1-indexed chunk id to token span:
30
+ # {1: (16, 32), 2: (32, 48)}
31
+ # 2. sequence of spans (auto-indexed from 1):
32
+ # [(16, 32), (32, 48)]
33
+
34
+ ChunkSpans = (
35
+ dict[ChunkId, Span] |
36
+ Sequence[Span]
37
+ )
38
+
39
+ TokenizerDecode = Callable[[int], str]
40
+
41
+ # general streaming tag parser
42
+
43
+ class StreamingTagParser(HTMLParser):
44
+ def __init__(self, on_tag: Callable[[bool, str, dict[str, str]], None]):
45
+ super().__init__()
46
+ self.on_tag = on_tag
47
+
48
+ def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]):
49
+ self.on_tag(False, tag, dict(attrs))
50
+
51
+ def handle_endtag(self, tag: str):
52
+ self.on_tag(True, tag, {})
53
+
54
+ # declarative attention state machine (Ho et al., 2026)
55
+
56
+ class DeclarativeStateMachine(StateMachine):
57
+ """
58
+ Declarative Attention State Machine
59
+
60
+ - <global> : attends to all context chunks (default)
61
+ - <focus chunks="1,2"> : attends only to specified chunk(s)
62
+ - <local> : attends to 0 context chunks
63
+ """
64
+
65
+ global_mode = State(initial = True)
66
+ focus_mode = State()
67
+ local_mode = State()
68
+
69
+ focus = focus_mode.from_(global_mode, local_mode)
70
+ local = local_mode.from_(global_mode, focus_mode)
71
+ revert = global_mode.from_(focus_mode, local_mode)
72
+
73
+ def __init__(
74
+ self,
75
+ chunk_spans: ChunkSpans,
76
+ tokenizer_decode: TokenizerDecode | None = None
77
+ ):
78
+ if isinstance(chunk_spans, (list, tuple)):
79
+ chunk_spans = {i + 1: span for i, span in enumerate(chunk_spans)}
80
+
81
+ self.chunk_spans = chunk_spans
82
+ self.active_chunks = set(chunk_spans.keys())
83
+ self.parser = StreamingTagParser(self._handle_tag)
84
+ self.tokenizer_decode = default(tokenizer_decode, default_decode_fn)
85
+ self.handlers: dict[str, Callable] = dict()
86
+
87
+ super().__init__()
88
+
89
+ # register custom tag handlers for arbitrary researcher logic
90
+
91
+ def on(self, tag: str, fn: Callable | None = None):
92
+ def decorator(handler: Callable):
93
+ self.handlers[tag.lower()] = handler
94
+ return handler
95
+
96
+ return maybe(decorator, default = decorator)(fn)
97
+
98
+ # dispatch streaming tags to transitions or custom handlers
99
+
100
+ def _handle_tag(self, is_closing: bool, tag: str, attrs: dict[str, str]):
101
+ tag = tag.lower()
102
+
103
+ if is_closing:
104
+ if hasattr(self, f"revert_{tag}"):
105
+ self.send(f"revert_{tag}")
106
+ elif not self.is_global:
107
+ try: self.revert()
108
+ except Exception: pass
109
+ return
110
+
111
+ # opening tag: try transition event first, then custom handler
112
+
113
+ if tag in [e.id for e in self.events]:
114
+ try:
115
+ self.send(tag, **attrs)
116
+ return
117
+ except Exception:
118
+ pass
119
+
120
+ if tag in self.handlers:
121
+ self.handlers[tag](self, **attrs)
122
+
123
+ # transition callbacks
124
+
125
+ @focus.on
126
+ def _on_focus(
127
+ self,
128
+ chunks: str | Sequence[int] | None = None,
129
+ chunk: str | int | None = None,
130
+ magic_chunks: str | None = None
131
+ ):
132
+ target = default(chunks, default(chunk, magic_chunks))
133
+
134
+ if isinstance(target, str):
135
+ self.active_chunks = {int(c) for c in target.split(',') if c.strip().isdigit()}
136
+ elif isinstance(target, int):
137
+ self.active_chunks = {target}
138
+ elif exists(target):
139
+ self.active_chunks = set(target)
140
+ else:
141
+ self.active_chunks = set()
142
+
143
+ @local.on
144
+ def _on_local(self):
145
+ self.active_chunks = set()
146
+
147
+ @revert.on
148
+ def _on_revert(self):
149
+ self.active_chunks = set(self.chunk_spans.keys())
150
+
151
+ # properties
152
+
153
+ @property
154
+ def is_global(self):
155
+ return self.global_mode.is_active
156
+
157
+ @property
158
+ def is_focus(self):
159
+ return self.focus_mode.is_active
160
+
161
+ @property
162
+ def is_local(self):
163
+ return self.local_mode.is_active
164
+
165
+ @property
166
+ def mode(self):
167
+ return "global" if self.is_global else ("focus" if self.is_focus else "local")
168
+
169
+ @property
170
+ def is_streaming_tag(self) -> bool:
171
+ """True if the parser is currently mid-tag (e.g. between '<' and '>')"""
172
+ return "<" in self.parser.rawdata
173
+
174
+ # token streaming (accepts int, torch LongTensor, or str)
175
+
176
+ def step(self, token: int | Tensor | str):
177
+ if isinstance(token, Tensor):
178
+ token = token.item()
179
+
180
+ text = self.tokenizer_decode(token) if isinstance(token, int) else str(token)
181
+ self.parser.feed(text)
182
+
183
+ # attention mask generation
184
+
185
+ def get_mask(self, total_len: int, device = None) -> Tensor:
186
+ mask = torch.ones(total_len, dtype = torch.bool, device = device)
187
+
188
+ if self.is_global:
189
+ return mask
190
+
191
+ for chunk_id, (start, end) in self.chunk_spans.items():
192
+ if chunk_id not in self.active_chunks:
193
+ mask[start:end] = False
194
+
195
+ return mask
196
+
197
+ def __call__(self, total_len: int, device = None) -> Tensor:
198
+ return self.get_mask(total_len, device = device)
@@ -0,0 +1,67 @@
1
+ [project]
2
+ name = "declarative-attention"
3
+ version = "0.0.4"
4
+ description = "Declarative Attention - Pytorch"
5
+ authors = [
6
+ { name = "Phil Wang", email = "lucidrains@gmail.com" }
7
+ ]
8
+ readme = "README.md"
9
+ requires-python = ">= 3.10"
10
+ license = { file = "LICENSE" }
11
+ keywords = [
12
+ 'artificial intelligence',
13
+ 'deep learning',
14
+ 'attention mechanisms'
15
+ ]
16
+
17
+ classifiers=[
18
+ 'Development Status :: 4 - Beta',
19
+ 'Intended Audience :: Developers',
20
+ 'Topic :: Scientific/Engineering :: Artificial Intelligence',
21
+ 'License :: OSI Approved :: MIT License',
22
+ 'Programming Language :: Python :: 3.10',
23
+ ]
24
+
25
+ dependencies = [
26
+ "einops>=0.8.1",
27
+ "python-statemachine>=2.5",
28
+ "torch-einops-utils",
29
+ "torch>=2.5",
30
+ ]
31
+
32
+ [project.urls]
33
+ Homepage = "https://pypi.org/project/declarative-attention/"
34
+ Repository = "https://github.com/lucidrains/declarative-attention"
35
+
36
+ [project.optional-dependencies]
37
+ examples = []
38
+ test = [
39
+ "pytest"
40
+ ]
41
+
42
+ [tool.pytest.ini_options]
43
+ pythonpath = [
44
+ "."
45
+ ]
46
+
47
+ [build-system]
48
+ requires = ["hatchling"]
49
+ build-backend = "hatchling.build"
50
+
51
+ [tool.rye]
52
+ managed = true
53
+ dev-dependencies = []
54
+
55
+ [tool.hatch.metadata]
56
+ allow-direct-references = true
57
+
58
+ [tool.hatch.build]
59
+ include = [
60
+ "declarative_attention",
61
+ "pyproject.toml",
62
+ "README.md",
63
+ "LICENSE"
64
+ ]
65
+
66
+ [tool.hatch.build.targets.wheel]
67
+ packages = ["declarative_attention"]