Update app.py

app.py

@@ -4,6 +4,7 @@ from transformers import pipeline
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 import pandas as pd
 import numpy as np
+import string
 
 
 tokenizer = AutoTokenizer.from_pretrained("moussaKam/AraBART", max_length=128, padding=True, pad_to_max_length = True, truncation=True)
@@ -15,9 +16,13 @@ def next_word(text, pipe):
   'Word':[],
   'Score':[],
   }
+   
+    
+  
     for e in pipe(text):
-        res_dict['Word'].append(e['token_str'])
-        res_dict['Score'].append(e['score'])
+        if e['token_str'] not in string.punctuation:
+            res_dict['Word'].append(e['token_str'])
+            res_dict['Score'].append(e['score'])
     return res_dict
 
 st.title("المساعدة اللغوية في التنبؤ بالمتلازمات والمتصاحبات والتعبيرات الاصطلاحية وتصحيحها")
@@ -29,7 +34,7 @@ st.checkbox('استعمال الرسم البياني', value=False)
 
 text_st = sent+ ' <mask>'
 
-pipe = pipeline("fill-mask", tokenizer=tokenizer, model=model)
+pipe = pipeline("fill-mask", tokenizer=tokenizer, model=model, top_k=10)
 dict_next_words = next_word(text_st, pipe)
 df = pd.DataFrame.from_dict(dict_next_words)
 df.reset_index(drop=True, inplace=True)