Update app.py

app.py

@@ -1,43 +1,37 @@
 import streamlit as st
 from transformers import pipeline
 
-unmasker = pipeline('fill-mask', model='dsfsi/zabantu-sot-ven-170m')
+unmasker = pipeline('fill-mask', model='dsfsi/zabantu-ven-120m'')
 
-def fill_mask(sentences):
-    results = {}
-    for sentence in sentences:
-        unmasked = unmasker(sentence)
-        results[sentence] = unmasked
-    return results
+def fill_mask(text):
+    unmasked = unmasker(text)
+    output = {}
+    for unmask in unmasked:
+        output[unmask["token_str"]] = unmask["score"]
+    return output
 
-def replace_mask(sentence, predicted_word):
-    return sentence.replace("<mask>", predicted_word)
+st.title("Fill Mask | Zabantu-Bantu-250m")
+st.write("This app predicts the missing word in a sentence using a BERT model.")
 
-st.title("Fill Mask | Zabantu-sot-ven-170m")
-st.write("This app predicts the missing word in a sentence using a Zabantu-sot-ven-170m model.")
+col1, col2 = st.columns(2)
 
-sample_sentences = ["Rabulasi wa <mask> u khou bvelela nga u lima,",
+with col1:
+    sample_sentences = ["Rabulasi wa <mask> u khou bvelela nga u lima,",
                     "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 (one per line):", 
+    text_input = st.text_area("Enter sentences with <mask> token (one per line):", 
                           "\n".join(sample_sentences))
 
-input_sentences = text_input.split(",")
+    input_sentences = text_input.split(",")
 
-if st.button("Submit"):
-    result = fill_mask(input_sentences)
+with col2:
+    if st.button("Submit"):
+        result = fill_mask(text_input)
+        
+        st.write("Predicted words and their probabilities:")
+        for token, score in result.items():
+            st.write(f"**{token}**: {score * 100:.2f}%")
 
-    if result:
-        for sentence, predictions in result.items():
-            st.write(f"**Original sentence**: {sentence}")
-            for prediction in predictions:
-                predicted_word = prediction['token_str']
-                score = prediction['score'] * 100
-                full_sentence = replace_mask(sentence, predicted_word)
-                
-                st.write(f"Predicted word: {predicted_word} - Score: {score:.2f}%")
-                st.write(f"Full sentence: {full_sentence}")
-                st.write("=" * 80)
 
 css = """
 <style>
@@ -62,7 +56,7 @@ footer {display:none !important}
     background-color: #3c4a6b;
 }
 
-.stTextInput, .stTextArea {
+.stTextInput {
     border: 1px solid #e6e6e6;
     padding: 0.75em;
     border-radius: 10px;
@@ -70,7 +64,7 @@ footer {display:none !important}
     width: 100%;
 }
 
-.stTextInput:focus, .stTextArea:focus {
+.stTextInput:focus {
     border-color: #17152e;
     outline: none;
     box-shadow: 0px 0px 5px rgba(23, 21, 46, 0.5);
@@ -84,6 +78,32 @@ div[data-testid="stMarkdownContainer"] p {
     padding: 2em;
     font-family: 'Poppins', sans-serif;
 }
+
+.examples {
+    margin-top: 2em;
+    margin-bottom: 2em;
+}
+
+.examples button {
+    background-color: #f0f0f0;
+    border: 1px solid #e6e6e6;
+    padding: 0.5em 1em;
+    margin: 5px;
+    border-radius: 8px;
+    cursor: pointer;
+}
+
+.examples button:hover {
+    background-color: #e0e0e0;
+}
+
+.stMarkdown h3 {
+    margin-bottom: 0.75em;
+    font-size: 24px;
+    font-weight: 500;
+}
 </style>
 """
+
+# Inject the custom CSS into the Streamlit app
 st.markdown(css, unsafe_allow_html=True)