Update utils/vulnerability_classifier.py

utils/vulnerability_classifier.py

@@ -9,7 +9,51 @@ import streamlit as st
 from transformers import pipeline
 from setfit import SetFitModel
 
+label_dict= {0: 'Agricultural communities',
+ 1: 'Children',
+ 2: 'Coastal communities',
+ 3: 'Ethnic, racial or other minorities',
+ 4: 'Fishery communities',
+ 5: 'Informal sector workers',
+ 6: 'Members of indigenous and local communities',
+ 7: 'Migrants and displaced persons',
+ 8: 'Older persons',
+ 9: 'Other',
+ 10: 'Persons living in poverty',
+ 11: 'Persons with disabilities',
+ 12: 'Persons with pre-existing health conditions',
+ 13: 'Residents of drought-prone regions',
+ 14: 'Rural populations',
+ 15: 'Sexual minorities (LGBTQI+)',
+ 16: 'Urban populations',
+ 17: 'Women and other genders'}
 
+def getlabels(preds):
+    # Get label names
+    preds_list = preds.tolist()
+    
+    predictions_names=[]
+    
+    # loop through each prediction
+    for ele in preds_list:
+    
+      # see if there is a value 1 and retrieve index
+      try:
+        index_of_one = ele.index(1)
+      except ValueError:
+        index_of_one = "NA"
+    
+      # Retrieve the name of the label (if no prediction made = NA)
+      if index_of_one != "NA":
+        name  = label_dict[index_of_one]
+      else:
+        name = "NA"
+    
+      # Append name to list
+      predictions_names.append(name)
+    
+      return predictions_names
+        
 @st.cache_resource
 def load_vulnerabilityClassifier(config_file:str = None, classifier_name:str = None):
     """
@@ -89,12 +133,11 @@ def vulnerability_classification(haystack_doc:pd.DataFrame,
     
         predictions = classifier_model(list(haystack_doc.text))
 
-        haystack_doc['Vulnerability Label'] = predictions
+        
 
-    # list_ = []
-    # for i in range(len(predictions)):
+        pred_labels = getlabels(predictions)
 
-    #   temp = predictions[i]
+        haystack_doc['Vulnerability Label'] = pred_labels
     #   placeholder = {}
     #   for j in range(len(temp)):
     #     placeholder[temp[j]['label']] = temp[j]['score']