plot

app.py

@@ -1,8 +1,8 @@
 import numpy as np
 import pandas as pd
 import streamlit as st
-#import matplotlib.pyplot as plt
-#import seaborn as sns
+import matplotlib.pyplot as plt
+import seaborn as sns
 
 #import jax
 #import jax.numpy as jnp
@@ -169,7 +169,6 @@ if __name__=='__main__':
     load_css('style.css')
     tokenizer,model = load_model()
     num_layers, num_heads = model.config.num_hidden_layers, model.config.num_attention_heads
-    st.write(num_layers,num_heads)
     mask_id = tokenizer('[MASK]').input_ids[1:-1][0]
 
     main_area = st.empty()
@@ -260,11 +259,20 @@ if __name__=='__main__':
             st.dataframe(df.style.highlight_max(axis=1))
 
             multihead = True
-            for layer_id in range(num_layers):
-                interventions = [create_interventions(15,['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)
-                effect = ((probs_original-probs)[0,0] + (probs_original-probs)[1,1] + (probs-probs_original)[0,1] + (probs-probs_original)[1,0])/4
-                st.write(effect)
+            effect_array = []
+            for token_id in range(1,len(masked_ids_option_1['sent_1'])-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)]
+                    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)
+                    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)
+            effects = np.array(effect_array)
+
+            fig,ax = plt.subplots(1,1,figsize=(8,6))
+            ax.imshow(effects.T)
+            st.pyplot(fig)