Update app.py

app.py

@@ -6,7 +6,7 @@ from transformers import BertTokenizer, BertForMaskedLM, BertForSequenceClassifi
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
-@st.cache
+@st.cache(allow_output_mutation=True)
 def load_models():
     tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
     bert_mlm_positive = BertForMaskedLM.from_pretrained('any0019/text_style_mlm_positive', return_dict=True).to(device).train(True)
@@ -90,14 +90,16 @@ def get_best_hypotheses(sentence, beam_size, max_steps, epsilon=1e-3, pretty_out
     used_poss = []
     
     st.write(f"step #0:")
-    st.write(f"-- 1: (positive probability ~ {round(current_beam[sentence], 5)})\n      {sentence}")
+    st.write(f"-- 1: (positive probability ~ {round(current_beam[sentence], 5)})"
+    st.write(f"      {sentence}")
     
     for step in range(max_steps):
         current_beam, used_pos = beam_get_replacements(current_beam, beam_size, epsilon, used_poss)
         
         st.write(f"\nstep #{step+1}:")
         for i, (sent, prob) in enumerate(current_beam.items()):
-            st.write(f"-- {i+1}: (positive probability ~ {round(prob, 5)})\n      {highlight_diff(sent, sentence) if pretty_output else sent}")
+            st.write(f"-- {i+1}: (positive probability ~ {round(prob, 5)})"
+            st.write(f"      {highlight_diff(sent, sentence) if pretty_output else sent}")
         
         if used_pos is None:
             return current_beam, used_poss