Spaces:
Sleeping
Sleeping
#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() |