debug

app.py

@@ -211,10 +211,19 @@ 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]))
 
     if st.session_state['page_status'] == 'finish_debug':
+            option_1_tokens_1 = np.array(input_ids_dict['sent_1'])[np.array(option_1_locs['sent_1'])]
+            option_1_tokens_2 = np.array(input_ids_dict['sent_2'])[np.array(option_1_locs['sent_2'])]
+            option_2_tokens_1 = np.array(input_ids_dict['sent_1'])[np.array(option_2_locs['sent_1'])]
+            option_2_tokens_2 = np.array(input_ids_dict['sent_2'])[np.array(option_2_locs['sent_2'])]
+            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
+
             for layer_id in range(num_layers):
                 interventions = [create_interventions(16,['lay','qry','key','val'],num_heads) if i==layer_id else {'lay':[],'qry':[],'key':[],'val':[]} for i in range(num_layers)]
                 for masked_ids in [masked_ids_option_1, masked_ids_option_2]: