A minimal replication of finding and evaluating contextual n-grams in Pythia series models.
Use the included Dockerfile, or alternatively install PyTorch then run:
pip install nltk kaleido tqdm einops seaborn plotly-express fancy-einsum scikit-learn torchmetrics ipykernel ipywidgets nbformat git+https://github.com/neelnanda-io/TransformerLens git+https://github.com/callummcdougall/CircuitsVis.git#subdirectory=python git+https://github.com/neelnanda-io/neelutils.git git+https://github.com/neelnanda-io/neel-plotly.git
Generate data by running each script from the command line:
python generate_foo.py --model pythia-70m
Some scripts are extremely slow because they run over hundreds of model checkpoints. We advise using an A6000 with 100GB of RAM or equivalent.
Then replicate figures by running figures.py
python figures.py --model pythia-70m