Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,8 +7,8 @@ import os
 token = os.environ.get("token")
 
 
-def instantiate_gpt2(max_length_ : int, num_return_sequences : int, text : str) -> dict:
-    pipe = pipeline(task='text-generation', model='Iscte-Sintra/GPT2-Kriolu', tokenizer='Iscte-Sintra/GPT2-Kriolu', token=token, truncation=True)
+def instantiate_gpt2(model_name: str,max_length_ : int, num_return_sequences : int, text : str) -> dict:
+    pipe = pipeline(task='text-generation', model=f'Iscte-Sintra/{model_name}', tokenizer=f'Iscte-Sintra/{model_name}', token=token, truncation=True)
     results = pipe(text, max_length=max_length_, num_return_sequences=num_return_sequences)
     return results
 
@@ -21,15 +21,15 @@ def instantiate_albertina(top_k : int, text : str) -> dict:
     return pipe(text, top_k=top_k)
 
 
-def build_gpt2_page():
+def build_decoder_page(model_name):
     try:
-        st.title("GPT-2 : Decoder")
+        st.title(f"{model_name} : Decoder - Text Generation Task")
         max_length : int = st.sidebar.slider("Max Length", 10, 200)
         num_return_sequences : int = st.sidebar.number_input('Number of Sequences to be Generated', min_value=1, max_value=10, value=1, step=1)
         text : str = st.text_area("Text", "Katxor sta trás di pórta.", height=75)
 
         if st.button("Submit"):
-            results = instantiate_gpt2(max_length, num_return_sequences, text)
+            results = instantiate_gpt2(model_name ,max_length, num_return_sequences, text)
             if results:
                 for result in results:
                     st.write(f"**Generated Text**: {result['generated_text']}")
@@ -114,4 +114,4 @@ if model_dict[selected_model] == 1:
 elif model_dict[selected_model] == 2:
     build_albertina_page()
 else:
-    build_gpt2_page()
+    build_gpt2_page('GPT2-Kriolu')