Merge branch 'main' into olmo2

Author: aryamanarora

.github/workflows/documentation.yml

@@ -14,7 +14,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install -r docs/requirements.txt
-          pip install -r requirements.txt
+          pip install -e ".[dev]"
       - name: Sphinx build
         run: |
           sphinx-build docs/source docs/build

.github/workflows/python-app.yml

@@ -23,12 +23,11 @@ jobs:
       uses: actions/setup-python@v3
       with:
         python-version: "3.10"
-        cache: 'pip'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install flake8 pytest
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+        if [ -f pyproject.toml ]; then pip install -e ".[dev]"; fi
     - name: Lint with flake8
       run: |
         # stop the build if there are Python syntax errors or undefined names

MANIFEST.in

@@ -1,6 +1,6 @@
 <br />
 <div align="center">
-  <h1 align="center"><img src="https://i.ibb.co/BNkhQH3/pyvene-logo.png"></h1>
+  <h1 align="center"><img src="https://i.ibb.co/BNkhQH3/pyvene-logo.png" height="100"></h1>
   <a href="https://arxiv.org/abs/2403.07809"><strong>Read our paper »</strong></a> | <a href="https://stanfordnlp.github.io/pyvene/"><strong>Read the docs »</strong></a>
 </div>     
 

README.md

@@ -0,0 +1,78 @@
+[project]
+name = "pyvene"
+version = "0.1.8"
+description = "Use Activation Intervention to Interpret Causal Mechanism of Model"
+readme = "README.md"
+authors = [
+    { name="Zhengxuan Wu", email="[email protected]" }
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: POSIX :: Linux",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+]
+license = { text = "Apache License 2.0" }
+requires-python = ">=3.9"
+dependencies = [
+    "torch>=2.0.0",
+    "transformers>=4.55.0.dev0",
+    "tokenizers>=0.20.0",
+    "datasets>=3.0.1",
+    "protobuf>=3.20.0",
+    "matplotlib>=3.7.4",
+    "ipywidgets>=8.1.1",
+    "plotnine>=0.12.4",
+    "huggingface-hub>=0.25.1",
+    "numpy>1.24.4",
+    "fsspec>=2023.6.0",
+    "accelerate>=0.34.2",
+    "sentencepiece>=0.2.0",
+]
+
+[dependency-groups]
+dev = [
+    "flake8>=7.1.1",
+    "pytest>=8.3.4",
+    "ipykernel>=6.29.5",
+]
+
+[project.urls]
+Documentation = "https://stanfordnlp.github.io/pyvene"
+Homepage = "https://github.com/stanfordnlp/pyvene"
+Repository = "https://github.com/stanfordnlp/pyvene.git"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["pyvene"]
+
+[tool.hatch.build]
+include = [
+    "pyvene/**/*.py",
+    "pyvene/**/*.json",
+    "pyvene/**/*.yaml",
+    "pyvene/**/*.yml",
+    "pyvene/**/*.txt",
+    "pyvene/**/*.md",
+    "pyvene/**/*.ipynb",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["*Test", "Test*"]
+python_functions = ["test_*"]
+addopts = "-v"
+
+[tool.flake8]
+max-line-length = 127
+extend-ignore = ["E203"]
+exclude = [".git", "__pycache__", "build", "dist", "*.egg-info"]
+per-file-ignores = [
+    "__init__.py: F401"
+]

pyproject.toml

@@ -2,7 +2,11 @@
 
 # Generic APIs
 from .data_generators.causal_model import CausalModel
-from .models.intervenable_base import IntervenableModel, IntervenableNdifModel, build_intervenable_model
+from .models.intervenable_base import (
+    IntervenableModel,
+    IntervenableNdifModel,
+    build_intervenable_model,
+)
 from .models.configuration_intervenable_model import IntervenableConfig
 from .models.configuration_intervenable_model import RepresentationConfig
 
@@ -37,7 +41,10 @@
 # Utils
 from .models.basic_utils import *
 from .models.intervention_utils import _do_intervention_by_swap
-from .models.intervenable_modelcard import type_to_module_mapping, type_to_dimension_mapping
+from .models.intervenable_modelcard import (
+    type_to_module_mapping,
+    type_to_dimension_mapping,
+)
 from .models.gpt2.modelings_intervenable_gpt2 import create_gpt2
 from .models.gpt2.modelings_intervenable_gpt2 import create_gpt2_lm
 from .models.blip.modelings_intervenable_blip import create_blip
@@ -51,5 +58,8 @@
 from .models.gru.modelings_gru import GRUConfig
 from .models.llama.modelings_intervenable_llama import create_llama
 from .models.mlp.modelings_intervenable_mlp import create_mlp_classifier
-from .models.backpack_gpt2.modelings_intervenable_backpack_gpt2 import create_backpack_gpt2
-
+from .models.backpack_gpt2.modelings_intervenable_backpack_gpt2 import (
+    create_backpack_gpt2,
+)
+from .models.olmo.modelings_intervenable_olmo import create_olmo
+from .models.gpt_oss.modelings_intervenable_gpt_oss import create_gpt_oss

pyvene/__init__.py

@@ -0,0 +1,24 @@
+## How to add new models?
+
+You can prompt a LM to generate files, or modifying existing ones in this folder by simply following these steps:
+
+- Get the relevent implementation file from `https://github.com/huggingface/transformers/blob/main/src/transformers/models/` (e.g., the implementation for `gpt-oss` [here](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_oss/modeling_gpt_oss.py)).
+
+- Copy the whole transformer model src file.
+
+- Create a new folder for your new model.
+
+- Move one of the existing model file to your new folder (e.g., `/gpt2/modelings_intervenable_gpt2.py` along with the default `__init__.py` file).
+
+- Prompt a language model with the following template:
+
+```text
+[YOUR_EXAMPLE_PYVENE_MODEL_FILE_COPY]
+
+Generate a new mapping file based on the existing one above for the following new model:
+
+[HF_TRANSFORMER_MODEL_SRC_FILE_COPY]
+
+You also need to pay attention to these details:
+- [OTHER_REQ_GOES_HERE] (e.g., you need to take care of the MoE strcuture)
+```

pyvene/models/README.md

@@ -22,7 +22,7 @@
     # 'vis.attention_output': ("vision_model.encoder.layers[%s].self_attn", CONST_OUTPUT_HOOK),
     # 'vis.attention_input': ("vision_model.encoder.layers[%s].self_attn", CONST_INPUT_HOOK),
     "block_input": ("text_encoder.encoder.layer[%s]", CONST_INPUT_HOOK),
-    "block_output": ("text_encoder.encoder.layer[%s]", CONST_INPUT_HOOK),
+    "block_output": ("text_encoder.encoder.layer[%s]", CONST_OUTPUT_HOOK),
     "mlp_activation": (
         "text_encoder.encoder.layer[%s].intermediate.dense",
         CONST_OUTPUT_HOOK,

pyvene/models/blip/modelings_intervenable_blip.py

@@ -22,7 +22,7 @@
     # 'vis.attention_output': ("vision_model.encoder.layers[%s].self_attn", CONST_OUTPUT_HOOK),
     # 'vis.attention_input': ("vision_model.encoder.layers[%s].self_attn", CONST_INPUT_HOOK),
     "block_input": ("text_encoder.encoder.layer[%s]", CONST_INPUT_HOOK),
-    "block_output": ("text_encoder.encoder.layer[%s]", CONST_INPUT_HOOK),
+    "block_output": ("text_encoder.encoder.layer[%s]", CONST_OUTPUT_HOOK),
     "mlp_activation": (
         "text_encoder.encoder.layer[%s].intermediate.dense",
         CONST_OUTPUT_HOOK,