fix

app.py

@@ -72,8 +72,7 @@ def annotate_mask(sent_id,sent):
                     mask_locs.append(word_id)
                 else:
                     mask_locs.remove(word_id)
-    st.markdown(show_annotated_sentence(decoded_sent,mask_locs=mask_locs), unsafe_allow_html = True)
-    return mask_locs
+    return decoded_sent, mask_locs
 
 def annotate_options(sent_id,sent):
     st.write(f'Sentence {sent_id}')
@@ -136,10 +135,16 @@ if __name__=='__main__':
             sent_2 = st.session_state['sent_2']
 
             st.write('2. Select sites to mask out and click "Confirm"')
-            st.session_state[f'mask_locs_1'] = annotate_mask(1,sent_1)
-            st.session_state[f'mask_locs_2'] = annotate_mask(2,sent_2)
+            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)
             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':