slide

app.py

@@ -222,102 +222,101 @@ if __name__=='__main__':
             st.experimental_rerun()
 
     if st.session_state['page_status']=='analysis':
-        with main_area.container():
-            sent_1 = st.session_state['sent_1']
-            sent_2 = st.session_state['sent_2']
-            #show_annotated_sentence(st.session_state['decoded_sent_1'],
-            #                            option_locs=st.session_state['option_locs_1'],
-            #                            mask_locs=st.session_state['mask_locs_1'])
-            #show_annotated_sentence(st.session_state['decoded_sent_2'],
-            #                            option_locs=st.session_state['option_locs_2'],
-            #                            mask_locs=st.session_state['mask_locs_2'])
-
-            option_1_locs, option_2_locs = {}, {}
-            pron_locs = {}
-            input_ids_dict = {}
-            masked_ids_option_1 = {}
-            masked_ids_option_2 = {}
-            for sent_id in [1,2]:
-                option_1_locs[f'sent_{sent_id}'], option_2_locs[f'sent_{sent_id}'] = separate_options(st.session_state[f'option_locs_{sent_id}'])
-                pron_locs[f'sent_{sent_id}'] = st.session_state[f'mask_locs_{sent_id}']
-                input_ids_dict[f'sent_{sent_id}'] = tokenizer(st.session_state[f'sent_{sent_id}']).input_ids
-
-                masked_ids_option_1[f'sent_{sent_id}'] = mask_out(input_ids_dict[f'sent_{sent_id}'],
-                                                                        pron_locs[f'sent_{sent_id}'],
-                                                                        option_1_locs[f'sent_{sent_id}'],mask_id)
-                masked_ids_option_2[f'sent_{sent_id}'] = mask_out(input_ids_dict[f'sent_{sent_id}'],
-                                                                        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)
-            #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]
-            option_2_tokens_1 = np.array(input_ids_dict['sent_1'])[np.array(option_2_locs['sent_1'])+1]
-            option_2_tokens_2 = np.array(input_ids_dict['sent_2'])[np.array(option_2_locs['sent_2'])+1]
-            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
-
-            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)
-            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=[tokenizer.decode(option_1_tokens),tokenizer.decode(option_2_tokens)],
-                                    index=['Sentence 1','Sentence 2'])
-            cols = st.columns(3)
-            with cols[1]:
-                show_instruction('Probability of predicting each option in each sentence',fontsize=12)
-                st.dataframe(df.style.highlight_max(axis=1),use_container_width=True)
-
-            compare_1 = np.array(masked_ids_option_1['sent_1'])!=np.array(masked_ids_option_1['sent_2'])
-            compare_2 = np.array(masked_ids_option_2['sent_1'])!=np.array(masked_ids_option_2['sent_2'])
-            assert np.all(compare_1.astype(int)==compare_2.astype(int))
-            context_locs = list(np.arange(len(masked_ids_option_1['sent_1']))[compare_1]-1)  # match the indexing for annotation
-
-            multihead = True
-            assert np.all(np.array(pron_locs['sent_1'])==np.array(pron_locs['sent_2']))
-            assert np.all(np.array(option_1_locs['sent_1'])==np.array(option_1_locs['sent_2']))
-            assert np.all(np.array(option_2_locs['sent_1'])==np.array(option_2_locs['sent_2']))
-            token_id_list = pron_locs['sent_1'] + option_1_locs['sent_1'] + option_2_locs['sent_1'] + context_locs
-            #st.write(token_id_list)
-
-            effect_array = []
-            for token_id in token_id_list:
-                token_id += 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)
-            effect_array = np.transpose(np.array(effect_array),(1,0,2))
-
-            cols = st.columns(len(masked_ids_option_1['sent_1'])-2)
-            token_id = 0
-            for col_id,col in enumerate(cols):
-                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())
-                        ax.set_xticks([])
-                        ax.set_xticklabels([])
-                        ax.set_yticks([])
-                        ax.set_yticklabels([])
-                        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)
+        sent_1 = st.session_state['sent_1']
+        sent_2 = st.session_state['sent_2']
+        #show_annotated_sentence(st.session_state['decoded_sent_1'],
+        #                            option_locs=st.session_state['option_locs_1'],
+        #                            mask_locs=st.session_state['mask_locs_1'])
+        #show_annotated_sentence(st.session_state['decoded_sent_2'],
+        #                            option_locs=st.session_state['option_locs_2'],
+        #                            mask_locs=st.session_state['mask_locs_2'])
+
+        option_1_locs, option_2_locs = {}, {}
+        pron_locs = {}
+        input_ids_dict = {}
+        masked_ids_option_1 = {}
+        masked_ids_option_2 = {}
+        for sent_id in [1,2]:
+            option_1_locs[f'sent_{sent_id}'], option_2_locs[f'sent_{sent_id}'] = separate_options(st.session_state[f'option_locs_{sent_id}'])
+            pron_locs[f'sent_{sent_id}'] = st.session_state[f'mask_locs_{sent_id}']
+            input_ids_dict[f'sent_{sent_id}'] = tokenizer(st.session_state[f'sent_{sent_id}']).input_ids
+
+            masked_ids_option_1[f'sent_{sent_id}'] = mask_out(input_ids_dict[f'sent_{sent_id}'],
+                                                                    pron_locs[f'sent_{sent_id}'],
+                                                                    option_1_locs[f'sent_{sent_id}'],mask_id)
+            masked_ids_option_2[f'sent_{sent_id}'] = mask_out(input_ids_dict[f'sent_{sent_id}'],
+                                                                    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)
+        #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]
+        option_2_tokens_1 = np.array(input_ids_dict['sent_1'])[np.array(option_2_locs['sent_1'])+1]
+        option_2_tokens_2 = np.array(input_ids_dict['sent_2'])[np.array(option_2_locs['sent_2'])+1]
+        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
+
+        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)
+        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=[tokenizer.decode(option_1_tokens),tokenizer.decode(option_2_tokens)],
+                                index=['Sentence 1','Sentence 2'])
+        cols = st.columns(3)
+        with cols[1]:
+            show_instruction('Probability of predicting each option in each sentence',fontsize=12)
+            st.dataframe(df.style.highlight_max(axis=1),use_container_width=True)
+
+        compare_1 = np.array(masked_ids_option_1['sent_1'])!=np.array(masked_ids_option_1['sent_2'])
+        compare_2 = np.array(masked_ids_option_2['sent_1'])!=np.array(masked_ids_option_2['sent_2'])
+        assert np.all(compare_1.astype(int)==compare_2.astype(int))
+        context_locs = list(np.arange(len(masked_ids_option_1['sent_1']))[compare_1]-1)  # match the indexing for annotation
+
+        multihead = True
+        assert np.all(np.array(pron_locs['sent_1'])==np.array(pron_locs['sent_2']))
+        assert np.all(np.array(option_1_locs['sent_1'])==np.array(option_1_locs['sent_2']))
+        assert np.all(np.array(option_2_locs['sent_1'])==np.array(option_2_locs['sent_2']))
+        token_id_list = pron_locs['sent_1'] + option_1_locs['sent_1'] + option_2_locs['sent_1'] + context_locs
+        #st.write(token_id_list)
+
+        effect_array = []
+        for token_id in token_id_list:
+            token_id += 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)
+        effect_array = np.transpose(np.array(effect_array),(1,0,2))
+
+        cols = st.columns(len(masked_ids_option_1['sent_1'])-2)
+        token_id = 0
+        for col_id,col in enumerate(cols):
+            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())
+                    ax.set_xticks([])
+                    ax.set_xticklabels([])
+                    ax.set_yticks([])
+                    ax.set_yticklabels([])
+                    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)