improved demo

app.py

@@ -12,52 +12,78 @@ img {
 th {
     text-align: left!important
 }
+
+td {
+    font-size:
+}
 </style>
 """
+
 MASK_TOKEN = "<mask>"
 
+EMOTION_MAP = {
+    "anger": "😡",
+    "fear": "😱",
+    "happy": "😄",
+    "love": "😍",
+    "sadness": "😭",
+}
 
-def display_table(df):
-    st.subheader("Top 5 Prediction.")
-    df.drop(columns=["token", "token_str"], inplace=True)
-    df = df.style.set_properties(subset=["sequence", "score"], **{"text-align": "left"})
+
+def display_table(df: pd.DataFrame, subheader: str):
+    st.subheader(subheader)
     st.table(df)
 
 
-def main():
+def setup():
     st.markdown(STYLE, unsafe_allow_html=True)
-    st.title("Indonesian RoBERTa Base")
-    user_input = st.text_input("Insert a sentence to predict with a mask token: <mask>")
-    mask_api = InferenceApi("flax-community/indonesian-roberta-base")
+    st.title("🇮🇩 Indonesian RoBERTa Base")
+
 
-    emot_name = "StevenLimcorn/indonesian-roberta-base-emotion-classifier"
-    emot_pipeline = pipeline("sentiment-analysis", model=emot_name, tokenizer=emot_name)
+def main():
+    setup()
+
+    user_input = st.text_input(
+        f"Insert a sentence to predict with a {MASK_TOKEN} token // Masukkan kalimat untuk diisi dengan token {MASK_TOKEN}",
+        value=f"Aduh... gimana nih.. hari ini {MASK_TOKEN} banget...",
+    )
+
+    mlm_model = "flax-community/indonesian-roberta-base"
+    mask_api = InferenceApi(mlm_model)
+
+    sa_model = "StevenLimcorn/indonesian-roberta-base-emotion-classifier"
+    sa_pipeline = pipeline("sentiment-analysis", model=sa_model, tokenizer=sa_model)
 
     if len(user_input) > 0:
         try:
             user_input.index(MASK_TOKEN)
         except ValueError:
-            st.error("Please enter a sentence with the correct mask token: <mask>")
+            st.error(
+                f"Please enter a sentence with the correct {MASK_TOKEN} token // Harap masukkan kalimat dengan token {MASK_TOKEN} yang benar"
+            )
         else:
-            # A List of dict with keys: sequence, score, token, token_str
-            result = mask_api(inputs=user_input)
-            df = pd.DataFrame(result)
-            display_table(df)
-
-            # emot
-            st.subheader("Emotion Analysis of the Top 5 Prediction")
-            emot_df = pd.DataFrame(columns=["sequence", "label", "score"])
-            for sequence in df["sequence"].values:
-                emot_output = emot_pipeline(sequence)
+            # render masked language modeling table
+            mlm_result = mask_api(inputs=user_input)
+            mlm_df = pd.DataFrame(mlm_result)
+            mlm_df.drop(columns=["token", "token_str"], inplace=True)
+            mlm_df_styled = mlm_df.style.set_properties(
+                subset=["sequence", "score"], **{"text-align": "left"}
+            )
+            display_table(mlm_df_styled, "🎈 Top 5 Predictions")
 
+            # render sentiment analysis table
+            sa_df = pd.DataFrame(columns=["sequence", "label", "score"])
+            for sequence in mlm_df["sequence"].values:
+                sa_output = sa_pipeline(sequence)  # predict for every mlm output
                 result_dict = {"sequence": sequence}
-                result_dict.update(emot_output[0])
-                emot_df = emot_df.append(result_dict, ignore_index=True)
+                result_dict.update(sa_output[0])
+                sa_df = sa_df.append(result_dict, ignore_index=True)
 
-            emot_df = emot_df.style.set_properties(
+            sa_df["label"] = sa_df["label"].apply(lambda x: x + " " + EMOTION_MAP[x])
+            sa_df_styled = sa_df.style.set_properties(
                 subset=["sequence", "label", "score"], **{"text-align": "left"}
             )
-            st.table(emot_df)
+            display_table(sa_df_styled, "🤔 By saying that, I guess you are feeling..")
 
 
 main()