Spaces:
Build error
Build error
| 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() | |