MaskedLM_App / app.py
MUmairAB's picture
Update app.py
2d6b756
raw
history blame
1.75 kB
# 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()