Update app.py

app.py

@@ -3,53 +3,83 @@ from transformers import pipeline
 
 unmasker = pipeline('fill-mask', model='dsfsi/zabantu-nso-ven-170m')
 
+st.set_page_config(layout="wide")
+
 def fill_mask(sentences):
     results = {}
+    warnings = []
     for sentence in sentences:
-        unmasked = unmasker(sentence)
-        results[sentence] = unmasked
-    return results
+        if "<mask>" in sentence:
+            unmasked = unmasker(sentence)
+            results[sentence] = unmasked
+        else:
+            warnings.append(f"Warning: No <mask> token found in sentence: {sentence}")
+    return results, warnings
 
 def replace_mask(sentence, predicted_word):
     return sentence.replace("<mask>", f"**{predicted_word}**")
 
-st.title("Fill Mask | Zabantu-nso-ven-170m")
+st.title("Fill Mask | Zabantu-sot-ven-170m")
 st.write(f"")
 
 col1, col2 = st.columns(2)
 
-with col1:
-    sample_sentences = [
-        "Le nna ke gona, ke a <mask>"
-    ]
+if 'text_input' not in st.session_state:
+    st.session_state['text_input'] = ""
 
-    text_input = st.text_area(
-        "Enter sentences with <mask> token (one per line):",
-        "\n".join(sample_sentences)
-    )
+if 'warnings' not in st.session_state:
+    st.session_state['warnings'] = []
+
+with col1:
+    with st.container(border=True):
+        st.markdown("Input :clipboard:")
+        sample_sentence = "Vhana vhane vha kha ḓi bva u bebwa vha kha khombo ya u <mask> nga Listeriosis."
+    
+        text_input = st.text_area(
+            "Enter sentences with <mask> token:",
+            value=st.session_state['text_input']
+        )
+    
+        input_sentences = text_input.split("\n")
 
-    input_sentences = text_input.split("\n")
+        button1, button2, _ = st.columns([1.25, 1, 5.7])
+        with button1:
+            if st.button("Test Example"):
+                # st.rerun()
+                result, warnings = fill_mask(sample_sentence.split("\n"))
+                # st.session_state['text_input'] = sample_sentence
+        
+        with button2:
+            if st.button("Submit"):
+                result, warnings = fill_mask(input_sentences)
+                st.session_state['warnings'] = warnings 
+    
+        if st.session_state['warnings']:
+            for warning in st.session_state['warnings']:
+                st.warning(warning)
 
-    if st.button("Submit"):
-        result = fill_mask(input_sentences)
+        st.markdown("Example")
+        st.code(sample_sentence, wrap_lines=True)
 
 with col2:
-    if 'result' in locals():  
-        if result:
-            for sentence, predictions in result.items():
-                for prediction in predictions:
-                    predicted_word = prediction['token_str']
-                    score = 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}</div>
-                        <div style="align-items: right;">{score:.2f}%</div>
-                    </div>
-                    """, unsafe_allow_html=True)
+    with st.container(border=True):
+        st.markdown("Output :bar_chart:")
+        if 'result' in locals() and result:  
+            if result:
+                for sentence, predictions in result.items():
+                    for prediction in predictions:
+                        predicted_word = prediction['token_str']
+                        score = 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}</div>
+                            <div style="align-items: center;">{score:.2f}%</div>
+                        </div>
+                        """, unsafe_allow_html=True)
 
 if 'result' in locals():  
         if result:
@@ -123,6 +153,7 @@ footer {display:none !important;}
     --tw-text-opacity: 1 !important;
     color:rgb(37 56 133 / var(--tw-text-opacity)) !important;
 }
+
 .container {
     display: flex;
     justify-content: space-between;
@@ -131,7 +162,7 @@ footer {display:none !important;}
     width: 100%;
 }
 .bar {
-    width: 70%;
+    # width: 70%;
     background-color: #e6e6e6;
     border-radius: 12px;
     overflow: hidden;
@@ -143,6 +174,7 @@ footer {display:none !important;}
     height: 100%;
     border-radius: 12px;
 }
+
 </style>
 """