fix: bug in confidences on output (#419)

capitalone · Oct 16, 2021 · 2730d34 · 2730d34
1 parent 7d9a1e2
commit 2730d34
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/dataprofiler/labelers/character_level_cnn_model.py b/dataprofiler/labelers/character_level_cnn_model.py
@@ -966,7 +966,7 @@ def predict(self, data, batch_size=32, show_confidences=False,
                 in enumerate(sentence_lengths[:allocation_index]):
             predictions_list[index] = list(predictions[index][:sentence_length])
             if show_confidences:
-                confidences_list = list(confidences[index][:sentence_length])
+                confidences_list[index] = list(confidences[index][:sentence_length])
 
         if show_confidences:
             return {'pred': predictions_list, 'conf': confidences_list} 

diff --git a/dataprofiler/tests/labelers/test_data_processing.py b/dataprofiler/tests/labelers/test_data_processing.py
@@ -1209,6 +1209,7 @@ def test_match_sentence_lengths(self):
                                              inplace=True)
         self.assertEqual(results, post_process_results)
 
+
 class TestPreandPostCharacterProcessorConnection(unittest.TestCase):
 
     def test_flatten_convert(self):

diff --git a/dataprofiler/tests/labelers/test_integration_unstructured_data_labeler.py b/dataprofiler/tests/labelers/test_integration_unstructured_data_labeler.py
@@ -115,8 +115,16 @@ def test_default_confidences(self):
         model_predictions_char_level, model_confidences_char_level = \
             results["pred"], results["conf"]
 
-        # for now just checking that it's not empty
+        # for now just checking that it's not empty and appropriate size
+        num_labels = max(default.label_mapping.values()) + 1
+        len_text = len(sample[0])
         self.assertIsNotNone(model_confidences_char_level)
+        self.assertEqual((len_text, num_labels),
+                         model_confidences_char_level[0].shape)
+
+        len_text = len(sample[1])
+        self.assertEqual((len_text, num_labels),
+                         model_confidences_char_level[1].shape)
 
     def test_default_edge_cases(self):
         """more complicated test for edge cases for the default model"""