Metadata-Version: 2.4
Name: mechir
Version: 0.2.0
Summary: A package for mechanistic interpretability in Neural IR
Author-email: Andrew Parry <0andrewparry@gmail.com>, Catherine Chen <catherine_s_chen@brown.edu>
Maintainer-email: Andrew Parry <0andrewparry@gmail.com>
License-Expression: Apache-2.0
Project-URL: Repository, https://github.com/Parry-Parry/MechIR.git
Keywords: pytorch,transformers,ranking,information retrieval
Classifier: Programming Language :: Python
Classifier: Operating System :: OS Independent
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: transformer_lens<3,>=2.15
Requires-Dist: sae_lens>=6
Requires-Dist: torch
Requires-Dist: transformers<5,>=4.40
Requires-Dist: ir_datasets
Requires-Dist: numpy<2,>=1.26
Requires-Dist: datasets<4
Requires-Dist: pyarrow<21
Requires-Dist: jaxtyping
Requires-Dist: pandas
Provides-Extra: rag
Requires-Dist: datasets<4; extra == "rag"
Requires-Dist: ir_datasets; extra == "rag"
Provides-Extra: plotting
Requires-Dist: seaborn; extra == "plotting"
Requires-Dist: matplotlib; extra == "plotting"
Requires-Dist: plotly; extra == "plotting"
Requires-Dist: streamlit; extra == "plotting"
Provides-Extra: examples
Requires-Dist: fire; extra == "examples"
Requires-Dist: tqdm; extra == "examples"
Requires-Dist: python-terrier; extra == "examples"
Requires-Dist: pyterrier-dr; extra == "examples"
Provides-Extra: axiomatic
Requires-Dist: python-terrier; extra == "axiomatic"
Requires-Dist: ir_axioms; extra == "axiomatic"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: flake8; extra == "dev"
Requires-Dist: black==26.10.1; extra == "dev"
Dynamic: license-file

# MechIR

## An IR Library for Mechanistic Interpretability

This library exists to make encoders first-class citizens in Mechanistic Interpretability, particularly with a focus on Information Retrieval and representation learning. We presented this work as a demo at ECIR 2025: https://arxiv.org/pdf/2501.10165

We have done our best to support many of the features present in [transformer_lens](https://github.com/TransformerLensOrg/TransformerLens/tree/main). Additionally, we provide extensive functionality for activation patching (more details in our demo!).

We have not released any SAEs. `mechir.modelling.sae.SAEMixin` lets you attach an existing [sae_lens](https://github.com/jbloomAus/SAELens) 6 SAE to a `Dot`, `Cat`, `MonoT5` or `GenerativeRanker` model and run with it hooked in, but that path is currently untested — treat it as experimental.


Heavily inspired by the paper **Axiomatic Causal Interventions for Reverse Engineering Relevance Computation in Neural Retrieval Models** by Catherine Chen, Jack Merullo, and Carsten Eickhoff (Brown University, University of Tübingen), their original code can be found [here](https://github.com/catherineschen/axiomatic-ir-interventions/tree/main)

## Demonstration

Primary files can be found in the /notebooks folder. If you use our live versions and want to run our experiments make sure to choose a GPU instance of Colab. You can easily change our notebook to observe different behaviour so try your own experiments!

* experiment.ipynb: This notebook demonstrates how to use the MechIR library to perform activation patching on a simple neural retrieval model. Here is a live version [on Colab](https://drive.google.com/file/d/1ZR1ZD2bwcLh5mQG_x561PCpydHK4mhap/view?usp=sharing)
* activation_patching_considerations.ipynb: This notebook provides a more in-depth look at the activation patching process and the considerations that go into it. Here is a live version [on Colab](https://drive.google.com/file/d/1TSRCvixMo4YPRwDqsjf7kuci6fGwwQG2/view?usp=sharing)

## Installation

### Latest Release (Unstable)
```
pip install git+https://github.com/Parry-Parry/MechIR.git
```

### PyPI (Stable)
```
pip install mechir
```

### Extras

```
pip install "mechir[plotting]"    # seaborn / plotly / streamlit, for mechir.plotting
pip install "mechir[axiomatic]"   # pyterrier / ir_axioms, for index-backed perturbations
pip install "mechir[examples]"    # fire / pyterrier / pyterrier_dr, for examples/ and scripts/
```

Plotting dependencies are an extra as of 0.2.0. `mechir.plotting` imports
them lazily, so a missing one raises an error naming the extra rather than a
bare `ModuleNotFoundError`.

## Usage

### Models 

Currently we support common bi- and cross-encoder models for neural retrieval. The following models are available:

* Dot: A bi-encoder scored by dot product, with `cls`, `mean` or `last` pooling
* Cat: A cross-encoder with joint query-document encoding and a classification head
* MonoT5: An encoder-decoder cross-encoder scored by contrasting `true`/`false` token logits
* GenerativeRanker: A decoder-backbone reranker scored by contrasting two label-token logits, e.g. Qwen3-Reranker

Encoder backbones (BERT, DistilBERT, ELECTRA) and decoder backbones
(Qwen2/Qwen3, Llama, Mistral) are both supported; `Dot` picks the right one
from the checkpoint's architecture. Decoder bi-encoders such as
Qwen3-Embedding use `pooling_type="last"`.

To load a model, for example TAS-B, you can use the following code:

```python
from mechir import Dot

model = Dot('sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco')
```

### Datasets and Perturbations 

To probe model behaviour we need queries and documents, we retrieve these from ir-datasets though you can use your own (MechDataset). To load an IR dataset, for example MS MARCO, you can use the following code:

```python
from mechir.data import MechIRDataset

dataset = MechIRDataset('msmarco-passage/dev')
```

The second step of probing is to create a perturbation of text to observe how model behaviour changes, we can do this simply with the perturbation decorator:

```python
from mechir.perturb import perturbation

@perturbation
def my_perturbation(text):
    return text + "MechIR"
```

We can then apply this perturbation efficiently using our dataset and a torch dataloader

```python
from torch.utils.data import DataLoader
from mechir.data import DotDataCollator

collate_fn = DotDataCollator(model.tokenizer, transformation_func=my_perturbation)
dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
```

### Activation Patching

Activation patching is a method to isolate model behaviour to particular components of the model commonly attention heads. There are several ways to perform activation patching, a simple case is to patch all heads:

```python
import torch

patch_output = []
for batch in dataloader:
    # patch() returns (result, cache); cache is None unless the model was
    # built with return_cache=True.
    result, _ = model.patch(**batch, patch_type="head_all")
    patch_output.append(result)

patch_output = torch.mean(torch.stack(patch_output), dim=0)
```

We can then easily visualise the attention heads which activate strongly for our perturbation:

```python
from mechir.plotting import plot_components

plot_components(patch_output)
```

## Citation
If you use our work in research please let us know! We will showcase the work in this repository to encourage more encoder interpretability. Please use the following citation if you find our work useful!
```
@inproceedings{mechir:2025,
  author       = {Andrew Parry and
                  Catherine Chen and
                  Carsten Eickhoff and
                  Sean MacAvaney},
  title        = {MechIR: {A} Mechanistic Interpretability Framework for Information
                  Retrieval},
  booktitle    = {Advances in Information Retrieval - 47th European Conference on Information
                  Retrieval, {ECIR} 2025, Lucca, Italy, April 6-10, 2025, Proceedings,
                  Part {V}},
  series       = {Lecture Notes in Computer Science},
  volume       = {15576},
  pages        = {89--95},
  publisher    = {Springer},
  year         = {2025},
  url          = {https://doi.org/10.1007/978-3-031-88720-8\_16},
  doi          = {10.1007/978-3-031-88720-8\_16},
  biburl       = {https://dblp.org/rec/conf/ecir/ParryCEM25.bib},
  bibsource    = {dblp computer science bibliography, https://dblp.org}
}
```
