Update src/streamlit_app.py

src/streamlit_app.py

@@ -16,6 +16,50 @@ def instantiate_encoder(model_name: str, top_k : int, text : str) -> dict:
     pipe = pipeline("fill-mask", model=f"Iscte-Sintra/{model_name}", tokenizer=f"Iscte-Sintra/{model_name}", token=token)
     return pipe(text, top_k=top_k)
 
+def instantiate_encoder_decoder(model_name: str, top_k : int, text : str) -> dict:
+    pipe = pipeline("translation", model=f"Iscte-Sintra/{model_name}", tokenizer=f"Iscte-Sintra/{model_name}", token=token)
+    return pipe(text, top_k=top_k)
+
+
+def build_translation_page(model_name):
+    try:
+        st.title(f"{model_name} : Translation Task")
+
+        # Sidebar options
+        src_lang = st.sidebar.selectbox(
+            "Source Language",
+            ["auto", "en", "fr", "de", "es", "pt", "it", "nl", "zh", "ar"],
+            index=0,
+        )
+
+        tgt_lang = st.sidebar.selectbox(
+            "Target Language",
+            ["en", "fr", "de", "es", "pt", "it", "nl", "zh", "ar"],
+            index=4,  # default to Portuguese in this example
+        )
+
+        text = st.text_area("Enter text to translate", "Katxor sta trás di pórta.", height=100)
+
+        if st.button("Translate"):
+            if not text.strip():
+                st.warning("Please enter some text to translate!", icon="⚠️")
+                return
+
+            # Call translation model (you can define your own implementation)
+            results = instantiate_translation_model(model_name, text, src_lang, tgt_lang)
+
+            if results:
+                st.subheader("Translation Result")
+                if isinstance(results, list):
+                    for i, result in enumerate(results, start=1):
+                        st.write(f"**Translation {i}:** {result}")
+                else:
+                    st.write(f"**Translated Text:** {results}")
+    except Exception as e:
+        st.warning("An error occurred during translation!", icon="⚠️")
+        st.warning(e)
+
+
 
 def build_decoder_page(model_name):
     try:
@@ -72,12 +116,15 @@ def build_encoder_page(model_name:str):
 
 # Your dictionary of models
 model_dict = {'RoBERTa-Kriolu': "Encoder",'Albertina-Kriolu':"Encoder" ,
-              'GPT2-Kriolu': "Decoder", "GPT2-Small": "Decoder","GPT2_v1.15": "Decoder", "GPT2_v1.17":"Decoder","GPT2_v1.18":"Decoder" }
+              'GPT2-Kriolu': "Decoder", "GPT2_v1.15": "Decoder", "GPT2_v1.17":"Decoder","GPT2_v1.18":"Decoder", "GPT2-small":"Encoder-Decoder"}
 
 # Always appears at the top of the sidebar
 selected_model = st.sidebar.selectbox("Architecture", list(model_dict.keys()))
 
 if model_dict[selected_model] == "Encoder":
     build_encoder_page(selected_model)
+
+elif model_dict[selected_model] == "Encoder-Decoder":
+    build_translation_page(selected_model)
 else:
     build_decoder_page(selected_model)