Update utils/vulnerability_classifier.py

utils/vulnerability_classifier.py

@@ -37,13 +37,12 @@ def get_vulnerability_labels(preds):
     
     # Get label names
     preds_list = preds.tolist()
-    st.write(preds_list)
     
     # Get the name of the group where the prediction is equal to "1"
     result = []
 
-    for sublist in my_list_of_lists:
-        result.extend([my_dict[key] for key, value in enumerate(sublist) if value == 1])
+    for sublist in preds_list:
+        result.extend([label_dict[key] for key, value in enumerate(sublist) if value == 1])
 
     st.write("Here are the results")
     st.write(result)