import gradio as gr
from transformers import pipeline

def fill_mask(text):
    model_name = "AventIQ-AI/RoBERTa"
    mask_filler = pipeline("fill-mask", model=model_name)
    
    results = mask_filler(text)
    
    output = "\n".join([f"{i+1}. {res['sequence']} (Confidence: {res['score']:.4f})" for i, res in enumerate(results)])
    return output

iface = gr.Interface(
    fn=fill_mask,
    inputs=gr.Textbox(label="Input Text (Use <mask> for missing word)"),
    outputs=gr.Textbox(label="Predictions"),
    title="RoBERTa Fill Mask",
    description="Enter a sentence with <mask>, and the RoBERTa model will predict the missing word."
)

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