multihead

app.py

@@ -116,11 +116,14 @@ def show_instruction(sent,fontsize=20):
     suffix = '</span></p>'
     return st.markdown(prefix + sent + suffix, unsafe_allow_html = True)
 
-def create_interventions(token_id,interv_types,num_heads):
+def create_interventions(token_id,interv_types,num_heads,multihead=False):
     interventions = {}
     for rep in ['lay','qry','key','val']:
         if rep in interv_types:
-            interventions[rep] = [(head_id,token_id,[head_id,head_id+num_heads]) for head_id in range(num_heads)]
+            if multihead:
+                interventions[rep] = [(head_id,token_id,[0,1]) for head_id in range(num_heads)]
+            else:
+                interventions[rep] = [(head_id,token_id,[head_id,head_id+num_heads]) for head_id in range(num_heads)]
         else:
             interventions[rep] = []
     return interventions
@@ -251,10 +254,17 @@ if __name__=='__main__':
             interventions = [{'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
             probs_original = run_intervention(interventions,1,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
             df = pd.DataFrame(data=[[probs_original[0,0][0],probs_original[1,0][0]],
-                                    [probs_original[0,1][0],probs_original[1,1][0]]],columns=['Option 1','Option 2'],index=['Sentence 1','Sentence 2'])
+                                    [probs_original[0,1][0],probs_original[1,1][0]]],
+                                    columns=[tokenizer.decode(option_1_tokens),tokenizer.decode(option_2_tokens)],
+                                    index=['Sentence 1','Sentence 2'])
             st.dataframe(df.style.highlight_max(axis=1))
 
-    if st.session_state['page_status'] == 'finish_debug':
-            for layer_id in range(num_layers):
-                interventions = [create_interventions(16,['lay','qry','key','val'],num_heads) if i==layer_id else {'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
-                probs = run_intervention(interventions,num_heads,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
+            multihead = True
+            for layer_id in range(num_layers)[:1]:
+                interventions = [create_interventions(16,['lay','qry','key','val'],num_heads,multihead) if i==layer_id else {'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
+                if multihead:
+                    probs = run_intervention(interventions,1,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
+                else
+                    probs = run_intervention(interventions,num_heads,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
+
+                st.write(probs)