Update app.py

app.py

@@ -1,5 +1,6 @@
 from transformers import pipeline
 import streamlit as st
+import re
 
 unmasker = pipeline('fill-mask', model='distilbert-base-uncased')
 
@@ -8,10 +9,15 @@ def main():
     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')
     
-    if '[MASK]' is in masked.split(' '):
+    mask_str = '[MASK]'
+    mask_check = re.search(masked, mask_str)
+    
+    if mask_check:
         output = unmasker(masked)
     else:
         print('Sentence is not masked!!!')
   
 if __name__ == "__main__":
-    main()
+    main()
+
+x = re.search("^The.*Spain$", txt)