Update caption translation logic

apps/mlm.py

@@ -109,21 +109,21 @@ def app(state):
 
     # Display caption
     new_col2.write("Write your text with exactly one [MASK] token.")
-    caption = new_col2.text_input(
+    mlm_state.caption = new_col2.text_input(
         label="Text",
         value=mlm_state.caption,
         help="Type your masked caption regarding the image above in one of the four languages.",
     )
 
-    if caption == mlm_state.caption:
+    if mlm_state.unmasked_caption == mlm_state.caption.replace("[MASK]", mlm_state.currently_masked_token):
         new_col2.markdown("**Masked Token**: "+mlm_state.currently_masked_token)
         new_col2.markdown("**English Translation: " + mlm_state.unmasked_caption if mlm_state.caption_lang_id == "en" else translate(mlm_state.unmasked_caption, 'en'))
 
     else:
         new_col2.markdown(
-            f"""**English Translation**: {caption if mlm_state.caption_lang_id == "en" else translate(caption, 'en')}"""
+            f"""**English Translation**: {mlm_state.caption if mlm_state.caption_lang_id == "en" else translate(mlm_state.caption, 'en')}"""
         )
-    caption_inputs = get_text_attributes(caption)
+    caption_inputs = get_text_attributes(mlm_state.caption)
 
     # Display Top-5 Predictions
     with st.spinner("Predicting..."):