Update public API docstrings

.github/workflows/ci.yml

@@ -107,7 +107,7 @@ jobs:
           python-version: 3.8
 
       - name: Install tools
-        run: pip install mkdocs mkdocstrings mkdocs-exclude mkdocs-material taskipy
+        run: pip install mkdocs mkdocstrings-python mkdocs-exclude mkdocs-material taskipy
 
       - name: Generate
         run: task docs

docs/reference/data.md

@@ -0,0 +1,3 @@
+# Dataset Ingestion and Adaptation
+
+::: charmory.data

docs/reference/engine.md

@@ -0,0 +1,5 @@
+# Armory Engines
+
+::: charmory.engine.EvaluationEngine
+
+::: charmory.engine.AdversarialDatasetEngine

docs/reference/evaluation.md

@@ -0,0 +1,3 @@
+# Evaluation Configuration
+
+::: charmory.evaluation

docs/reference/labels.md

@@ -0,0 +1,3 @@
+# Label Targeters
+
+::: charmory.labels

docs/reference/model.md

@@ -0,0 +1,7 @@
+# Model Ingestion and Adaptation
+
+::: charmory.model.base
+
+::: charmory.model.image_classification
+
+::: charmory.model.object_detection

docs/reference/profiler.md

@@ -0,0 +1,3 @@
+# Computational Performance Profilers
+
+::: armory.metrics.compute

docs/reference/task.md

@@ -0,0 +1,7 @@
+# Evaluation Tasks
+
+::: charmory.tasks.base.BaseEvaluationTask
+
+::: charmory.tasks.image_classification.ImageClassificationTask
+
+::: charmory.tasks.object_detection.ObjectDetectionTask

docs/reference/track.md

@@ -0,0 +1,3 @@
+# Evaluation Tracking
+
+::: charmory.track

docs/reference/transforms.md

@@ -0,0 +1,3 @@
+# Dataset Transforms
+
+::: charmory.experimental.transforms

docs/reference/utils.md

@@ -0,0 +1,3 @@
+# Utilities
+
+::: charmory.utils

library/src/armory/metrics/compute.py

@@ -1,5 +1,5 @@
 """
-Computational metrics
+Profilers to collect computational metrics
 """
 
 import cProfile
@@ -23,9 +23,7 @@ def results(self) -> Mapping[str, float]:
 
 
 class NullProfiler:
-    """
-    Measures computational resource use
-    """
+    """Profiler that does nothing (no-op)"""
 
     def __init__(self):
         self.measurement_dict = {}
@@ -39,6 +37,8 @@ def results(self):
 
 
 class BasicProfiler(NullProfiler):
+    """Profiler using `time.perf_counter`"""
+
     @contextlib.contextmanager
     def measure(self, name):
         startTime = time.perf_counter()
@@ -72,6 +72,8 @@ def results(self):
 
 
 class DeterministicProfiler(NullProfiler):
+    """Profiler using cProfile for deterministic profiling"""
+
     def __init__(self):
         super().__init__()
         log.warning(

library/src/charmory/data.py

@@ -16,9 +16,32 @@
 
 
 class ArmoryDataset(Dataset):
-    """Wrapper around a dataset to apply an adapter to all samples obtained from the dataset"""
+    """
+    Wrapper around a PyTorch dataset to apply an adapter to all samples obtained
+    from the dataset.
+
+    Example::
+
+        from charmory.data import ArmoryDataset
+
+        def rename_fields(sample):
+            # Rename the 'data' field in the sample to 'image'
+            sample["image"] = sample.pop("data")
+            return sample
+
+        # assuming `dataset` has been defined elsewhere
+        renamed_dataset = ArmoryDataset(dataset, rename_fields)
+    """
 
     def __init__(self, dataset, adapter: DatasetOutputAdapter):
+        """
+        Initializes the dataset.
+
+        Args:
+            dataset: Source dataset to be wrapped. It must be subscriptable and
+                support the `len` operator.
+            adapter: Dataset sample adapter
+        """
         self._dataset = dataset
         self._adapter = adapter
 
@@ -30,14 +53,38 @@ def __getitem__(self, index):
 
 
 class TupleDataset(ArmoryDataset):
-    """Dataset wrapper with a pre-applied adapter to adapt tuples to map-like samples"""
+    """
+    Dataset wrapper with a pre-applied adapter to adapt tuples to map-like
+    samples.
+
+    Example::
+
+        from charmory.data import TupleDataset
+
+        # assuming `dataset` has been defined elsewhere
+        print(dataset[0])
+        # output: [[0, 0, 0], [0, 0, 0]], [5]
+
+        tuple_ds = TupleDataset(dataset, x_key="image", y_key="label")
+        print(tuple_ds[0])
+        # output: {'image': [[0, 0, 0], [0, 0, 0]], 'label': [5]}
+    """
 
     def __init__(
         self,
         dataset,
         x_key: str,
         y_key: str,
     ):
+        """
+        Initializes the dataset.
+
+        Args:
+            dataset: Source dataset where samples are a two-entry tuple of data,
+                or x, and target, or y.
+            x_key: Key name to use for x data in the adapted sample dictionary
+            y_key: Key name to use for y data in the adapted sample dictionary
+        """
         super().__init__(dataset, self._adapt)
         self._x_key = x_key
         self._y_key = y_key

library/src/charmory/engine/__init__.py

@@ -1,2 +1,3 @@
+"""This package contains the Armory runtime engines."""
 from charmory.engine.adversarial_dataset import AdversarialDatasetEngine
 from charmory.engine.evaluation import EvaluationEngine

library/src/charmory/engine/adversarial_dataset.py

@@ -1,12 +1,14 @@
 """Armory engine to create adversarial datasets"""
-from typing import Any, Callable, Mapping, Optional
+from typing import TYPE_CHECKING, Any, Callable, Generator, Mapping, Optional
 
-import datasets
 import numpy as np
 
 from charmory.tasks.base import BaseEvaluationTask
 from charmory.track import get_current_params
 
+if TYPE_CHECKING:
+    import datasets
+
 SampleAdapter = Callable[[Mapping[str, Any]], Mapping[str, Any]]
 """
 An adapter for generated samples. The input argument and return types are a
@@ -19,7 +21,7 @@ class AdversarialDatasetEngine:
     Armory engine to create adversarial datasets. An adversarial dataset has
     an adversarial attack already applied to every sample in the dataset.
 
-    Example:
+    Example::
 
         from charmory.engine import AdversarialDatasetEngine
 
@@ -40,14 +42,14 @@ def __init__(
         task: BaseEvaluationTask,
         output_dir: Optional[str] = None,
         adapter: Optional[SampleAdapter] = None,
-        features: Optional[datasets.Features] = None,
+        features: Optional["datasets.Features"] = None,
         num_batches: Optional[int] = None,
     ):
         """
         Initializes the engine.
 
         Args:
-            evaluation: Armory evaluation from which to generate the dataset
+            task: Armory evaluation task from which to generate the dataset
             output_dir: Optional, directory to which to write the generated dataset
             adapter: Optional, adapter to perform additional modifications to samples
             features: Optional, dataset features
@@ -63,12 +65,14 @@ def __init__(
         self.num_batches = num_batches
 
     @staticmethod
-    def _default_adapter(sample: Mapping[str, Any]):
+    def _default_adapter(sample: Mapping[str, Any]) -> Mapping[str, Any]:
         # do nothing
         return sample
 
-    def generate(self) -> datasets.Dataset:
+    def generate(self) -> "datasets.Dataset":
         """Create the adversarial dataset"""
+        import datasets
+
         dataset = datasets.Dataset.from_generator(
             self._generator, features=self.features
         )
@@ -80,7 +84,7 @@ def generate(self) -> datasets.Dataset:
 
         return dataset
 
-    def _generator(self):
+    def _generator(self) -> Generator[Mapping[str, Any], None, None]:
         """
         Iterates over every batch in the source dataset, applies the adversarial
         attack, and yields the pre-attacked samples.

library/src/charmory/engine/evaluation.py

@@ -1,39 +1,73 @@
-from typing import Optional, Union
+"""Armory engine to perform model robustness evaluations"""
+from typing import Mapping, Optional, TypedDict
 
 import lightning.pytorch as pl
 import lightning.pytorch.loggers as pl_loggers
 from lightning.pytorch.utilities import rank_zero_only
+from torch import Tensor
 
 from charmory.tasks.base import BaseEvaluationTask
 from charmory.track import get_current_params, init_tracking_uri, track_system_metrics
 
 
+class EvaluationResults(TypedDict):
+    """Robustness evaluation results"""
+
+    compute: Mapping[str, float]
+    """Computational metrics"""
+    metrics: Mapping[str, Tensor]
+    """Task-specific evaluation metrics"""
+
+
 class EvaluationEngine:
+    """
+    Armory engine to perform model robustness evaluations.
+
+    Example::
+
+        from charmory.engine import EvaluationEngine
+
+        # assuming `task` has been defined using a `charmory.tasks` class
+        engine = EvaluationEngine(task)
+        results = engine.run()
+    """
+
     def __init__(
-        self,
-        task: BaseEvaluationTask,
-        limit_test_batches: Optional[Union[int, float]] = None,
+        self, task: BaseEvaluationTask, run_id: Optional[str] = None, **kwargs
     ):
+        """
+        Initializes the engine.
+
+        Args:
+            task: Armory evaluation task to perform model inference and
+                application of adversarial attacks
+            run_id: Optional, MLflow run ID to which to record evaluation results
+            **kwargs: All other keyword arguments will be forwarded to the
+                `lightning.pytorch.Trainer` class.
+        """
         self.task = task
         self._logger = pl_loggers.MLFlowLogger(
             experiment_name=self.task.evaluation.name,
             tags={"mlflow.note.content": self.task.evaluation.description},
             tracking_uri=init_tracking_uri(self.task.evaluation.sysconfig.armory_home),
+            run_id=run_id,
         )
         self.trainer = pl.Trainer(
             inference_mode=False,
-            limit_test_batches=limit_test_batches,
             logger=self._logger,
+            **kwargs,
         )
-        self.run_id: Optional[str] = None
+        self.run_id = run_id
         self._was_run = False
 
     @rank_zero_only
     def _log_params(self):
+        """Log tracked params with MLflow"""
         self.run_id = self._logger.run_id
         self._logger.log_hyperparams(get_current_params())
 
-    def run(self):
+    def run(self) -> EvaluationResults:
+        """Perform the evaluation"""
         if self._was_run:
             raise RuntimeError(
                 "Evaluation engine has already been run. Create a new EvaluationEngine "
@@ -47,7 +81,7 @@ def run(self):
             self.trainer.test(
                 self.task, dataloaders=self.task.evaluation.dataset.test_dataloader
             )
-        return dict(
+        return EvaluationResults(
             compute=self.task.evaluation.metric.profiler.results(),
             metrics=self.trainer.callback_metrics,
         )