Update app.py

app.py

@@ -3,7 +3,10 @@ import streamlit as st
 from termcolor import colored
 import torch
 from transformers import BertTokenizer, BertForMaskedLM, BertForSequenceClassification
+
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
 @st.cache
 def load_models():
     tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
@@ -11,7 +14,11 @@ def load_models():
     bert_mlm_negative = BertForMaskedLM.from_pretrained('text_style_mlm_negative', return_dict=True).to(device).train(True)
     bert_classifier = BertForSequenceClassification.from_pretrained('text_style_classifier', num_labels=2).to(device).train(True)
     return tokenizer, bert_mlm_positive, bert_mlm_negative, bert_classifier
+
+
 tokenizer, bert_mlm_positive, bert_mlm_negative, bert_classifier = load_models()
+
+
 def highlight_diff(sent, sent_main):
     tokens = tokenizer.tokenize(sent)
     tokens_main = tokenizer.tokenize(sent_main)
@@ -24,11 +31,14 @@ def highlight_diff(sent, sent_main):
             new_toks.append(tok)
     
     return ' '.join(new_toks)
-    
+
+
 def get_classifier_prob(sent):
     bert_classifier.eval()
     with torch.no_grad():
         return bert_classifier(**{k: v.to(device) for k, v in tokenizer(sent, return_tensors='pt').items()}).logits.softmax(dim=-1)[0].cpu().numpy()
+
+
 def beam_get_replacements(current_beam, beam_size, epsilon=1e-3, used_positions=[]):
     """
     - for each sentence in :current_beam: - split the sentence into tokens using the INGSOC-approved BERT tokenizer
@@ -74,6 +84,8 @@ def beam_get_replacements(current_beam, beam_size, epsilon=1e-3, used_positions=
         else:
             st.write("No more new hypotheses")
             return current_beam, None
+
+
 def get_best_hypotheses(sentence, beam_size, max_steps, epsilon=1e-3, pretty_output=False):
     current_beam = {sentence: get_classifier_prob(sentence)[1]}
     used_poss = []
@@ -94,10 +106,14 @@ def get_best_hypotheses(sentence, beam_size, max_steps, epsilon=1e-3, pretty_out
             used_poss.append(used_pos)
     
     return current_beam, used_poss
+
+
 st.title("Correcting opinions")
+
 default_value = "write your review here (in lower case - vocab reasons)"
 sentence = st.text_area("Text", default_value, height = 275)
 beam_size = st.sidebar.slider("Beam size", value = 3, min_value = 1, max_value=20, step=1)
 max_steps = st.sidebar.slider("Max steps", value = 3, min_value = 1, max_value=10, step=1)
 prettyfy = st.sidebar.slider("Higlight changes", value = 0, min_value = 0, max_value=1, step=1)
+
 beam, used_poss = get_best_hypotheses(sentence, beam_size=beam_size, max_steps=max_steps, pretty_output=bool(prettyfy))