Update app.py

app.py

@@ -19,21 +19,29 @@ default_value = "أستاذ التعليم"
 sent = st.text_area('المدخل',default_value)
 
 tokenizer = AutoTokenizer.from_pretrained("moussaKam/AraBART", max_length=128, padding=True, pad_to_max_length = True, truncation=True)
+
 model = AutoModelForMaskedLM.from_pretrained("Hamda/test-1-finetuned-AraBART")
 pipe = pipeline("fill-mask", tokenizer=tokenizer, model=model, top_k=10)
+
 def next_word(text, pipe):
-        res_dict= {  
-          'الكلمة المقترحة':[],
-          'العلامة':[],
-        }
-        for e in pipe(text):
-            if all(c not in list(string.punctuation) for c in e['token_str']):
-                res_dict['الكلمة المقترحة'].append(e['token_str'])
-                res_dict['العلامة'].append(e['score'])
-        return res_dict
+    
+    filter_words = list(string.punctuation).extend(sent)
+    res_dict= {  
+      'الكلمة المقترحة':[],
+      'العلامة':[],
+    }
+    for e in pipe(text):
+        if all(c not in filter_words for c in e['token_str']):
+            res_dict['الكلمة المقترحة'].append(e['token_str'])
+            res_dict['العلامة'].append(e['score'])
+    return res_dict
+
 if (st.button('بحث', disabled=False)):
+    
     text_st = sent+ ' <mask>'
+    
     dict_next_words = next_word(text_st, pipe)
+    
     df = pd.DataFrame.from_dict(dict_next_words)
     df.reset_index(drop=True, inplace=True)
     st.dataframe(df)
@@ -57,7 +65,12 @@ if (st.checkbox('الاستعانة بالرسم البياني المعرفي
     
     #@st.cache
     def Query2id(voc, query):
-        return [voc.index[voc['word'] == word].values[0] for word in query.split()]
+        try:
+            return [voc.index[voc['word'] == word].values[0] for word in query.split()]
+        except KeyError:
+            st.markdown("""---""")
+            st.write('لم يتم التعرف على الكلمة')
+            st.markdown("""---""")      
     
     id_list = Query2id(df3, sent)
     #@st.cache
@@ -87,13 +100,19 @@ if (st.checkbox('الاستعانة بالرسم البياني المعرفي
             key  = df3.loc[nc].values[0]
             tmt[key] = score
         return tmt
+    
     tmt = compute_score(L, id_list)    
+    
+    for e in sent:
+        del tmt[e]
+    
     exp_terms = []
     t_li = tmt.values()
     tmexp = sorted(tmt.items(), key=lambda x: x[1], reverse=True)
     i = 0
     dict_res = {'الكلمة المقترحة':[], 
     'العلامة':[]}
+    
     for key, value in tmexp:
         new_score=((value-min(t_li))/(max(t_li)-min(t_li)))-0.0001
         dict_res['العلامة'].append(str(new_score)[:6])
@@ -101,6 +120,7 @@ if (st.checkbox('الاستعانة بالرسم البياني المعرفي
         i+=1
         if (i==10):
             break
+    
     res_df = pd.DataFrame.from_dict(dict_res)
     res_df.index += 1