def test_confusion_matrix_correct_metrics(): from collections import Counter evaluated = [ EvaluationResult(results=Counter({ ('O', 'O'): 150, ('O', 'PERSON'): 30, ('O', 'COMPANY'): 30, ('PERSON', 'PERSON'): 40, ('COMPANY', 'COMPANY'): 40, ('PERSON', 'COMPANY'): 10, ('COMPANY', 'PERSON'): 10, ('PERSON', 'O'): 30, ('COMPANY', 'O'): 30 }), model_errors=None, text=None) ] model = MockTokensModel(prediction=None, entities_to_keep=['PERSON', 'COMPANY']) scores = model.calculate_score(evaluated, beta=2.5) assert scores.pii_precision == 0.625 assert scores.pii_recall == 0.625 assert scores.entity_recall_dict['PERSON'] == 0.5 assert scores.entity_precision_dict['PERSON'] == 0.5 assert scores.entity_recall_dict['COMPANY'] == 0.5 assert scores.entity_precision_dict['COMPANY'] == 0.5
def test_confusion_matrix_2_correct_metrics(): from collections import Counter evaluated = [ EvaluationResult(results=Counter({ ('O', 'O'): 65467, ('O', 'ORG'): 4189, ('GPE', 'O'): 3370, ('PERSON', 'PERSON'): 2024, ('GPE', 'PERSON'): 1488, ('GPE', 'GPE'): 1033, ('O', 'GPE'): 964, ('ORG', 'ORG'): 914, ('O', 'PERSON'): 834, ('GPE', 'ORG'): 401, ('PERSON', 'ORG'): 35, ('PERSON', 'O'): 33, ('ORG', 'O'): 8, ('PERSON', 'GPE'): 5, ('ORG', 'PERSON'): 1 }), model_errors=None, text=None) ] model = MockTokensModel(prediction=None) scores = model.calculate_score(evaluated, beta=2.5) pii_tp = evaluated[0].results[('PERSON', 'PERSON')] + \ evaluated[0].results[('ORG', 'ORG')] + \ evaluated[0].results[('GPE', 'GPE')] + \ evaluated[0].results[('ORG', 'GPE')] + \ evaluated[0].results[('ORG', 'PERSON')] + \ evaluated[0].results[('GPE', 'ORG')] + \ evaluated[0].results[('GPE', 'PERSON')] + \ evaluated[0].results[('PERSON', 'GPE')] + \ evaluated[0].results[('PERSON', 'ORG')] pii_fp = evaluated[0].results[('O', 'PERSON')] + \ evaluated[0].results[('O', 'GPE')] + \ evaluated[0].results[('O', 'ORG')] pii_fn = evaluated[0].results[('PERSON', 'O')] + \ evaluated[0].results[('GPE', 'O')] + \ evaluated[0].results[('ORG', 'O')] assert scores.pii_precision == pii_tp / (pii_tp + pii_fp) assert scores.pii_recall == pii_tp / (pii_tp + pii_fn)
def test_evaluate_multiple_examples_ignore_entity_correct_statistics(): prediction = ["O", "O", "O", "U-PERSON", "O", "U-TENNIS_PLAYER"] model = MockTokensModel(prediction=prediction, labeling_scheme='BILOU', entities_to_keep=['PERSON', 'TENNIS_PLAYER']) input_sample = InputSample("My name is Raphael or David", masked=None, spans=None) input_sample.tokens = ["My", "name", "is", "Raphael", "or", "David"] input_sample.tags = ["O", "O", "O", "U-PERSON", "O", "U-PERSON"] evaluated = model.evaluate_all( [input_sample, input_sample, input_sample, input_sample]) scores = model.calculate_score(evaluated) assert scores.pii_precision == 1 assert scores.pii_recall == 1