import streamlit as st
from transformers import pipeline

# Page config
st.set_page_config(page_title="Blank Space Filling Model", page_icon="📝", layout="centered")

st.title("📝 Blank Space Filling Model")
st.write("Type a sentence with a blank using **____** or **[MASK]**.")

# Load Hugging Face model
@st.cache_resource
def load_model():
    return pipeline("fill-mask", model="bert-base-uncased")

fill_mask = load_model()

# Input box
user_input = st.text_input(
    "Enter your sentence:",
    "India is a ____ country."
)

# Prediction button
if st.button("Fill Blank"):
    sentence = user_input.replace("____", "[MASK]")

    if "[MASK]" not in sentence:
        st.error("Please include a blank like ____ or [MASK].")
    else:
        with st.spinner("Predicting..."):
            results = fill_mask(sentence)

        st.success("Prediction completed!")

        st.subheader("Top Predictions")
        for i, result in enumerate(results[:5], start=1):
            word = result["token_str"].strip()
            sentence_output = result["sequence"]
            confidence = round(result["score"] * 100, 2)

            st.write(f"### {i}. {word}")
            st.write(f"**Sentence:** {sentence_output}")
            st.write(f"**Confidence:** {confidence}%")
            st.markdown("---")