show table

app.py

@@ -233,9 +233,9 @@ if __name__=='__main__':
                                                                         pron_locs[f'sent_{sent_id}'],
                                                                         option_2_locs[f'sent_{sent_id}'],mask_id)
 
-            st.write(option_1_locs)
-            st.write(option_2_locs)
-            st.write(pron_locs)
+            #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]))
 
@@ -246,10 +246,14 @@ 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)
-            st.write(probs_original)
+            df = pd.DataFrame(data=[[probs_original[0,0][0],probs_original[1,0][0]],
+                                    [probs_original[0,1][0],probs_original[1,1][0]]],columns=['Option 1','Option 2'],index=['Sentence 1','Sentence 2'])
+            st.dataframe(df.style.highlight_max(axis=1))
 
     if st.session_state['page_status'] == 'finish_debug':
             for layer_id in range(num_layers):