MUmairAB's picture
Update app.py
ff26108
raw
history blame contribute delete
No virus
1.89 kB
#import necessary libraries
import gradio as gr
from transformers import pipeline
#Instantiate the model
model = pipeline(task="fill-mask",
model="MUmairAB/bert-based-MaskedLM")
#Typically, the model should be imported within a function. However, in this case, we are downloading it outside the function to avoid a significant delay that could annoy the user when downloading it inside the main function. By loading the model at this point, it will be downloaded when the app runs, and the user will overlook this initial loading time, as opposed to experiencing a delay after entering the input.
def fill_the_mask(text):
if "[MASK]" not in text:
return "You did not enter \"[MASK]\" in the text. Please write your text again!"
else:
#Apply the model
model_out = model(text)
#First sort the list of dictionaries according to the score
model_out = sorted(model_out, key=lambda x: x['score'],reverse=True)
#Create a dictionary to store the model output
out_dict = {}
#Iterate over the list of dictionaries and get the required ouput
for sub_dict in model_out:
out_dict[sub_dict["sequence"]] = round(sub_dict["score"], 3)
return out_dict
#Create a Gradio user interface
my_interface = gr.Interface(title="Masked Language Model APP\n(by Umair Akram)",
description="This App uses a fine-tuned DistilBERT-Base-Uncased Masked Language Model to predict the missed word in a sentence.\nEnter your text and put \"[MASK]\" at the word which you want to predict, as shown in the following example: Can we [MASK] to Paris?",
fn=fill_the_mask,
inputs="text",
outputs="label")
#Define the main function
if __name__ == "__main__":
#Launch the Gradio interface
my_interface.launch()