import streamlit as st
from transformers import pipeline

unmasker = pipeline('fill-mask', model='dsfsi/zabantu-ven-120m')

def fill_mask(sentences):
    results = {}
    for sentence in sentences:
        unmasked = unmasker(sentence)
        results[sentence] = unmasked
    return results

def replace_mask(sentence, predicted_word):
    return sentence.replace("<mask>", predicted_word)

st.title("Fill Mask | Zabantu-nso-ven-120m")
st.write("This app predicts the missing word in a sentence using a Zabantu-nso-ven-120m model.")

sample_sentences = ["Rabulasi wa <mask> u khou bvelela nga u lima,",
                    "Vhana vhane vha kha ḓi bva u bebwa vha kha khombo ya u <mask> nga Listeriosis"]

text_input = st.text_area("Enter sentences with <mask> token (one per line):", 
                          "\n".join(sample_sentences))

input_sentences = text_input.split(",")

if st.button("Submit"):
    result = fill_mask(input_sentences)

    if result:
        for sentence, predictions in result.items():
            st.write(f"**Original sentence**: {sentence}")
            for prediction in predictions:
                predicted_word = prediction['token_str']
                score = prediction['score'] * 100
                full_sentence = replace_mask(sentence, predicted_word)
                
                st.write(f"Predicted word: {predicted_word} - Score: {score:.2f}%")
                st.write(f"Full sentence: {full_sentence}")
            st.write("=" * 80)

css = """
<style>
footer {display:none !important}

.stButton > button {
    background-color: #17152e;
    color: white;
    border: none;
    padding: 0.75em 2em;
    text-align: center;
    text-decoration: none;
    display: inline-block;
    font-size: 16px;
    margin: 4px 2px;
    cursor: pointer;
    border-radius: 12px;
    transition: background-color 0.3s ease;
}

.stButton > button:hover {
    background-color: #3c4a6b;
}

.stTextInput, .stTextArea {
    border: 1px solid #e6e6e6;
    padding: 0.75em;
    border-radius: 10px;
    font-size: 16px;
    width: 100%;
}

.stTextInput:focus, .stTextArea:focus {
    border-color: #17152e;
    outline: none;
    box-shadow: 0px 0px 5px rgba(23, 21, 46, 0.5);
}

div[data-testid="stMarkdownContainer"] p {
    font-size: 16px;
}

.stApp {
    padding: 2em;
    font-family: 'Poppins', sans-serif;
}
</style>
"""
st.markdown(css, unsafe_allow_html=True)