Added more sanity checking.

app.py

@@ -134,14 +134,13 @@ scrambler = pipeline("temp-scale", model="anferico/bert-for-patents")
 
 def unmask(text, temp, rounds):
     sampling = 'multi'
-    successful_iters = 0
     for round in range(rounds):
         text = add_mask(text, size=1)
         split_text = text.split()
         res = scrambler(text, temp=temp, top_k=10)
         mask_pos = [i for i, t in enumerate(split_text) if 'MASK' in t][0]
         out = {item["token_str"]: item["score"] for item in res}
-        score_to_str = {out[k]:k for k in out.keys()}
+        score_to_str = {out[k] : k for k in out.keys()}
         score_list = list(score_to_str.keys())
         if sampling == 'multi':
             idx = np.argmax(np.random.multinomial(1, score_list, 1))
@@ -151,12 +150,11 @@ def unmask(text, temp, rounds):
         new_token = score_to_str[score]
         if len(list(new_token)) < 2:
             continue
-        split_text[mask_pos] = new_token
+        split_text[mask_pos] = '*' + new_token + '*'
         text = ' '.join(split_text)
-        successful_iters += 1
     text = list(text)
     text[0] = text[0].upper()
-    return ''.join(text) + str(successful_iters)
+    return ''.join(text)
 
 
 textbox = gr.Textbox(label="Example prompts", lines=5)