Metadata-Version: 2.1
Name: tms-kit
Version: 0.4.0
Summary: Toolkit for running toy models of superposition experiments
Author-Email: Daniel Tan <dtch1997@users.noreply.github.com>
License: MIT
Requires-Python: >=3.11
Requires-Dist: torch>=2.4.1
Requires-Dist: numpy>=2.1.1
Requires-Dist: jaxtyping>=0.2.34
Requires-Dist: tqdm>=4.66.5
Requires-Dist: einops>=0.8.0
Requires-Dist: plotly>=5.24.1
Requires-Dist: pandas>=2.2.3
Requires-Dist: matplotlib>=3.9.2
Requires-Dist: torchmetrics>=1.4.2
Requires-Dist: rich>=13.9.2
Provides-Extra: dev
Requires-Dist: sphinx>=8.0.2; extra == "dev"
Requires-Dist: ruff>=0.6.9; extra == "dev"
Requires-Dist: pytest>=8.3.3; extra == "dev"
Requires-Dist: furo>=2024.8.6; extra == "dev"
Provides-Extra: tutorials
Requires-Dist: kaleido==0.2.1; extra == "tutorials"
Requires-Dist: rich>=13.9.2; extra == "tutorials"
Requires-Dist: simple-parsing>=0.1.6; extra == "tutorials"
Description-Content-Type: text/markdown

# TMS

![tests](https://github.com/dtch1997/tms/actions/workflows/tests.yaml/badge.svg)

![5_2_superposition](experiments/bottleneck_tms/5_2_superposition.png)

Codebase for quickly implementing experiments with toy models of superposition. 

## Quickstart

To set up the environment, see [setup docs](docs/setup.md) for detailed instructions. 

To reproduce figures from Anthropic's [Toy Models of Superposition](https://transformer-circuits.pub/2022/toy_model/index.html), see [sample script](experiments/bottleneck_tms/run.py) 

## Acknowledgements

This codebase is heavily adapted from the [ARENA 3.0 codebase](https://github.com/callummcdougall/ARENA_3.0/tree/main), designed and maintained by Callum McDougall. Many thanks to Callum and the ARENA team!