add heads

app.py

@@ -130,14 +130,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,multihead=False):
+def create_interventions(token_id,interv_types,num_heads,multihead=False,heads=[]):
     interventions = {}
     for rep in ['lay','qry','key','val']:
         if rep in interv_types:
             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)]
+                interventions[rep] = [(head_id,token_id,[i,i+len(heads)]) for i,head_id in enumerate(heads)]
         else:
             interventions[rep] = []
     return interventions
@@ -176,6 +176,27 @@ def run_intervention(interventions,batch_size,skeleton_model,model,masked_ids_op
     assert probs.shape[0]==2 and probs.shape[1]==2 and probs.shape[2]==batch_size
     return probs
 
+def show_results(effect_array,masked_sent,token_id_list,num_layers):
+    cols = st.columns(len(masked_sent)-2)
+    for col_id,col in enumerate(cols):
+        with col:
+            st.write(tokenizer.decode([masked_sent[col_id+1]]))
+            if col_id in token_id_list:
+                interv_id = token_id_list.index(col_id)
+                fig,ax = plt.subplots()
+                ax.set_box_aspect(num_layers)
+                ax.imshow(effect_array[:,interv_id:interv_id+1],cmap=sns.color_palette("light:r", as_cmap=True),
+                                vmin=effect_array.min(),vmax=effect_array.max())
+                ax.set_xticks([])
+                ax.set_xticklabels([])
+                ax.set_yticks([])
+                ax.set_yticklabels([])
+                ax.spines['top'].set_visible(False)
+                ax.spines['bottom'].set_visible(False)
+                ax.spines['right'].set_visible(False)
+                ax.spines['left'].set_visible(False)
+                st.pyplot(fig)
+
 if __name__=='__main__':
     wide_setup()
     load_css('style.css')
@@ -217,7 +238,7 @@ if __name__=='__main__':
         show_instruction('2. Select sites to mask out and click "Confirm"',fontsize=16)
         #show_instruction('------------------------------',fontsize=32)
         annotate_mask(1,sent_1)
-        show_instruction('------------------------------',fontsize=32)
+        show_instruction('------------------------------',fontsize=24)
         annotate_mask(2,sent_2)
         if st.button('Confirm',key='confirm_mask'):
             st.session_state['page_status'] = 'annotate_options'
@@ -230,21 +251,34 @@ if __name__=='__main__':
         show_instruction('3. Select options and click "Confirm"',fontsize=16)
         #show_instruction('------------------------------',fontsize=32)
         annotate_options(1,sent_1)
-        show_instruction('------------------------------',fontsize=32)
+        show_instruction('------------------------------',fontsize=24)
         annotate_options(2,sent_2)
         if st.button('Confirm',key='confirm_option'):
             st.session_state['page_status'] = 'analysis'
             st.experimental_rerun()
 
     if st.session_state['page_status']=='analysis':
+        interv_reps = st.multiselect('Select the types of representations to intervene.',['layer','query','key','value'])
+        rep_dict = {'layer':'lay','query':'qry','key':'key','value':'val'}
+        multihead = not st.checkbox('Perform individual head analysis (takes time)')
+        if not multihead:
+            heads = st.multiselect('Select heads to intervene.',list(np.arange(1,num_heads+1)))
+        else:
+            heads = []
+
+        if st.button('Run',key='run'):
+            st.session_state['reps'] = [rep_dict[rep] for rep in interv_reps]
+            st.session_state['multihead'] = multihead
+            st.session_state['heads'] = heads
+            st.session_state['page_status'] = 'results'
+            st.experimental_rerun()
+
+    if st.session_state['page_status']=='results':
         sent_1 = st.session_state['sent_1']
         sent_2 = st.session_state['sent_2']
-        #show_annotated_sentence(st.session_state['decoded_sent_1'],
-        #                            option_locs=st.session_state['option_locs_1'],
-        #                            mask_locs=st.session_state['mask_locs_1'])
-        #show_annotated_sentence(st.session_state['decoded_sent_2'],
-        #                            option_locs=st.session_state['option_locs_2'],
-        #                            mask_locs=st.session_state['mask_locs_2'])
+        multihead = st.session_state['multihead']
+        heads = st.session_state['heads']
+        reps = st.session_state['reps']
 
         option_1_locs, option_2_locs = {}, {}
         pron_locs = {}
@@ -263,12 +297,6 @@ if __name__=='__main__':
                                                                     pron_locs[f'sent_{sent_id}'],
                                                                     option_2_locs[f'sent_{sent_id}'],mask_id)
 
-        #st.write(option_1_locs)
-        #st.write(option_2_locs)
-        #st.write(pron_locs)
-        #for token_ids in [masked_ids_option_1['sent_1'],masked_ids_option_1['sent_2'],masked_ids_option_2['sent_1'],masked_ids_option_2['sent_2']]:
-        #    st.write(' '.join([tokenizer.decode([token]) for token in token_ids]))
-
         option_1_tokens_1 = np.array(input_ids_dict['sent_1'])[np.array(option_1_locs['sent_1'])+1]
         option_1_tokens_2 = np.array(input_ids_dict['sent_2'])[np.array(option_1_locs['sent_2'])+1]
         option_2_tokens_1 = np.array(input_ids_dict['sent_1'])[np.array(option_2_locs['sent_1'])+1]
@@ -293,45 +321,31 @@ if __name__=='__main__':
         assert np.all(compare_1.astype(int)==compare_2.astype(int))
         context_locs = list(np.arange(len(masked_ids_option_1['sent_1']))[compare_1]-1)  # match the indexing for annotation
 
-        multihead = True
         assert np.all(np.array(pron_locs['sent_1'])==np.array(pron_locs['sent_2']))
         assert np.all(np.array(option_1_locs['sent_1'])==np.array(option_1_locs['sent_2']))
         assert np.all(np.array(option_2_locs['sent_1'])==np.array(option_2_locs['sent_2']))
         token_id_list = pron_locs['sent_1'] + option_1_locs['sent_1'] + option_2_locs['sent_1'] + context_locs
-        #st.write(token_id_list)
 
         effect_array = []
         for token_id in token_id_list:
             token_id += 1
             effect_list = []
             for layer_id in range(num_layers):
-                interventions = [create_interventions(token_id,['lay','qry','key','val'],num_heads,multihead) if i==layer_id else {'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
+                interventions = [create_interventions(token_id,reps,num_heads,multihead,[head_id-1 for head_id in heads])
+                                    if i==layer_id else {'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
                 if multihead:
                     probs = run_intervention(interventions,1,skeleton_model,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
                 else:
-                    probs = run_intervention(interventions,num_heads,skeleton_model,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
+                    probs = run_intervention(interventions,len(heads),skeleton_model,model,masked_ids_option_1,masked_ids_option_2,option_1_tokens,option_2_tokens,pron_locs)
                 effect = ((probs_original-probs)[0,0] + (probs_original-probs)[1,1] + (probs-probs_original)[0,1] + (probs-probs_original)[1,0])/4
                 effect_list.append(effect)
             effect_array.append(effect_list)
         effect_array = np.transpose(np.array(effect_array),(1,0,2))
 
-        cols = st.columns(len(masked_ids_option_1['sent_1'])-2)
-        token_id = 0
-        for col_id,col in enumerate(cols):
-            with col:
-                st.write(tokenizer.decode([masked_ids_option_1['sent_1'][col_id+1]]))
-                if col_id in token_id_list:
-                    interv_id = token_id_list.index(col_id)
-                    fig,ax = plt.subplots()
-                    ax.set_box_aspect(num_layers)
-                    ax.imshow(effect_array[:,interv_id:interv_id+1,0],cmap=sns.color_palette("light:r", as_cmap=True),
-                                    vmin=effect_array[:,:,0].min(),vmax=effect_array[:,:,0].max())
-                    ax.set_xticks([])
-                    ax.set_xticklabels([])
-                    ax.set_yticks([])
-                    ax.set_yticklabels([])
-                    ax.spines['top'].set_visible(False)
-                    ax.spines['bottom'].set_visible(False)
-                    ax.spines['right'].set_visible(False)
-                    ax.spines['left'].set_visible(False)
-                    st.pyplot(fig)
+        if multihead:
+            show_results(effect_array[:,:,0],masked_ids_option_1['sent_1'],token_id_list,num_layers)
+        else:
+            tabs = st.tabs(heads)
+            for i,tab in enumerate(tabs):
+                with tab:
+                    show_results(effect_array[:,:,i],masked_ids_option_1['sent_1'],token_id_list,num_layers)