test

app.py

@@ -240,8 +240,8 @@ if __name__=='__main__':
             #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]))
+            #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]
@@ -250,8 +250,6 @@ if __name__=='__main__':
             assert np.all(option_1_tokens_1==option_1_tokens_2) and np.all(option_2_tokens_1==option_2_tokens_2)
             option_1_tokens = option_1_tokens_1
             option_2_tokens = option_2_tokens_1
-            st.write(option_1_tokens)
-            st.write(option_2_tokens)
 
             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)
@@ -262,11 +260,10 @@ if __name__=='__main__':
             st.dataframe(df.style.highlight_max(axis=1))
 
             multihead = True
-            for layer_id in range(num_layers)[:1]:
+            for layer_id in range(num_layers):
                 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_original-probs)