First try entity_ratio option

FairEval.py

@@ -147,21 +147,24 @@ class FairEvaluation(evaluate.Metric):
         true_spans = seq_to_fair(true_spans)
         pred_spans = seq_to_fair(pred_spans)
 
-        # (3) COUNT ERRORS AND CALCULATE SCORES
+        # (3) COUNT ERRORS AND CALCULATE SCORES (counting total ground truth entities too)
         total_errors = compare_spans([], [])
+        total_ref_entities = 0
         for i in range(len(true_spans)):
+            total_ref_entities += len(true_spans[i])
             sentence_errors = compare_spans(true_spans[i], pred_spans[i])
             total_errors = add_dict(total_errors, sentence_errors)
 
         if weights is None and mode == 'weighted':
-            print("The chosen mode is \'weighted\', but no weights are given. Setting weights to:\n")
             weights = {"TP": {"TP": 1},
-                 "FP": {"FP": 1},
-                 "FN": {"FN": 1},
-                 "LE": {"TP": 0, "FP": 0.5, "FN": 0.5},
-                 "BE": {"TP": 0.5, "FP": 0.25, "FN": 0.25},
-                 "LBE": {"TP": 0, "FP": 0.5, "FN": 0.5}}
-            print(weights)
+                       "FP": {"FP": 1},
+                       "FN": {"FN": 1},
+                       "LE": {"TP": 0, "FP": 0.5, "FN": 0.5},
+                       "BE": {"TP": 0.5, "FP": 0.25, "FN": 0.25},
+                       "LBE": {"TP": 0, "FP": 0.5, "FN": 0.5}}
+            print("The chosen mode is \'weighted\', but no weights are given. Setting weights to:\n")
+            for k in weights:
+                print(k, ":", weights[k])
 
         config = {"labels": "all", "eval_method": [mode], "weights": weights,}
         results = calculate_results(total_errors, config)
@@ -170,34 +173,36 @@ class FairEvaluation(evaluate.Metric):
         # (4) SELECT OUTPUT MODE AND REFORMAT AS SEQEVAL-HUGGINGFACE OUTPUT
         # initialize empty dictionary and count errors
         output = {}
-        total_trad_errors = results['overall']['traditional']['FP'] + results['overall']['traditional']['FN']
-        total_fair_errors = results['overall']['fair']['FP'] + results['overall']['fair']['FN'] + \
-                            results['overall']['fair']['LE'] + results['overall']['fair']['BE'] + \
-                            results['overall']['fair']['LBE']
+        # control the divider for the error_format (count, proportion over total errors or over total entities)
+        if error_format == 'count':
+            trad_divider = 1,
+            fair_divider = 1,
+        elif error_format == 'entity_ratio':
+            trad_divider = total_ref_entities
+            fair_divider = total_ref_entities
+        elif error_format == 'error_ratio':
+            trad_divider = results['overall']['traditional']['FP'] + results['overall']['traditional']['FN']
+            fair_divider = results['overall']['fair']['FP'] + results['overall']['fair']['FN'] + \
+                                results['overall']['fair']['LE'] + results['overall']['fair']['BE'] + \
+                                results['overall']['fair']['LBE']
+
 
         # assert valid options
         assert mode in ['traditional', 'fair', 'weighted'], 'mode must be \'traditional\', \'fair\' or \'weighted\''
-        assert error_format in ['count', 'proportion'], 'error_format must be \'count\' or \'proportion\''
+        assert error_format in ['count', 'error_ratio', 'entity_ratio'], 'error_format must be \'count\', \'error_ratio\' or \'entity_ratio\''
 
         # append entity-level errors and scores
         if mode == 'traditional':
             for k, v in results['per_label'][mode].items():
-                if error_format == 'count':
-                    output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'], 'TP': v['TP'],
-                                 'FP': v['FP'], 'FN': v['FN']}
-                elif error_format == 'proportion':
-                    output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'], 'TP': v['TP'],
-                                 'FP': v['FP'] / total_trad_errors, 'FN': v['FN'] / total_trad_errors}
+                output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'],
+                             'TP': v['TP'] / trad_divider if error_format == 'entity_ratio' else v['TP'],
+                             'FP': v['FP'] / trad_divider, 'FN': v['FN'] / trad_divider}
         elif mode == 'fair' or mode == 'weighted':
             for k, v in results['per_label'][mode].items():
-                if error_format == 'count':
-                    output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'], 'TP': v['TP'],
-                                 'FP': v['FP'], 'FN': v['FN'], 'LE': v['LE'], 'BE': v['BE'], 'LBE': v['LBE']}
-                elif error_format == 'proportion':
-                    output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'], 'TP': v['TP'],
-                                 'FP': v['FP'] / total_fair_errors, 'FN': v['FN'] / total_fair_errors,
-                                 'LE': v['LE'] / total_fair_errors, 'BE': v['BE'] / total_fair_errors,
-                                 'LBE': v['LBE'] / total_fair_errors}
+                output[k] = {'precision': v['Prec'], 'recall': v['Rec'], 'f1': v['F1'],
+                             'TP': v['TP'] / fair_divider if error_format == 'entity_ratio' else v['TP'],
+                             'FP': v['FP'] / fair_divider, 'FN': v['FN'] / fair_divider,
+                             'LE': v['LE'] / fair_divider, 'BE': v['BE'] / fair_divider, 'LBE': v['LBE'] / fair_divider}
 
         # append overall scores
         output['overall_precision'] = results['overall'][mode]['Prec']
@@ -206,25 +211,16 @@ class FairEvaluation(evaluate.Metric):
 
         # append overall error counts
         if mode == 'traditional':
-            output['TP'] = results['overall'][mode]['TP']
-            output['FP'] = results['overall'][mode]['FP']
-            output['FN'] = results['overall'][mode]['FN']
-            if error_format == 'proportion':
-                output['FP'] = output['FP'] / total_trad_errors
-                output['FN'] = output['FN'] / total_trad_errors
+            output['TP'] = results['overall'][mode]['TP'] / trad_divider if error_format == 'entity_ratio' else results['overall'][mode]['TP']
+            output['FP'] = results['overall'][mode]['FP'] / trad_divider
+            output['FN'] = results['overall'][mode]['FN'] / trad_divider
         elif mode == 'fair' or 'weighted':
-            output['TP'] = results['overall'][mode]['TP']
-            output['FP'] = results['overall'][mode]['FP']
-            output['FN'] = results['overall'][mode]['FN']
-            output['LE'] = results['overall'][mode]['LE']
-            output['BE'] = results['overall'][mode]['BE']
-            output['LBE'] = results['overall'][mode]['LBE']
-            if error_format == 'proportion':
-                output['FP'] = output['FP'] / total_fair_errors
-                output['FN'] = output['FN'] / total_fair_errors
-                output['LE'] = output['LE'] / total_fair_errors
-                output['BE'] = output['BE'] / total_fair_errors
-                output['LBE'] = output['LBE'] / total_fair_errors
+            output['TP'] = results['overall'][mode]['TP'] / fair_divider if error_format == 'entity_ratio' else results['overall'][mode]['TP']
+            output['FP'] = results['overall'][mode]['FP'] / fair_divider
+            output['FN'] = results['overall'][mode]['FN'] / fair_divider
+            output['LE'] = results['overall'][mode]['LE'] / fair_divider
+            output['BE'] = results['overall'][mode]['BE'] / fair_divider
+            output['LBE'] = results['overall'][mode]['LBE'] / fair_divider
 
         return output