Update app.py

app.py

@@ -21,8 +21,6 @@ def replace_mask(sentence, predicted_word):
     return sentence.replace("<mask>", f"**{predicted_word}**")
 
 st.title("Fill Mask | Zabantu-XLM-Roberta")
-st.write("")
-
 st.markdown("Zabantu-XLMR refers to a fleet of models trained on different combinations of South African Bantu languages.")
 
 col1, col2 = st.columns(2)
@@ -33,56 +31,63 @@ if 'text_input' not in st.session_state:
 if 'warnings' not in st.session_state:
     st.session_state['warnings'] = []
 
+language_options = ['Zulu', 'Tshivenda', 'Sepedi', 'Tswana', 'Tsonga']
+
 with col1:
-    with st.container(border=True): 
-        st.markdown("Input :clipboard:")
-        sample_sentence = {
-            '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."  
-        }
-
-        text_input = st.text_area(
-            "Enter sentences with <mask> token:",
-            value=st.session_state['text_input']
-        )
-    
-        input_sentences = {f"sentence_{i}": sentence for i, sentence in enumerate(text_input.split("\n")) if sentence}
+    with st.container():
+        st.markdown("### Input :clipboard:")
+        
+        input_sentences = {}
+        for i in range(5):
+            language = st.selectbox(f"Select language for sentence {i+1}:", language_options, key=f'language_{i}')
+            sentence = st.text_area(f"Enter sentence for {language} (with <mask>):", key=f'text_input_{i}')
+            if sentence:
+                input_sentences[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': "Mufana uyo <mask> vhukuma.",
+                    'sepedi': "Mosadi o <mask> pheka.",
+                    'tswana': "Monna o <mask> tsamaya.",
+                    'tsonga': "N'wana wa xisati u <mask> ku tsaka."
+                }
                 input_sentences = sample_sentence
-                result, warnings = fill_mask(input_sentences) 
-                # st.session_state['text_input'] = "\n".join([f"{lang}: {sentence}" for lang, sentence in sample_sentence.items()])
+                result, warnings = fill_mask(input_sentences)
         
         with button2:
             if st.button("Submit"):
                 result, warnings = fill_mask(input_sentences)
-                st.session_state['warnings'] = warnings  
-    
+                st.session_state['warnings'] = warnings
+
         if st.session_state['warnings']:
             for warning in st.session_state['warnings']:
                 st.warning(warning)
 
-        st.markdown("Example")
-        st.code(sample_sentence, wrap_lines=True)
+        st.markdown("### Example")
+        st.code({
+            '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."
+        }, wrap_lines=True)
 
 with col2:
-    with st.container(border=True): 
-        st.markdown("Output :bar_chart:")
+    with st.container():
+        st.markdown("### Output :bar_chart:")
         if 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>
@@ -96,8 +101,8 @@ with col2:
 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'])  
+            original_sentence = input_sentences[language]
+            predicted_sentence = replace_mask(original_sentence, language_predictions[0]['token_str'])
             st.write(f"{language}: {predicted_sentence}\n")
 
 css = """
@@ -125,5 +130,4 @@ footer {display:none !important;}
 }
 </style>
 """
-
 st.markdown(css, unsafe_allow_html=True)