Shangkhonil's picture
Update app.py
f0a3481 verified
raw
history blame contribute delete
959 Bytes
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()