Added edit slider and changed sampling back to multinomial.

app.py

@@ -6,12 +6,16 @@ from transformers.pipelines import PIPELINE_REGISTRY, FillMaskPipeline
 from transformers import AutoModelForMaskedLM
 
 # unmasker = pipeline("temp-scale", model="anferico/bert-for-patents")
-example = 'A crustless [MASK] made from two slices of baked bread.'
-example = 'The invention provides a method for altering or modifying [MASK] of one or more gene products.'
-example = 'The graphite [MASK] is composed of a two-dimensional hexagonal lattice of carbon atoms.'
+examples = [['A crustless [MASK] made from two slices of baked bread.', 1.2],
+           ['The invention provides a method for altering or modifying [MASK] of one or more gene products.', 1.1],
+           ['The graphite [MASK] is composed of a two-dimensional hexagonal lattice of carbon atoms.', 1.4]]
 
 def add_mask(text, size=1):
     split_text = text.split()
+
+    # If the user supplies a mask, don't add more
+    if '[MASK]' in split_text:
+        return text
     idx = np.random.randint(len(split_text), size=size)
     for i in idx:
         split_text[i] = '[MASK]'
@@ -114,32 +118,38 @@ PIPELINE_REGISTRY.register_pipeline(
 )
 scrambler = pipeline("temp-scale", model="anferico/bert-for-patents")
 
-def unmask(text, temp, sampling='uniform'):
-    # text = add_mask(text)
-    split_text = text.split()
-    res = scrambler(text, temp=temp, top_k=3)
-    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_list = list(score_to_str.keys())
-    if sampling == 'multi':
-        idx = np.argmax(np.random.multinomial(1, score_list, 1))
-    else:
-        idx = np.random.randint(0, len(score_list))
-    score = score_list[idx]
-    new_token = score_to_str[score]
-    split_text[mask_pos] = new_token
-    return out, ' '.join(split_text)
+
+def unmask(text, temp, rounds):
+    sampling = 'multi'
+
+    for _ 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_list = list(score_to_str.keys())
+        if sampling == 'multi':
+            idx = np.argmax(np.random.multinomial(1, score_list, 1))
+        else:
+            idx = np.random.randint(0, len(score_list))
+        score = score_list[idx]
+        new_token = score_to_str[score]
+        split_text[mask_pos] = new_token
+        text = ' '.join(split_text)
+    return text
 
 textbox = gr.Textbox(label="Type language here", lines=5)
-textbox2 = gr.Textbox(placeholder="Type here...", lines=4)
-temp_slider = gr.Slider(1.0, 2.0, value=1.0, label='Temperature Scale')
+textbox2 = gr.Textbox(placeholder="", lines=4)
+temp_slider = gr.Slider(1.0, 2.0, value=1.0, label='Creativity')
+edit_slider = gr.Slider(1, 50, step=1, value=1.0, label='Number of edits')
 
 demo = gr.Interface(
     fn=unmask,
-    inputs=[textbox, temp_slider],
-    outputs=["label", textbox2],
-    examples=[[example, 1.2]],
+    inputs=[textbox, temp_slider, edit_slider],
+    outputs=[textbox2],
+    examples=examples,
 )
 
 demo.launch()