app update

app.py

@@ -6,6 +6,8 @@ from transformers import pipeline
 import os
 import json
 import random
+import numpy as np
+
 
 @st.cache(show_spinner=False,persist=True)
 def load_model(masked_text,model_name):
@@ -19,7 +21,7 @@ def load_model(masked_text,model_name):
     masked_text = masked_text.replace("<mask>",MASK_TOKEN)
     result_sentence = nlp(masked_text)
 
-    return result_sentence[0]['sequence']
+    return result_sentence[0]['sequence'], result_sentence[0]['token_str']
 
 def main():
 
@@ -30,7 +32,7 @@ def main():
 
     models = st.multiselect(
         "Choose models",                                         
-        ['flax-community/roberta-hindi','mrm8488/HindiBERTa','ai4bharat/indic-bert',\
+        ['flax-community/roberta-hindi','mrm8488/HindiBERTa',\
         'neuralspace-reverie/indic-transformers-hi-bert', 
           'surajp/RoBERTa-hindi-guj-san'],
         ["flax-community/roberta-hindi"]
@@ -44,14 +46,12 @@ def main():
     st.sidebar.title("Hindi MLM")
      
     pick_random = st.sidebar.checkbox("Pick any random text")
-
-    #st.write('You selected:', masked_text)
     
-    results_df = pd.DataFrame(columns = ['Model Name','Masked Text','Filled Masked Text'])
+    results_df = pd.DataFrame(columns = ['Model Name','Filled Token','Filled Text'])
      
     model_names = []
-    masked_texts = []
     filled_masked_texts = []
+    filled_tokens = []
     
     if pick_random:
         random_text = texts[random.randint(0,texts.shape[0]-1)]
@@ -60,22 +60,24 @@ def main():
         select_text = st.sidebar.selectbox('Select any of the following text',\
      texts)
         masked_text = st.text_area("Please type a masked sentence to fill",select_text)
-
+    
+    #pd.set_option('max_colwidth',30)
     if st.button('Fill the Mask!'):
         with st.spinner("Filling the Mask..."):
 
             for selected_model in models:
 
-                filled_sentence = load_model(masked_text,selected_model)
+                filled_sentence,filled_token = load_model(masked_text,selected_model)
                 model_names.append(selected_model)
-                masked_texts.append(masked_text)
+                filled_tokens.append(filled_token)
                 filled_masked_texts.append(filled_sentence)
 
             results_df['Model Name'] = model_names
-            results_df['Masked Text'] = masked_texts
-            results_df['Filled Masked Text'] = filled_masked_texts
+            results_df['Filled Token'] = filled_tokens
+            results_df['Filled Text'] = filled_masked_texts
+        
+            #st.table(results_df)
+            st.write(results_df)
 
-            st.table(results_df)
-            
 if __name__ == "__main__":
     main()