import gradio as gr
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-mpnet-base-v2")
model = AutoModelForMaskedLM.from_pretrained("sentence-transformers/all-mpnet-base-v2")

def fill_mask(text):
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    predictions = torch.topk(outputs.logits, k=1).indices.squeeze(0)
    decoded_output = tokenizer.decode(predictions)
    return decoded_output

title = "Masked Language Model (MPNet)"
description = "Provide input text with [MASK] and the model will predict the masked token."

interface = gr.Interface(
    fn=fill_mask,
    inputs=gr.Textbox(label="Input Text", placeholder="Type something with [MASK]..."),
    outputs=gr.Textbox(label="Predicted Text"),
    title=title,
    description=description,
)

if __name__ == "__main__":
    interface.launch()