Stefan Heimersheim commited on
Commit
8b9fdad
1 Parent(s): 40f75e8

Added code

Browse files
Files changed (3) hide show
  1. Home.py +245 -0
  2. INFO.md +8 -0
  3. requirements.txt +7 -0
Home.py ADDED
@@ -0,0 +1,245 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformer_lens import HookedTransformer, utils
3
+ from io import StringIO
4
+ import sys
5
+ import torch
6
+ from functools import partial
7
+ import plotly.offline as pyo
8
+ import plotly.graph_objs as go
9
+ import numpy as np
10
+ import plotly.express as px
11
+ import circuitsvis as cv
12
+
13
+ # Little bit of front end for model selector
14
+
15
+ # Radio buttons
16
+ model_name = st.sidebar.radio("Model", [
17
+ "gelu-1l",
18
+ "gelu-2l",
19
+ #"gelu-3l",
20
+ #"gelu-4l",
21
+ #"attn-only-1l",
22
+ #"attn-only-2l",
23
+ #"attn-only-3l",
24
+ #"attn-only-4l",
25
+ #"solu-1l",
26
+ #"solu-2l",
27
+ #"solu-3l",
28
+ #"solu-4l",
29
+ #"solu-6l",
30
+ #"solu-8l",
31
+ #"solu-10l",
32
+ #"solu-12l",
33
+ #"gpt2-small",
34
+ #"gpt2-medium",
35
+ #"gpt2-large",
36
+ #"gpt2-xl",
37
+ ], index=1)
38
+
39
+
40
+ # Backend code
41
+
42
+ model = HookedTransformer.from_pretrained(model_name)
43
+
44
+ def predict_next_token(prompt):
45
+ logits = model(prompt)[0,-1]
46
+ answer_index = logits.argmax()
47
+ answer = model.tokenizer.decode(answer_index)
48
+ answer = f"<b>|{answer}|</b> (answer by {model.cfg.model_name})"
49
+ return answer
50
+
51
+ def test_prompt(prompt, answer):
52
+ output = StringIO()
53
+ sys.stdout = output
54
+ utils.test_prompt(prompt, answer, model)
55
+ output = output.getvalue()
56
+ return output
57
+
58
+ def compute_residual_stream_patch(clean_prompt=None, answer=None, corrupt_prompt=None, corrupt_answer=None, layers=None):
59
+ model.reset_hooks()
60
+ clean_answer_index = model.tokenizer.encode(answer)[0]
61
+ corrupt_answer_index = model.tokenizer.encode(corrupt_answer)[0]
62
+ clean_tokens = model.to_str_tokens(clean_prompt)
63
+ _, corrupt_cache = model.run_with_cache(corrupt_prompt)
64
+ # Patching function
65
+ def patch_residual_stream(activations, hook, layer="blocks.6.hook_resid_post", pos=5):
66
+ activations[:, pos, :] = corrupt_cache[layer][:, pos, :]
67
+ return activations
68
+ # Compute logit diffs
69
+ n_layers = len(layers)
70
+ n_pos = len(clean_tokens)
71
+ patching_effect = torch.zeros(n_layers, n_pos)
72
+ for l, layer in enumerate(layers):
73
+ for pos in range(n_pos):
74
+ fwd_hooks = [(layer, partial(patch_residual_stream, layer=layer, pos=pos))]
75
+ prediction_logits = model.run_with_hooks(clean_prompt, fwd_hooks=fwd_hooks)[0, -1]
76
+ patching_effect[l, pos] = prediction_logits[clean_answer_index] - prediction_logits[corrupt_answer_index]
77
+ return patching_effect
78
+
79
+ def compute_attn_patch(clean_prompt=None, answer=None, corrupt_prompt=None, corrupt_answer=None):
80
+ use_attn_result_prev = model.cfg.use_attn_result
81
+ model.cfg.use_attn_result = True
82
+ clean_answer_index = model.tokenizer.encode(answer)[0]
83
+ corrupt_answer_index = model.tokenizer.encode(corrupt_answer)[0]
84
+ clean_tokens = model.to_str_tokens(clean_prompt)
85
+ _, corrupt_cache = model.run_with_cache(corrupt_prompt)
86
+ # Patching function
87
+ def patch_head_result(activations, hook, head=None, pos=None):
88
+ activations[:, pos, head, :] = corrupt_cache[hook.name][:, pos, head, :]
89
+ return activations
90
+
91
+ n_layers = model.cfg.n_layers
92
+ n_heads = model.cfg.n_heads
93
+ n_pos = len(clean_tokens)
94
+ patching_effect = torch.zeros(n_layers*n_heads, n_pos)
95
+ for layer in range(n_layers):
96
+ for head in range(n_heads):
97
+ for pos in range(n_pos):
98
+ fwd_hooks = [(f"blocks.{layer}.attn.hook_result", partial(patch_head_result, head=head, pos=pos))]
99
+ prediction_logits = model.run_with_hooks(clean_prompt, fwd_hooks=fwd_hooks)[0, -1]
100
+ patching_effect[n_heads*layer+head, pos] = prediction_logits[clean_answer_index] - prediction_logits[corrupt_answer_index]
101
+ model.cfg.use_attn_result = use_attn_result_prev
102
+ return patching_effect
103
+
104
+ def imshow(tensor, xlabel="X", ylabel="Y", zlabel=None, xticks=None, yticks=None, c_midpoint=0.0, c_scale="RdBu", **kwargs):
105
+ tensor = utils.to_numpy(tensor)
106
+ xticks = [str(x) for x in xticks]
107
+ yticks = [str(y) for y in yticks]
108
+ labels = {"x": xlabel, "y": ylabel}
109
+ if zlabel is not None:
110
+ labels["color"] = zlabel
111
+ fig = px.imshow(tensor, x=xticks, y=yticks, labels=labels, color_continuous_midpoint=c_midpoint,
112
+ color_continuous_scale=c_scale, **kwargs)
113
+ return fig
114
+
115
+ def plot_residual_stream_patch(clean_prompt=None, answer=None, corrupt_prompt=None, corrupt_answer=None):
116
+ layers = ["blocks.0.hook_resid_pre", *[f"blocks.{i}.hook_resid_post" for i in range(model.cfg.n_layers)]]
117
+ token_labels = model.to_str_tokens(clean_prompt)
118
+ patching_effect = compute_residual_stream_patch(clean_prompt=clean_prompt, answer=answer, corrupt_prompt=corrupt_prompt, corrupt_answer=corrupt_answer, layers=layers)
119
+ fig = imshow(patching_effect, xticks=token_labels, yticks=layers, xlabel="Position", ylabel="Layer",
120
+ zlabel="Logit Difference", title="Patching residual stream at specific layer and position")
121
+ return fig
122
+
123
+ def plot_attn_patch(clean_prompt=None, answer=None, corrupt_prompt=None, corrupt_answer=None):
124
+ clean_tokens = model.to_str_tokens(clean_prompt)
125
+ n_layers = model.cfg.n_layers
126
+ n_heads = model.cfg.n_heads
127
+ layerhead_labels = [f"{l}.{h}" for l in range(n_layers) for h in range(n_heads)]
128
+ token_labels = [f"(pos {i:2}) {t}" for i, t in enumerate(clean_tokens)]
129
+ patching_effect = compute_attn_patch(clean_prompt=clean_prompt, answer=answer, corrupt_prompt=corrupt_prompt, corrupt_answer=corrupt_answer)
130
+ return imshow(patching_effect, xticks=token_labels, yticks=layerhead_labels, xlabel="Position", ylabel="Layer.Head",
131
+ zlabel="Logit Difference", title=f"Patching attention outputs for specific layer, head, and position", width=600, height=300+200*n_layers)
132
+
133
+
134
+ # Frontend code
135
+ st.title("Simple Trafo Mech Int")
136
+ st.subheader("Transformer Mechanistic Interpretability")
137
+ st.markdown("Powered by [TransformerLens](https://github.com/neelnanda-io/TransformerLens/)")
138
+ st.markdown("For _what_ these plots are, and _why_, see this [tutorial](https://docs.google.com/document/d/1e6cs8d9QNretWvOLsv_KaMp6kSPWpJEW0GWc0nwjqxo/).")
139
+
140
+ # Predict next token
141
+ st.header("Predict the next token")
142
+ st.markdown("Just a simple test UI, enter a prompt and the model will predict the next token")
143
+ prompt_simple = st.text_input("Prompt:", "Today, the weather is", key="prompt_simple")
144
+
145
+ if "prompt_simple_output" not in st.session_state:
146
+ st.session_state.prompt_simple_output = None
147
+
148
+ if st.button("Run model", key="key_button_prompt_simple"):
149
+ res = predict_next_token(prompt_simple)
150
+ st.session_state.prompt_simple_output = res
151
+
152
+ if st.session_state.prompt_simple_output:
153
+ st.markdown(st.session_state.prompt_simple_output, unsafe_allow_html=True)
154
+
155
+
156
+ # Test prompt
157
+ st.header("Verbose test prompt")
158
+ st.markdown("Enter a prompt and the correct answer, the model will run the prompt and print the results")
159
+
160
+ prompt = st.text_input("Prompt:", "The most popular programming language is", key="prompt")
161
+ answer = st.text_input("Answer:", " Java", key="answer")
162
+
163
+ if "test_prompt_output" not in st.session_state:
164
+ st.session_state.test_prompt_output = None
165
+
166
+ if st.button("Run model", key="key_button_test_prompt"):
167
+ res = test_prompt(prompt, answer)
168
+ st.session_state.test_prompt_output = res
169
+
170
+ if st.session_state.test_prompt_output:
171
+ st.code(st.session_state.test_prompt_output)
172
+
173
+
174
+ # Residual stream patching
175
+
176
+ st.header("Residual stream patching")
177
+ st.markdown("Enter a clean prompt, correct answer, corrupt prompt and corrupt answer, the model will compute the patching effect")
178
+
179
+ default_clean_prompt = "Her name was Alex Hart. Tomorrow at lunch time Alex"
180
+ default_clean_answer = "Hart"
181
+ default_corrupt_prompt = "Her name was Alex Carroll. Tomorrow at lunch time Alex"
182
+ default_corrupt_answer = "Carroll"
183
+
184
+ clean_prompt = st.text_input("Clean Prompt:", default_clean_prompt)
185
+ clean_answer = st.text_input("Correct Answer:", default_clean_answer)
186
+ corrupt_prompt = st.text_input("Corrupt Prompt:", default_corrupt_prompt)
187
+ corrupt_answer = st.text_input("Corrupt Answer:", default_corrupt_answer)
188
+
189
+ if "residual_stream_patch_out" not in st.session_state:
190
+ st.session_state.residual_stream_patch_out = None
191
+
192
+ if st.button("Run model", key="key_button_residual_stream_patch"):
193
+ fig = plot_residual_stream_patch(clean_prompt=clean_prompt, answer=clean_answer, corrupt_prompt=corrupt_prompt, corrupt_answer=corrupt_answer)
194
+ st.session_state.residual_stream_patch_out = fig
195
+
196
+ if st.session_state.residual_stream_patch_out:
197
+ st.plotly_chart(st.session_state.residual_stream_patch_out)
198
+
199
+
200
+ # Attention head output
201
+
202
+ st.header("Attention head output patching")
203
+ st.markdown("Enter a clean prompt, correct answer, corrupt prompt and corrupt answer, the model will compute the patching effect")
204
+
205
+ clean_prompt_attn = st.text_input("Clean Prompt:", default_clean_prompt, key="key2_clean_prompt_attn")
206
+ clean_answer_attn = st.text_input("Correct Answer:", default_clean_answer, key="key2_clean_answer_attn")
207
+ corrupt_prompt_attn = st.text_input("Corrupt Prompt:", default_corrupt_prompt, key="key2_corrupt_prompt_attn")
208
+ corrupt_answer_attn = st.text_input("Corrupt Answer:", default_corrupt_answer, key="key2_corrupt_answer_attn")
209
+
210
+ if "attn_head_patch_out" not in st.session_state:
211
+ st.session_state.attn_head_patch_out = None
212
+
213
+ if st.button("Run model", key="key_button_attn_head_patch"):
214
+ fig = plot_attn_patch(clean_prompt=clean_prompt_attn, answer=clean_answer_attn, corrupt_prompt=corrupt_prompt_attn, corrupt_answer=corrupt_answer_attn)
215
+ st.session_state.attn_head_patch_out = fig
216
+
217
+ if st.session_state.attn_head_patch_out:
218
+ st.plotly_chart(st.session_state.attn_head_patch_out)
219
+
220
+
221
+ # Attention Head Visualization
222
+
223
+ st.header("Attention Pattern Visualization")
224
+ st.markdown("Powered by [CircuitsVis](https://github.com/alan-cooney/CircuitsVis)")
225
+ st.markdown("Enter a prompt, show attention patterns")
226
+
227
+ default_prompt_attn = "Her name was Alex Hart. Tomorrow at lunch time Alex"
228
+ prompt_attn = st.text_input("Prompt:", default_prompt_attn)
229
+
230
+ if "attn_html" not in st.session_state:
231
+ st.session_state.attn_html = None
232
+
233
+ if st.button("Run model", key="key_button_attention_head"):
234
+ _, cache = model.run_with_cache(prompt_attn)
235
+ st.session_state.attn_html = []
236
+ for layer in range(model.cfg.n_layers):
237
+ html = cv.attention.attention_patterns(tokens=model.to_str_tokens(prompt_attn),
238
+ attention=cache[f'blocks.{layer}.attn.hook_pattern'][0])
239
+ st.session_state.attn_html.append(html.show_code())
240
+
241
+ if st.session_state.attn_html:
242
+ for layer in range(len(st.session_state.attn_html)):
243
+ st.write(f"Attention patterns Layer {layer}:")
244
+ st.components.v1.html(st.session_state.attn_html[layer], height=500)
245
+
INFO.md ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ # Trafo Mech Int playground
2
+
3
+ Mechanistic Interpretability for everyone!
4
+ Website to visualise Transformer internals
5
+
6
+ By [Stefan Heimersheim](https://github.com/Stefan-Heimersheim/) and [Jonathan Ng](https://github.com/derpyplops).
7
+
8
+ [Mechanistic Interpretability Hackathon](https://itch.io/jam/mechint) submission.
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ git+https://github.com/neelnanda-io/TransformerLens/
2
+ torch
3
+ flask
4
+ gunicorn
5
+ plotly
6
+ circuitsvis
7
+ streamlit