Update src/streamlit_app.py

src/streamlit_app.py

@@ -12,7 +12,7 @@ def instantiate_gpt2(model_name: str,max_length_ : int, num_return_sequences : i
     results = pipe(text, max_length=max_length_, num_return_sequences=num_return_sequences)
     return results
 
-def instantiate_roberta(top_k : int, text : str) -> dict:
+def instantiate_encoder(model_name: str, 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)
 
@@ -37,9 +37,9 @@ def build_decoder_page(model_name):
         st.warning('Max length must be greater than default sentence number of tokens!', icon="⚠️")
         st.warning(e)
 
-def build_roberta_page():
+def build_encoder_page(model_name:str):
 
-    st.title("RoBERTa : Encoder")
+    st.title(f"{model_name} : Encoder - Fill-Mask Task")
 
     top_k = st.sidebar.number_input('Number of predictions to return', min_value=1, max_value=5, value=1, step=1)
 
@@ -56,7 +56,7 @@ def build_roberta_page():
         submit = st.button("Submit")
         try:
             if submit and input_text:
-                results = instantiate_roberta(top_k, input_text)
+                results = instantiate_encoder(model_name, top_k, input_text)
         except Exception as e:
             st.warning('There must be a special token "<mask>" in sentence!', icon="⚠️")
             st.warning(e)
@@ -70,48 +70,14 @@ 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_albertina(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,'Albertina':2 ,'GPT-2': 3}
+model_dict = {'RoBERTa-Kriolu': "Encoder",'Albertina-Kriolu':"Encoder" ,'GPT2-Kriolu': "Decoder"}
 
 # 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()
+if model_dict[selected_model] == "Encoder":
+    build_encoder_page(selected_model)
 else:
-    build_decoder_page('GPT2-Kriolu')
+    build_decoder_page(selected_model)