import streamlit as st
import transformers
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForMaskedLM
import pandas as pd
import numpy as np


tokenizer = AutoTokenizer.from_pretrained("moussaKam/AraBART", padding= True, truncation=True, max_length=128)
model = AutoModelForMaskedLM.from_pretrained("moussaKam/AraBART")

#@st.cache
def next_word(text, pipe):
    res_dict= {  
  'Word':[],
  'Score':[],
  }
    for e in pipe(text):
        res_dict['Word'].append(e['token_str'])
        res_dict['Score'].append(e['score'])
    return res_dict

st.title("Predict Next Word")
st.write("Expand your query by leveraging various models")
default_value = "بيعت الأسلحة في السوق"
# sent is the variable holding the user's input
sent = st.text_area("Input", default_value, height=30)

st.checkbox('Use Graph', value=False)

text_st = sent+ ' <mask>'

pipe = pipeline("fill-mask", tokenizer=tokenizer, model=model)
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)


#st.table(df)