test

app.py

@@ -91,7 +91,6 @@ def annotate_options(sent_id,sent):
                     st.session_state[f'option_locs_{sent_id}'].append(word_id)
                 else:
                     st.session_state[f'option_locs_{sent_id}'].remove(word_id)
-    st.write([st.session_state[f'option_locs_{sent_id}']])
     st.markdown(show_annotated_sentence(decoded_sent,
                                         option_locs=st.session_state[f'option_locs_{sent_id}'],
                                         mask_locs=st.session_state[f'mask_locs_{sent_id}']), unsafe_allow_html = True)
@@ -123,41 +122,36 @@ if __name__=='__main__':
         st.session_state['page_status'] = 'type_in'
 
     if st.session_state['page_status']=='type_in':
-        with main_area.container():
-            st.write('1. Type in the sentences and click "Tokenize"')
-            sent_1 = st.text_input('Sentence 1',value='It is better to play a prank on Samuel than Craig because he gets angry less often.')
-            sent_2 = st.text_input('Sentence 2',value='It is better to play a prank on Samuel than Craig because he gets angry more often.')
-            if st.button('Tokenize'):
-                st.session_state['page_status'] = 'annotate_mask'
-                st.session_state['sent_1'] = sent_1
-                st.session_state['sent_2'] = sent_2
-                main_area.empty()
+        st.write('1. Type in the sentences and click "Tokenize"')
+        sent_1 = st.text_input('Sentence 1',value='It is better to play a prank on Samuel than Craig because he gets angry less often.')
+        sent_2 = st.text_input('Sentence 2',value='It is better to play a prank on Samuel than Craig because he gets angry more often.')
+        if st.button('Tokenize'):
+            st.session_state['page_status'] = 'annotate_mask'
+            st.session_state['sent_1'] = sent_1
+            st.session_state['sent_2'] = sent_2
+            st.experimental_rerun()
 
     if st.session_state['page_status']=='annotate_mask':
-        with main_area.container():
-            sent_1 = st.session_state['sent_1']
-            sent_2 = st.session_state['sent_2']
+        sent_1 = st.session_state['sent_1']
+        sent_2 = st.session_state['sent_2']
 
-            st.write('2. Select sites to mask out and click "Confirm"')
-            annotate_mask(1,sent_1)
-            annotate_mask(2,sent_2)
-            st.write(st.session_state['mask_locs_1'])
-            st.write(st.session_state['mask_locs_2'])
-            if st.button('Confirm',key='mask'):
-                st.session_state['page_status'] = 'annotate_options'
-                main_area.empty()
+        st.write('2. Select sites to mask out and click "Confirm"')
+        annotate_mask(1,sent_1)
+        annotate_mask(2,sent_2)
+        if st.button('Confirm',key='mask'):
+            st.session_state['page_status'] = 'annotate_options'
+            st.experimental_rerun()
 
     if st.session_state['page_status'] == 'annotate_options':
-        with main_area.container():
-            sent_1 = st.session_state['sent_1']
-            sent_2 = st.session_state['sent_2']
+        sent_1 = st.session_state['sent_1']
+        sent_2 = st.session_state['sent_2']
 
-            st.write('2. Select options and click "Confirm"')
-            annotate_options(1,sent_1)
-            annotate_options(2,sent_2)
-            if st.button('Confirm',key='option'):
-                st.session_state['page_status'] = 'analysis'
-                main_area.empty()
+        st.write('2. Select options and click "Confirm"')
+        annotate_options(1,sent_1)
+        annotate_options(2,sent_2)
+        if st.button('Confirm',key='option'):
+            st.session_state['page_status'] = 'analysis'
+            st.experimental_rerun()
 
     if st.session_state['page_status']=='analysis':
         with main_area.container():
@@ -168,7 +162,8 @@ if __name__=='__main__':
             input_ids_2 = tokenizer(sent_2).input_ids
             input_ids = torch.tensor([input_ids_1,input_ids_2])
 
-            outputs = SkeletonAlbertForMaskedLM(model,input_ids,interventions = {0:{'lay':[(8,1,[0,1])]}})
+            outputs = SkeletonAlbertForMaskedLM(model,input_ids,
+                                                    interventions = {0:{'lay':[(head_id,17,[0,1]) for head_id in range(64)]}})
             logprobs = F.log_softmax(outputs['logits'], dim = -1)
-            preds = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[0]]
+            preds = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[0][1:-1]]
             st.write([tokenizer.decode([token]) for token in preds])