add checkbox

app.py

@@ -3,12 +3,27 @@ from transformers import pipeline
 import utils
 
 st.title("Indo LEGO-ABSA")
-example_text = "EXAMPLE: Pizza nya enak sekali| aspect: <extra_id_0>, opinion: <extra_id_1>, sentiment: <extra_id_2>"
+example_text = "EXAMPLE: Pizza nya enak sekali"
 text = st.text_area("Enter your text here", "", placeholder=example_text)
+aspect = st.checkbox("Aspect", value=True)
+opinion = st.checkbox("Opinion")
+sentiment = st.checkbox("Sentiment")
 pipe = pipeline(model="rdyzakya/IndoLEGO-ABSA")
+prompt = ""
+se_order = ""
 
 if text:
-    out = pipe(text)
-    se_order = utils.get_se_order(text)
-    result = utils.catch_answer(out[0]["generated_text"], se_order)
-    st.json(result)
+    if aspect:
+        se_order += 'a'
+    if opinion:
+        se_order += 'o'
+    if sentiment:
+        se_order += 's'
+    se_order = set(se_order)
+    if not se_order.issubset("acos"):
+        st.json([])
+    else:
+        inputs = utils.add_prompt(text, se_order)
+        out = pipe(inputs)
+        result = utils.catch_answer(out[0]["generated_text"], str(se_order))
+        st.json(result)

utils.py

@@ -1,13 +1,12 @@
 import re
 import constant
-mask_pattern = r"<extra_id_\d+>"
-def get_se_order(text):
-    prompt = text.split('|')[-1].strip()
-    sentiment_elements = prompt.split(',')
-    for i in range(len(sentiment_elements)):
-        sentiment_elements[i] = re.sub(mask_pattern, '', sentiment_elements[i]).strip()
-    se_order = ''.join([se[0] for se in sentiment_elements])
-    return se_order
+def add_prompt(text, se_order="aos"):
+    prompt = []
+    for counter, se in enumerate(se_order):
+        prompt.append(constant.SENTIMENT_ELEMENT[se] + " : " + f"<extra_id_{counter}>")
+    prompt = " ,".join(prompt)
+    result = text + "| " + prompt
+    return result
 
 def catch_answer(out, se_order):
         if out == constant.NO_TARGET:
@@ -22,10 +21,4 @@ def catch_answer(out, se_order):
         for i in range(len(result)):
             for k, v in result[i].items():
                 result[i][k] = v.strip()
-        return result
-
-if __name__ == "__main__":
-    print(catch_answer(
-        '<extra_id_0> Pizza nya <extra_id_1> enak sekali <extra_id_2> positive',
-        'aos'
-    ))
+        return result