Reduced suggested edits and max edits for faster testing. Changed s

app.py

@@ -18,17 +18,24 @@ ex_str3 = "The graphite plane is composed of a two-dimensional hexagonal lattice
           "length and a width parallel to the graphite plane and a thickness orthogonal to the graphite plane with at " \
           "least one of the length, width, and thickness values being 100 nanometers or smaller. "
 
-examples = [[ex_str1, 1.2, 1],
-           [ex_str2, 1.5, 10],
-           [ex_str3, 1.4, 5]]
+tab_two_examples = [[ex_str1, 1.2, 1],
+                    [ex_str2, 1.5, 10],
+                    [ex_str3, 1.4, 5]]
+
+tab_one_examples = [['A crustless _ made from two slices of baked bread.'],
+                    ['The present disclosure provides a DNA-targeting RNA that comprises a targeting _.'],
+                    ['The _ plane is composed of a two-dimensional hexagonal lattice of carbon atoms.']
+                    ]
 
 
 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
+    if '_' in split_text:
+        u_pos = [i for i, s in enumerate(split_text) if '_' in s][0]
+        split_text[u_pos] = '[MASK]'
+        return ' '.join(split_text), '[MASK]'
     idx = np.random.randint(len(split_text), size=size)
     masked_strings = []
     for i in idx:
@@ -146,6 +153,14 @@ def sample_output(out, sampling):
     return score_to_str[score]
 
 
+def unmask_single(text, temp=1):
+    tp = add_mask(text, size=1)
+    masked_text, masked = tp[0], tp[1]
+    res = scrambler(masked_text, temp=temp, top_k=10)
+    out = {item["token_str"]: item["score"] for item in res}
+    return out
+
+
 def unmask(text, temp, rounds):
     sampling = 'multi'
     for _ in range(rounds):
@@ -161,7 +176,7 @@ def unmask(text, temp, rounds):
             if unsuccessful_iters > 5:
                 break
             print('skipped', new_token)
-            new_token = sample_output(out, sampling='uniform')
+            new_token = sample_output(out, sampling=sampling)
             unsuccessful_iters += 1
         if new_token == masked[0]:
             split_text[mask_pos] = new_token
@@ -173,18 +188,43 @@ def unmask(text, temp, rounds):
     text[0] = text[0].upper()
     return ''.join(text)
 
+textbox1 = gr.Textbox(label="Input Sentence", lines=5)
+output_textbox1 = gr.Textbox(placeholder="Output will appear here", lines=4)
+
+textbox2 = gr.Textbox(label="Input Sentences", lines=5)
+output_textbox2 = gr.Textbox(placeholder="Output will appear here", lines=4)
+temp_slider2 = gr.Slider(1.0, 3.0, value=1.0, label='Creativity')
+edit_slider2 = gr.Slider(1, 20, step=1, value=1.0, label='Number of edits')
+
+title1 = "Patent-BERT Sentence Remix-er: Single Edit"
+description1 = """<p>Try inserting a '_' where you want the model to generate a list of likely words.
+<br/>
+<p/>"""
+title2 = "Patent-BERT Sentence Remix-er: Multiple Edits"
+description2 = """<p>Try typing in a sentence for the model to remix. Adjust the 'creativity' scale bar to change the 
+the model's confidence in its likely substitutions and the 'number of edits' for the number of edits you want 
+the model to attempt to make. <br/> <p/> """
+
+demo1 = gr.Interface(
+    fn=unmask_single,
+    inputs=[textbox1],
+    outputs='label',
+    examples=tab_one_examples,
+    allow_flagging='never',
+    title=title1,
+    description=description1
+)
 
-textbox = gr.Textbox(label="Example prompts", lines=5)
-output_textbox = gr.Textbox(placeholder="Output will appear here", lines=4)
-temp_slider = gr.Slider(1.0, 3.0, value=1.0, label='Creativity')
-edit_slider = gr.Slider(1, 20, step=5, value=1.0, label='Number of edits')
-
-demo = gr.Interface(
+demo2 = gr.Interface(
     fn=unmask,
-    inputs=[textbox, temp_slider, edit_slider],
-    outputs=[output_textbox],
-    examples=examples,
-    allow_flagging='never'
+    inputs=[textbox2, temp_slider2, edit_slider2],
+    outputs=[output_textbox2],
+    examples=tab_two_examples,
+    allow_flagging='never',
+    title=title2,
+    description=description2
 )
 
-demo.launch()
+gr.TabbedInterface(
+    [demo1, demo2], ["Single edit", "Multiple Edits"]
+).launch()