fix: type confidence as a float in label parser (#1586)

kili-technology · Nov 13, 2023 · 266b6df · 266b6df
1 parent fcdb86a
commit 266b6df
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 9 deletions.
diff --git a/src/kili/services/label_data_parsing/annotation.py b/src/kili/services/label_data_parsing/annotation.py
@@ -119,7 +119,7 @@ def category(self) -> "category_module.Category":
         return self._json_data["categories"][0]
 
     @typechecked
-    def add_category(self, name: str, confidence: Optional[int] = None) -> None:
+    def add_category(self, name: str, confidence: Optional[float] = None) -> None:
         """Add a category to an annotation job with categories."""
         if "categories" not in self._json_data:
             category_list = category_module.CategoryList(

diff --git a/src/kili/services/label_data_parsing/category.py b/src/kili/services/label_data_parsing/category.py
@@ -145,13 +145,13 @@ def display_name(self, name: str) -> None:
             ) from err
 
     @property
-    def confidence(self) -> int:
+    def confidence(self) -> float:
         """Returns the confidence of the category label."""
         return self._json_data["confidence"]
 
     @confidence.setter
     @typechecked
-    def confidence(self, confidence: int) -> None:
+    def confidence(self, confidence: float) -> None:
         """Set the confidence of the category label."""
         if not 0 <= confidence <= 100:
             raise ValueError(f"Confidence must be between 0 and 100, got {confidence}")
@@ -226,7 +226,10 @@ def _check_can_append_category(self, category: Category) -> None:
 
     @typechecked
     def add_category(
-        self, name: str, confidence: Optional[int] = None, children: Optional[Dict] = None
+        self,
+        name: str,
+        confidence: Optional[float] = None,
+        children: Optional[Dict] = None,
     ) -> None:
         """Add a category object to the CategoryList object."""
         category_dict: Dict[str, object] = {"name": name}

diff --git a/src/kili/services/label_data_parsing/job_response.py b/src/kili/services/label_data_parsing/job_response.py
@@ -141,7 +141,7 @@ def children(self, children: Dict) -> None:
         self._json_data["children"] = parsed_children_job
 
     @typechecked
-    def add_category(self, name: str, confidence: Optional[int] = None) -> None:
+    def add_category(self, name: str, confidence: Optional[float] = None) -> None:
         """Adds a category to a job with categories."""
         if self._job_interface["mlTask"] != "CLASSIFICATION":
             raise AttributeNotCompatibleWithJobError("add_category")

diff --git a/tests/unit/services/label_data_parsing/parsing/test_resp_parsing.py b/tests/unit/services/label_data_parsing/parsing/test_resp_parsing.py
@@ -76,8 +76,8 @@ def test_attribute_categories_multiclass():
     json_response_dict = {
         "JOB_0": {
             "categories": [
-                {"name": "CATEGORY_A", "confidence": 42},
-                {"name": "CATEGORY_B", "confidence": 100},
+                {"name": "CATEGORY_A", "confidence": 42.289},
+                {"name": "CATEGORY_B", "confidence": 88.75},
             ]
         }
     }
@@ -86,10 +86,10 @@ def test_attribute_categories_multiclass():
     parsed_jobs = ParsedJobs(json_response=json_response_dict, project_info=project_info)
     categories = parsed_jobs["JOB_0"].categories
 
-    assert categories[0].confidence == 42
+    assert categories[0].confidence == 42.289
     assert categories[0].display_name == "Category A"
     assert categories[0].name == "CATEGORY_A"
-    assert categories[1].confidence == 100
+    assert categories[1].confidence == 88.75
     assert categories[1].display_name == "Category B"
     assert categories[1].name == "CATEGORY_B"