from transformers import pipeline
import streamlit as st
import re

unmasker = pipeline('fill-mask', model='distilbert-base-cased')

def main():
    with st.form("Mask_filling_form"):
        st.title('EN sentence mask filling')
        st.markdown('Mask a word with [MASK]. For example, **I [MASK] dog for a walk**')
        masked = st.text_input('Masked sentence','Please insert a sentence')
        
        submitted = st.form_submit_button("Submit")
        
        if submitted:
            mask_str = '[MASK]'
            mask_check = bool(re.search(mask_str, masked))
            
            if mask_check:
                st.success("Submitted")
                output = unmasker(masked)
                
                st.write('Filled sentence:',output[0]['sequence'])
            else:
                st.error('Sentence is not masked!!!')
  
if __name__ == "__main__":
    main()