Update src/streamlit_app.py

src/streamlit_app.py

@@ -16,6 +16,10 @@ def instantiate_roberta(top_k : int, text : str) -> dict:
     pipe = pipeline("fill-mask", model="Iscte-Sintra/RoBERTa-Kriolu", tokenizer="Iscte-Sintra/RoBERTa-Kriolu", token=token)
     return pipe(text, top_k=top_k)
 
+def instatiate_albertina(top_k : int, text : str) -> dict:
+    pipe = pipeline("fill-mask", model="Iscte-Sintra/Albertina-Kriolu", tokenizer="Iscte-Sintra/Albertina-Kriolu", token=token)
+    return pipe(text, top_k=top_k)
+
 
 def build_gpt2_page():
     try:
@@ -66,13 +70,48 @@ def build_roberta_page():
         else:
             predicted_text = st.text_input("Predicted Token", disabled=True)
 
+def build_albertina_page():
+    st.title("Albertina : Encoder")
+
+    top_k = st.sidebar.number_input('Number of predictions to return', min_value=1, max_value=5, value=1, step=1)
+
+    st.write("Enter a sentence with a **<mask>** token, and the model will predict the missing word.")
+
+    results = None
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        st.subheader("Input")
+        input_text = st.text_input("Input Sentence", "Katxor sta trás di <mask>.")
+
+        submit = st.button("Submit")
+        try:
+            if submit and input_text:
+                results = instantiate_roberta(top_k, input_text)
+        except Exception as e:
+            st.warning('There must be a special token "<mask>" in sentence!', icon="⚠️")
+            st.warning(e)
+
+    with col2:
+        st.subheader("Prediction")
+        if results:
+            predicted_text = st.text_input("Predicted Token", value=results[0]['sequence'], disabled=True)
+            for result in results:
+                st.write(f"**Prediction**: {result['token_str']} | **Confidence**: {round(result['score'], 4)}")
+        else:
+            predicted_text = st.text_input("Predicted Token", disabled=True)
+
+
 # Your dictionary of models
-model_dict = {'RoBERTa': 1, 'GPT-2': 2}
+model_dict = {'RoBERTa': 1,'Albertina':2 ,'GPT-2': 3}
 
 # Always appears at the top of the sidebar
 selected_model = st.sidebar.selectbox("Architecture", list(model_dict.keys()))
 
 if model_dict[selected_model] == 1:
     build_roberta_page()
+elif model_dict[selected_model] == 2:
+    build_albertina_page()
 else:
     build_gpt2_page()