from pandas.io.formats.format import return_docstring
import streamlit as st
import pandas as pd
from transformers import AutoTokenizer,AutoModelForMaskedLM
from transformers import pipeline
import os
import json

@st.cache(show_spinner=False,persist=True)
def load_model(masked_text,model_name):

    model = AutoModelForMaskedLM.from_pretrained(model_name, from_flax=True)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    # tokenizer.save_pretrained('exported_pytorch_model')
    # model.save_pretrained('exported_pytorch_model')
    nlp = pipeline('fill-mask', model=model, tokenizer=tokenizer)

    result_sentence = nlp(masked_text)

    return result_sentence


def main():

    st.title("RoBERTa Hindi")
    st.markdown(
    "This demo uses pretrained RoBERTa variants for Mask Language Modeling (MLM)"
    )

    models = st.multiselect(
        "Choose models",
        ['flax-community/roberta-hindi','mrm8488/HindiBERTa','ai4bharat/indic-bert',\
        'neuralspace-reverie/indic-transformers-hi-bert', 
          'surajp/RoBERTa-hindi-guj-san'],
        ["flax-community/roberta-hindi"]
    )   
    
    target_text_path = './mlm_custom/mlm_targeted_text.csv'
    target_text_df = pd.read_csv(target_text_path)
    
    texts = target_text_df['text']
    
    st.sidebar.title("Hindi MLM")
    masked_text = st.sidebar.selectbox('Select any of the following text',
     texts)

    st.write('You selected:', masked_text)
    
    selected_model = models[0]

    if st.button('Fill the Mask!'):
         with st.spinner("Filling the Mask..."):
            filled_sentence = load_model(masked_text,selected_model)
            st.write(filled_sentence)

    
if __name__ == "__main__":
    main()