Update app.py

app.py

@@ -5,7 +5,6 @@ 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):
@@ -47,33 +46,29 @@ with col1:
 
         input1, input2 = st.columns(2)
         
-        # 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}', index=0)
             with input2:
                 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])
 
-        # 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.")
+                'zulu': "Le ndoda ithi izo <mask> ukudla.",
+                'tshivenda': "Mufana uyo <mask> vhukuma.",
+                'sepedi': "Mosadi o <mask> pheka.",
+                'tswana': "Monna o <mask> tsamaya.",
+                '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)
@@ -90,7 +85,6 @@ with col1:
 with col2:
     with st.container():
         st.markdown("Output :bar_chart:")
-        # 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 ""
@@ -109,8 +103,12 @@ with col2:
                     </div>
                     """, unsafe_allow_html=True)
 
-                    predicted_sentence = replace_mask(original_sentence, predicted_word)
-                    st.write(f"{language}: {predicted_sentence}\n")
+if 'predictions' in locals():  
+    if result:
+        for language, language_predictions in result.items(): 
+            original_sentence = sample_sentence[language]
+            predicted_sentence = replace_mask(original_sentence, language_predictions[0]['token_str'])  
+            st.write(f"{language}: {predicted_sentence}\n")
 
 css = """
 <style>
@@ -137,4 +135,4 @@ footer {display:none !important;}
 }
 </style>
 """
-st.markdown(css, unsafe_allow_html=True)
+st.markdown(css, unsafe_allow_html=True)