13 | return (
14 |
15 | {tokens.map(({ token, activation, normalized_activation }, i) => {
16 | const color = getInterpolatedColor(colors, boundaries, normalized_activation || activation);
17 | return
27 | {token}
28 |
29 | })}
30 |
31 | )
32 | }
33 |
--------------------------------------------------------------------------------
/neuron-viewer/package.json:
--------------------------------------------------------------------------------
1 | {
2 | "name": "neuron-viewer",
3 | "version": "0.1.67",
4 | "homepage": "https://openaipublic.blob.core.windows.net/neuron-explainer/neuron-viewer",
5 | "dependencies": {
6 | "@headlessui/react": "^1.7.8",
7 | "@headlessui/tailwindcss": "^0.1.2",
8 | "@types/d3-scale": "^4.0.3",
9 | "@types/lodash": "^4.14.194",
10 | "@types/react": "^18.0.37",
11 | "@types/react-dom": "^18.0.11",
12 | "d3-scale": "^4.0.2",
13 | "lodash": "^4.17.21",
14 | "react": "^18.2.0",
15 | "react-dom": "^18.2.0",
16 | "react-router-dom": "^6.10.0",
17 | "web-vitals": "^3.0.3"
18 | },
19 | "scripts": {
20 | "startpy": "nodemon python/server.py",
21 | "start": "parcel src/index.html",
22 | "build": "parcel build src/index.html",
23 | "serve": "parcel serve src/index.html",
24 | "typecheck": "tsc -p ."
25 | },
26 | "eslintConfig": {
27 | "extends": [
28 | "react-app"
29 | ]
30 | },
31 | "alias": {
32 | "preact/jsx-dev-runtime": "preact/jsx-runtime"
33 | },
34 | "devDependencies": {
35 | "@observablehq/plot": "^0.6.5",
36 | "@parcel/transformer-typescript-tsc": "^2.8.3",
37 | "@parcel/validator-typescript": "^2.8.3",
38 | "nodemon": "^2.0.22",
39 | "parcel": "^2.8.3",
40 | "preact": "^10.13.2",
41 | "process": "^0.11.10",
42 | "react-refresh": "0.10.0",
43 | "tailwindcss": "^3.2.4",
44 | "typescript": "^5.0.4"
45 | }
46 | }
47 |
--------------------------------------------------------------------------------
/neuron-viewer/python/server.py:
--------------------------------------------------------------------------------
1 | # %%
2 | import logging
3 |
4 | from flask import Flask, request
5 | from flask_cors import CORS
6 |
7 | import json
8 |
9 | import urllib.request
10 |
11 | def load_az_json(url):
12 | with urllib.request.urlopen(url) as f:
13 | return json.load(f)
14 |
15 | def start(
16 | dev: bool = False,
17 | host_name: str = "0.0.0.0",
18 | port: int = 80,
19 | ):
20 | app = Flask("interpretability chat")
21 | app.logger.setLevel(logging.INFO)
22 | # app.logger.disabled = True
23 | CORS(app)
24 |
25 | @app.after_request
26 | def after_request(response):
27 | response.headers.add("Access-Control-Allow-Origin", "*")
28 | response.headers.add(
29 | "Access-Control-Allow-Headers", "Content-Type,Authorization"
30 | )
31 | response.headers.add(
32 | "Access-Control-Allow-Methods", "GET,PUT,POST,DELETE,OPTIONS"
33 | )
34 | return response
35 |
36 | @app.route("/load_az", methods=["GET", "POST"])
37 | async def load_az():
38 | args = request.get_json()
39 | path = args["path"]
40 | result = load_az_json(path)
41 | return result
42 |
43 | app.run(debug=dev, host=host_name, port=port, use_reloader=False)
44 |
45 |
46 | def main(dev: bool = True, host_name: str = "0.0.0.0", port: int = 8000):
47 | start(dev=dev, host_name=host_name, port=port)
48 |
49 |
50 | if __name__ == "__main__":
51 | main()
52 |
--------------------------------------------------------------------------------
/neuron-viewer/src/feed.jsx:
--------------------------------------------------------------------------------
1 | import * as Panes from "./panes"
2 | import React, { useEffect } from "react"
3 | import Welcome from "./welcome"
4 | import { useState } from "react"
5 | import { useParams, Link } from "react-router-dom"
6 |
7 | export default function Feed() {
8 | const params = useParams()
9 | // If params is missing either index, there's no neuron selected.
10 | let activeNeuron;
11 | if (params.layer === undefined || params.neuron === undefined) {
12 | activeNeuron = null
13 | } else {
14 | // Grab the layer and neuron indices from the params, casting them to ints.
15 | activeNeuron = {
16 | "layer": parseInt(params.layer),
17 | "neuron": parseInt(params.neuron),
18 | }
19 | }
20 |
21 | const Pane = ({ children }) => (
22 |
{children}
23 | )
24 |
25 | return (
26 |
27 |
28 |
29 | Neuron Viewer
30 |
31 | {activeNeuron && (
32 |
33 | Neuron {activeNeuron.layer}:{activeNeuron.neuron}
34 |
35 | )}
36 |
37 |
38 |
41 |
42 | {activeNeuron ?
43 | <>
44 |
45 | {React.createElement(Panes["Explanation"], { activeNeuron })}
46 |
47 |
48 | {React.createElement(Panes["DatasetList"], { activeNeuron })}
49 |
50 |
51 | {React.createElement(Panes["TopTokens"], { activeNeuron })}
52 |
53 |
54 | {React.createElement(Panes["SimilarNeurons"], { activeNeuron })}
55 |
56 | :
57 |
58 | }
59 |
60 |
61 |
62 |
63 | )
64 | }
65 |
--------------------------------------------------------------------------------
/neuron-explainer/neuron_explainer/explanations/puzzles.py:
--------------------------------------------------------------------------------
1 | import json
2 | import os
3 | from dataclasses import dataclass
4 |
5 | from neuron_explainer.activations.activations import ActivationRecord
6 |
7 |
8 | @dataclass(frozen=True)
9 | class Puzzle:
10 | """A puzzle is a ground truth explanation, a collection of sentences (stored as ActivationRecords) with activations
11 | according to that explanation, and a collection of false explanations"""
12 |
13 | name: str
14 | explanation: str
15 | activation_records: list[ActivationRecord]
16 | false_explanations: list[str]
17 |
18 |
19 | def convert_puzzle_to_tokenized_sentences(puzzle: Puzzle) -> list[list[str]]:
20 | """Converts a puzzle to a list of tokenized sentences."""
21 | return [record.tokens for record in puzzle.activation_records]
22 |
23 |
24 | def convert_puzzle_dict_to_puzzle(puzzle_dict: dict) -> Puzzle:
25 | """Converts a json dictionary representation of a puzzle to the Puzzle class."""
26 | puzzle_activation_records = []
27 | for sentence in puzzle_dict["sentences"]:
28 | # Token-activation pairs are listed as either a string or a list of a string and a float. If it is a list, the float is the activation.
29 | # If it is only a string, the activation is assumed to be 0. This is useful for readability and reducing redundancy in the data.
30 | tokens = [t[0] if type(t) is list else t for t in sentence]
31 | assert all([type(t) is str for t in tokens]), "All tokens must be strings"
32 | activations = [float(t[1]) if type(t) is list else 0.0 for t in sentence]
33 | assert all([type(t) is float for t in activations]), "All activations must be floats"
34 |
35 | puzzle_activation_records.append(ActivationRecord(tokens=tokens, activations=activations))
36 |
37 | return Puzzle(
38 | name=puzzle_dict["name"],
39 | explanation=puzzle_dict["explanation"],
40 | activation_records=puzzle_activation_records,
41 | false_explanations=puzzle_dict["false_explanations"],
42 | )
43 |
44 |
45 | PUZZLES_BY_NAME: dict[str, Puzzle] = dict()
46 | script_dir = os.path.dirname(os.path.abspath(__file__))
47 | with open(os.path.join(script_dir, "puzzles.json"), "r") as f:
48 | puzzle_dicts = json.loads(f.read())
49 | for name in puzzle_dicts.keys():
50 | PUZZLES_BY_NAME[name] = convert_puzzle_dict_to_puzzle(puzzle_dicts[name])
51 |
--------------------------------------------------------------------------------
/neuron-explainer/neuron_explainer/activations/token_connections.py:
--------------------------------------------------------------------------------
1 | from dataclasses import dataclass
2 | from typing import List, Union
3 |
4 | import blobfile as bf
5 | from neuron_explainer.fast_dataclasses import FastDataclass, loads, register_dataclass
6 | from neuron_explainer.azure import standardize_azure_url
7 | import urllib.request
8 |
9 |
10 | @register_dataclass
11 | @dataclass
12 | class TokensAndWeights(FastDataclass):
13 | tokens: List[str]
14 | strengths: List[float]
15 |
16 |
17 | @register_dataclass
18 | @dataclass
19 | class WeightBasedSummaryOfNeuron(FastDataclass):
20 | input_positive: TokensAndWeights
21 | input_negative: TokensAndWeights
22 | output_positive: TokensAndWeights
23 | output_negative: TokensAndWeights
24 |
25 |
26 | def load_token_weight_connections_of_neuron(
27 | layer_index: Union[str, int],
28 | neuron_index: Union[str, int],
29 | dataset_path: str = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/related-tokens/weight-based",
30 | ) -> WeightBasedSummaryOfNeuron:
31 | """Load the TokenLookupTableSummaryOfNeuron for the specified neuron."""
32 | url = "/".join([dataset_path, str(layer_index), f"{neuron_index}.json"])
33 | url = standardize_azure_url(url)
34 | with urllib.request.urlopen(url) as f:
35 | return loads(f.read(), backwards_compatible=False)
36 |
37 |
38 | @register_dataclass
39 | @dataclass
40 | class TokenLookupTableSummaryOfNeuron(FastDataclass):
41 | """List of tokens and the average activations of a given neuron in response to each
42 | respective token. These are selected from among the tokens in the vocabulary with the
43 | highest average activations across an internet text dataset, with the highest activations
44 | first."""
45 |
46 | tokens: List[str]
47 | average_activations: List[float]
48 |
49 |
50 | def load_token_lookup_table_connections_of_neuron(
51 | layer_index: Union[str, int],
52 | neuron_index: Union[str, int],
53 | dataset_path: str = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/related-tokens/activation-based",
54 | ) -> TokenLookupTableSummaryOfNeuron:
55 | """Load the TokenLookupTableSummaryOfNeuron for the specified neuron."""
56 | url = "/".join([dataset_path, str(layer_index), f"{neuron_index}.json"])
57 | url = standardize_azure_url(url)
58 | with urllib.request.urlopen(url) as f:
59 | return loads(f.read(), backwards_compatible=False)
60 |
--------------------------------------------------------------------------------
/neuron-viewer/src/index.html:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
11 |
15 |
24 |
25 |
26 |
Neuron viewer
27 |
28 |
29 |
30 |
31 |
42 |
43 |
44 |
45 |
46 |
56 |
57 |
58 |
59 |
60 |
--------------------------------------------------------------------------------
/neuron-explainer/demos/explain_puzzles.ipynb:
--------------------------------------------------------------------------------
1 | {
2 | "cells": [
3 | {
4 | "cell_type": "code",
5 | "execution_count": null,
6 | "metadata": {},
7 | "outputs": [],
8 | "source": [
9 | "%load_ext autoreload\n",
10 | "%autoreload 2"
11 | ]
12 | },
13 | {
14 | "cell_type": "code",
15 | "execution_count": null,
16 | "metadata": {},
17 | "outputs": [],
18 | "source": [
19 | "import os\n",
20 | "\n",
21 | "os.environ[\"OPENAI_API_KEY\"] = \"put-key-here\"\n",
22 | "\n",
23 | "from neuron_explainer.activations.activation_records import calculate_max_activation\n",
24 | "from neuron_explainer.explanations.explainer import TokenActivationPairExplainer\n",
25 | "from neuron_explainer.explanations.prompt_builder import PromptFormat\n",
26 | "from neuron_explainer.explanations.puzzles import PUZZLES_BY_NAME\n",
27 | "\n",
28 | "\n",
29 | "EXPLAINER_MODEL_NAME = \"gpt-4\"\n",
30 | "\n",
31 | "explainer = TokenActivationPairExplainer(\n",
32 | " model_name=EXPLAINER_MODEL_NAME,\n",
33 | " prompt_format=PromptFormat.HARMONY_V4,\n",
34 | " max_concurrent=1,\n",
35 | ")\n",
36 | "\n",
37 | "for puzzle_name, puzzle in PUZZLES_BY_NAME.items():\n",
38 | " print(f\"{puzzle_name=}\")\n",
39 | " puzzle_answer = puzzle.explanation\n",
40 | " # Generate an explanation for the puzzle.\n",
41 | " explanations = await explainer.generate_explanations(\n",
42 | " all_activation_records=puzzle.activation_records,\n",
43 | " max_activation=calculate_max_activation(puzzle.activation_records),\n",
44 | " num_samples=1,\n",
45 | " )\n",
46 | " assert len(explanations) == 1\n",
47 | " model_generated_explanation = explanations[0]\n",
48 | " print(f\"{model_generated_explanation=}\")\n",
49 | " print(f\"{puzzle_answer=}\\n\")\n",
50 | "\n"
51 | ]
52 | }
53 | ],
54 | "metadata": {
55 | "kernelspec": {
56 | "display_name": "openai",
57 | "language": "python",
58 | "name": "openai"
59 | },
60 | "language_info": {
61 | "codemirror_mode": {
62 | "name": "ipython",
63 | "version": 3
64 | },
65 | "file_extension": ".py",
66 | "mimetype": "text/x-python",
67 | "name": "python",
68 | "nbconvert_exporter": "python",
69 | "pygments_lexer": "ipython3",
70 | "version": "3.9.9"
71 | },
72 | "orig_nbformat": 4
73 | },
74 | "nbformat": 4,
75 | "nbformat_minor": 2
76 | }
77 |
--------------------------------------------------------------------------------
/neuron-viewer/src/types.ts:
--------------------------------------------------------------------------------
1 | import { scaleLinear } from "d3-scale"
2 | import { min, max, flatten } from "lodash"
3 |
4 | export type Neuron = {
5 | layer: number;
6 | neuron: number;
7 | }
8 |
9 | export type TokenAndActivation = {
10 | token: string,
11 | activation: number
12 | normalized_activation?: number
13 | }
14 |
15 | export type TokenSequence = TokenAndActivation[]
16 |
17 | export const normalizeTokenActs = (...sequences: TokenSequence[][]) => {
18 | // console.log('sequences', sequences)
19 | let flattened: TokenAndActivation[] = flatten(flatten(sequences))
20 | // Replace all activations less than 0 in data.tokens with 0. This matches the format in the
21 | // top + random activation records displayed in the main grid.
22 | flattened = flattened.map(({token, activation}) => {
23 | return {
24 | token,
25 | activation: Math.max(activation, 0)
26 | }
27 | })
28 | const maxActivation = max(flattened.map((ta) => ta.activation)) || 0;
29 | const neuronScale = scaleLinear()
30 | // Even though we're only displaying positive activations, we still need to scale in a way that
31 | // accounts for the existence of negative activations, since our color scale includes them.
32 | .domain([0, maxActivation])
33 | .range([0, 1])
34 |
35 | return sequences.map((seq) => seq.map((tas) => tas.map(({ token, activation }) => ({
36 | token,
37 | activation,
38 | normalized_activation: neuronScale(activation),
39 | }))))
40 | }
41 |
42 | export type Color = {r: number, g: number, b: number};
43 | export function interpolateColor(color_l: Color, color_r: Color, value: number) {
44 | const color = {
45 | r: Math.round(color_l.r + (color_r.r - color_l.r) * value),
46 | g: Math.round(color_l.g + (color_r.g - color_l.g) * value),
47 | b: Math.round(color_l.b + (color_r.b - color_l.b) * value),
48 | }
49 | return color
50 | }
51 |
52 | export function getInterpolatedColor(colors: Color[], boundaries: number[], value: number) {
53 | const index = boundaries.findIndex((boundary) => boundary >= value)
54 | const colorIndex = Math.max(0, index - 1)
55 | const color_left = colors[colorIndex]
56 | const color_right = colors[colorIndex + 1]
57 | const boundary_left = boundaries[colorIndex]
58 | const boundary_right = boundaries[colorIndex + 1]
59 | const ratio = (value - boundary_left) / (boundary_right - boundary_left)
60 | const color = interpolateColor(color_left, color_right, ratio)
61 | return color
62 | }
63 |
64 | export const DEFAULT_COLORS = [
65 | // { r: 255, g: 0, b: 105 },
66 | { r: 255, g: 255, b: 255 },
67 | { r: 0, g: 255, b: 0 },
68 | ]
69 | export const DEFAULT_BOUNDARIES = [
70 | // 0, 0.5, 1
71 | 0, 1
72 | ]
73 |
74 |
--------------------------------------------------------------------------------
/neuron-explainer/neuron_explainer/fast_dataclasses/test_fast_dataclasses.py:
--------------------------------------------------------------------------------
1 | from dataclasses import dataclass
2 |
3 | import pytest
4 |
5 | from .fast_dataclasses import FastDataclass, dumps, loads, register_dataclass
6 |
7 |
8 | # Inheritance is a bit tricky with our setup. dataclass_name must be set for instances of these
9 | # classes to serialize and deserialize correctly, but if it's given a default value, then subclasses
10 | # can't have any fields that don't have default values, because of how constructors are generated
11 | # for dataclasses (fields with no default value can't follow those with default values). To work
12 | # around this, we set dataclass_name in __post_init__ on the base class, which is called after the
13 | # constructor. The implementation does the right thing for both the base class and the subclass.
14 | @register_dataclass
15 | @dataclass
16 | class DataclassC(FastDataclass):
17 | ints: list[int]
18 |
19 |
20 | @register_dataclass
21 | @dataclass
22 | class DataclassC_ext(DataclassC):
23 | s: str
24 |
25 |
26 | @register_dataclass
27 | @dataclass
28 | class DataclassB(FastDataclass):
29 | str_to_c: dict[str, DataclassC]
30 | cs: list[DataclassC]
31 |
32 |
33 | @register_dataclass
34 | @dataclass
35 | class DataclassA(FastDataclass):
36 | floats: list[float]
37 | strings: list[str]
38 | bs: list[DataclassB]
39 |
40 |
41 | @register_dataclass
42 | @dataclass
43 | class DataclassD(FastDataclass):
44 | s1: str
45 | s2: str = "default"
46 |
47 |
48 | def test_dataclasses() -> None:
49 | a = DataclassA(
50 | floats=[1.0, 2.0],
51 | strings=["a", "b"],
52 | bs=[
53 | DataclassB(
54 | str_to_c={"a": DataclassC(ints=[1, 2]), "b": DataclassC(ints=[3, 4])},
55 | cs=[DataclassC(ints=[5, 6]), DataclassC_ext(ints=[7, 8], s="s")],
56 | ),
57 | DataclassB(
58 | str_to_c={"c": DataclassC_ext(ints=[9, 10], s="t"), "d": DataclassC(ints=[11, 12])},
59 | cs=[DataclassC(ints=[13, 14]), DataclassC(ints=[15, 16])],
60 | ),
61 | ],
62 | )
63 | assert loads(dumps(a)) == a
64 |
65 |
66 | def test_c_and_c_ext() -> None:
67 | c_ext = DataclassC_ext(ints=[3, 4], s="s")
68 | assert loads(dumps(c_ext)) == c_ext
69 |
70 | c = DataclassC(ints=[1, 2])
71 | assert loads(dumps(c)) == c
72 |
73 |
74 | def test_bad_serialized_data() -> None:
75 | assert type(loads(dumps(DataclassC(ints=[3, 4])))) == DataclassC
76 | assert type(loads('{"ints": [3, 4]}', backwards_compatible=False)) == dict
77 | assert type(loads('{"ints": [3, 4], "dataclass_name": "DataclassC"}')) == DataclassC
78 | with pytest.raises(TypeError):
79 | loads('{"ints": [3, 4], "bogus_extra_field": "foo", "dataclass_name": "DataclassC"}')
80 | with pytest.raises(TypeError):
81 | loads('{"ints_field_is_missing": [3, 4], "dataclass_name": "DataclassC"}')
82 | assert type(loads('{"s1": "test"}', backwards_compatible=False)) == dict
83 | assert type(loads('{"s1": "test"}', backwards_compatible=True)) == DataclassD
84 |
--------------------------------------------------------------------------------
/neuron-explainer/neuron_explainer/fast_dataclasses/fast_dataclasses.py:
--------------------------------------------------------------------------------
1 | # Utilities for dataclasses that are very fast to serialize and deserialize, with limited data
2 | # validation. Fields must not be tuples, since they get serialized and then deserialized as lists.
3 | #
4 | # The unit tests for this library show how to use it.
5 |
6 | import json
7 | from dataclasses import dataclass, field, fields, is_dataclass
8 | from functools import partial
9 | from typing import Any, Union
10 |
11 | import orjson
12 |
13 | dataclasses_by_name = {}
14 | dataclasses_by_fieldnames = {}
15 |
16 |
17 | @dataclass
18 | class FastDataclass:
19 | dataclass_name: str = field(init=False)
20 |
21 | def __post_init__(self) -> None:
22 | self.dataclass_name = self.__class__.__name__
23 |
24 |
25 | def register_dataclass(cls): # type: ignore
26 | assert is_dataclass(cls), "Only dataclasses can be registered."
27 | dataclasses_by_name[cls.__name__] = cls
28 | name_set = frozenset(f.name for f in fields(cls) if f.name != "dataclass_name")
29 | dataclasses_by_fieldnames[name_set] = cls
30 | return cls
31 |
32 |
33 | def dumps(obj: Any) -> bytes:
34 | return orjson.dumps(obj, option=orjson.OPT_SERIALIZE_NUMPY)
35 |
36 |
37 | def _object_hook(d: Any, backwards_compatible: bool = True) -> Any:
38 | # If d is a list, recurse.
39 | if isinstance(d, list):
40 | return [_object_hook(x, backwards_compatible=backwards_compatible) for x in d]
41 | # If d is not a dict, return it as is.
42 | if not isinstance(d, dict):
43 | return d
44 | cls = None
45 | if "dataclass_name" in d:
46 | if d["dataclass_name"] in dataclasses_by_name:
47 | cls = dataclasses_by_name[d["dataclass_name"]]
48 | else:
49 | assert backwards_compatible, (
50 | f"Dataclass {d['dataclass_name']} not found, set backwards_compatible=True if you "
51 | f"are okay with that."
52 | )
53 | # Load objects created without dataclass_name set.
54 | else:
55 | # Try our best to find a dataclass if backwards_compatible is True.
56 | if backwards_compatible:
57 | d_fields = frozenset(d.keys())
58 | if d_fields in dataclasses_by_fieldnames:
59 | cls = dataclasses_by_fieldnames[d_fields]
60 | elif len(d_fields) > 0:
61 | # Check if the fields are a subset of a dataclass (if the dataclass had extra fields
62 | # added since the data was created). Note that this will fail if fields were removed
63 | # from the dataclass.
64 | for key, possible_cls in dataclasses_by_fieldnames.items():
65 | if d_fields.issubset(key):
66 | cls = possible_cls
67 | break
68 | else:
69 | print(f"Could not find dataclass for {d_fields} {cls}")
70 | new_d = {
71 | k: _object_hook(v, backwards_compatible=backwards_compatible)
72 | for k, v in d.items()
73 | if k != "dataclass_name"
74 | }
75 | if cls is not None:
76 | return cls(**new_d)
77 | else:
78 | return new_d
79 |
80 |
81 | def loads(s: Union[str, bytes], backwards_compatible: bool = True) -> Any:
82 | return json.loads(
83 | s,
84 | object_hook=partial(_object_hook, backwards_compatible=backwards_compatible),
85 | )
86 |
--------------------------------------------------------------------------------
/neuron-explainer/demos/generate_and_score_explanation.ipynb:
--------------------------------------------------------------------------------
1 | {
2 | "cells": [
3 | {
4 | "cell_type": "code",
5 | "execution_count": null,
6 | "metadata": {},
7 | "outputs": [],
8 | "source": [
9 | "%load_ext autoreload\n",
10 | "%autoreload 2"
11 | ]
12 | },
13 | {
14 | "cell_type": "code",
15 | "execution_count": null,
16 | "metadata": {},
17 | "outputs": [],
18 | "source": [
19 | "import os\n",
20 | "\n",
21 | "os.environ[\"OPENAI_API_KEY\"] = \"put-key-here\"\n",
22 | "\n",
23 | "from neuron_explainer.activations.activation_records import calculate_max_activation\n",
24 | "from neuron_explainer.activations.activations import ActivationRecordSliceParams, load_neuron\n",
25 | "from neuron_explainer.explanations.calibrated_simulator import UncalibratedNeuronSimulator\n",
26 | "from neuron_explainer.explanations.explainer import TokenActivationPairExplainer\n",
27 | "from neuron_explainer.explanations.prompt_builder import PromptFormat\n",
28 | "from neuron_explainer.explanations.scoring import simulate_and_score\n",
29 | "from neuron_explainer.explanations.simulator import ExplanationNeuronSimulator\n",
30 | "\n",
31 | "EXPLAINER_MODEL_NAME = \"gpt-4\"\n",
32 | "SIMULATOR_MODEL_NAME = \"text-davinci-003\"\n",
33 | "\n",
34 | "\n",
35 | "# test_response = await client.make_request(prompt=\"test 123<|endofprompt|>\", max_tokens=2)\n",
36 | "# print(\"Response:\", test_response[\"choices\"][0][\"text\"])\n",
37 | "\n",
38 | "# Load a neuron record.\n",
39 | "neuron_record = load_neuron(9, 6236)\n",
40 | "\n",
41 | "# Grab the activation records we'll need.\n",
42 | "slice_params = ActivationRecordSliceParams(n_examples_per_split=5)\n",
43 | "train_activation_records = neuron_record.train_activation_records(\n",
44 | " activation_record_slice_params=slice_params\n",
45 | ")\n",
46 | "valid_activation_records = neuron_record.valid_activation_records(\n",
47 | " activation_record_slice_params=slice_params\n",
48 | ")\n",
49 | "\n",
50 | "# Generate an explanation for the neuron.\n",
51 | "explainer = TokenActivationPairExplainer(\n",
52 | " model_name=EXPLAINER_MODEL_NAME,\n",
53 | " prompt_format=PromptFormat.HARMONY_V4,\n",
54 | " max_concurrent=1,\n",
55 | ")\n",
56 | "explanations = await explainer.generate_explanations(\n",
57 | " all_activation_records=train_activation_records,\n",
58 | " max_activation=calculate_max_activation(train_activation_records),\n",
59 | " num_samples=1,\n",
60 | ")\n",
61 | "assert len(explanations) == 1\n",
62 | "explanation = explanations[0]\n",
63 | "print(f\"{explanation=}\")\n",
64 | "\n",
65 | "# Simulate and score the explanation.\n",
66 | "simulator = UncalibratedNeuronSimulator(\n",
67 | " ExplanationNeuronSimulator(\n",
68 | " SIMULATOR_MODEL_NAME,\n",
69 | " explanation,\n",
70 | " max_concurrent=1,\n",
71 | " prompt_format=PromptFormat.INSTRUCTION_FOLLOWING,\n",
72 | " )\n",
73 | ")\n",
74 | "scored_simulation = await simulate_and_score(simulator, valid_activation_records)\n",
75 | "print(f\"score={scored_simulation.get_preferred_score():.2f}\")\n"
76 | ]
77 | }
78 | ],
79 | "metadata": {
80 | "kernelspec": {
81 | "display_name": "openai",
82 | "language": "python",
83 | "name": "python3"
84 | },
85 | "language_info": {
86 | "codemirror_mode": {
87 | "name": "ipython",
88 | "version": 3
89 | },
90 | "file_extension": ".py",
91 | "mimetype": "text/x-python",
92 | "name": "python",
93 | "nbconvert_exporter": "python",
94 | "pygments_lexer": "ipython3",
95 | "version": "3.9.9"
96 | },
97 | "orig_nbformat": 4
98 | },
99 | "nbformat": 4,
100 | "nbformat_minor": 2
101 | }
102 |
--------------------------------------------------------------------------------
/neuron-viewer/src/panes/datasetList.jsx:
--------------------------------------------------------------------------------
1 | import HeatmapGrid from "../heatmapGrid"
2 | import React, { useEffect, useState } from "react"
3 | import { normalizeTokenActs } from "../types"
4 |
5 | import {get_neuron_record} from "../interpAPI"
6 |
7 | function zip_sequences(sequences) {
8 | return sequences.map(({ activations, tokens }) => {
9 | return tokens.map((token, idx) => ({
10 | token,
11 | activation: activations[idx],
12 | }))
13 | })
14 | }
15 |
16 | export default ({ activeNeuron }) => {
17 | const [data, setData] = useState(null)
18 | const [showingMore, setShowingMore] = useState({})
19 | const [isLoading, setIsLoading] = useState(true)
20 |
21 | useEffect(() => {
22 | async function fetchData() {
23 | if (data) {
24 | return
25 | }
26 | const result = await get_neuron_record(activeNeuron)
27 | console.log(result)
28 | const all_sequences = []
29 | all_sequences.push({
30 | // label: '[0.999, 1] (Top quantile, sorted. 50 of 50000)',
31 | label: 'Top',
32 | sequences: zip_sequences(result.most_positive_activation_records),
33 | default_show: 4,
34 | })
35 | all_sequences.push({
36 | label: 'Quantile range [0.99, 0.999] sample',
37 | sequences: zip_sequences(result.random_sample_by_quantile[3]),
38 | default_show: 1,
39 | })
40 | all_sequences.push({
41 | label: 'Quantile range [0.9, 0.99] sample',
42 | sequences: zip_sequences(result.random_sample_by_quantile[2]),
43 | default_show: 1,
44 | })
45 | all_sequences.push({
46 | label: 'Quantile range [0.5, 0.9] sample',
47 | sequences: zip_sequences(result.random_sample_by_quantile[1]),
48 | default_show: 1,
49 | })
50 | all_sequences.push({
51 | label: 'Quantile range [0, 0.5] sample',
52 | sequences: zip_sequences(result.random_sample_by_quantile[0]),
53 | default_show: 1,
54 | })
55 | all_sequences.push({
56 | // label: '[0, 1] (Random)',
57 | label: 'Random sample',
58 | sequences: zip_sequences(result.random_sample),
59 | default_show: 2,
60 | })
61 | // for reference
62 | // intervals = [(0, 1), (0, 0.5), (0.5, 0.9), (0.9, 0.99), (0.99, 0.999), (0.999, 1)]
63 | // saved_activations_by_interval = [neuron_record.random_sample] + neuron_record.random_sample_by_decile[:-1] + [neuron_record.top_activations]
64 | setData(all_sequences)
65 | setIsLoading(false)
66 | }
67 | fetchData()
68 | }, [activeNeuron])
69 |
70 | if (isLoading) {
71 | return (
72 |
73 |
74 |
loading top dataset examples
75 |
76 | )
77 | }
78 |
79 | // const activations = data.top_activations;
80 | const all_normalized_sequences = normalizeTokenActs(...data.map(({sequences}) => sequences))
81 |
82 | return (
83 |
84 |
Activations
85 | {
86 | data.map(({label, default_show}, idx) => {
87 | const n_show = showingMore[label] ? all_normalized_sequences[idx].length : default_show;
88 | return (
89 |
90 |
91 | {label}
92 |
96 |
97 |
98 |
99 | )
100 | })
101 | }
102 |
103 | )
104 | }
105 |
--------------------------------------------------------------------------------
/neuron-explainer/demos/generate_and_score_token_look_up_table_explanation.ipynb:
--------------------------------------------------------------------------------
1 | {
2 | "cells": [
3 | {
4 | "cell_type": "code",
5 | "execution_count": null,
6 | "metadata": {},
7 | "outputs": [],
8 | "source": [
9 | "%load_ext autoreload\n",
10 | "%autoreload 2"
11 | ]
12 | },
13 | {
14 | "cell_type": "code",
15 | "execution_count": null,
16 | "metadata": {},
17 | "outputs": [],
18 | "source": [
19 | "import os\n",
20 | "\n",
21 | "os.environ[\"OPENAI_API_KEY\"] = \"put-key-here\"\n",
22 | "\n",
23 | "from neuron_explainer.activations.activations import ActivationRecordSliceParams, load_neuron\n",
24 | "from neuron_explainer.activations.token_connections import load_token_lookup_table_connections_of_neuron\n",
25 | "from neuron_explainer.explanations.calibrated_simulator import UncalibratedNeuronSimulator\n",
26 | "from neuron_explainer.explanations.explainer import TokenSpaceRepresentationExplainer\n",
27 | "from neuron_explainer.explanations.prompt_builder import PromptFormat\n",
28 | "from neuron_explainer.explanations.scoring import simulate_and_score\n",
29 | "from neuron_explainer.explanations.simulator import ExplanationNeuronSimulator\n",
30 | "\n",
31 | "EXPLAINER_MODEL_NAME = \"gpt-4\"\n",
32 | "SIMULATOR_MODEL_NAME = \"text-davinci-003\"\n",
33 | "\n",
34 | "\n",
35 | "# test_response = await client.make_request(prompt=\"test 123<|endofprompt|>\", max_tokens=2)\n",
36 | "# print(\"Response:\", test_response[\"choices\"][0][\"text\"])\n",
37 | "\n",
38 | "layer_index = 9\n",
39 | "neuron_index = 6236\n",
40 | "\n",
41 | "# Load a token lookup table.\n",
42 | "token_lookup_table = load_token_lookup_table_connections_of_neuron(layer_index, neuron_index)\n",
43 | "\n",
44 | "# Load a neuron record.\n",
45 | "neuron_record = load_neuron(layer_index, neuron_index)\n",
46 | "\n",
47 | "# Grab the activation records we'll need.\n",
48 | "slice_params = ActivationRecordSliceParams(n_examples_per_split=5)\n",
49 | "valid_activation_records = neuron_record.valid_activation_records(\n",
50 | " activation_record_slice_params=slice_params\n",
51 | ")\n",
52 | "\n",
53 | "# Generate an explanation for the neuron.\n",
54 | "explainer = TokenSpaceRepresentationExplainer(\n",
55 | " model_name=EXPLAINER_MODEL_NAME,\n",
56 | " prompt_format=PromptFormat.HARMONY_V4,\n",
57 | " max_concurrent=1,\n",
58 | ")\n",
59 | "explanations = await explainer.generate_explanations(\n",
60 | " tokens=token_lookup_table.tokens,\n",
61 | " num_samples=1,\n",
62 | ")\n",
63 | "assert len(explanations) == 1\n",
64 | "explanation = explanations[0]\n",
65 | "print(f\"{explanation=}\")\n",
66 | "\n",
67 | "# Simulate and score the explanation.\n",
68 | "simulator = UncalibratedNeuronSimulator(\n",
69 | " ExplanationNeuronSimulator(\n",
70 | " SIMULATOR_MODEL_NAME,\n",
71 | " explanation,\n",
72 | " max_concurrent=1,\n",
73 | " prompt_format=PromptFormat.INSTRUCTION_FOLLOWING,\n",
74 | " )\n",
75 | ")\n",
76 | "scored_simulation = await simulate_and_score(simulator, valid_activation_records)\n",
77 | "print(f\"score={scored_simulation.get_preferred_score():.2f}\")\n"
78 | ]
79 | }
80 | ],
81 | "metadata": {
82 | "kernelspec": {
83 | "display_name": "Python 3",
84 | "language": "python",
85 | "name": "python3"
86 | },
87 | "language_info": {
88 | "codemirror_mode": {
89 | "name": "ipython",
90 | "version": 3
91 | },
92 | "file_extension": ".py",
93 | "mimetype": "text/x-python",
94 | "name": "python",
95 | "nbconvert_exporter": "python",
96 | "pygments_lexer": "ipython3",
97 | "version": "3.9.8"
98 | },
99 | "vscode": {
100 | "interpreter": {
101 | "hash": "fd71fb58b1ad02dde67c8ac595a52586dd87d3465221a699fc288aa2c48d5565"
102 | }
103 | }
104 | },
105 | "nbformat": 4,
106 | "nbformat_minor": 2
107 | }
108 |
--------------------------------------------------------------------------------
/neuron-viewer/src/panes/similarNeurons.jsx:
--------------------------------------------------------------------------------
1 | import React, { useEffect, useState } from "react"
2 | import _ from "lodash"
3 | import { Link } from "react-router-dom"
4 |
5 | import { get_explanations, get_top_neuron_connections } from "../interpAPI"
6 |
7 | function NeuronInfo({ neuron, strength }) {
8 | const [info, setInfo] = useState(null)
9 |
10 | useEffect(() => {
11 | async function fetchInfo() {
12 | const result = (await get_explanations({
13 | layer: neuron.layer,
14 | neuron: neuron.neuron,
15 | }))
16 | setInfo(result)
17 | }
18 |
19 | if (!info) {
20 | fetchInfo()
21 | }
22 | }, [])
23 |
24 | if (!info) {
25 | return (
26 |
27 |
28 | Loading neuron {neuron.layer}:{neuron.neuron}...
29 |
30 |
31 |
32 | )
33 | }
34 |
35 | return (
36 |
37 |
38 |
40 |
41 | Neuron {neuron.layer}:{neuron.neuron}
42 |
43 |
44 |
45 | Connection strength: {strength.toFixed(2)}
46 |
47 |
48 | {info.scored_explanations.map((explanation, i) => (
49 |
50 |
51 | {explanation.explanation}
52 |
53 |
54 | score: {explanation.scored_simulation.ev_correlation_score.toFixed(2)}
55 |
56 |
57 | ))}
58 |
59 |
60 |
61 | )
62 | }
63 |
64 | export default function SimilarNeurons({ activeNeuron: neuron }) {
65 | const [similarNeurons, setSimilarNeurons] = useState([])
66 | const [isLoading, setIsLoading] = useState(true)
67 |
68 | useEffect(() => {
69 | async function fetchSimilarNeurons() {
70 | const result = await get_top_neuron_connections(neuron)
71 | setSimilarNeurons(result)
72 | setIsLoading(false)
73 | }
74 |
75 | fetchSimilarNeurons()
76 | }, [neuron])
77 |
78 | if (isLoading) {
79 | return (
80 |
83 | )
84 | }
85 |
86 | const n_show = 3;
87 | return (
88 |
89 |
Related neurons
90 |
91 |
92 | {
93 | similarNeurons.input ?
94 |
95 |
Upstream
96 |
97 | {similarNeurons.input.slice(0, n_show).map(([layer, neuron, strength]) => (
98 |
99 | ))}
100 |
101 |
: null
102 | }
103 | {
104 | similarNeurons.output ?
105 |
106 |
Downstream
107 |
108 | {similarNeurons.output.slice(0, n_show).map(([layer, neuron, strength]) => (
109 |
110 | ))}
111 |
112 |
: null
113 | }
114 |
115 |
116 |
117 | )
118 | }
119 |
--------------------------------------------------------------------------------
/neuron-viewer/src/simulationHeatmap.tsx:
--------------------------------------------------------------------------------
1 | import React, { useState } from 'react';
2 |
3 | import { interpolateColor, Color, getInterpolatedColor, DEFAULT_COLORS, DEFAULT_BOUNDARIES, TokenAndActivation } from './types'
4 |
5 | type Props = {
6 | sequences: TokenAndActivation[][],
7 | simulated_sequences: TokenAndActivation[][],
8 | overlay_activations: boolean,
9 | colors?: Color[],
10 | boundaries?: number[],
11 | }
12 | export default function SimulationSequences({ sequences, simulated_sequences, overlay_activations, colors = DEFAULT_COLORS, boundaries = DEFAULT_BOUNDARIES }: Props) {
13 | return <>
14 | {
15 | sequences.map((tokens, i) => {
16 | let simulated_tokens = simulated_sequences[i];
17 | if (overlay_activations) {
18 | return (
19 |
20 | {tokens.map(({ token, activation, normalized_activation }, j) => {
21 | const { token: simulated_token, activation: simulated_activation, normalized_activation: simulated_normalized_activation } = simulated_tokens[j];
22 | if (simulated_token !== token) {
23 | throw new Error('simulated tokens not matching')
24 | }
25 | const color = getInterpolatedColor(colors, boundaries, normalized_activation || activation);
26 | const simcolor = getInterpolatedColor(colors, boundaries, simulated_normalized_activation || simulated_activation);
27 |
28 | return
29 |
30 | {token}
37 | {token}
44 |
45 |
46 | })}
47 |
48 | )
49 | } else {
50 | return (
51 |
52 |
53 | Real activations:
59 | {tokens.map(({ token, activation, normalized_activation }, j) => {
60 | const color = getInterpolatedColor(colors, boundaries, normalized_activation || activation);
61 | return {token}
68 | })}
69 |
70 |
71 |
72 | Simulated activations:
78 | {simulated_tokens.map(({ token, activation, normalized_activation }, j) => {
79 | const color = getInterpolatedColor(colors, boundaries, normalized_activation || activation);
80 | return {token}
87 | })}
88 |
89 |
90 | )
91 | }
92 | })
93 | }
94 |
95 | }
96 |
--------------------------------------------------------------------------------
/neuron-viewer/src/panes/topTokens.jsx:
--------------------------------------------------------------------------------
1 | import React, { useState, useEffect } from "react"
2 | import { get_top_tokens } from "../interpAPI"
3 |
4 |
5 | const TokenDisplay = ({ activeNeuron }) => {
6 | const [isLoading, setIsLoading] = useState(true)
7 | const [data, setData] = useState(null)
8 |
9 | const loadTokens = async () => {
10 | setIsLoading(true)
11 | const weightStrengths = await get_top_tokens(activeNeuron, 'weight')
12 | const activationStrengths = await get_top_tokens(activeNeuron, 'activation')
13 |
14 | const data = {
15 | activeNeuron,
16 | weightStrengths,
17 | activationStrengths,
18 | }
19 |
20 | setData(data)
21 | setIsLoading(false)
22 | }
23 |
24 | useEffect(() => {
25 | if (!data) {
26 | loadTokens()
27 | }
28 | }, [])
29 |
30 |
31 | return (
32 |
33 |
Related tokens
34 | {isLoading ? (
35 |
38 | ) : (
39 | <>
40 |
Mean-activation-based
41 |
42 | {data.activationStrengths.tokens.map((token, idx) => {
43 | return (
44 | data.activationStrengths.average_activations[idx] === null ? null :
45 |
50 | {token}
51 |
52 | )
53 | })}
54 |
55 |
Weight-based
56 |
57 |
Input tokens:
58 | {data.weightStrengths.input_positive.tokens.slice(0, 20).map((token, idx) => {
59 | return (
60 | data.weightStrengths.input_positive.strengths[idx] === null ? null :
61 |
66 | {token}
67 |
68 | )
69 | })}
70 |
71 | {
72 |
73 |
Input tokens negative:
74 | {data.weightStrengths.input_negative.tokens.slice(0, 20).map((token, idx) => {
75 | return (
76 | data.weightStrengths.input_negative.strengths[idx] === null ? null :
77 |
82 | {token}
83 |
84 | )
85 | })}
86 |
87 | }
88 |
89 |
Output tokens:
90 | {data.weightStrengths.output_positive.tokens.slice(0, 20).map((token, idx) => {
91 | return (
92 | data.weightStrengths.output_positive.strengths[idx] === null ? null :
93 |
98 | {token}
99 |
100 | )
101 | })}
102 |
103 | {
104 |
105 |
Output tokens negative:
106 | {data.weightStrengths.output_negative.tokens.slice(0, 20).map((token, idx) => {
107 | return (
108 |
113 | {token}
114 |
115 | )
116 | })}
117 |
118 | }
119 |
120 | )}
121 |
122 | )
123 | }
124 | export default TokenDisplay
125 |
--------------------------------------------------------------------------------
/neuron-explainer/neuron_explainer/explanations/prompt_builder.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | from enum import Enum
4 | from typing import TypedDict, Union
5 |
6 | import tiktoken
7 |
8 | HarmonyMessage = TypedDict(
9 | "HarmonyMessage",
10 | {
11 | "role": str,
12 | "content": str,
13 | },
14 | )
15 |
16 |
17 | class PromptFormat(str, Enum):
18 | """
19 | Different ways of formatting the components of a prompt into the format accepted by the relevant
20 | API server endpoint.
21 | """
22 |
23 | NONE = "none"
24 | """Suitable for use with models that don't use special tokens for instructions."""
25 | INSTRUCTION_FOLLOWING = "instruction_following"
26 | """Suitable for IF models that use <|endofprompt|>."""
27 | HARMONY_V4 = "harmony_v4"
28 | """
29 | Suitable for Harmony models that use a structured turn-taking role+content format. Generates a
30 | list of HarmonyMessage dicts that can be sent to the /chat/completions endpoint.
31 | """
32 |
33 | @classmethod
34 | def from_string(cls, s: str) -> PromptFormat:
35 | for prompt_format in cls:
36 | if prompt_format.value == s:
37 | return prompt_format
38 | raise ValueError(f"{s} is not a valid PromptFormat")
39 |
40 |
41 | class Role(str, Enum):
42 | """See https://platform.openai.com/docs/guides/chat"""
43 |
44 | SYSTEM = "system"
45 | USER = "user"
46 | ASSISTANT = "assistant"
47 |
48 |
49 | class PromptBuilder:
50 | """Class for accumulating components of a prompt and then formatting them into an output."""
51 |
52 | def __init__(self) -> None:
53 | self._messages: list[HarmonyMessage] = []
54 |
55 | def add_message(self, role: Role, message: str) -> None:
56 | self._messages.append(HarmonyMessage(role=role, content=message))
57 |
58 | def prompt_length_in_tokens(self, prompt_format: PromptFormat) -> int:
59 | # TODO(sbills): Make the model/encoding configurable. This implementation assumes GPT-4.
60 | encoding = tiktoken.get_encoding("cl100k_base")
61 | if prompt_format == PromptFormat.HARMONY_V4:
62 | # Approximately-correct implementation adapted from this documentation:
63 | # https://platform.openai.com/docs/guides/chat/introduction
64 | num_tokens = 0
65 | for message in self._messages:
66 | num_tokens += (
67 | 4 # every message follows <|im_start|>{role/name}\n{content}<|im_end|>\n
68 | )
69 | num_tokens += len(encoding.encode(message["content"], allowed_special="all"))
70 | num_tokens += 2 # every reply is primed with <|im_start|>assistant
71 | return num_tokens
72 | else:
73 | prompt_str = self.build(prompt_format)
74 | assert isinstance(prompt_str, str)
75 | return len(encoding.encode(prompt_str, allowed_special="all"))
76 |
77 | def build(
78 | self, prompt_format: PromptFormat, *, allow_extra_system_messages: bool = False
79 | ) -> Union[str, list[HarmonyMessage]]:
80 | """
81 | Validates the messages added so far (reasonable alternation of assistant vs. user, etc.)
82 | and returns either a regular string (maybe with <|endofprompt|> tokens) or a list of
83 | HarmonyMessages suitable for use with the /chat/completions endpoint.
84 |
85 | The `allow_extra_system_messages` parameter allows the caller to specify that the prompt
86 | should be allowed to contain system messages after the very first one.
87 | """
88 | # Create a deep copy of the messages so we can modify it and so that the caller can't
89 | # modify the internal state of this object.
90 | messages = [message.copy() for message in self._messages]
91 |
92 | expected_next_role = Role.SYSTEM
93 | for message in messages:
94 | role = message["role"]
95 | assert role == expected_next_role or (
96 | allow_extra_system_messages and role == Role.SYSTEM
97 | ), f"Expected message from {expected_next_role} but got message from {role}"
98 | if role == Role.SYSTEM:
99 | expected_next_role = Role.USER
100 | elif role == Role.USER:
101 | expected_next_role = Role.ASSISTANT
102 | elif role == Role.ASSISTANT:
103 | expected_next_role = Role.USER
104 |
105 | if prompt_format == PromptFormat.INSTRUCTION_FOLLOWING:
106 | last_user_message = None
107 | for message in messages:
108 | if message["role"] == Role.USER:
109 | last_user_message = message
110 | assert last_user_message is not None
111 | last_user_message["content"] += "<|endofprompt|>"
112 |
113 | if prompt_format == PromptFormat.HARMONY_V4:
114 | return messages
115 | elif prompt_format in [PromptFormat.NONE, PromptFormat.INSTRUCTION_FOLLOWING]:
116 | return "".join(message["content"] for message in messages)
117 | else:
118 | raise ValueError(f"Unknown prompt format: {prompt_format}")
119 |
--------------------------------------------------------------------------------
/neuron-explainer/neuron_explainer/activations/activation_records.py:
--------------------------------------------------------------------------------
1 | """Utilities for formatting activation records into prompts."""
2 |
3 | import math
4 | from typing import Optional, Sequence
5 |
6 | from neuron_explainer.activations.activations import ActivationRecord
7 |
8 | UNKNOWN_ACTIVATION_STRING = "unknown"
9 |
10 |
11 | def relu(x: float) -> float:
12 | return max(0.0, x)
13 |
14 |
15 | def calculate_max_activation(activation_records: Sequence[ActivationRecord]) -> float:
16 | """Return the maximum activation value of the neuron across all the activation records."""
17 | flattened = [
18 | # Relu is used to assume any values less than 0 are indicating the neuron is in the resting
19 | # state. This is a simplifying assumption that works with relu/gelu.
20 | max(relu(x) for x in activation_record.activations)
21 | for activation_record in activation_records
22 | ]
23 | return max(flattened)
24 |
25 |
26 | def normalize_activations(activation_record: list[float], max_activation: float) -> list[int]:
27 | """Convert raw neuron activations to integers on the range [0, 10]."""
28 | if max_activation <= 0:
29 | return [0 for x in activation_record]
30 | # Relu is used to assume any values less than 0 are indicating the neuron is in the resting
31 | # state. This is a simplifying assumption that works with relu/gelu.
32 | return [min(10, math.floor(10 * relu(x) / max_activation)) for x in activation_record]
33 |
34 |
35 | def _format_activation_record(
36 | activation_record: ActivationRecord,
37 | max_activation: float,
38 | omit_zeros: bool,
39 | hide_activations: bool = False,
40 | start_index: int = 0,
41 | ) -> str:
42 | """Format neuron activations into a string, suitable for use in prompts."""
43 | tokens = activation_record.tokens
44 | normalized_activations = normalize_activations(activation_record.activations, max_activation)
45 | if omit_zeros:
46 | assert (not hide_activations) and start_index == 0, "Can't hide activations and omit zeros"
47 | tokens = [
48 | token for token, activation in zip(tokens, normalized_activations) if activation > 0
49 | ]
50 | normalized_activations = [x for x in normalized_activations if x > 0]
51 |
52 | entries = []
53 | assert len(tokens) == len(normalized_activations)
54 | for index, token, activation in zip(range(len(tokens)), tokens, normalized_activations):
55 | activation_string = str(int(activation))
56 | if hide_activations or index < start_index:
57 | activation_string = UNKNOWN_ACTIVATION_STRING
58 | entries.append(f"{token}\t{activation_string}")
59 | return "\n".join(entries)
60 |
61 |
62 | def format_activation_records(
63 | activation_records: Sequence[ActivationRecord],
64 | max_activation: float,
65 | *,
66 | omit_zeros: bool = False,
67 | start_indices: Optional[list[int]] = None,
68 | hide_activations: bool = False,
69 | ) -> str:
70 | """Format a list of activation records into a string."""
71 | return (
72 | "\n
\n"
73 | + "\n\n\n".join(
74 | [
75 | _format_activation_record(
76 | activation_record,
77 | max_activation,
78 | omit_zeros=omit_zeros,
79 | hide_activations=hide_activations,
80 | start_index=0 if start_indices is None else start_indices[i],
81 | )
82 | for i, activation_record in enumerate(activation_records)
83 | ]
84 | )
85 | + "\n\n"
86 | )
87 |
88 |
89 | def _format_tokens_for_simulation(tokens: Sequence[str]) -> str:
90 | """
91 | Format tokens into a string with each token marked as having an "unknown" activation, suitable
92 | for use in prompts.
93 | """
94 | entries = []
95 | for token in tokens:
96 | entries.append(f"{token}\t{UNKNOWN_ACTIVATION_STRING}")
97 | return "\n".join(entries)
98 |
99 |
100 | def format_sequences_for_simulation(
101 | all_tokens: Sequence[Sequence[str]],
102 | ) -> str:
103 | """
104 | Format a list of lists of tokens into a string with each token marked as having an "unknown"
105 | activation, suitable for use in prompts.
106 | """
107 | return (
108 | "\n\n"
109 | + "\n\n\n".join(
110 | [_format_tokens_for_simulation(tokens) for tokens in all_tokens]
111 | )
112 | + "\n\n"
113 | )
114 |
115 |
116 | def non_zero_activation_proportion(
117 | activation_records: Sequence[ActivationRecord], max_activation: float
118 | ) -> float:
119 | """Return the proportion of activation values that aren't zero."""
120 | total_activations_count = sum(
121 | [len(activation_record.activations) for activation_record in activation_records]
122 | )
123 | normalized_activations = [
124 | normalize_activations(activation_record.activations, max_activation)
125 | for activation_record in activation_records
126 | ]
127 | non_zero_activations_count = sum(
128 | [len([x for x in activations if x != 0]) for activations in normalized_activations]
129 | )
130 | return non_zero_activations_count / total_activations_count
131 |
--------------------------------------------------------------------------------
/neuron-viewer/src/interpAPI.ts:
--------------------------------------------------------------------------------
1 | import {Neuron} from './types';
2 | import {memoizeAsync} from "./utils"
3 |
4 | export const load_file_no_cache = async(path: string) => {
5 | const data = {
6 | path: path
7 | }
8 | const url = new URL("/load_az", window.location.href)
9 | url.port = '8000';
10 | return await (
11 | await fetch(url, {
12 | method: "POST", // or 'PUT'
13 | headers: {
14 | "Content-Type": "application/json",
15 | },
16 | body: JSON.stringify(data),
17 | })
18 | ).json()
19 |
20 | }
21 |
22 | export const load_file_az = async(path: string) => {
23 | const res = (
24 | await fetch(path, {
25 | method: "GET",
26 | mode: "cors",
27 | headers: {
28 | "Content-Type": "application/json",
29 | },
30 | })
31 | )
32 | if (!res.ok) {
33 | console.error(`HTTP error: ${res.status} - ${res.statusText}`);
34 | return;
35 | }
36 | return await res.json()
37 | }
38 |
39 |
40 | // export const load_file = memoizeAsync('load_file', load_file_no_cache)
41 | export const load_file = window.location.host.indexOf('localhost:') === -1 ? load_file_az : load_file_no_cache;
42 |
43 |
44 | // # (derived from az://oaialignment/datasets/interp/gpt2_xl/v1/webtext1/len_nomax/n_50000/mlp_post_act/ranked_by_max_activation)
45 | // const NEURON_RECORDS_PATH = "az://oaisbills/rcall/oss/migrated_make_crow_datasets/gpt2_xl_n_50000_64_token/neurons"
46 | const NEURON_RECORDS_PATH = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/collated-activations"
47 |
48 | // # (derived from az://oaialignment/datasets/interp/gpt2_xl/v1/webtext1/len_nomax/n_50000/mlp_post_act/ranked_by_max_activation/neurons/explanations/canonical-run-v1)
49 | // const EXPLANATIONS_PATH = "az://oaisbills/rcall/oss/migrated_explanation_datasets/canonical_gpt2_xl_all_neurons"
50 | const EXPLANATIONS_PATH = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/explanations"
51 |
52 | // weight-based
53 | // const WHOLE_LAYER_WEIGHT_TOKENS_PATH = "az://oaidan/rcall/data/interpretability/connections/gpt2-xl/mlp/unnorm_token_representations_uncommon_vanilla"
54 | // const WEIGHT_TOKENS_PATH = "az://oaijeffwu/jeffwu-data/interpretability/neuron-connections/gpt2-xl/weight-based"
55 | const WEIGHT_TOKENS_PATH = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/related-tokens/weight-based"
56 | // lookup table
57 | // const WHOLE_LAYER_ACTIVATION_TOKENS_PATH = "az://oaidan/rcall/data/interpretability/connections/gpt2_xl/mlp/unnorm_token_representations_vanilla_and_common_in_colangv2_unigram"
58 | // const ACTIVATION_TOKENS_PATH = "az://oaijeffwu/jeffwu-data/interpretability/neuron-connections/gpt2-xl/lookup-table"
59 | const ACTIVATION_TOKENS_PATH = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/related-tokens/activation-based"
60 |
61 | // const CONNECTIONS_PATH = "az://oaialignment/datasets/interp/connections/gpt2/neuron_space/incl_attn_False"
62 | const CONNECTIONS_PATH = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/related-neurons/weight-based"
63 |
64 |
65 | export const get_explanations = async (activeNeuron: Neuron) => {
66 | const result = await load_file(`${EXPLANATIONS_PATH}/${activeNeuron.layer}/${activeNeuron.neuron}.jsonl`)
67 | return result
68 | }
69 |
70 | export const get_top_tokens = async (activeNeuron: Neuron, weightType: string) => {
71 | let TOKENS_PATH;
72 | if (weightType === 'weight') {
73 | TOKENS_PATH = WEIGHT_TOKENS_PATH;
74 | } else if (weightType === 'activation') {
75 | TOKENS_PATH = ACTIVATION_TOKENS_PATH;
76 | } else {
77 | throw new Error(`Invalid weightType: ${weightType}`)
78 | }
79 | const result = await load_file(`${TOKENS_PATH}/${activeNeuron.layer}/${activeNeuron.neuron}.json`)
80 | return result
81 | // const result = await load_file_no_cache(`${ORIG_TOKENS_PATH}/${activeNeuron.layer}.json`)
82 | // return result.neuron_summaries[activeNeuron.neuron]
83 | }
84 |
85 | export const get_top_neuron_connections = async (activeNeuron: Neuron) => {
86 | const result = await load_file(`${CONNECTIONS_PATH}/${activeNeuron.layer}/${activeNeuron.neuron}.json`)
87 |
88 | const res: {[key: string]: [number, number]} = {};
89 | ["input", "output"].forEach((direction) => {
90 | const sign = "positive" // "negative"
91 | const weight_name: string = {output: "c_proj", input: "c_fc"}[direction] as string;
92 | const res_for_dir = result[weight_name];
93 | if (res_for_dir === null) {
94 | return
95 | }
96 | // let key = 'top_negative_neurons'
97 | const top_neuron_strs = res_for_dir[`top_${sign}_neurons`] // {layer}_{neuron} strings for each top-connected neuron
98 | const top_weights = res_for_dir[`top_${sign}_weights`]
99 | const top_layer_neuron_tuples = top_neuron_strs.map((neuron_str: string, i: number) => {
100 | const [layer, neuron] = neuron_str.split("_").map((x: string) => parseInt(x))
101 | return [layer, neuron, top_weights[i]] as [number, number, number]
102 | })
103 | res[direction] = top_layer_neuron_tuples.slice(0, 10)
104 | })
105 |
106 | return res
107 | }
108 |
109 | export const get_neuron_record = async(activeNeuron: Neuron) => {
110 | const result = await load_file(`${NEURON_RECORDS_PATH}/${activeNeuron.layer}/${activeNeuron.neuron}.json`)
111 | return result
112 | }
113 |
114 |
--------------------------------------------------------------------------------
/neuron-explainer/neuron_explainer/explanations/token_space_few_shot_examples.py:
--------------------------------------------------------------------------------
1 | from dataclasses import dataclass
2 | from enum import Enum
3 | from typing import List
4 |
5 | from neuron_explainer.fast_dataclasses import FastDataclass
6 |
7 |
8 | @dataclass
9 | class Example(FastDataclass):
10 | """
11 | An example list of tokens as strings corresponding to top token space inputs of a neuron, with a
12 | string explanation of the neuron's behavior on these tokens.
13 | """
14 |
15 | tokens: List[str]
16 | explanation: str
17 |
18 |
19 | class TokenSpaceFewShotExampleSet(Enum):
20 | """Determines which few-shot examples to use when sampling explanations."""
21 |
22 | ORIGINAL = "original"
23 | TEST = "test"
24 |
25 | def get_examples(self) -> list[Example]:
26 | """Returns regular examples for use in a few-shot prompt."""
27 | if self is TokenSpaceFewShotExampleSet.ORIGINAL:
28 | return ORIGINAL_EXAMPLES
29 | elif self is TokenSpaceFewShotExampleSet.TEST:
30 | return TEST_EXAMPLES
31 | else:
32 | raise ValueError(f"Unhandled example set: {self}")
33 |
34 |
35 | ORIGINAL_EXAMPLES = [
36 | Example(
37 | tokens=[
38 | "actual",
39 | " literal",
40 | " actual",
41 | " hyper",
42 | " real",
43 | " EX",
44 | " Real",
45 | "^",
46 | "Full",
47 | " full",
48 | " optical",
49 | " style",
50 | "any",
51 | "ALL",
52 | "extreme",
53 | " miniature",
54 | " Optical",
55 | " faint",
56 | "~",
57 | " Physical",
58 | " REAL",
59 | "*",
60 | "virtual",
61 | "TYPE",
62 | " technical",
63 | "otally",
64 | " physic",
65 | "Type",
66 | "<",
67 | "images",
68 | "atic",
69 | " sheer",
70 | " Style",
71 | " partial",
72 | " natural",
73 | "Hyper",
74 | " Any",
75 | " theoretical",
76 | "|",
77 | " ultimate",
78 | "oing",
79 | " constant",
80 | "ANY",
81 | "antically",
82 | "ishly",
83 | " ex",
84 | " visual",
85 | "special",
86 | "omorphic",
87 | "visual",
88 | ],
89 | explanation=" adjectives related to being real, or to physical properties and evidence",
90 | ),
91 | Example(
92 | tokens=[
93 | "cephal",
94 | "aeus",
95 | " coma",
96 | "bered",
97 | "abetes",
98 | "inflamm",
99 | "rugged",
100 | "alysed",
101 | "azine",
102 | "hered",
103 | "cells",
104 | "aneously",
105 | "fml",
106 | "igm",
107 | "culosis",
108 | "iani",
109 | "CTV",
110 | "disabled",
111 | "heric",
112 | "ulo",
113 | "geoning",
114 | "awi",
115 | "translation",
116 | "iral",
117 | "govtrack",
118 | "mson",
119 | "cloth",
120 | "nesota",
121 | " Dise",
122 | " Lyme",
123 | " dementia",
124 | "agn",
125 | " reversible",
126 | " susceptibility",
127 | "esthesia",
128 | "orf",
129 | " inflamm",
130 | " Obesity",
131 | " tox",
132 | " Disorders",
133 | "uberty",
134 | "blind",
135 | "ALTH",
136 | "avier",
137 | " Immunity",
138 | " Hurt",
139 | "ulet",
140 | "ueless",
141 | " sluggish",
142 | "rosis",
143 | ],
144 | explanation=" words related to physical medical conditions",
145 | ),
146 | Example(
147 | tokens=[
148 | " January",
149 | "terday",
150 | "cember",
151 | " April",
152 | " July",
153 | "September",
154 | "December",
155 | "Thursday",
156 | "quished",
157 | "November",
158 | "Tuesday",
159 | "uesday",
160 | " Sept",
161 | "ruary",
162 | " March",
163 | ";;;;;;;;;;;;",
164 | " Monday",
165 | "Wednesday",
166 | " Saturday",
167 | " Wednesday",
168 | "Reloaded",
169 | "aturday",
170 | " August",
171 | "Feb",
172 | "Sunday",
173 | "Reviewed",
174 | "uggest",
175 | " Dhabi",
176 | "ACTED",
177 | "tten",
178 | "Year",
179 | "August",
180 | "alogue",
181 | "MX",
182 | " Janeiro",
183 | "yss",
184 | " Leilan",
185 | " Fiscal",
186 | " referen",
187 | "semb",
188 | "eele",
189 | "wcs",
190 | "detail",
191 | "ertation",
192 | " Reborn",
193 | " Sunday",
194 | "itially",
195 | "aturdays",
196 | " Dise",
197 | "essage",
198 | ],
199 | explanation=" nouns related to time and dates",
200 | ),
201 | ]
202 |
203 | TEST_EXAMPLES = [
204 | Example(
205 | tokens=[
206 | "these",
207 | " are",
208 | " tokens",
209 | ],
210 | explanation=" this is a test explanation",
211 | ),
212 | ]
213 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Automated interpretability
2 |
3 | ## Code and tools
4 |
5 | This repository contains code and tools associated with the [Language models can explain neurons in
6 | language models](https://openaipublic.blob.core.windows.net/neuron-explainer/paper/index.html) paper, specifically:
7 |
8 | * Code for automatically generating, simulating, and scoring explanations of neuron behavior using
9 | the methodology described in the paper. See the
10 | [neuron-explainer README](neuron-explainer/README.md) for more information.
11 |
12 | Note: if you run into errors of the form "Error: Could not find any credentials that grant access to storage account: 'openaipublic' and container: 'neuron-explainer'"." you might be able to fix this by signing up for an azure account and specifying the credentials as described in the error message.
13 |
14 | * A tool for viewing neuron activations and explanations, accessible
15 | [here](https://openaipublic.blob.core.windows.net/neuron-explainer/neuron-viewer/index.html). See
16 | the [neuron-viewer README](neuron-viewer/README.md) for more information.
17 |
18 | ## Public datasets
19 |
20 | Together with this code, we're also releasing public datasets of GPT-2 XL neurons and explanations.
21 | Here's an overview of those datasets.
22 |
23 | * Neuron activations: `az://openaipublic/neuron-explainer/data/collated-activations/{layer_index}/{neuron_index}.json`
24 | - Tokenized text sequences and their activations for the neuron. We
25 | provide multiple sets of tokens and activations: top-activating ones, random
26 | samples from several quantiles; and a completely random sample. We also provide
27 | some basic statistics for the activations.
28 | - Each file contains a JSON-formatted
29 | [`NeuronRecord`](neuron-explainer/neuron_explainer/activations/activations.py#L89) dataclass.
30 | * Neuron explanations: `az://openaipublic/neuron-explainer/data/explanations/{layer_index}/{neuron_index}.jsonl`
31 | - Scored model-generated explanations of the behavior of the neuron, including simulation results.
32 | - Each file contains a JSON-formatted
33 | [`NeuronSimulationResults`](neuron-explainer/neuron_explainer/explanations/explanations.py#L146)
34 | dataclass.
35 | * Related neurons: `az://openaipublic/neuron-explainer/data/related-neurons/weight-based/{layer_index}/{neuron_index}.json`
36 | - Lists of the upstream and downstream neurons with the most positive and negative connections (see below for definition).
37 | - Each file contains a JSON-formatted dataclass whose definition is not included in this repo.
38 | * Tokens with high average activations:
39 | `az://openaipublic/neuron-explainer/data/related-tokens/activation-based/{layer_index}/{neuron_index}.json`
40 | - Lists of tokens with the highest average activations for individual neurons, and their average activations.
41 | - Each file contains a JSON-formatted [`TokenLookupTableSummaryOfNeuron`](neuron-explainer/neuron_explainer/activations/token_connections.py#L36)
42 | dataclass.
43 | * Tokens with large inbound and outbound weights:
44 | `az://openaipublic/neuron-explainer/data/related-tokens/weight-based/{layer_index}/{neuron_index}.json`
45 | - List of the most-positive and most-negative input and output tokens for individual neurons,
46 | as well as the associated weight (see below for definition).
47 | - Each file contains a JSON-formatted [`WeightBasedSummaryOfNeuron`](neuron-explainer/neuron_explainer/activations/token_connections.py#L17)
48 | dataclass.
49 |
50 | Update (July 5, 2023):
51 | We also released a set of explanations for GPT-2 Small. The methodology is slightly different from the methodology used for GPT-2 XL so the results aren't directly comparable.
52 | * Neuron activations: `az://openaipublic/neuron-explainer/gpt2_small_data/collated-activations/{layer_index}/{neuron_index}.json`
53 | * Neuron explanations: `az://openaipublic/neuron-explainer/gpt2_small_data/explanations/{layer_index}/{neuron_index}.jsonl`
54 |
55 | Update (August 30, 2023): We recently discovered a bug in how we performed inference on the GPT-2 series models used for the paper and for these datasets. Specifically, we used an optimized GELU implementation rather than the original GELU implementation associated with GPT-2. While the model’s behavior is very similar across these two configurations, the post-MLP activation values we used to generate and simulate explanations differ from the correct values by the following amounts for GPT-2 small:
56 |
57 | - Median: 0.0090
58 | - 90th percentile: 0.0252
59 | - 99th percentile: 0.0839
60 | - 99.9th percentile: 0.1736
61 |
62 | ### Definition of connection weights
63 |
64 | Refer to [GPT-2 model code](https://github.com/openai/gpt-2/blob/master/src/model.py) for
65 | understanding of model weight conventions.
66 |
67 | *Neuron-neuron*: For two neurons `(l1, n1)` and `(l2, n2)` with `l1 < l2`, the connection strength is defined as
68 | `h{l1}.mlp.c_proj.w[:, n1, :] @ diag(h{l2}.ln_2.g) @ h{l2}.mlp.c_fc.w[:, :, n2]`.
69 |
70 | *Neuron-token*: For token `t` and neuron `(l, n)`, the input weight is computed as
71 | `wte[t, :] @ diag(h{l}.ln_2.g) @ h{l}.mlp.c_fc.w[:, :, n]`
72 | and the output weight is computed as
73 | `h{l}.mlp.c_proj.w[:, n, :] @ diag(ln_f.g) @ wte[t, :]`.
74 |
75 | ### Misc Lists of Interesting Neurons
76 | Lists of neurons we thought were interesting according to different criteria, with some preliminary descriptions.
77 | * [Interesting Neurons (external)](https://docs.google.com/spreadsheets/d/1p7fYs31NU8sJoeKyUx4Mn2laGx8xXfHg_KcIvYiKPpg/edit#gid=0)
78 | * [Neurons that score high on random, possibly monosemantic? (external)](https://docs.google.com/spreadsheets/d/1TqKFcz-84jyIHLU7VRoTc8BoFBMpbgac-iNBnxVurQ8/edit?usp=sharing)
79 | * [Clusters of neurons well explained by activation explanation but not by tokens](https://docs.google.com/document/d/1lWhKowpKDdwTMALD_K541cdwgGoQx8DFUSuEe1U2AGE/edit?usp=sharing)
80 | * [Neurons sensitive to truncation](https://docs.google.com/document/d/1x89TWBvuHcyC2t01EDbJZJ5LQYHozlcS-VUmr5shf_A/edit?usp=sharing)
81 |
--------------------------------------------------------------------------------
/neuron-explainer/neuron_explainer/api_client.py:
--------------------------------------------------------------------------------
1 | import asyncio
2 | import contextlib
3 | import os
4 | import random
5 | import traceback
6 | from asyncio import Semaphore
7 | from functools import wraps
8 | from typing import Any, Callable, Optional
9 |
10 | import httpx
11 | import orjson
12 |
13 |
14 | def is_api_error(err: Exception) -> bool:
15 | if isinstance(err, httpx.HTTPStatusError):
16 | response = err.response
17 | error_data = response.json().get("error", {})
18 | error_message = error_data.get("message")
19 | if response.status_code in [400, 404, 415]:
20 | if error_data.get("type") == "idempotency_error":
21 | print(f"Retrying after idempotency error: {error_message} ({response.url})")
22 | return True
23 | else:
24 | # Invalid request
25 | return False
26 | else:
27 | print(f"Retrying after API error: {error_message} ({response.url})")
28 | return True
29 |
30 | elif isinstance(err, httpx.ConnectError):
31 | print(f"Retrying after connection error... ({err.request.url})")
32 | return True
33 |
34 | elif isinstance(err, httpx.TimeoutException):
35 | print(f"Retrying after a timeout error... ({err.request.url})")
36 | return True
37 |
38 | elif isinstance(err, httpx.ReadError):
39 | print(f"Retrying after a read error... ({err.request.url})")
40 | return True
41 |
42 | print(f"Retrying after an unexpected error: {repr(err)}")
43 | traceback.print_tb(err.__traceback__)
44 | return True
45 |
46 |
47 | def exponential_backoff(
48 | retry_on: Callable[[Exception], bool] = lambda err: True
49 | ) -> Callable[[Callable], Callable]:
50 | """
51 | Returns a decorator which retries the wrapped function as long as the specified retry_on
52 | function returns True for the exception, applying exponential backoff with jitter after
53 | failures, up to a retry limit.
54 | """
55 | init_delay_s = 1.0
56 | max_delay_s = 10.0
57 | # Roughly 30 minutes before we give up.
58 | max_tries = 200
59 | backoff_multiplier = 2.0
60 | jitter = 0.2
61 |
62 | def decorate(f: Callable) -> Callable:
63 | assert asyncio.iscoroutinefunction(f)
64 |
65 | @wraps(f)
66 | async def f_retry(*args: Any, **kwargs: Any) -> None:
67 | delay_s = init_delay_s
68 | for i in range(max_tries):
69 | try:
70 | return await f(*args, **kwargs)
71 | except Exception as err:
72 | if not retry_on(err) or i == max_tries - 1:
73 | raise
74 | jittered_delay = random.uniform(delay_s * (1 - jitter), delay_s * (1 + jitter))
75 | await asyncio.sleep(jittered_delay)
76 | delay_s = min(delay_s * backoff_multiplier, max_delay_s)
77 |
78 | return f_retry
79 |
80 | return decorate
81 |
82 |
83 | API_KEY = os.getenv("OPENAI_API_KEY")
84 | assert API_KEY, "Please set the OPENAI_API_KEY environment variable"
85 | API_HTTP_HEADERS = {
86 | "Content-Type": "application/json",
87 | "Authorization": "Bearer " + API_KEY,
88 | }
89 | BASE_API_URL = "https://api.openai.com/v1"
90 |
91 |
92 | class ApiClient:
93 | """Performs inference using the OpenAI API. Supports response caching and concurrency limits."""
94 |
95 | def __init__(
96 | self,
97 | model_name: str,
98 | # If set, no more than this number of HTTP requests will be made concurrently.
99 | max_concurrent: Optional[int] = None,
100 | # Whether to cache request/response pairs in memory to avoid duplicating requests.
101 | cache: bool = False,
102 | ):
103 | self.model_name = model_name
104 |
105 | if max_concurrent is not None:
106 | self._concurrency_check: Optional[Semaphore] = Semaphore(max_concurrent)
107 | else:
108 | self._concurrency_check = None
109 |
110 | if cache:
111 | self._cache: Optional[dict[str, Any]] = {}
112 | else:
113 | self._cache = None
114 |
115 | @exponential_backoff(retry_on=is_api_error)
116 | async def make_request(
117 | self, timeout_seconds: Optional[int] = None, **kwargs: Any
118 | ) -> dict[str, Any]:
119 | if self._cache is not None:
120 | key = orjson.dumps(kwargs)
121 | if key in self._cache:
122 | return self._cache[key]
123 | async with contextlib.AsyncExitStack() as stack:
124 | if self._concurrency_check is not None:
125 | await stack.enter_async_context(self._concurrency_check)
126 | http_client = await stack.enter_async_context(
127 | httpx.AsyncClient(timeout=timeout_seconds)
128 | )
129 | # If the request has a "messages" key, it should be sent to the /chat/completions
130 | # endpoint. Otherwise, it should be sent to the /completions endpoint.
131 | url = BASE_API_URL + ("/chat/completions" if "messages" in kwargs else "/completions")
132 | kwargs["model"] = self.model_name
133 | response = await http_client.post(url, headers=API_HTTP_HEADERS, json=kwargs)
134 | # The response json has useful information but the exception doesn't include it, so print it
135 | # out then reraise.
136 | try:
137 | response.raise_for_status()
138 | except Exception as e:
139 | print(response.json())
140 | raise e
141 | if self._cache is not None:
142 | self._cache[key] = response.json()
143 | return response.json()
144 |
145 |
146 | if __name__ == "__main__":
147 |
148 | async def main() -> None:
149 | client = ApiClient(model_name="gpt-3.5-turbo", max_concurrent=1)
150 | print(await client.make_request(prompt="Why did the chicken cross the road?", max_tokens=9))
151 |
152 | asyncio.run(main())
153 |
--------------------------------------------------------------------------------
/neuron-explainer/neuron_explainer/explanations/scoring.py:
--------------------------------------------------------------------------------
1 | from __future__ import annotations
2 |
3 | import asyncio
4 | import logging
5 | from typing import Any, Callable, Coroutine, Sequence
6 |
7 | import numpy as np
8 | from neuron_explainer.activations.activations import ActivationRecord
9 | from neuron_explainer.explanations.calibrated_simulator import (
10 | CalibratedNeuronSimulator,
11 | LinearCalibratedNeuronSimulator,
12 | )
13 | from neuron_explainer.explanations.explanations import (
14 | ScoredSequenceSimulation,
15 | ScoredSimulation,
16 | SequenceSimulation,
17 | )
18 | from neuron_explainer.explanations.simulator import ExplanationNeuronSimulator, NeuronSimulator
19 |
20 |
21 | def flatten_list(list_of_lists: Sequence[Sequence[Any]]) -> list[Any]:
22 | return [item for sublist in list_of_lists for item in sublist]
23 |
24 |
25 | def correlation_score(
26 | real_activations: Sequence[float] | np.ndarray,
27 | predicted_activations: Sequence[float] | np.ndarray,
28 | ) -> float:
29 | return np.corrcoef(real_activations, predicted_activations)[0, 1]
30 |
31 |
32 | def score_from_simulation(
33 | real_activations: ActivationRecord,
34 | simulation: SequenceSimulation,
35 | score_function: Callable[[Sequence[float] | np.ndarray, Sequence[float] | np.ndarray], float],
36 | ) -> float:
37 | return score_function(real_activations.activations, simulation.expected_activations)
38 |
39 |
40 | def rsquared_score_from_sequences(
41 | real_activations: Sequence[float] | np.ndarray,
42 | predicted_activations: Sequence[float] | np.ndarray,
43 | ) -> float:
44 | return float(
45 | 1
46 | - np.mean(np.square(np.array(real_activations) - np.array(predicted_activations)))
47 | / np.mean(np.square(np.array(real_activations)))
48 | )
49 |
50 |
51 | def absolute_dev_explained_score_from_sequences(
52 | real_activations: Sequence[float] | np.ndarray,
53 | predicted_activations: Sequence[float] | np.ndarray,
54 | ) -> float:
55 | return float(
56 | 1
57 | - np.mean(np.abs(np.array(real_activations) - np.array(predicted_activations)))
58 | / np.mean(np.abs(np.array(real_activations)))
59 | )
60 |
61 |
62 | async def make_explanation_simulator(
63 | explanation: str,
64 | calibration_activation_records: Sequence[ActivationRecord],
65 | model_name: str,
66 | calibrated_simulator_class: type[CalibratedNeuronSimulator] = LinearCalibratedNeuronSimulator,
67 | ) -> CalibratedNeuronSimulator:
68 | """
69 | Make a simulator that uses an explanation to predict activations and calibrates it on the given
70 | activation records.
71 | """
72 | simulator = ExplanationNeuronSimulator(model_name, explanation)
73 | calibrated_simulator = calibrated_simulator_class(simulator)
74 | await calibrated_simulator.calibrate(calibration_activation_records)
75 | return calibrated_simulator
76 |
77 |
78 | async def _simulate_and_score_sequence(
79 | simulator: NeuronSimulator, activations: ActivationRecord
80 | ) -> ScoredSequenceSimulation:
81 | """Score an explanation of a neuron by how well it predicts activations on a sentence."""
82 | simulation = await simulator.simulate(activations.tokens)
83 | logging.debug(simulation)
84 | rsquared_score = score_from_simulation(activations, simulation, rsquared_score_from_sequences)
85 | absolute_dev_explained_score = score_from_simulation(
86 | activations, simulation, absolute_dev_explained_score_from_sequences
87 | )
88 | scored_sequence_simulation = ScoredSequenceSimulation(
89 | simulation=simulation,
90 | true_activations=activations.activations,
91 | ev_correlation_score=score_from_simulation(activations, simulation, correlation_score),
92 | rsquared_score=rsquared_score,
93 | absolute_dev_explained_score=absolute_dev_explained_score,
94 | )
95 | return scored_sequence_simulation
96 |
97 |
98 | def aggregate_scored_sequence_simulations(
99 | scored_sequence_simulations: list[ScoredSequenceSimulation],
100 | ) -> ScoredSimulation:
101 | """
102 | Aggregate a list of scored sequence simulations. The logic for doing this is non-trivial for EV
103 | scores, since we want to calculate the correlation over all activations from all sequences at
104 | once rather than simply averaging per-sequence correlations.
105 | """
106 | all_true_activations: list[float] = []
107 | all_expected_values: list[float] = []
108 | for scored_sequence_simulation in scored_sequence_simulations:
109 | all_true_activations.extend(scored_sequence_simulation.true_activations or [])
110 | all_expected_values.extend(scored_sequence_simulation.simulation.expected_activations)
111 | ev_correlation_score = (
112 | correlation_score(all_true_activations, all_expected_values)
113 | if len(all_true_activations) > 0
114 | else None
115 | )
116 | rsquared_score = rsquared_score_from_sequences(all_true_activations, all_expected_values)
117 | absolute_dev_explained_score = absolute_dev_explained_score_from_sequences(
118 | all_true_activations, all_expected_values
119 | )
120 |
121 | return ScoredSimulation(
122 | scored_sequence_simulations=scored_sequence_simulations,
123 | ev_correlation_score=ev_correlation_score,
124 | rsquared_score=rsquared_score,
125 | absolute_dev_explained_score=absolute_dev_explained_score,
126 | )
127 |
128 |
129 | async def simulate_and_score(
130 | simulator: NeuronSimulator,
131 | activation_records: Sequence[ActivationRecord],
132 | ) -> ScoredSimulation:
133 | """
134 | Score an explanation of a neuron by how well it predicts activations on the given text
135 | sequences.
136 | """
137 | scored_sequence_simulations = await asyncio.gather(
138 | *[
139 | _simulate_and_score_sequence(
140 | simulator,
141 | activation_record,
142 | )
143 | for activation_record in activation_records
144 | ]
145 | )
146 | return aggregate_scored_sequence_simulations(scored_sequence_simulations)
147 |
148 |
149 | async def make_simulator_and_score(
150 | make_simulator: Coroutine[None, None, NeuronSimulator],
151 | activation_records: Sequence[ActivationRecord],
152 | ) -> ScoredSimulation:
153 | """Chain together creating the simulator and using it to score activation records."""
154 | simulator = await make_simulator
155 | return await simulate_and_score(simulator, activation_records)
156 |
--------------------------------------------------------------------------------
/neuron-viewer/src/panes/explanation.jsx:
--------------------------------------------------------------------------------
1 | import React, { useState, useEffect } from "react"
2 | import { get_explanations } from "../interpAPI"
3 | // import HeatmapGrid from "../heatmapGrid"
4 | import SimulationHeatmap from "../simulationHeatmap"
5 | import { normalizeTokenActs } from "../types"
6 |
7 |
8 | function zip_simulated_sequences(sequences) {
9 | return sequences.map(({ simulation }) => {
10 | return simulation.tokens.map((token, idx) => ({
11 | token,
12 | activation: simulation.expected_activations[idx],
13 | }))
14 | })
15 | }
16 |
17 | function zip_real_sequences(sequences) {
18 | return sequences.map(({ simulation, true_activations }) => {
19 | return simulation.tokens.map((token, idx) => ({
20 | token,
21 | activation: true_activations[idx],
22 | }))
23 | })
24 | }
25 |
26 | const ExplanationDisplay = ({ activeNeuron }) => {
27 | const [isLoading, setIsLoading] = useState(true)
28 | const [data, setData] = useState(null)
29 | const [showingScoringDetails, setShowingScoringDetails] = useState(false)
30 | const [toggle, setToggle] = useState(false);
31 |
32 | const loadExplanation = async () => {
33 | const result = await get_explanations(activeNeuron);
34 | setData(result.scored_explanations[0])
35 | setIsLoading(false)
36 | }
37 |
38 | useEffect(() => {
39 | if (!data) {
40 | loadExplanation()
41 | }
42 | }, [])
43 |
44 | const handleToggleChange = () => {
45 | setToggle(!toggle);
46 | };
47 |
48 | let sim_sequences;
49 | if (data) {
50 | sim_sequences = zip_simulated_sequences(data.scored_simulation.scored_sequence_simulations);
51 | [sim_sequences] = normalizeTokenActs(sim_sequences)
52 | } else {
53 | sim_sequences = []
54 | }
55 |
56 | let real_sequences;
57 | if (data) {
58 | real_sequences = zip_real_sequences(data.scored_simulation.scored_sequence_simulations);
59 | [real_sequences] = normalizeTokenActs(real_sequences)
60 | } else {
61 | real_sequences = []
62 | }
63 |
64 | const suggest_explanation_link = "https://docs.google.com/forms/d/e/1FAIpQLSckMyDQedGhdISIqaqn0YGUtd2xqEWgPu7ehoPUTT2pTge_-g/viewform?"
65 | + `usp=pp_url&entry.541490611=${activeNeuron.layer}`
66 | + `&entry.1688855196=${activeNeuron.neuron}`
67 | + `&entry.495312202=https://openaipublic.blob.core.windows.net/neuron-explainer/neuron-viewer/index.html%23/layers/${activeNeuron.layer}/neurons/${activeNeuron.neuron}`;
68 |
69 | return (
70 | <>
71 |
72 |
Explanation
73 | {isLoading ? (
74 |
77 | ) : (
78 | <>
79 |
80 |
81 | {data.explanation}
82 |
83 |
84 | score: {data.scored_simulation.ev_correlation_score.toFixed(2)}
85 |
86 |
87 | Suggest Better Explanation
88 |
89 |
90 |
93 | {
94 | showingScoringDetails ?
95 | <>
96 |
101 |
115 |
128 |
142 |
143 |
144 | {toggle ? 'Activations overlaid (top = real, bottom = simulated)' : 'Activations not overlaid'}
145 |
146 |
Top
147 |
152 |
Random
153 |
158 | : null
159 | }
160 |
161 | )}
162 |
93 |
Welcome! Pick a neuron:
94 |
126 |
132 |
133 |
Interesting neurons:
134 |
135 |
138 | {knownGoodNeurons.map(({ heading, layer, neuron, label, description }, j) => (
139 | heading ?
140 | {heading}
141 |
:
150 | ))}
151 |
152 |
153 |
154 |