Metadata-Version: 2.4
Name: attnread
Version: 0.1.2
Summary: Recover attention weights without changing the attention backend.
License-Expression: MIT
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: nnsight==0.7.0
Requires-Dist: torch<2.15,>=2.14
Requires-Dist: transformers<5.20,>=5.19
Provides-Extra: dev
Requires-Dist: pytest<10,>=9; extra == "dev"
Requires-Dist: build<2,>=1.2; extra == "dev"
Requires-Dist: twine<7,>=6; extra == "dev"
Dynamic: license-file

# AttnRead

**Recover attention weights without changing the attention backend.**

AttnRead is a lightweight Python toolkit for extracting attention patterns from optimized transformer attention implementations, including PyTorch SDPA and FlashAttention 2.

Using identity-value passes and activation restoration, AttnRead recovers attention patterns while preserving the model's downstream computation. No fallback to eager attention required.

Version **0.1.2**, distributed under the MIT License. The distribution and
Python import name are both `attnread`.

## Installation

Use Python 3.11 or newer in a venv. Install your preferred CPU/CUDA PyTorch
build first, then install AttnRead:

```bash
python -m pip install attnread==0.1.2
```

## Usage

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from attnread import extract_attention

checkpoint = "Qwen/Qwen3-4B"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval().requires_grad_(False)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer("Describe the colors of an autumn leaf.",
                   return_tensors="pt").to(model.device)

patterns = extract_attention(model, inputs, heads=[(0, 3), (12, 7)])
head = patterns[("text", 0, 3)]  # CPU tensor: [batch, query tokens, key tokens].
```

Use the same call with `AutoModelForImageTextToText` and processor outputs:

```python
# (layer, head) means text attention; qualify a vision head with its stream.
patterns = extract_attention(
    vlm, processed_inputs, heads=[(12, 3), ("vision", 4, 2)],
)
text = patterns[("text", 12, 3)]
vision = patterns[("vision", 4, 2)]
```

**Selection.** `heads=None` extracts every discovered text head. A sequence of
`(layer, head)` or `(stream, layer, head)` addresses extracts only those heads;
returned keys always use the full triple, in request order. Layers and heads are
zero-based within each stream, following discovered module order. Use explicit paths for sparse/hybrid stacks or unusual
module naming. Different streams can have different sequence lengths, so results
are a dictionary. Probes are grouped by layer and save/transfer only selected
heads. The underlying model still computes all heads; head selection reduces
saved memory, not the number of identity probes. Start with one short prefix and
one head. Unknown, duplicate or out-of-range addresses raise `ValueError`.

**Model interface.** Pass an ordinary HF model in evaluation mode and its
preprocessed tensor input mapping. Discovery recognizes separate `q_proj`,
`k_proj`, `v_proj` with `o_proj`/`out_proj`, GPT-2's fused `c_attn`/`c_proj`, and
Qwen3-VL vision `qkv`/`proj`. Vision module paths identify the `vision` stream;
others belong to `text`. For unusual or ambiguous naming, supply explicit paths:

```python
patterns = extract_attention(model, inputs, heads=[("text", 0, 1)],
    attention_modules={"text": ["model.layers.7.self_attn"]})
# With explicit lists, layer 0 denotes the first entry in that list.
```

**Dependencies.** AttnRead requires Python 3.11 or newer, NNsight 0.7.0,
Transformers 5.19.x, and PyTorch 2.14.x. Install a PyTorch build suitable for your
CPU or GPU. FlashAttention 2 is optional and installed separately.

## Development

From the development checkout:

```bash
python -m venv .venv
source .venv/bin/activate
# Install your preferred CPU/CUDA PyTorch build first.
python -m pip install -e '.[dev]'
python -m pytest -q test_attnread.py
python -m build
python -m twine check --strict dist/*
```

Tests and the example notebook are development-only and excluded from release
archives.

## License

MIT License. Copyright (c) 2026 Etha Tianze Hua. The full license text is included in the
source archive and wheel.
