Update app.py

app.py

@@ -15,7 +15,6 @@ 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)
-#@st.cache
 def next_word(text, pipe):
         res_dict= {  
           'الكلمة المقترحة':[],
@@ -35,6 +34,7 @@ if (st.button('بحث', disabled=False)):
 #using Graph   
 if (st.checkbox('الاستعانة بالرسم البياني', value=False)):
     a = time()
+    state_cb = True
     VocMap = './voc.csv'
     ScoreMap = './BM25.csv'
     
@@ -57,8 +57,13 @@ if (st.checkbox('الاستعانة بالرسم البياني', value=False)):
     #@st.cache
     def setQueriesVoc(df, id_list):
         res = []
+        
         for e in id_list:
-            res.extend(list(df.loc[e]['ID2'].values))   
+            try:
+                res.extend(list(df.loc[e]['ID2'].values)) 
+            except KeyError:
+                st.write('Token Unfound')
+                continue
         return list(set(res))
     
     L = setQueriesVoc(df_in, id_list)