Update app.py

app.py

@@ -33,15 +33,17 @@ def id_replace_complex(s, threshold=0.4):
   # if score >= threshold select tokens[idx]
   compl_tok = [tokens[idx] for idx, x in enumerate(
       cwi_pipe(cands)) if x['score'] >= threshold]
-
+  
+  replacements = []
   # potentially parallelizable, depends on desired behaviour
   for t in compl_tok:
     idx = s.index(t)
     s = s[:idx] + '<mask>' + s[idx+len(t):]
     # get top candidate for mask fill in complex token
-    s = fill_pipe(s)[0]['sequence']
-
-  return s, compl_tok
+    top_result = fill_pipe(s)[0]
+    s = top_result['sequence']
+    replacements.append(top_result['token_str'])
+  return s, compl_tok, replacements
 
 
 def generate_candidate_text(s, model, tokenizer, tokenized=False):
@@ -72,10 +74,10 @@ def rank_candidate_text(sentences):
 
 def full_pipeline(source, simpl_model, simpl_tok, tokens, lexical=False):
   
-  modified, complex_words  = id_replace_complex(source, threshold=0.2) if lexical else (source, None)
+  modified, complex_words, replacements  = id_replace_complex(source, threshold=0.2) if lexical else (source, None, None)
   cands = generate_candidate_text(tokens+modified, simpl_model, simpl_tok)
   output = rank_candidate_text(cands)
-  return output, complex_words
+  return output, complex_words, replacements
   
 def main():
 
@@ -117,10 +119,10 @@ def main():
         if (submit):
             
             tokens = " ".join([t+str(v) for t, v in zip(aug_tok, user_values)]) + " "
-            output, words = full_pipeline(input_sentence, simpl_model, simpl_tok, tokens, lexical)
-            print(tokens, lexical)
-           
-            c1, c2 = st.columns([1,2])
+            output, words, replacements = full_pipeline(input_sentence, simpl_model, simpl_tok, tokens, lexical)
+            
+    
+            c1, c2, c3 = st.columns([1,1,2])
 
             with c1:
                 st.markdown("#### Words identified as complex")
@@ -132,6 +134,15 @@ def main():
                     st.markdown("None :smile:")
 
             with c2:
+                st.markdown("#### Their mask-predicted replacement")
+                if replacements:
+                    for w in replacements:
+                        st.markdown(f"* {w}")
+
+                else:
+                    st.markdown("None :smile:")
+
+            with c3:
                 st.markdown(f"#### Original Sentence:\n > {input_sentence}") 
                 st.markdown(f"#### Output Sentence:\n > {output}")