Update app.py

app.py

@@ -6,9 +6,9 @@ import pandas as pd
 import numpy as np
 import string
 
-
+device = 'cuda'
 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")
+model = AutoModelForMaskedLM.from_pretrained("Hamda/test-1-finetuned-AraBART").to(device)
 
 #@st.cache
 def next_word(text, pipe):
@@ -31,7 +31,7 @@ st.checkbox('استعمال الرسم البياني', value=False)
 
 text_st = sent+ ' <mask>'
 
-pipe = pipeline("fill-mask", tokenizer=tokenizer, model=model, top_k=6)
+pipe = pipeline("fill-mask", tokenizer=tokenizer, model=model, top_k=15)
 dict_next_words = next_word(text_st, pipe)
 df = pd.DataFrame.from_dict(dict_next_words)
 df.reset_index(drop=True, inplace=True)