Update app.py

app.py

@@ -33,21 +33,34 @@ infill: sports teams are profitable for owners. ( accumulating vast sums / stock
 
 original:"""
 
+def prefix_format(sentence):
+  words = sentence.split()
+  if "[MASK]" in sentence:
+    words2 = words.index("[MASK]")
+    #print(words2)
+    output = ("<|SUF|> " + ' '.join(words[words2+1:]) + " <|PRE|> " + ' '.join(words[:words2]) + " <|MID|>")
+    st.write(output)
+  else:
+    st.write("Add [MASK] to sentence")
+    
 with st.form(key='my_form'):
     prompt = st.text_area(label='Enter sentence', value=g)
     submit_button = st.form_submit_button(label='Submit')
+    submit_button6 = st.form_submit_button(label='Turn Into Infill Format. Just add [MASK] where you want it infilled')
     if submit_button:
       with torch.no_grad():
         text = tokenizer.encode(prompt)
         myinput, past_key_values = torch.tensor([text]), None
         myinput = myinput
-        myinput= myinput.to(device)
+        myinput= myinput
         logits, past_key_values = model(myinput, past_key_values = past_key_values, return_dict=False)
         logits = logits[0,-1]
         probabilities = torch.nn.functional.softmax(logits)
-        best_logits, best_indices = logits.topk(300)
+        best_logits, best_indices = logits.topk(250)
         best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]
         text.append(best_indices[0].item())
         best_probabilities = probabilities[best_indices].tolist()
         words = []              
-        st.write(best_words)
+        st.write(best_words)
+    if submit_button6:
+        prefix_format(prompt)