Update utils/target_classifier.py

utils/target_classifier.py

@@ -69,21 +69,40 @@ def target_classification(haystack_doc:pd.DataFrame,
     x: Series object with the unique SDG covered in the document uploaded and 
     the number of times it is covered/discussed/count_of_paragraphs. 
     """
-    logging.info("Working on action/target extraction")
+                            
+    logging.info("Working on target/action identification")
+                            
+    haystack_doc['Vulnerability Label'] = 'NA'
+
     if not classifier_model:
+
         classifier_model = st.session_state['target_classifier']
+
+        # Get predictions
+        predictions = classifier_model(list(haystack_doc.text))
+
+        # Get labels for predictions
+        pred_labels = getlabels(predictions)
+
+        # Save labels
+        haystack_doc['Target Label'] = pred_labels
+
+                            
+    # logging.info("Working on action/target extraction")
+    # if not classifier_model:
+    #     classifier_model = st.session_state['target_classifier']
     
-    results = classifier_model(list(haystack_doc.text))
-    labels_= [(l[0]['label'],
-               l[0]['score']) for l in results]
+    # results = classifier_model(list(haystack_doc.text))
+    # labels_= [(l[0]['label'],
+    #            l[0]['score']) for l in results]
            
 
-    df1 = DataFrame(labels_, columns=["Target Label","Target Score"])
-    df = pd.concat([haystack_doc,df1],axis=1)
+    # df1 = DataFrame(labels_, columns=["Target Label","Target Score"])
+    # df = pd.concat([haystack_doc,df1],axis=1)
     
-    df = df.sort_values(by="Target Score", ascending=False).reset_index(drop=True)
-    df['Target Score'] = df['Target Score'].round(2)
-    df.index += 1
-    # df['Label_def'] = df['Target Label'].apply(lambda i: _lab_dict[i])
+    # df = df.sort_values(by="Target Score", ascending=False).reset_index(drop=True)
+    # df['Target Score'] = df['Target Score'].round(2)
+    # df.index += 1
+    # # df['Label_def'] = df['Target Label'].apply(lambda i: _lab_dict[i])
 
     return df