Move functions to main class

isco_hierarchical_accuracy.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 """ISCO-08 Hierarchical Accuracy Measure."""
 
+from typing import List, Set, Dict, Tuple
 import evaluate
 import datasets
-import ham
-import isco
+
+# import ham
+# import isco
 
 
 # TODO: Add BibTeX citation
@@ -110,11 +112,144 @@ class ISCO_Hierarchical_Accuracy(evaluate.Metric):
             reference_urls=["http://path.to.reference.url/new_module"],
         )
 
+    def create_hierarchy_dict(file: str) -> dict:
+        """
+        Creates a dictionary where keys are nodes and values are sets of parent nodes representing the group level hierarchy of the ISCO-08 structure.
+        The function assumes that the input CSV file has a column named 'unit' with the 4-digit ISCO-08 codes.
+        A csv file with the ISCO-08 structure can be downloaded from the International Labour Organization (ILO) at [https://www.ilo.org/ilostat-files/ISCO/newdocs-08-2021/ISCO-08/ISCO-08 EN.csv](https://www.ilo.org/ilostat-files/ISCO/newdocs-08-2021/ISCO-08/ISCO-08%20EN.csv)
+
+        Args:
+        - file: A string representing the path to the CSV file containing the 4-digit ISCO-08 codes. It can be a local path or a web URL.
+
+        Returns:
+        - A dictionary where keys are ISCO-08 unit codes and values are sets of their parent codes.
+        """
+
+        try:
+            import requests
+            import csv
+        except ImportError as error:
+            raise error
+
+        isco_hierarchy = {}
+
+        if file.startswith("http://") or file.startswith("https://"):
+            response = requests.get(file)
+            lines = response.text.splitlines()
+        else:
+            with open(file, newline="") as csvfile:
+                lines = csvfile.readlines()
+
+        reader = csv.DictReader(lines)
+        for row in reader:
+            unit_code = row["unit"].zfill(4)
+            minor_code = unit_code[0:3]
+            sub_major_code = unit_code[0:2]
+            major_code = unit_code[0]
+            isco_hierarchy[unit_code] = {minor_code, major_code, sub_major_code}
+
+        return isco_hierarchy
+
+    def find_ancestors(node: str, hierarchy: dict) -> set:
+        """
+        Find the ancestors of a given node in a hierarchy.
+
+        Args:
+            node (str): The node for which to find ancestors.
+            hierarchy (dict): A dictionary representing the hierarchy, where the keys are nodes and the values are their parents.
+
+        Returns:
+            set: A set of ancestors of the given node.
+        """
+        ancestors = set()
+        nodes_to_visit = [node]
+        while nodes_to_visit:
+            current_node = nodes_to_visit.pop()
+            if current_node in hierarchy:
+                parents = hierarchy[current_node]
+                ancestors.update(parents)
+                nodes_to_visit.extend(parents)
+        return ancestors
+
+    def extend_with_ancestors(self, classes: set, hierarchy: dict) -> set:
+        """
+        Extend the given set of classes with their ancestors from the hierarchy.
+
+        Args:
+            classes (set): The set of classes to extend.
+            hierarchy (dict): The hierarchy of classes.
+
+        Returns:
+            set: The extended set of classes including their ancestors.
+        """
+        extended_classes = set(classes)
+        for cls in classes:
+            ancestors = self.find_ancestors(cls, hierarchy)
+            extended_classes.update(ancestors)
+        return extended_classes
+
+    def calculate_hierarchical_precision_recall(
+        reference_codes: List[str],
+        predicted_codes: List[str],
+        hierarchy: Dict[str, Set[str]],
+    ) -> Tuple[float, float]:
+        """
+        Calculates the hierarchical precision and recall given the reference codes, predicted codes, and hierarchy definition.
+
+        Args:
+            real_codes (List[str]): The list of reference codes.
+            predicted_codes (List[str]): The list of predicted codes.
+            hierarchy (Dict[str, Set[str]]): The hierarchy definition where keys are nodes and values are sets of parent nodes.
+
+        Returns:
+            Tuple[float, float]: A tuple containing the hierarchical precision and recall floating point values.
+        """
+        # Extend the sets of real and predicted codes with their ancestors
+        extended_real = set()
+        for code in reference_codes:
+            extended_real.add(code)
+            extended_real.update(hierarchy.get(code, set()))
+
+        extended_predicted = set()
+        for code in predicted_codes:
+            extended_predicted.add(code)
+            extended_predicted.update(hierarchy.get(code, set()))
+
+        # Calculate the intersection
+        correct_predictions = extended_real.intersection(extended_predicted)
+
+        # Calculate hierarchical precision and recall
+        hP = (
+            len(correct_predictions) / len(extended_predicted)
+            if extended_predicted
+            else 0
+        )
+        hR = len(correct_predictions) / len(extended_real) if extended_real else 0
+
+        return hP, hR
+
+    def hierarchical_f_measure(hP, hR, beta=1.0):
+        """
+        Calculate the hierarchical F-measure.
+
+        Parameters:
+        hP (float): The hierarchical precision.
+        hR (float): The hierarchical recall.
+        beta (float, optional): The beta value for F-measure calculation. Default is 1.0.
+
+        Returns:
+        float: The hierarchical F-measure.
+        """
+        if hP + hR == 0:
+            return 0
+        return (beta**2 + 1) * hP * hR / (beta**2 * hP + hR)
+
     def _download_and_prepare(self, dl_manager):
         """Download external ISCO-08 csv file from the ILO website for creating the hierarchy dictionary."""
         isco_csv = dl_manager.download_and_extract(ISCO_CSV_MIRROR_URL)
         print(f"ISCO CSV file downloaded")
-        self.isco_hierarchy = isco.create_hierarchy_dict(isco_csv)
+        # self.isco_hierarchy = isco.create_hierarchy_dict(isco_csv)
+        self.isco_hierarchy = self.create_hierarchy_dict(isco_csv)
         print("ISCO hierarchy dictionary created")
         print(self.isco_hierarchy)
 
@@ -132,10 +267,10 @@ class ISCO_Hierarchical_Accuracy(evaluate.Metric):
 
         # Calculate hierarchical precision, recall and f-measure
         hierarchy = self.isco_hierarchy
-        hP, hR = ham.calculate_hierarchical_precision_recall(
+        hP, hR = self.calculate_hierarchical_precision_recall(
             references, predictions, hierarchy
         )
-        hF = ham.hierarchical_f_measure(hP, hR)
+        hF = self.hierarchical_f_measure(hP, hR)
         print(
             f"Hierarchical Precision: {hP}, Hierarchical Recall: {hR}, Hierarchical F-measure: {hF}"
         )