Update utils/vulnerability_classifier.py

utils/vulnerability_classifier.py

@@ -84,24 +84,26 @@ def vulnerability_classification(haystack_doc:pd.DataFrame,
         classifier_model = st.session_state['vulnerability_classifier']
     
         predictions = classifier_model(list(haystack_doc.text))
-        
-    list_ = []
-    for i in range(len(predictions)):
 
-      temp = predictions[i]
-      placeholder = {}
-      for j in range(len(temp)):
-        placeholder[temp[j]['label']] = temp[j]['score']
-      list_.append(placeholder)
-    labels_ = [{**list_[l]} for l in range(len(predictions))]
-    truth_df = DataFrame.from_dict(labels_)
-    truth_df = truth_df.round(2)
-    truth_df = truth_df.astype(float) >= threshold
-    truth_df = truth_df.astype(str)
-    categories = list(truth_df.columns)
-    truth_df['Vulnerability Label'] = truth_df.apply(lambda x: {i if x[i]=='True' else 
-                                              None for i in categories}, axis=1)
-    truth_df['Vulnerability Label'] = truth_df.apply(lambda x: list(x['Vulnerability Label'] 
-                                                            -{None}),axis=1)
-    haystack_doc['Vulnerability Label'] = list(truth_df['Vulnerability Label'])
+        haystack_doc['Vulnerability Label'] = predictions
+
+    # list_ = []
+    # for i in range(len(predictions)):
+
+    #   temp = predictions[i]
+    #   placeholder = {}
+    #   for j in range(len(temp)):
+    #     placeholder[temp[j]['label']] = temp[j]['score']
+    #   list_.append(placeholder)
+    # labels_ = [{**list_[l]} for l in range(len(predictions))]
+    # truth_df = DataFrame.from_dict(labels_)
+    # truth_df = truth_df.round(2)
+    # truth_df = truth_df.astype(float) >= threshold
+    # truth_df = truth_df.astype(str)
+    # categories = list(truth_df.columns)
+    # truth_df['Vulnerability Label'] = truth_df.apply(lambda x: {i if x[i]=='True' else 
+    #                                           None for i in categories}, axis=1)
+    # truth_df['Vulnerability Label'] = truth_df.apply(lambda x: list(x['Vulnerability Label'] 
+    #                                                         -{None}),axis=1)
+    # haystack_doc['Vulnerability Label'] = list(truth_df['Vulnerability Label'])
     return haystack_doc