Update appStore/vulnerability_analysis.py

appStore/vulnerability_analysis.py

@@ -72,23 +72,21 @@ def app():
                 df = vulnerability_classification(haystack_doc=df,
                                             threshold= params['threshold'])
 
-                # Filter the dataframe to only show the paragraphs with references
-                df_filtered = df[df['Vulnerability Label'].apply(lambda x: len(x) > 0 and 'Other' not in x)]
-
-                # Rename column 
-                df_filtered.rename(columns={'Vulnerability Label': 'Group identified'}, inplace=True)
-
 
                 # Store df in session state with key1
-                st.session_state.key1 = df_filtered
+                st.session_state.key1 = df
 
 
 def vulnerability_display(): 
     
-    # Assign dataframe a name
-    df_vul = st.session_state['key0']
+    # Get the vulnerability df
+    df = st.session_state['key1']
 
-    #st.write(df_vul)
+    # Filter the dataframe to only show the paragraphs with references
+    df_filtered = df[df['Vulnerability Label'].apply(lambda x: len(x) > 0 and 'Other' not in x)]
+
+    # Rename column 
+    df_filtered.rename(columns={'Vulnerability Label': 'Group(s)'}, inplace=True)
 
     # Header
     st.subheader("Explore references to vulnerable groups:")
@@ -99,8 +97,8 @@ def vulnerability_display():
         
 
         # Text 
-        num_paragraphs = len(df_vul['Vulnerability Label'])
-        num_references = df_vul['Vulnerability Label'].apply(lambda x: 'Other' not in x).sum()
+        num_paragraphs = len(df['Vulnerability Label'])
+        num_references = len(df_filtered['Group(s)'])
        
         st.markdown(f"""<div style="text-align: justify;"> The document contains a
                 total of <span style="color: red;">{num_paragraphs}</span> paragraphs.
@@ -118,14 +116,14 @@ def vulnerability_display():
         # # Create a df that stores all the labels
         df_labels = pd.DataFrame(list(label_dict.items()), columns=['Label ID', 'Label'])
 
-        # Count how often each label appears in the "Vulnerability Labels" column
+        # Count how often each label appears in the "Group identified" column
         group_counts = {}
 
         # Iterate through each sublist
-        for index, row in df_vul.iterrows():
+        for index, row in df_filtered.iterrows():
             
             # Iterate through each group in the sublist
-            for sublist in row['Vulnerability Label']:
+            for sublist in row['Group(s)']:
                 
                 # Update the count in the dictionary
                 group_counts[sublist] = group_counts.get(sublist, 0) + 1
@@ -147,8 +145,4 @@ def vulnerability_display():
                      labels={'Count': 'Frequency'})
 
         #Show plot
-        st.plotly_chart(fig, use_container_width=True)
-
-    # ### Table 
-    #st.write(df_vul[df_vul['Vulnerability Label'].apply(lambda x: 'Other' not in x)])
-              
+        st.plotly_chart(fig, use_container_width=True)