Update appStore/vulnerability_analysis.py

appStore/vulnerability_analysis.py

@@ -56,8 +56,11 @@ def app():
 
     ### Main app code ###
     with st.container():
-                   
+
+            # If a document has been processed                   
             if 'key0' in st.session_state:
+
+                # Run vulnerability classifier
                 df = st.session_state.key0
                 classifier = load_vulnerabilityClassifier(classifier_name=params['model_name'])
                 st.session_state['{}_classifier'.format(classifier_identifier)] = classifier
@@ -67,8 +70,8 @@ def app():
                 df = vulnerability_classification(haystack_doc=df,
                                             threshold= params['threshold'])
 
-                # Store df in session state
-                st.session_state.key0 = df
+                # Store df in session state with key1
+                st.session_state.key1 = df