Update utils/vulnerability_classifier.py

utils/vulnerability_classifier.py

@@ -75,16 +75,16 @@ def vulnerability_classification(haystack_doc:pd.DataFrame,
     the number of times it is covered/discussed/count_of_paragraphs. 
     """
     logging.info("Working on vulnerability Identification")
-    haystack_doc['Indicator Label'] = 'NA'
-    haystack_doc['PA_check'] = haystack_doc['Policy-Action Label'].apply(lambda x: True if len(x) != 0 else False)
+    haystack_doc['Sector Label'] = 'NA'
+    # haystack_doc['PA_check'] = haystack_doc['Policy-Action Label'].apply(lambda x: True if len(x) != 0 else False)
                             
-    df1 = haystack_doc[haystack_doc['PA_check'] == True]
-    df = haystack_doc[haystack_doc['PA_check'] == False]
+    # df1 = haystack_doc[haystack_doc['PA_check'] == True]
+    # df = haystack_doc[haystack_doc['PA_check'] == False]
     if not classifier_model:
         classifier_model = st.session_state['vulnerability_classifier']
     
-        predictions = classifier_model(list(df1.text))
-
+        predictions = classifier_model(list(haystack_doc.text))
+        
     list_ = []
     for i in range(len(predictions)):
 
@@ -99,11 +99,9 @@ def vulnerability_classification(haystack_doc:pd.DataFrame,
     truth_df = truth_df.astype(float) >= threshold
     truth_df = truth_df.astype(str)
     categories = list(truth_df.columns)
-    truth_df['Indicator Label'] = truth_df.apply(lambda x: {i if x[i]=='True' else 
+    truth_df['Vulnerability Label'] = truth_df.apply(lambda x: {i if x[i]=='True' else 
                                               None for i in categories}, axis=1)
-    truth_df['Indicator Label'] = truth_df.apply(lambda x: list(x['Indicator Label'] 
+    truth_df['Vulnerability Label'] = truth_df.apply(lambda x: list(x['Vulnerability Label'] 
                                                             -{None}),axis=1)
-    df1['Indicator Label'] = list(truth_df['Indicator Label'])
-    df = pd.concat([df,df1])
-    df = df.drop(columns = ['PA_check'])
-    return df
+    haystack_doc['Vulnerability Label'] = list(truth_df['Vulnerability Label'])
+    return haystack_doc