Update app.py

app.py

@@ -8,11 +8,11 @@ st.set_page_config(layout="wide")
 def fill_mask(sentences):
     results = {}
     warnings = []
-    for key, (language, sentence) in sentences.items():
+    for language, sentence in sentences.items():
         if "<mask>" in sentence:
             masked_sentence = sentence.replace('<mask>', unmasker.tokenizer.mask_token)
             unmasked = unmasker(masked_sentence)
-            results[key] = (language, unmasked)
+            results[language] = unmasked
         else:
             warnings.append(f"Warning: No <mask> token found in sentence: {sentence}")
     return results, warnings
@@ -25,48 +25,51 @@ st.markdown("Zabantu-XLMR refers to a fleet of models trained on different combi
 
 col1, col2 = st.columns(2)
 
-if 'text_input' not in st.session_state:
-    st.session_state['text_input'] = ""
-
+# Initialize session states
+if 'submit_clicked' not in st.session_state:
+    st.session_state['submit_clicked'] = False
 if 'warnings' not in st.session_state:
     st.session_state['warnings'] = []
 
-if 'result' not in st.session_state:
-    st.session_state['result'] = {}
-
 language_options = ['Choose language', 'Zulu', 'Tshivenda', 'Sepedi', 'Tswana', 'Tsonga']
 
-input_sentences = {}
-
 with col1:
     with st.container():
         st.markdown("Input :clipboard:")
 
         input1, input2 = st.columns(2)
         
+        input_sentences = {}
         for i in range(5):
-             with input1:
+            with input1:
                 language = st.selectbox(f"Select language for sentence {i+1}:", language_options, key=f'language_{i}')
             with input2:
+                # Disable text input if language is not selected
                 disabled = True if language == "Choose language" else False
                 sentence = st.text_input(f"Enter sentence for {language} (with <mask>):", key=f'text_input_{i}', disabled=disabled)
                 if not disabled and sentence:
                     input_sentences[language.lower()] = sentence
 
         button1, button2, _ = st.columns([2, 2, 4])
-
-        if st.button("Test Example"):
-            sample_sentences = {
-                'zulu_1': ('zulu', "Le ndoda ithi izo <mask> ukudla."),
-                'tshivenda_2': ('tshivenda', "Vhana vhane vha kha ḓi bva u bebwa vha kha khombo ya u <mask> nga Listeriosis."),
-                'tshivenda_3': ('tshivenda', "Rabulasi wa <mask> u khou bvelela nga u lima"),
-                'tswana_4': ('tswana', "Monna o <mask> tsamaya."),
-                'tsonga_5': ('tsonga', "N'wana wa xisati u <mask> ku tsaka.")
-            }
-            st.session_state['result'], st.session_state['warnings'] = fill_mask(sample_sentences)
-
-        if st.button("Submit"):
-            st.session_state['result'], st.session_state['warnings'] = fill_mask(input_sentences)
+        
+        with button1:
+            if st.button("Test Example"):
+                sample_sentence = {
+                    'zulu': "Le ndoda ithi izo <mask> ukudla.",
+                    'tshivenda': "Vhana vhane vha kha ḓi bva u bebwa vha kha khombo ya u <mask> nga Listeriosis.",
+                    'tshivenda': "Rabulasi wa <mask> u khou bvelela nga u lima",
+                    'tswana': "Monna o <mask> tsamaya.",
+                    'tsonga': "N'wana wa xisati u <mask> ku tsaka."
+                }
+                input_sentences = sample_sentence
+                result, warnings = fill_mask(input_sentences)
+        
+        with button2:
+            # Set session state when "Submit" is clicked
+            if st.button("Submit"):
+                st.session_state['submit_clicked'] = True
+                result, warnings = fill_mask(input_sentences)
+                st.session_state['warnings'] = warnings
 
         if st.session_state['warnings']:
             for warning in st.session_state['warnings']:
@@ -84,25 +87,29 @@ with col1:
 with col2:
     with st.container():
         st.markdown("Output :bar_chart:")
-        if st.session_state['result']:
-            for key, (language, predictions) in st.session_state['result'].items():
-                original_sentence = input_sentences[key][1]
-                predicted_word = predictions[0]['token_str']
-                score = predictions[0]['score'] * 100
-
-                st.markdown(f"""
-                <div class="bar">
-                    <div class="bar-fill" style="width: {score}%;"></div>
-                </div>
-                <div class="container">
-                    <div style="align-items: left;">{predicted_word} ({language})</div>
-                    <div style="align-items: right;">{score:.2f}%</div>
-                </div>
-                """, unsafe_allow_html=True)
-
-                predicted_sentence = replace_mask(original_sentence, predicted_word)
-                st.write(f"{language}: {predicted_sentence}\n")
 
+        # Ensure output only runs after "Submit" is clicked
+        if st.session_state['submit_clicked'] and input_sentences:
+            for language, sentence in input_sentences.items():
+                masked_sentence = sentence.replace('<mask>', unmasker.tokenizer.mask_token)
+                predictions = unmasker(masked_sentence)
+
+                if predictions:
+                    top_prediction = predictions[0]
+                    predicted_word = top_prediction['token_str']
+                    score = top_prediction['score'] * 100
+
+                    st.markdown(f"""
+                    <div class="bar">
+                        <div class="bar-fill" style="width: {score}%;"></div>
+                    </div>
+                    <div class="container">
+                        <div style="align-items: left;">{predicted_word} ({language})</div>
+                        <div style="align-items: right;">{score:.2f}%</div>
+                    </div>
+                    """, unsafe_allow_html=True)
+
+# CSS to hide footer and style the output
 css = """
 <style>
 footer {display:none !important;}