feat: docs and convenience methods to explain category map schema

Author: mihow

ami/ml/models/algorithm.py

@@ -20,6 +20,41 @@
 class AlgorithmCategoryMap(BaseModel):
     """
     A list of classification labels for a given algorithm version
+
+    Expected schema for `data` field. This is the primary "category map" used by the model
+    to map from the category index in the model output to a human-readable label and other metadata.
+
+    IMPORTANT: Currently only `label` & `taxon_rank` are imported to the Taxon model if the taxon does
+    not already exist in the Antenna database. But the Taxon model can store any metadata, so this is
+    extensible in the future.
+    [
+        {
+            "index": 0,
+            "gbif_key": 123456,
+            "label": "Vanessa atalanta",
+            "taxon_rank": "SPECIES",
+        },
+        {
+            "index": 1,
+            "gbif_key": 789012,
+            "label": "Limenitis",
+            "taxon_rank": "GENUS",
+        },
+        {
+            "id": 3,
+            "gbif_key": 345678,
+            "label": "Nymphalis californica",
+            "taxon_rank": "SPECIES",
+        }
+    ]
+
+    The labels field is a simple list of string labels the correct index order used by the model.
+    [
+        "Vanessa atalanta",
+        "Limenitis",
+        "Nymphalis californica",
+    ]
+
     """
 
     data = models.JSONField(
@@ -55,6 +90,14 @@ def make_labels_hash(cls, labels):
         """
         return hash("".join(labels))
 
+    @classmethod
+    def labels_from_data(cls, data, label_field="label"):
+        return [category[label_field] for category in data]
+
+    @classmethod
+    def data_from_labels(cls, labels, label_field="label"):
+        return [{"index": i, label_field: label} for i, label in enumerate(labels)]
+
     def get_category(self, label, label_field="label"):
         # Can use JSON containment operators
         return self.data.index(next(category for category in self.data if category[label_field] == label))

ami/ml/tests.py

@@ -675,12 +675,12 @@ def test_labels_hash_auto_generation(self):
         from ami.ml.models import AlgorithmCategoryMap
 
         # Test data
-        test_labels = ["coleoptera", "diptera", "lepidoptera"]
         test_data = [
             {"index": 0, "label": "coleoptera"},
             {"index": 1, "label": "diptera"},
             {"index": 2, "label": "lepidoptera"},
         ]
+        test_labels = AlgorithmCategoryMap.labels_from_data(test_data)
 
         # Create instance using objects.create()
         category_map = AlgorithmCategoryMap.objects.create(labels=test_labels, data=test_data, version="test-v1")
@@ -696,3 +696,22 @@ def test_labels_hash_auto_generation(self):
         category_map2 = AlgorithmCategoryMap.objects.create(labels=test_labels, data=test_data, version="test-v2")
 
         self.assertEqual(category_map.labels_hash, category_map2.labels_hash)
+
+    def test_labels_data_conversion_methods(self):
+        from ami.ml.models import AlgorithmCategoryMap
+
+        # Test data
+        test_data = [
+            {"index": 0, "label": "coleoptera"},
+            {"index": 1, "label": "diptera"},
+            {"index": 2, "label": "lepidoptera"},
+        ]
+        test_labels = AlgorithmCategoryMap.labels_from_data(test_data)
+
+        # Convert labels to data and back
+        converted_data = AlgorithmCategoryMap.data_from_labels(test_labels)
+        converted_labels = AlgorithmCategoryMap.labels_from_data(converted_data)
+
+        # Verify conversions are correct
+        self.assertEqual(test_data, converted_data)
+        self.assertEqual(test_labels, converted_labels)