Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,75 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import transformers
+from transformers import pipeline
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+def infer(sent, max_length, num_return_sequences):
+    generator = pipeline('text-generation', model='gpt2')
+
+    return generator(sent, max_length=max_length, num_return_sequences=num_return_sequences)
+
+def instantiate_gpt2(max_length_ : int, num_return_sequences : int, text : str) -> dict:
+    pipe = pipeline(task='text-generation', model='Iscte-Sintra/GPT2-Kriolu')
+    results = pipe(text, max_length=max_length_, num_return_sequences=num_return_sequences)
+    return results
+
+def instantiate_roberta(top_k : int, text : str) -> dict:
+    pipe = pipeline("fill-mask", model="Iscte-Sintra/RoBERTa-Kriolu")
+    return pipe(text, top_k=top_k)
+
+def build_gpt2_page():
+    try:
+        st.title("GPT-2 : Decoder")
+        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)
+            if results:
+                for result in results:
+                    st.write(f"**Generated Text**: {result['generated_text']}")
+    except Exception as e:
+        st.warning('Max length must be greater than default sentence number of tokens!', icon="⚠️")
+
+def build_roberta_page():
+
+    st.title("RoBERTa : 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", "Nha ca<mask>, nha regra!")
+
+        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="⚠️")
+
+    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}
+
+# 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()
+else:
+    build_gpt2_page()