Components as decorators (#84)

* Make components and tasks decorators rather than classes.

* Expose CLI.

* Bug fix.

* Fix another bug.

* Make experiment loss optional.

* Fix yet another bug.

* Bug fix once more.

* Another bug fix.

* Update defaults.

* WIP

* Tidy-up.

* Export `configure` correctly.

* Update examples.

* Make linting happy.

* Update zookeeper/core/component.py

Co-Authored-By: Lukas Geiger <[email protected]>

* Remove erroneous prints.

* Review suggestion.

* Move `colorama` import to top.

examples/larq_experiment.py

@@ -7,19 +7,20 @@
 
 import larq as lq
 import tensorflow as tf
-import tensorflow_datasets as tfds
 
-from zookeeper import Dataset, Experiment, Model, Preprocessing, TFDSDataset
-from zookeeper.cli import add_task_to_cli, cli
+from zookeeper import cli, component, task
+from zookeeper.tf import Dataset, Experiment, Model, Preprocessing, TFDSDataset
 
 
+@component
 class Cifar10(TFDSDataset):
     name = "cifar10"
     # CIFAR-10 has only train and test, so validate on test.
-    train_split = tfds.Split.TRAIN
-    validation_split = tfds.Split.TEST
+    train_split = "train"
+    validation_split = "test"
 
 
+@component
 class PadCropAndFlip(Preprocessing):
     pad_size: int
     output_size: int
@@ -44,6 +45,7 @@ def output(self, data):
         return data["label"]
 
 
+@component
 class BinaryNet(Model):
     dataset: Dataset
     preprocessing: Preprocessing
@@ -107,7 +109,7 @@ def build(self, input_shape):
         )
 
 
-@add_task_to_cli
+@task
 class BinaryNetCifar10(Experiment):
     dataset = Cifar10()
     preprocessing = PadCropAndFlip(pad_size=40, output_size=32)

zookeeper/__init__.py

@@ -1,16 +1,3 @@
-from zookeeper.component import Component
-from zookeeper.dataset import Dataset, TFDSDataset
-from zookeeper.experiment import Experiment
-from zookeeper.model import Model
-from zookeeper.preprocessing import Preprocessing
-from zookeeper.task import Task
+from zookeeper.core import cli, component, configure, task
 
-__all__ = [
-    "Component",
-    "Dataset",
-    "Experiment",
-    "Model",
-    "Preprocessing",
-    "Task",
-    "TFDSDataset",
-]
+__all__ = ["cli", "component", "configure", "task"]