# app.py
import torch
from transformers import pipeline
import gradio as gr

# Initialize the model pipeline
pipe = pipeline(
    "fill-mask",
    model="answerdotai/ModernBERT-base",
    torch_dtype=torch.bfloat16,
)


# Function to run the model
def fill_mask(input_text):
    return pipe(input_text)[0]["sequence"]


# Set up the Gradio interface
iface = gr.Interface(
    fn=fill_mask,
    inputs="text",
    outputs="text",
    title="Mask Filling with ModernBERT",
    description="Enter a sentence with a [MASK] token, and the model will predict the missing word.",
)

# Launch the interface
iface.launch()