Update app.py

app.py

@@ -1,17 +1,117 @@
 import streamlit as st
-q1 = st.text_area(placeholder='Question 1 ...', key='q1', label='Question 1')
-a1 = st.text_area(placeholder='Answer for question 1 ...', key='a1', label='Answer for question 1')
+import pandas as pd
+import os
+import requests
 
-q2 = st.text_area(placeholder='Question 2 ...', key='q2', label='Question 2')
-a2 = st.text_area(placeholder='Answer for question 2 ...', key='a2', label='Answer for question 2')
+st.set_page_config(page_icon='🦜', page_title='Text Generation Labeling Tool', layout='wide', initial_sidebar_state="collapsed")
+st.markdown("<h1 style='text-align: center;'>Text Generation Labeling Tool</h1>", unsafe_allow_html=True)
 
-q3 = st.text_area(placeholder='Question 3 ...', key='q3', label='Question 3')
-a3 = st.text_area(placeholder='Answer for question 3 ...', key='a3', label='Answer for question 3')
+def file_selector(folder_path=r'./Datasets'):
+    filenames = os.listdir(folder_path)
+    return filenames, folder_path
 
-q4 = st.text_area(placeholder='Question 4 ...', key='q4', label='Question 4')
-a4 = st.text_area(placeholder='Answer for question 4 ...', key='a4', label='Answer for question 4')
+def revert_question_type_id(txt_question_type):
+    if txt_question_type == 'What':
+        return 0
+    elif txt_question_type == 'Who':
+        return 1
+    elif txt_question_type == 'When':
+        return 2
+    elif txt_question_type == 'Where':
+        return 3
+    elif txt_question_type == 'Why':
+        return 4
+    elif txt_question_type == 'How':
+        return 5
+    elif txt_question_type == 'Others':
+        return 6
 
-btn_gen_str = st.button('Generate')
+filenames, folder_path = file_selector()
+filename_input = st.sidebar.selectbox(label='Input dataset file:', options=filenames)
+df = pd.read_csv(f'./{folder_path}/{filename_input}')
 
-if btn_gen_str:
-    st.success("question: " + q1 + ", answer: " + a1 + "[SEP] " + "question: " + q2 + ", answer: " + a2 + "[SEP] " + "question: " + q3 + ", answer: " + a3 + "[SEP] " +"question: " + q4 + ", answer: " + a4 + "[SEP]")
+if 'idx' not in st.session_state:
+    st.session_state.idx = 0
+
+st.markdown(f"<h4 style='text-align: center;'>Sample {st.session_state.idx + 1}/{len(df)}</h4>", unsafe_allow_html=True)
+
+col_1, col_2, col_3, col_4, col_5, col_6, col_7, col_8, col_9, col_10 = st.columns([1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
+
+btn_previous = col_1.button(label=':arrow_backward: Previous sample', use_container_width=True)
+btn_next = col_2.button(label='Next sample :arrow_forward:', use_container_width=True)
+btn_save = col_3.button(label=':heavy_check_mark: Save change', use_container_width=True)
+txt_goto = col_4.selectbox(label='Sample', label_visibility='collapsed', options=list(range(1, len(df) + 1)))
+btn_goto = col_5.button(label=':fast_forward: Move to', use_container_width=True)
+
+if len(df) != 0:
+    col_1, col_2, col_3, col_4, col_5, col_6, col_7, col_8, col_9, col_10 = st.columns(spec=[1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
+    txt_context = st.text_area(height=200, label='Your context:', value=df['context'][st.session_state.idx])
+
+    col_11, col_12, col_13 = st.columns([4.5, 1, 4.5])
+    txt_question = col_11.text_area(height=90, label='Your question:', value=df['question'][st.session_state.idx])
+    txt_question_type = col_12.selectbox(label='Your question type:', options=['What', 'Who', 'When', 'Where', 'Why', 'How', 'Others'], index=int(df['question_type'][st.session_state.idx]))
+    txt_answer = col_13.text_area(height=90, label='Your answer:', value=df['answer'][st.session_state.idx])
+
+    st.markdown(f"<p style='text-align: left; font-weight: normal; font-size: 14px'>Your distractors:</p>", unsafe_allow_html=True)
+
+    col_21, col_22 = st.columns(spec=[9, 1])
+    txt_distractors = col_21.text_area(height=90, label='Your distractors:', label_visibility='collapsed', value=df['distract'][st.session_state.idx])  
+    btn_generate_distractor = col_22.button(label='Generate distractors', use_container_width=True)
+    
+    if btn_generate_distractor:
+        if filename_input == 'BiologyQA.csv':
+            expert = 'biologist'
+        elif filename_input == 'GeographyQA.csv':
+            expert = 'geographer'
+        elif filename_input == 'HistoryQA.csv':
+            expert = 'historian'
+        url = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent'
+        headers = {'Content-Type': 'application/json'}
+        data = {
+            'contents': [
+                {
+                    'parts': [
+                        {
+                            'text': f"You are a great {expert}, here is the following content: context: '{txt_context}', question: '{txt_question}', answer: '{txt_answer}' generate three distract answers. Distractor answers are separated by [SEP]. Example: Distract answer 1 [SEP] Distract answer 2 [SEP] Distract answer 3"
+                        }
+                    ]
+                }
+            ]
+        }
+        api_key = 'AIzaSyApFAbCUA1H-VHAidzqmyStHFe92ODeO1Y'
+        params = {'key': api_key}
+        response = requests.post(url, headers=headers, json=data, params=params)
+        if response.status_code == 200:
+            correct = response.json()['candidates'][0]['content']['parts'][0]['text']
+            st.success(f'3 distraction answers: {correct}')
+            st.cache_data.clear()
+        else:
+            st.error('Failed to generate distractors. Please check API and inputs.')
+            st.rerun()
+
+    if btn_previous:
+        if st.session_state.idx > 0:
+            st.session_state.idx -= 1
+            st.rerun()
+        else:
+            pass
+
+    if btn_next:
+        if st.session_state.idx < (len(df) - 1):
+            st.session_state.idx += 1
+            st.rerun()
+        else:
+            pass
+
+    if btn_save:
+        df['context'][st.session_state.idx] = txt_context
+        df['question'][st.session_state.idx] = txt_question
+        df['answer'][st.session_state.idx] = txt_answer
+        df['distract'][st.session_state.idx] = txt_distractors
+        df['question_type'][st.session_state.idx] = revert_question_type_id(txt_question_type)
+
+        df.to_csv(f'./Datasets/{filename_input}', index=None)
+
+    if btn_goto:
+        st.session_state.idx = txt_goto - 1
+        st.rerun()