Minor typo fix

app.py

@@ -30,10 +30,12 @@ def add_mask(text, size=1):
     if '[MASK]' in split_text:
         return text
     idx = np.random.randint(len(split_text), size=size)
-    masked = split_text[idx]
+    masked_strings = []
     for i in idx:
+        masked_strings.append(split_text[i])
         split_text[i] = '[MASK]'
-    return ' '.join(split_text), masked
+    masked_output = ' '.join(split_text)
+    return masked_output, masked_strings
 
 
 class TempScalePipe(FillMaskPipeline):
@@ -136,9 +138,10 @@ scrambler = pipeline("temp-scale", model="anferico/bert-for-patents")
 def unmask(text, temp, rounds):
     sampling = 'multi'
     for round in range(rounds):
-        text, masked = add_mask(text, size=1)
-        split_text = text.split()
-        res = scrambler(text, temp=temp, top_k=10)
+        tp = add_mask(text, size=1)
+        masked_text, masked = tp[0], tp[1]
+        split_text = masked_text.split()
+        res = scrambler(masked_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()}
@@ -149,7 +152,7 @@ def unmask(text, temp, rounds):
             idx = np.random.randint(0, len(score_list))
         score = score_list[idx]
         new_token = score_to_str[score]
-        if len(list(new_token)) < 2 or new_token == masked:
+        if len(list(new_token)) < 2 or new_token == masked[0]:
             continue
         split_text[mask_pos] = '*' + new_token + '*'
         text = ' '.join(split_text)