Update utils/ghg_classifier.py

utils/ghg_classifier.py

@@ -10,10 +10,9 @@ from transformers import pipeline
 
 # Labels dictionary ###
 _lab_dict = {
-                      'LABEL_0':'NEGATIVE',
-                      'LABEL_1':'NOT GHG',
-                      'LABEL_2':'GHG',
-                      'NA':'NA',
+                      'GHG':'GHG',
+                      'NOT_GHG':'NON GHG TRANSPORT TARGET',
+                      'NEGATIVE':'OTHERS',
                       }
         
 
@@ -74,9 +73,12 @@ def ghg_classification(haystack_doc:pd.DataFrame,
     """
     logging.info("Working on GHG Extraction")
     haystack_doc['GHG Label'] = 'NA'
-    haystack_doc['GHG Score'] = 'NA'
+    haystack_doc['GHG Score'] = 0.0
+    # applying GHG Identifier to only 'Target' paragraphs.
     temp = haystack_doc[haystack_doc['Target Label'] == 'TARGET']
+    temp = temp.reset_index(drop=True)
     df = haystack_doc[haystack_doc['Target Label'] == 'NEGATIVE']
+    df = df.reset_index(drop=True)
 
     if not classifier_model:
         classifier_model = st.session_state['ghg_classifier']
@@ -84,9 +86,11 @@ def ghg_classification(haystack_doc:pd.DataFrame,
     results = classifier_model(list(temp.text))
     labels_= [(l[0]['label'],l[0]['score']) for l in results]
     temp['GHG Label'],temp['GHG Score'] = zip(*labels_)
+    temp['GHG Label'] = temp['GHG Label'].apply(lambda x: _lab_dict[x])
+    # merge back Target and non-Target dataframe
     df = pd.concat([df,temp])
-    df['GHG Label'] = df['GHG Label'].apply(lambda i: _lab_dict[i])
     df = df.reset_index(drop =True)
+    df['GHG Score'] = df['GHG Score'].round(2)
     df.index += 1
 
     return df