Metadata-Version: 2.1
Name: token-trace
Version: 0.3.2
Summary: Transformer token flow visualizer
Author-Email: "Seongho Son, Liza Karmannaya, Daniel Tan, David Chanin" <chanindav@gmail.com>
License: MIT
Requires-Python: >=3.11
Requires-Dist: transformers>=4.39.3
Requires-Dist: transformer-lens>=1.15.0
Requires-Dist: sae-lens>=0.2.0
Requires-Dist: pandera>=0.18.3
Requires-Dist: streamlit>=1.33.0; extra == "client"
Requires-Dist: st-annotated-text>=4.0.1; extra == "client"
Requires-Dist: pytest>=8.1.1; extra == "client"
Requires-Dist: graphviz>=0.20.3; extra == "client"
Provides-Extra: client
Description-Content-Type: text/markdown

# Token-Trace

A tool and UI to construct prompt-centric views of SAE feature attributions.

Main functionality:
- We use this tool to identify which SAE features have the most 'attribution' towards decreasing the model loss. 
- In combination with Neuronpedia, we can identify what each SAE feature represents; this then gives us a rough idea of what computation the model is performing. 

This tool is a first step towards discovering information flow between the features / layers of a transformer


# Quickstart

## Installation

```
pip install token-trace
```
## Example Usage

```
from token_trace import compute_node_attribution

text = "When John and Mary went to the shops, John gave the bag to Mary"

df: pd.DataFrame = compute_node_attribution(
    model_name = "gpt2",
    text
)
```

Each row of `df` describes one node corresponding to an SAE feature or error term. 

## Visualizing SAE attribution statistics in frontend. 

We use [Streamlit](https://streamlit.io/) to create a UI. Start the app as follows:
```
streamlit run app/token_trace_app.py
```

## Methodology

Under the hood, we use attribution patching to compute indirect effect of the loss with respect to SAE features. The method is adapted heavily from [Sparse Feature Circuits](https://arxiv.org/abs/2403.19647). 

## Development

We use [PDM](https://github.com/pdm-project/pdm) to manage dependencies. Set up a development environment as follows:
```
pdm install # creates a .venv
source .venv/bin/activate
```

Once in the virtual environment, make sure to also install the pre-commit hooks
```
pre-commit install
```