{"id":27400846,"url":"https://github.com/butanium/nnterp","last_synced_at":"2025-04-14T03:43:29.952Z","repository":{"id":252387434,"uuid":"839907396","full_name":"Butanium/nnterp","owner":"Butanium","description":"A small package implementing some useful wrapping around nnsight","archived":false,"fork":false,"pushed_at":"2025-02-13T09:53:50.000Z","size":126,"stargazers_count":9,"open_issues_count":2,"forks_count":2,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-03-14T02:18:18.846Z","etag":null,"topics":["mechanistic-interpretability","nnsight","patchscopes"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/Butanium.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-08-08T15:06:39.000Z","updated_at":"2025-02-20T19:05:12.000Z","dependencies_parsed_at":"2024-09-05T18:20:37.365Z","dependency_job_id":"92cedbcc-cb48-4671-97ec-8f558edc958f","html_url":"https://github.com/Butanium/nnterp","commit_stats":null,"previous_names":["butanium/nnterp"],"tags_count":11,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Butanium%2Fnnterp","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Butanium%2Fnnterp/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Butanium%2Fnnterp/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Butanium%2Fnnterp/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/Butanium","download_url":"https://codeload.github.com/Butanium/nnterp/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248819114,"owners_count":21166470,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["mechanistic-interpretability","nnsight","patchscopes"],"created_at":"2025-04-14T03:43:29.367Z","updated_at":"2025-04-14T03:43:29.943Z","avatar_url":"https://github.com/Butanium.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# nnterp\n\n## Installation\n- `pip install nnterp`\n- `pip install nnterp[display]` if you want to use the `display` module for visualizations\n\n## Usage\n### Loading a Model\n\nFirst, let's load a model in `nnsight` using `nnterp`'s `load_model` function.\n\n```python\nfrom nnterp import load_model\n\nmodel_name = \"meta-llama/Llama-2-7b-hf\"\n# Load the model (float16 and gpu by default)\nnn_model = load_model(model_name)\ntokenizer = nn_model.tokenizer\n```\n\n### Collecting activations\n\nTo collect activations from a model using `nnterp`, you can use the `collect_activations` function. This function takes the following parameters:\n\n- `nn_model`: The NNSight model.\n- `prompts`: The prompts for which you want to collect activations.\n- `layers`: The layers for which you want to collect activations. If not specified, activations will be collected for all layers.\n- `get_activations`: A function to get the activations. By default, it will collect the layer output, but you can also collect other things like attention input/output\n- `remote`: Whether to run the model on the remote device.\n- `idx`: The index of the token to collect activations for.\n- `open_context`: Whether to open a context for the model trace. You can set to false if you want to collect activations in an already opened nnsight tracing context.\n\n```python\nfrom nnterp import collect_activations\n\n# Load the model\nnn_model = load_model(model_name)\n\n# Create a prompt\nprompt = \"The quick brown fox jumps over the lazy dog\"\n\n# Collect activations for all layers\nactivations = collect_activations(nn_model, [prompt])\n\n# Print the activations\nfor layer, activation in enumerate(activations):\n    print(f\"Layer {layer}: {activation.shape}\")\n```\n\n### Collecting activations in batches\n\nIf you have a large number of prompts and want to collect activations in batches to optimize memory usage, you can use the `collect_activations_batched` function. This function has similar parameters to `collect_activations`, but also takes a `batch_size` parameter to specify the batch size for collecting activations.\n\n```python\nfrom nnterp import collect_activations_batched\n# Load the model\nnn_model = load_model(model_name)\n\n# Create a list of prompts\nprompts = [\"The quick brown fox\", \"jumps over the lazy dog\"]\n\n# Collect activations in batches\nbatch_size = 2\nactivations = collect_activations_batched(nn_model, prompts, batch_size)\n\n# Print the activations\nfor layer, activation in enumerate(activations):\n    print(f\"Layer {layer}: {activation.shape}\")\n```\n\n### Creating and Running Prompts\n\nNext, we create some toy prompts and run them through the model to get the next token probabilities.\n\n```python\nfrom nnterp import Prompt, run_prompts\n\n# Create toy prompts\nprompts = [\n    Prompt.from_strings(\"The quick brown fox\", {\"target\": \"jumps\"}, tokenizer),\n    Prompt.from_strings(\"Hello, how are you\", {\"target\": \"doing\"}, tokenizer)\n]\n\n# Run prompts through the model and get the next token probabilities\ntarget_probs = run_prompts(nn_model, prompts, batch_size=2)\n\n# Print the results\nfor prompt, probs in zip(prompts, target_probs[\"target\"]):\n    print(f\"Prompt: {prompt.prompt}\")\n    print(f\"Target Probabilities: {probs}\")\n```\n\n### Using Interventions\n\nNow, let's use some interventions like `logit_lens`\n\n#### Logit Lens\n\n```python\nfrom nnterp import logit_lens\n\n# Create a toy prompt\nprompt = \"The quick brown fox jumps over the lazy dog\"\n\n# Get the logit lens probabilities\nlogit_probs = logit_lens(nn_model, prompt)\n\n# Print the results\nprint(f\"Logit Lens Probabilities: {logit_probs.shape}\")\n```\n\n#### Patchscope Lens\n\n```python\nfrom nnterp import patchscope_lens, TargetPrompt\n\n# Create source and target prompts\nsource_prompt = \"The quick brown fox\"\ntarget_prompt = TargetPrompt(prompt=\"jumps over the lazy dog\", index_to_patch=-1)\n\n# Get the patchscope lens probabilities\npatchscope_probs = patchscope_lens(nn_model, source_prompts=[source_prompt], target_patch_prompts=[target_prompt])\n\n# Print the results\nprint(f\"Patchscope Lens Probabilities: {patchscope_probs}\")\n```\n\n### Using the Display Module\n\n```python\nfrom nnterp import plot_topk_tokens\n\n# Plot Patchscope Lens Probabilities and save the figure to test.png and test.html\nfig = plot_topk_tokens(\n    patchscope_probs,\n    tokenizer,\n    k=5,\n    title=\"Patchscope Lens Probabilities\",\n    file=\"test.png\",\n    save_html=True,  # Default is True\n)\nfig.show()\n```\n\n### Full Example\n\nHere is a full example combining all the above functionalities:\n\n```python\nfrom nnterp import load_model\nfrom nnterp.prompt_utils import Prompt, run_prompts\nfrom nnterp.interventions import logit_lens, patchscope_lens, TargetPrompt\n\n# Load the model\nmodel_name = \"meta-llama/Llama-2-7b-hf\"\nnn_model = load_model(model_name, trust_remote_code=False, device_map=\"auto\")\ntokenizer = nn_model.tokenizer\n\n# Create toy prompts\nprompts = [\n    Prompt.from_strings(\"The quick brown fox\", {\"target\": \"jumps\"}, tokenizer),\n    Prompt.from_strings(\"Hello, how are you\", {\"target\": \"doing\"}, tokenizer)\n]\n\n# Run prompts through the model\ntarget_probs = run_prompts(nn_model, prompts, batch_size=2)\n\n# Print the results\nfor prompt, probs in zip(prompts, target_probs[\"target\"]):\n    print(f\"Prompt: {prompt.prompt}\")\n    print(f\"Target Probabilities: {probs}\")\n\n# Logit Lens\nprompt = \"The quick brown fox jumps over the lazy dog\"\nlogit_probs = logit_lens(nn_model, prompt)\nprint(f\"Logit Lens Probabilities: {logit_probs}\")\n\n# Patchscope Lens\nsource_prompt = \"The quick brown fox\"\ntarget_prompt = TargetPrompt(prompt=\"jumps over the lazy dog\", index_to_patch=-1)\npatchscope_probs = patchscope_lens(nn_model, source_prompts=[source_prompt], target_patch_prompts=[target_prompt])\nprint(f\"Patchscope Lens Probabilities: {patchscope_probs}\")\n```\n\n## Codebase Overview\n- `nnsight_utils.py` basically allows you to deal with TL and HF models in a similar way.\n- `interventions.py` is a module that contains tools like logit lens, patchscope lens and other interventions.\n- `prompt_utils.py` contains utils to create prompts for which you want to track specific tokens in the next token distribution and run interventions on them and collect the probabilities of the tokens you're interested in.\n\n# Contributing\n- Create a git tag with the version number `git tag vx.y.z; git push origin vx.y.z`\n- Build with `python -m build`\n- Publish with e.g. `twine upload dist/*x.y.z*`\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fbutanium%2Fnnterp","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fbutanium%2Fnnterp","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fbutanium%2Fnnterp/lists"}