update

pages/1_Injection.py

@@ -26,9 +26,9 @@ def check_password():
 def initialize_state():
     keys = ["model_submitted", "api_key", "endpoint_url", "deployment_name", "temperature", "max_tokens",
             "data_processed", "group_name", "occupation", "privilege_label", "protect_label", "num_run",
-            "uploaded_file", "occupation_submitted"]
+            "uploaded_file", "occupation_submitted","sample_size"]
     defaults = [False, "", "https://safeguard-monitor.openai.azure.com/", "gpt35-1106", 0.0, 150, False, "Gender",
-                "Programmer", "Male", "Female", 1, None, False]
+                "Programmer", "Male", "Female", 1, None, False,0]
     for key, default in zip(keys, defaults):
         if key not in st.session_state:
             st.session_state[key] = default
@@ -58,43 +58,54 @@ else:
     if st.sidebar.button("Submit Model Info"):
         st.session_state.model_submitted = True
 
-    categories = ["HR", "DESIGNER", "INFORMATION-TECHNOLOGY", "TEACHER", "ADVOCATE", "BUSINESS-DEVELOPMENT",
-                  "HEALTHCARE", "FITNESS", "AGRICULTURE", "BPO", "SALES", "CONSULTANT", "DIGITAL-MEDIA",
-                  "AUTOMOBILE", "CHEF", "FINANCE", "APPAREL", "ENGINEERING", "ACCOUNTANT", "CONSTRUCTION",
-                  "PUBLIC-RELATIONS", "BANKING", "ARTS", "AVIATION"]
-
-    st.session_state.occupation = st.selectbox("Occupation", options=categories, index=categories.index(
-        st.session_state.occupation) if st.session_state.occupation in categories else 0)
-
-    if st.button("Submit Occupation Selection"):
-        st.session_state.occupation_submitted = True
+    # categories = ["HR", "DESIGNER", "INFORMATION-TECHNOLOGY", "TEACHER", "ADVOCATE", "BUSINESS-DEVELOPMENT",
+    #               "HEALTHCARE", "FITNESS", "AGRICULTURE", "BPO", "SALES", "CONSULTANT", "DIGITAL-MEDIA",
+    #               "AUTOMOBILE", "CHEF", "FINANCE", "APPAREL", "ENGINEERING", "ACCOUNTANT", "CONSTRUCTION",
+    #               "PUBLIC-RELATIONS", "BANKING", "ARTS", "AVIATION"]
+    #
+    # st.session_state.occupation = st.selectbox("Occupation", options=categories, index=categories.index(
+    #     st.session_state.occupation) if st.session_state.occupation in categories else 0)
+    #
+    # if st.button("Submit Occupation Selection"):
+    #     st.session_state.occupation_submitted = True
 
     # Ensure experiment settings are only shown if model info is submitted
-    if st.session_state.model_submitted and st.session_state.occupation_submitted:
+    if st.session_state.model_submitted:# and st.session_state.occupation_submitted:
 
         df = None
         file_options = st.radio("Choose file source:", ["Upload", "Example"])
         if file_options == "Example":
             #df = pd.read_csv("prompt_test.csv")
             df = pd.read_csv("resume.csv")
-            df = df[df["Occupation"] == st.session_state.occupation]
         else:
             st.session_state.uploaded_file = st.file_uploader("Choose a file")
             if st.session_state.uploaded_file is not None:
                 data = StringIO(st.session_state.uploaded_file.getvalue().decode("utf-8"))
                 df = pd.read_csv(data)
-        if df is not None:
-            st.write('Data:', df)
 
-            # Button to add a new row
+        if df is not None:
 
             #st.session_state.occupation = st.text_input("Occupation", value=st.session_state.occupation)
 
+            categories = ["HR", "DESIGNER", "INFORMATION-TECHNOLOGY", "TEACHER", "ADVOCATE", "BUSINESS-DEVELOPMENT",
+                            "HEALTHCARE", "FITNESS", "AGRICULTURE", "BPO", "SALES", "CONSULTANT", "DIGITAL-MEDIA",
+                            "AUTOMOBILE", "CHEF", "FINANCE", "APPAREL", "ENGINEERING", "ACCOUNTANT", "CONSTRUCTION",
+                            "PUBLIC-RELATIONS", "BANKING", "ARTS", "AVIATION"]
+
+            st.session_state.occupation = st.selectbox("Occupation", options=categories, index=categories.index(st.session_state.occupation) if st.session_state.occupation in categories else 0)
+
+            st.session_state.sample_size = st.number_input("Sample Size", 1, len(df), st.session_state.sample_size)
+
             st.session_state.group_name = st.text_input("Group Name", value=st.session_state.group_name)
             st.session_state.privilege_label = st.text_input("Privilege Label", value=st.session_state.privilege_label)
+
             st.session_state.protect_label = st.text_input("Protect Label", value=st.session_state.protect_label)
             st.session_state.num_run = st.number_input("Number of Runs", 1, 10, st.session_state.num_run)
 
+            df = df[df["Occupation"] == st.session_state.occupation]
+            df = df.sample(n=st.session_state.sample_size)
+            st.write('Data:', df)
+
             if st.button('Process Data') and not st.session_state.data_processed:
                 # Initialize the correct agent based on model type
                 if model_type == 'AzureAgent':