Add entity percentages to unstructured labeler profile (#305)

* Add entity percentages to unstructured labeler profile * Update tests * Add detailed test for entity percentages
capitalone · Jul 6, 2021 · 897eac6 · 897eac6
1 parent eaa272c
commit 897eac6
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 0 deletions.
diff --git a/dataprofiler/profilers/unstructured_labeler_profile.py b/dataprofiler/profilers/unstructured_labeler_profile.py
@@ -138,6 +138,7 @@ def update(self, df_series):
     def profile(self):
         profile = {
             "entity_counts": self.entity_counts,
+            "entity_percentages": self.entity_percentages,
             "times": self.times
         }
         return profile

diff --git a/dataprofiler/tests/profilers/test_column_profile_compilers.py b/dataprofiler/tests/profilers/test_column_profile_compilers.py
@@ -268,6 +268,10 @@ def test_base(self, *mocks):
                     'postprocess_char_level': defaultdict(int),
                     'true_char_level': defaultdict(int),
                     'word_level': defaultdict(int)},
+                'entity_percentages': {
+                    'postprocess_char_level': defaultdict(int),
+                    'true_char_level': defaultdict(int),
+                    'word_level': defaultdict(int)},
                 'times': {'data_labeler_predict': 1.0}},
             'statistics': {
                 'times': {'vocab': 1.0, 'words': 1.0},

diff --git a/dataprofiler/tests/profilers/test_unstructured_labeler_profile.py b/dataprofiler/tests/profilers/test_unstructured_labeler_profile.py
@@ -121,6 +121,11 @@ def test_profile(self, processor_class_mock, model_class_mock):
                 'true_char_level': defaultdict(int, {'UNKNOWN': 1}),
                 'word_level': defaultdict(int)
             },
+            entity_percentages={
+                'postprocess_char_level': defaultdict(int, {'UNKNOWN': 1.0}),
+                'true_char_level': defaultdict(int, {'UNKNOWN': 1.0}),
+                'word_level': defaultdict(int)
+            },
             times=defaultdict(float, {'data_labeler_predict': 1.0})
         )
 
@@ -132,6 +137,39 @@ def test_profile(self, processor_class_mock, model_class_mock):
         # key and value populated correctly
         self.assertDictEqual(expected_profile, profile)
 
+    @mock.patch('dataprofiler.profilers.'
+                'unstructured_labeler_profile.DataLabeler')
+    @mock.patch('dataprofiler.profilers.'
+                'unstructured_labeler_profile.'
+                'CharPostprocessor')
+    def test_entity_percentages(self, mock1, mock2):
+        """
+        Tests to see that entity percentages match the counts given
+        """
+        profile = UnstructuredLabelerProfile()
+        profile.char_sample_size = 20
+        profile.word_sample_size = 10
+        profile.entity_counts["postprocess_char_level"]["UNKNOWN"] = 6
+        profile.entity_counts["postprocess_char_level"]["TEST"] = 14
+        profile.entity_counts["true_char_level"]["UNKNOWN"] = 4
+        profile.entity_counts["true_char_level"]["TEST"] = 16
+        profile.entity_counts["word_level"]["UNKNOWN"] = 5
+        profile.entity_counts["word_level"]["TEST"] = 5
+        profile.update(pd.Series(["a"]))
+
+        expected_percentages = {
+            'postprocess_char_level': defaultdict(int, {'UNKNOWN': 0.3,
+                                                        'TEST': 0.7}),
+            'true_char_level': defaultdict(int, {'UNKNOWN': 0.2,
+                                                 'TEST': 0.8}),
+            'word_level': defaultdict(int, {'UNKNOWN': 0.5,
+                                            'TEST': 0.5})
+        }
+
+        percentages = profile.profile['entity_percentages']
+
+        self.assertDictEqual(expected_percentages, percentages)
+
     @mock.patch('dataprofiler.profilers.'
                 'unstructured_labeler_profile.DataLabeler')
     def test_unstructured_labeler_profile_add(self, mock):