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>
 """