fix

app.py

@@ -72,7 +72,8 @@ def annotate_mask(sent_id,sent):
                     mask_locs.append(word_id)
                 else:
                     mask_locs.remove(word_id)
-    return decoded_sent, mask_locs
+    st.markdown(show_annotated_sentence(decoded_sent,mask_locs=mask_locs), unsafe_allow_html = True)
+    st.session_state[f'mask_locs_{sent_id}'] = mask_locs
 
 def annotate_options(sent_id,sent):
     st.write(f'Sentence {sent_id}')
@@ -135,16 +136,10 @@ if __name__=='__main__':
             sent_2 = st.session_state['sent_2']
 
             st.write('2. Select sites to mask out and click "Confirm"')
-            decoded_sent_1, mask_locs_1 = annotate_mask(1,sent_1)
-            container_1 = st.container()
-            decoded_sent_2, mask_locs_2 = annotate_mask(2,sent_2)
-            container_2 = st.container()
-            container_1.markdown(show_annotated_sentence(decoded_sent_1,mask_locs=mask_locs_1), unsafe_allow_html = True)
-            container_2.markdown(show_annotated_sentence(decoded_sent_2,mask_locs=mask_locs_2), unsafe_allow_html = True)
+            annotate_mask(1,sent_1)
+            annotate_mask(2,sent_2)
             if st.button('Confirm'):
                 st.session_state['page_status'] = 'annotate_options'
-                st.session_state['mask_locs_1'] = mask_locs_1
-                st.session_state['mask_locs_2'] = mask_locs_2
                 main_area.empty()
 
     if st.session_state['page_status'] == 'annotate_options':