tweak

app.py

@@ -87,7 +87,7 @@ if __name__=='__main__':
         st.write('2. Select sites to mask out and click "Confirm"')
         input_sent = tokenizer(sent_1).input_ids
         decoded_sent = [tokenizer.decode([token]) for token in input_sent[1:-1]]
-        char_nums = [len(word)+2 for word in decoded_sent]
+        char_nums = [len(word)+3 for word in decoded_sent]
         st.write(char_nums)
         cols = st.columns(char_nums)
         for word_id,(col,word) in enumerate(zip(cols,decoded_sent)):
@@ -95,6 +95,8 @@ if __name__=='__main__':
                 if st.button(word,key=f'word_{word_id}'):
                     if word_id not in st.session_state['masked_pos_1']:
                         st.session_state['masked_pos_1'].append(word_id)
+                    else:
+                        _ = st.session_state['masked_pos_1'].pop(word_id)
         st.write(f'Masked words: {", ".join([decoded_sent[word_id] for word_id in np.sort(st.session_state["masked_pos_1"])])}')