aes

app.py

@@ -288,10 +288,11 @@ if __name__=='__main__':
             cols = st.columns(len(masked_ids_option_1['sent_1'])-2)
             token_id = 0
             for col_id,col in enumerate(cols):
-                if col_id in token_id_list:
-                    interv_id = token_id_list.index(col_id)
-                    with col:
-                        fig,ax = plt.subplots(figsize=(1,6))
+                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())
@@ -299,5 +300,8 @@ if __name__=='__main__':
                         ax.set_xticklabels([])
                         ax.set_yticks([])
                         ax.set_yticklabels([])
-                        ax.set_title(tokenizer.decode([masked_ids_option_1['sent_1'][col_id+1]]))
+                        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)