interp-engine 0.0.24__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (27) hide show
  1. interp_engine-0.0.24/LICENSE +19 -0
  2. interp_engine-0.0.24/PKG-INFO +21 -0
  3. interp_engine-0.0.24/neuron_explainer/__init__.py +0 -0
  4. interp_engine-0.0.24/neuron_explainer/activations/__init__.py +0 -0
  5. interp_engine-0.0.24/neuron_explainer/activations/activation_records.py +130 -0
  6. interp_engine-0.0.24/neuron_explainer/activations/activations.py +311 -0
  7. interp_engine-0.0.24/neuron_explainer/activations/attention_utils.py +121 -0
  8. interp_engine-0.0.24/neuron_explainer/activations/token_connections.py +59 -0
  9. interp_engine-0.0.24/neuron_explainer/api_client.py +190 -0
  10. interp_engine-0.0.24/neuron_explainer/azure.py +5 -0
  11. interp_engine-0.0.24/neuron_explainer/explanations/__init__.py +0 -0
  12. interp_engine-0.0.24/neuron_explainer/explanations/calibrated_simulator.py +194 -0
  13. interp_engine-0.0.24/neuron_explainer/explanations/explainer.py +2585 -0
  14. interp_engine-0.0.24/neuron_explainer/explanations/explanations.py +230 -0
  15. interp_engine-0.0.24/neuron_explainer/explanations/few_shot_examples.py +3125 -0
  16. interp_engine-0.0.24/neuron_explainer/explanations/prompt_builder.py +118 -0
  17. interp_engine-0.0.24/neuron_explainer/explanations/puzzles.json +399 -0
  18. interp_engine-0.0.24/neuron_explainer/explanations/puzzles.py +50 -0
  19. interp_engine-0.0.24/neuron_explainer/explanations/scoring.py +155 -0
  20. interp_engine-0.0.24/neuron_explainer/explanations/simulator.py +1121 -0
  21. interp_engine-0.0.24/neuron_explainer/explanations/test_explainer.py +227 -0
  22. interp_engine-0.0.24/neuron_explainer/explanations/test_simulator.py +269 -0
  23. interp_engine-0.0.24/neuron_explainer/explanations/token_space_few_shot_examples.py +212 -0
  24. interp_engine-0.0.24/neuron_explainer/fast_dataclasses/__init__.py +3 -0
  25. interp_engine-0.0.24/neuron_explainer/fast_dataclasses/fast_dataclasses.py +85 -0
  26. interp_engine-0.0.24/neuron_explainer/fast_dataclasses/test_fast_dataclasses.py +83 -0
  27. interp_engine-0.0.24/pyproject.toml +19 -0
@@ -0,0 +1,19 @@
1
+ Copyright (c) 2023 Superalignment, OpenAI
2
+
3
+ Permission is hereby granted, free of charge, to any person obtaining a copy
4
+ of this software and associated documentation files (the "Software"), to deal
5
+ in the Software without restriction, including without limitation the rights
6
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7
+ copies of the Software, and to permit persons to whom the Software is
8
+ furnished to do so, subject to the following conditions:
9
+
10
+ The above copyright notice and this permission notice shall be included in all
11
+ copies or substantial portions of the Software.
12
+
13
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19
+ SOFTWARE.
@@ -0,0 +1,21 @@
1
+ Metadata-Version: 2.4
2
+ Name: interp-engine
3
+ Version: 0.0.24
4
+ Summary: OpenAI and Neuronpedia's implementation of automated-interpretability, with some updates. Not officially affiliated with OpenAI.
5
+ License-File: LICENSE
6
+ Author: OpenAI, Neuronpedia
7
+ Requires-Python: >=3.9,<4.0
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: Programming Language :: Python :: 3.9
10
+ Classifier: Programming Language :: Python :: 3.10
11
+ Classifier: Programming Language :: Python :: 3.11
12
+ Classifier: Programming Language :: Python :: 3.12
13
+ Classifier: Programming Language :: Python :: 3.13
14
+ Classifier: Programming Language :: Python :: 3.14
15
+ Requires-Dist: blobfile (>=2.1.1,<3.0.0)
16
+ Requires-Dist: boostedblob (>=0.15.3,<0.16.0)
17
+ Requires-Dist: httpx (>=0.27.0)
18
+ Requires-Dist: numpy (>=1.24.0,<2.0.0)
19
+ Requires-Dist: orjson (>=3.10.1,<4.0.0)
20
+ Requires-Dist: scikit-learn (>=1.2.0,<2.0.0)
21
+ Requires-Dist: tiktoken (>=0.6.0)
File without changes
@@ -0,0 +1,130 @@
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<start>\n"
73
+ + "\n<end>\n<start>\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<end>\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<start>\n"
109
+ + "\n<end>\n<start>\n".join(
110
+ [_format_tokens_for_simulation(tokens) for tokens in all_tokens]
111
+ )
112
+ + "\n<end>\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
@@ -0,0 +1,311 @@
1
+ # Dataclasses and enums for storing neuron-indexed information about activations. Also, related
2
+ # helper functions.
3
+
4
+ import math
5
+ from dataclasses import dataclass, field
6
+ from typing import List, Optional, Union
7
+
8
+ import urllib.request
9
+ import blobfile as bf
10
+ import boostedblob as bbb
11
+ from neuron_explainer.fast_dataclasses import FastDataclass, loads, register_dataclass
12
+ from neuron_explainer.azure import standardize_azure_url
13
+
14
+
15
+ @register_dataclass
16
+ @dataclass
17
+ class ActivationRecord(FastDataclass):
18
+ """Collated lists of tokens and their activations for a single neuron."""
19
+
20
+ tokens: List[str]
21
+ """Tokens in the text sequence, represented as strings."""
22
+ activations: List[float]
23
+ """Raw activation values for the neuron on each token in the text sequence."""
24
+ dfa_values: Optional[List[int]] = None
25
+ dfa_target_index: Optional[int] = None
26
+
27
+
28
+ @register_dataclass
29
+ @dataclass
30
+ class NeuronId(FastDataclass):
31
+ """Identifier for a neuron in an artificial neural network."""
32
+
33
+ layer_index: int
34
+ """The index of layer the neuron is in. The first layer used during inference has index 0."""
35
+ neuron_index: int
36
+ """The neuron's index within in its layer. Indices start from 0 in each layer."""
37
+
38
+
39
+ def _check_slices(
40
+ slices_by_split: dict[str, slice],
41
+ expected_num_values: int,
42
+ ) -> None:
43
+ """Assert that the slices are disjoint and fully cover the intended range."""
44
+ indices = set()
45
+ sum_of_slice_lengths = 0
46
+ n_splits = len(slices_by_split.keys())
47
+ for s in slices_by_split.values():
48
+ subrange = range(expected_num_values)[s]
49
+ sum_of_slice_lengths += len(subrange)
50
+ indices |= set(subrange)
51
+ assert (
52
+ sum_of_slice_lengths == expected_num_values
53
+ ), f"{sum_of_slice_lengths=} != {expected_num_values=}"
54
+ stride = n_splits
55
+ expected_indices = set.union(
56
+ *[
57
+ set(range(start_index, expected_num_values, stride))
58
+ for start_index in range(n_splits)
59
+ ]
60
+ )
61
+ assert indices == expected_indices, f"{indices=} != {expected_indices=}"
62
+
63
+
64
+ def get_slices_for_splits(
65
+ splits: list[str],
66
+ num_activation_records_per_split: int,
67
+ ) -> dict[str, slice]:
68
+ """
69
+ Get equal-sized interleaved subsets for each of a list of splits, given the number of elements
70
+ to include in each split.
71
+ """
72
+
73
+ stride = len(splits)
74
+ num_activation_records_for_even_splits = num_activation_records_per_split * stride
75
+ slices_by_split = {
76
+ split: slice(split_index, num_activation_records_for_even_splits, stride)
77
+ for split_index, split in enumerate(splits)
78
+ }
79
+ _check_slices(
80
+ slices_by_split=slices_by_split,
81
+ expected_num_values=num_activation_records_for_even_splits,
82
+ )
83
+ return slices_by_split
84
+
85
+
86
+ @dataclass
87
+ class ActivationRecordSliceParams:
88
+ """How to select splits (train, valid, etc.) of activation records."""
89
+
90
+ n_examples_per_split: Optional[int]
91
+ """The number of examples to include in each split."""
92
+
93
+
94
+ @register_dataclass
95
+ @dataclass
96
+ class NeuronRecord(FastDataclass):
97
+ """Neuron-indexed activation data, including summary stats and notable activation records."""
98
+
99
+ neuron_id: NeuronId
100
+ """Identifier for the neuron."""
101
+
102
+ random_sample: list[ActivationRecord] = field(default_factory=list)
103
+ """
104
+ Random activation records for this neuron. The random sample is independent from those used for
105
+ other neurons.
106
+ """
107
+ random_sample_by_quantile: Optional[list[list[ActivationRecord]]] = None
108
+ """
109
+ Random samples of activation records in each of the specified quantiles. None if quantile
110
+ tracking is disabled.
111
+ """
112
+ quantile_boundaries: Optional[list[float]] = None
113
+ """Boundaries of the quantiles used to generate the random_sample_by_quantile field."""
114
+
115
+ # Moments of activations
116
+ mean: Optional[float] = math.nan
117
+ variance: Optional[float] = math.nan
118
+ skewness: Optional[float] = math.nan
119
+ kurtosis: Optional[float] = math.nan
120
+
121
+ most_positive_activation_records: list[ActivationRecord] = field(
122
+ default_factory=list
123
+ )
124
+ """
125
+ Activation records with the most positive figure of merit value for this neuron over all dataset
126
+ examples.
127
+ """
128
+
129
+ @property
130
+ def max_activation(self) -> float:
131
+ """Return the maximum activation value over all top-activating activation records."""
132
+ return max(
133
+ [max(ar.activations) for ar in self.most_positive_activation_records]
134
+ )
135
+
136
+ def _get_top_activation_slices(
137
+ self, activation_record_slice_params: ActivationRecordSliceParams
138
+ ) -> dict[str, slice]:
139
+ splits = ["train", "calibration", "valid", "test"]
140
+ n_examples_per_split = activation_record_slice_params.n_examples_per_split
141
+ if n_examples_per_split is None:
142
+ n_examples_per_split = len(self.most_positive_activation_records) // len(
143
+ splits
144
+ )
145
+ assert len(self.most_positive_activation_records) >= n_examples_per_split * len(
146
+ splits
147
+ )
148
+ return get_slices_for_splits(splits, n_examples_per_split)
149
+
150
+ def _get_random_activation_slices(
151
+ self, activation_record_slice_params: ActivationRecordSliceParams
152
+ ) -> dict[str, slice]:
153
+ splits = ["calibration", "valid", "test"]
154
+ n_examples_per_split = activation_record_slice_params.n_examples_per_split
155
+ if n_examples_per_split is None:
156
+ n_examples_per_split = len(self.random_sample) // len(splits)
157
+ # NOTE: this assert could trigger on some old datasets with only 10 random samples, in which case you may have to remove "test" from the set of splits
158
+ assert len(self.random_sample) >= n_examples_per_split * len(splits)
159
+ return get_slices_for_splits(splits, n_examples_per_split)
160
+
161
+ def train_activation_records(
162
+ self,
163
+ activation_record_slice_params: ActivationRecordSliceParams,
164
+ ) -> list[ActivationRecord]:
165
+ """
166
+ Train split, typically used for generating explanations. Consists exclusively of
167
+ top-activating records since context window limitations make it difficult to include
168
+ random records.
169
+ """
170
+ return self.most_positive_activation_records[
171
+ self._get_top_activation_slices(activation_record_slice_params)["train"]
172
+ ]
173
+
174
+ def calibration_activation_records(
175
+ self,
176
+ activation_record_slice_params: ActivationRecordSliceParams,
177
+ ) -> list[ActivationRecord]:
178
+ """
179
+ Calibration split, typically used for calibrating neuron simulations. See
180
+ http://go/neuron_explanation_methodology for an explanation of calibration. Consists of
181
+ top-activating records and random records in a 1:1 ratio.
182
+ """
183
+ return (
184
+ self.most_positive_activation_records[
185
+ self._get_top_activation_slices(activation_record_slice_params)[
186
+ "calibration"
187
+ ]
188
+ ]
189
+ + self.random_sample[
190
+ self._get_random_activation_slices(activation_record_slice_params)[
191
+ "calibration"
192
+ ]
193
+ ]
194
+ )
195
+
196
+ def valid_activation_records(
197
+ self,
198
+ activation_record_slice_params: ActivationRecordSliceParams,
199
+ ) -> list[ActivationRecord]:
200
+ """
201
+ Validation split, typically used for evaluating explanations, either automatically with
202
+ simulation + correlation coefficient scoring, or manually by humans. Consists of
203
+ top-activating records and random records in a 1:1 ratio.
204
+ """
205
+ return (
206
+ self.most_positive_activation_records[
207
+ self._get_top_activation_slices(activation_record_slice_params)["valid"]
208
+ ]
209
+ + self.random_sample[
210
+ self._get_random_activation_slices(activation_record_slice_params)[
211
+ "valid"
212
+ ]
213
+ ]
214
+ )
215
+
216
+ def test_activation_records(
217
+ self,
218
+ activation_record_slice_params: ActivationRecordSliceParams,
219
+ ) -> list[ActivationRecord]:
220
+ """
221
+ Test split, typically used for explanation evaluations that can't use the validation split.
222
+ Consists of top-activating records and random records in a 1:1 ratio.
223
+ """
224
+ return (
225
+ self.most_positive_activation_records[
226
+ self._get_top_activation_slices(activation_record_slice_params)["test"]
227
+ ]
228
+ + self.random_sample[
229
+ self._get_random_activation_slices(activation_record_slice_params)[
230
+ "test"
231
+ ]
232
+ ]
233
+ )
234
+
235
+
236
+ def neuron_exists(
237
+ dataset_path: str, layer_index: Union[str, int], neuron_index: Union[str, int]
238
+ ) -> bool:
239
+ """Return whether the specified neuron exists."""
240
+ file = bf.join(dataset_path, "neurons", str(layer_index), f"{neuron_index}.json")
241
+ return bf.exists(file)
242
+
243
+
244
+ def load_neuron(
245
+ layer_index: Union[str, int],
246
+ neuron_index: Union[str, int],
247
+ dataset_path: str = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/collated-activations",
248
+ ) -> NeuronRecord:
249
+ """Load the NeuronRecord for the specified neuron."""
250
+ url = "/".join([dataset_path, str(layer_index), f"{neuron_index}.json"])
251
+ url = standardize_azure_url(url)
252
+ with urllib.request.urlopen(url) as f:
253
+ neuron_record = loads(f.read())
254
+ if not isinstance(neuron_record, NeuronRecord):
255
+ raise ValueError(
256
+ f"Stored data incompatible with current version of NeuronRecord dataclass."
257
+ )
258
+ return neuron_record
259
+
260
+
261
+ @bbb.ensure_session
262
+ async def load_neuron_async(
263
+ layer_index: Union[str, int],
264
+ neuron_index: Union[str, int],
265
+ dataset_path: str = "az://openaipublic/neuron-explainer/data/collated-activations",
266
+ ) -> NeuronRecord:
267
+ """Async version of load_neuron."""
268
+ file = bf.join(dataset_path, str(layer_index), f"{neuron_index}.json")
269
+ return await read_neuron_file(file)
270
+
271
+
272
+ @bbb.ensure_session
273
+ async def read_neuron_file(neuron_filename: str) -> NeuronRecord:
274
+ """Like load_neuron_async, but takes a raw neuron filename."""
275
+ raw_contents = await bbb.read.read_single(neuron_filename)
276
+ neuron_record = loads(raw_contents.decode("utf-8"))
277
+ if not isinstance(neuron_record, NeuronRecord):
278
+ raise ValueError(
279
+ f"Stored data incompatible with current version of NeuronRecord dataclass."
280
+ )
281
+ return neuron_record
282
+
283
+
284
+ def get_sorted_neuron_indices(
285
+ dataset_path: str, layer_index: Union[str, int]
286
+ ) -> List[int]:
287
+ """Returns the indices of all neurons in this layer, in ascending order."""
288
+ layer_dir = bf.join(dataset_path, "neurons", str(layer_index))
289
+ return sorted(
290
+ [
291
+ int(f.split(".")[0])
292
+ for f in bf.listdir(layer_dir)
293
+ if f.split(".")[0].isnumeric()
294
+ ]
295
+ )
296
+
297
+
298
+ def get_sorted_layers(dataset_path: str) -> List[str]:
299
+ """
300
+ Return the indices of all layers in this dataset, in ascending numerical order, as strings.
301
+ """
302
+ return [
303
+ str(x)
304
+ for x in sorted(
305
+ [
306
+ int(x)
307
+ for x in bf.listdir(bf.join(dataset_path, "neurons"))
308
+ if x.isnumeric()
309
+ ]
310
+ )
311
+ ]
@@ -0,0 +1,121 @@
1
+ """
2
+ Contains math utilities for converting from flattened representations of attention activations
3
+ (which are a scalar per token pair) to nested lists. The inner lists are attention activations
4
+ related to attention from the same token (to different tokens).
5
+
6
+ Tested in ./test_attention_utils.py.
7
+ """
8
+
9
+ import math
10
+
11
+ import numpy as np
12
+
13
+
14
+ def _inverse_triangular_number(n: int) -> int:
15
+ # the m'th triangular number t_m satisfies t_m = m(m+1)/2
16
+ # this function asserts that n is a triangular number, and returns the unique m such that t_m = n
17
+ # this is used to infer the number of sequence tokens from the number of activations
18
+ assert n >= 0
19
+ m: int = (
20
+ math.floor(math.sqrt(1 + 8 * n)) - 1
21
+ ) // 2 # from quadratic formula applied to m(m+1)/2 = n
22
+ assert m * (m + 1) // 2 == n
23
+ return m
24
+
25
+
26
+ def get_max_num_attended_to_sequence_tokens(num_sequence_tokens: int, num_activations: int) -> int:
27
+ # Attended to sequences are assumed to increase in length up to a maximum length, and then stay at that
28
+ # length for the remainder of the sequence. The maximum attended to sequence length is at most equal to the sequence length,
29
+ # but is permitted to be less
30
+ num_sequence_token_pairs = num_sequence_tokens * (num_sequence_tokens + 1) // 2
31
+ if num_activations == num_sequence_token_pairs:
32
+ # the maximum attended to sequence length is equal to the sequence length
33
+ return num_sequence_tokens
34
+ else:
35
+ # the maximum attended to sequence length is less than the sequence length, and
36
+ assert num_activations < num_sequence_token_pairs
37
+ num_missing_activations = num_sequence_token_pairs - num_activations
38
+ num_missing_sequence_tokens = _inverse_triangular_number(num_missing_activations)
39
+ max_num_attended_to_sequence_tokens = num_sequence_tokens - num_missing_sequence_tokens
40
+ assert max_num_attended_to_sequence_tokens > 0
41
+ return max_num_attended_to_sequence_tokens
42
+
43
+
44
+ def get_attended_to_sequence_length_per_sequence_token(
45
+ num_sequence_tokens: int, max_num_attended_to_sequence_tokens: int
46
+ ) -> list[int]:
47
+ # given a num_sequence_tokens and a max_num_attended_to_sequence_tokens, return a list of length num_sequence_tokens
48
+ # where the ith element is the length of the attended to sequence for the ith sequence token.
49
+ # The length of the attended to sequence starts at 1, increases up to max_num_attended_to_sequence_tokens, by 1 with each
50
+ # token, and then stays at max_num_attended_to_sequence_tokens for the remainder of the sequence
51
+ assert num_sequence_tokens >= max_num_attended_to_sequence_tokens
52
+ attended_to_sequence_lengths = list(range(1, max_num_attended_to_sequence_tokens + 1))
53
+ if num_sequence_tokens > max_num_attended_to_sequence_tokens:
54
+ attended_to_sequence_lengths.extend(
55
+ [
56
+ max_num_attended_to_sequence_tokens
57
+ for _ in range(num_sequence_tokens - max_num_attended_to_sequence_tokens)
58
+ ]
59
+ )
60
+ return attended_to_sequence_lengths
61
+
62
+
63
+ def get_attended_to_sequence_lengths(num_sequence_tokens: int, num_activations: int) -> list[int]:
64
+ max_num_attended_to_sequence_tokens = get_max_num_attended_to_sequence_tokens(
65
+ num_sequence_tokens, num_activations
66
+ )
67
+ return get_attended_to_sequence_length_per_sequence_token(
68
+ num_sequence_tokens, max_num_attended_to_sequence_tokens
69
+ )
70
+
71
+
72
+ def _convert_flattened_index_to_unflattened_index_assuming_square_matrix(
73
+ flat_index: int,
74
+ ) -> tuple[int, int]:
75
+ # this con
76
+ n = math.floor((-1 + math.sqrt(1 + 8 * flat_index)) / 2)
77
+ m = flat_index - n * (n + 1) // 2
78
+ return n, m
79
+
80
+
81
+ def convert_flattened_index_to_unflattened_index(
82
+ flattened_index: int,
83
+ num_sequence_tokens: int | None = None,
84
+ num_activations: int | None = None,
85
+ ) -> tuple[int, int]:
86
+ # given a flattened index, return the unflattened index
87
+ # if the attention matrix is square (most common), then the flattened_index uniquely determines the index within the square matrix
88
+ # if the attention matrix has more rows (sequence tokens) than columns (attended-to sequence tokens), then num_sequence_tokens
89
+ # and num_activations are required to determine the index within the matrix
90
+ # specify both num_sequence_tokens and num_activations, or neither
91
+ assert not (num_sequence_tokens is None) ^ (num_activations is None)
92
+
93
+ if (
94
+ num_sequence_tokens is None
95
+ or num_activations == num_sequence_tokens * (num_sequence_tokens + 1) // 2
96
+ ):
97
+ assume_square_matrix = True
98
+ else:
99
+ assume_square_matrix = False
100
+
101
+ if assume_square_matrix:
102
+ return _convert_flattened_index_to_unflattened_index_assuming_square_matrix(flattened_index)
103
+ else:
104
+ assert num_sequence_tokens is not None
105
+ assert num_activations is not None
106
+ assert flattened_index < num_activations
107
+ sequence_lengths = get_attended_to_sequence_lengths(num_sequence_tokens, num_activations)
108
+ sequence_lengths_cumsum = np.cumsum([0] + sequence_lengths)
109
+ sequence_index = int(
110
+ np.searchsorted(sequence_lengths_cumsum, flattened_index, side="right") - 1
111
+ )
112
+ assert sequence_lengths_cumsum[sequence_index] <= flattened_index, (
113
+ sequence_lengths_cumsum[sequence_index],
114
+ flattened_index,
115
+ )
116
+ assert sequence_lengths_cumsum[sequence_index + 1] >= flattened_index, (
117
+ sequence_lengths_cumsum[sequence_index + 1],
118
+ flattened_index,
119
+ )
120
+ index_within_sequence = flattened_index - sequence_lengths_cumsum[sequence_index]
121
+ return sequence_index, index_within_sequence
@@ -0,0 +1,59 @@
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)