# import the module
import streamlit as st
from transformers import pipeline


#This function accepts the masked text like: "How are [MASK]"
# and feeds this text to the model and prints the output in which [MASK] is filled with the appropriate word.
def print_the_mask(text):

    #Import the model
    model = pipeline(task="fill-mask",
                     model="MUmairAB/bert-based-MaskedLM")

    #Apply the model
    model_out = model("I want to [MASK]")

    #First sort the list of dictionaries according to the score
    model_out = sorted(model_out, key=lambda x: x['score'],reverse=True)
    for sub_dict in model_out:
        print(sub_dict["sequence"])


#The main function that will be executed when this file is executed
def main():
    # Set the title
    st.title("Masked Language Model App")
    st.write("Created by: [Umair Akram](https://www.linkedin.com/in/m-umair01/)")
   
    h1 = "This App uses a fine-tuned DistilBERT-Base-Uncased Masked Language Model to predict the missed word in a sentence."
    st.subheader(h1)

    st.write("Its code and other interesting projects are available on my [website](https://mumairab.github.io/)")
    h2 = "Enter your text and put \"[MASK]\" at the word which you want to predict, as shown in the following example: How are [MASK]"
    st.write(h2)

    text = st.text_input(label="Enter your text here:",
                 value="Type here ...")

    if(st.button('Submit')):
        # Perform the input validation
        if "[MASK]" not in text:
            st.write("You did not enter \"[MASK]\" in the text. Please write your text again!")
        else:
            print_the_mask(text)

#Call the main function
if __name__ == "__main__":
    #Launch the Gradio interface
    main()