Update app.py

app.py

@@ -5,14 +5,17 @@ unmasker = pipeline('fill-mask', model='dsfsi/zabantu-xlm-roberta')
 
 st.set_page_config(layout="wide")
 
+# Disable auto-rerun when selecting options or typing
+st.stop_rerun = True
+
 def fill_mask(sentences):
     results = {}
     warnings = []
-    for language, sentence in sentences.items():
+    for key, (language, sentence) in sentences.items():
         if "<mask>" in sentence:
             masked_sentence = sentence.replace('<mask>', unmasker.tokenizer.mask_token)
             unmasked = unmasker(masked_sentence)
-            results[language] = unmasked
+            results[key] = (language, unmasked)
         else:
             warnings.append(f"Warning: No <mask> token found in sentence: {sentence}")
     return results, warnings
@@ -25,52 +28,52 @@ st.markdown("Zabantu-XLMR refers to a fleet of models trained on different combi
 
 col1, col2 = st.columns(2)
 
-# Initialize session states
-if 'submit_clicked' not in st.session_state:
-    st.session_state['submit_clicked'] = False
+if 'text_input' not in st.session_state:
+    st.session_state['text_input'] = ""
+
 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 = {}
+        # Loop to gather input sentences
         for i in range(5):
             with input1:
-                language = st.selectbox(f"Select language for sentence {i+1}:", language_options, key=f'language_{i}')
+                language = st.selectbox(f"Select language for sentence {i+1}:", language_options, key=f'language_{i}', index=0)
             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
+                sentence = st.text_input(f"Enter sentence for {language} (with <mask>):", key=f'text_input_{i}')
+                if sentence:
+                    # Use a unique key for each sentence (even if languages are the same)
+                    input_sentences[f'{language.lower()}_{i+1}'] = (language.lower(), sentence)
 
         button1, button2, _ = st.columns([2, 2, 4])
-        
-        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
 
+        # Call fill_mask on button click, not on form input
+        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)
+
+        # Display warnings
         if st.session_state['warnings']:
             for warning in st.session_state['warnings']:
                 st.warning(warning)
@@ -87,13 +90,10 @@ with col1:
 with col2:
     with st.container():
         st.markdown("Output :bar_chart:")
-
-        # 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)
-
+        # Check for the result in session_state and display predictions
+        if st.session_state['result']:
+            for key, (language, predictions) in st.session_state['result'].items():
+                original_sentence = input_sentences[key][1] if key in input_sentences else ""
                 if predictions:
                     top_prediction = predictions[0]
                     predicted_word = top_prediction['token_str']
@@ -109,7 +109,9 @@ with col2:
                     </div>
                     """, unsafe_allow_html=True)
 
-# CSS to hide footer and style the output
+                    predicted_sentence = replace_mask(original_sentence, predicted_word)
+                    st.write(f"{language}: {predicted_sentence}\n")
+
 css = """
 <style>
 footer {display:none !important;}