Update app.py

app.py

@@ -31,15 +31,19 @@ if 'text_input' not in st.session_state:
 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:
                 language = st.selectbox(f"Select language for sentence {i+1}:", language_options, key=f'language_{i}')
@@ -50,22 +54,19 @@ with col1:
                     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_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.")
-                }
-                input_sentences = sample_sentences
-                result, warnings = fill_mask(input_sentences)
-        
-        with button2:
-            if st.button("Submit"):
-                result, warnings = fill_mask(input_sentences)
-                st.session_state['warnings'] = warnings
+
+        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)
 
         if st.session_state['warnings']:
             for warning in st.session_state['warnings']:
@@ -83,32 +84,24 @@ with col1:
 with col2:
     with st.container():
         st.markdown("Output :bar_chart:")
-        if input_sentences:
-            for key, (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)
-
-if 'predictions' in locals():  
-    if result:
-        for key, (language, language_predictions) in result.items(): 
-            original_sentence = input_sentences[key][1]
-            predicted_sentence = replace_mask(original_sentence, language_predictions[0]['token_str'])
-            st.write(f"{language}: {predicted_sentence}\n")
+        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")
 
 css = """
 <style>