Update app.py

app.py

@@ -16,6 +16,7 @@ from transformers import AutoTokenizer, AutoModelForMaskedLM
 artist_name = st.text_input("Model", "roberta-large")
 tokenizer = AutoTokenizer.from_pretrained("roberta-large")
 model = AutoModelForMaskedLM.from_pretrained(artist_name)
+model2 = AutoModelForMaskedLM.from_pretrained("BigSalmon/FormalRobertaa")
 
 
 
@@ -36,4 +37,13 @@ with st.form(key='my_form'):
         mask_hidden_state = last_hidden_state[mask_index]
         idx = torch.topk(mask_hidden_state, k=100, dim=0)[1]
         words = [tokenizer.decode(i.item()).strip() for i in idx]
-        st.text_area(label = 'Infill:', value=words)
+        a_list.append(words)
+      with torch.no_grad():
+        output = model2(token_ids)
+      last_hidden_state = output[0].squeeze()
+      for mask_index in masked_pos:
+        mask_hidden_state = last_hidden_state[mask_index]
+        idx = torch.topk(mask_hidden_state, k=100, dim=0)[1]
+        words2 = [tokenizer.decode(i.item()).strip() for i in idx]
+        a_list.append(words2)
+        st.text_area(label = 'Infill:', value=a_list)