reduce num runs

app.py

@@ -260,7 +260,13 @@ if __name__=='__main__':
 
             multihead = True
             effect_array = []
-            for token_id in range(1,len(masked_ids_option_1['sent_1'])-1):
+            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']
+            st.write(token_id_list)
+            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)]
@@ -274,15 +280,18 @@ if __name__=='__main__':
             effects = 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:
-                    fig,ax = plt.subplots()
-                    ax.set_box_aspect(effects.shape[0])
-                    ax.imshow(effects[:,col_id:col_id+1,0],cmap=sns.color_palette("light:r", as_cmap=True),
-                                    vmin=effects[:,:,0].min(),vmax=effects[:,:,0].max())
-                    ax.set_xticks([])
-                    ax.set_xticklabels([])
-                    ax.set_yticks([])
-                    ax.set_yticklabels([])
-                    ax.set_title(tokenizer.decode([masked_ids_option_1['sent_1'][col_id+1]]))
-                    st.pyplot(fig)
+                if col_id in token_id_list:
+                    with col:
+                        fig,ax = plt.subplots()
+                        ax.set_box_aspect(effects.shape[0])
+                        ax.imshow(effects[:,token_id:token_id+1,0],cmap=sns.color_palette("light:r", as_cmap=True),
+                                        vmin=effects[:,:,0].min(),vmax=effects[:,:,0].max())
+                        ax.set_xticks([])
+                        ax.set_xticklabels([])
+                        ax.set_yticks([])
+                        ax.set_yticklabels([])
+                        ax.set_title(tokenizer.decode([masked_ids_option_1['sent_1'][col_id+1]]))
+                        st.pyplot(fig)
+                        token_id += 1