CLN: Refactor model abstraction, fix #737 #726

src/vak/common/tensorboard.py

@@ -128,5 +128,5 @@ def events2df(
         ).set_index("step")
         if drop_wall_time:
             dfs[scalar_tag].drop("wall_time", axis=1, inplace=True)
-    df = pd.concat([v for k, v in dfs.items()], axis=1)
+    df = pd.concat([v for k, v in dfs.items() if k != "epoch"], axis=1)
     return df

src/vak/config/model.py

@@ -78,6 +78,8 @@ def from_config_dict(cls, config_dict: dict):
                 f"Model name not found in registry: {model_name}\n"
                 f"Model names in registry:\n{MODEL_NAMES}"
             )
+
+        # NOTE: we are getting model_config here
         model_config = config_dict[model_name]
         if not all(key in MODEL_TABLES for key in model_config.keys()):
             invalid_keys = (
@@ -89,7 +91,7 @@ def from_config_dict(cls, config_dict: dict):
             )
         # for any tables not specified, default to empty dict so we can still use ``**`` operator on it
         for model_table in MODEL_TABLES:
-            if model_table not in config_dict:
+            if model_table not in model_config:
                 model_config[model_table] = {}
         return cls(name=model_name, **model_config)
 

src/vak/models/__init__.py

@@ -1,5 +1,5 @@
-from . import base, decorator, definition, registry
-from .base import Model
+from . import decorator, definition, factory, registry
+from .factory import ModelFactory
 from .convencoder_umap import ConvEncoderUMAP
 from .decorator import model
 from .ed_tcn import ED_TCN
@@ -10,14 +10,14 @@
 from .tweetynet import TweetyNet
 
 __all__ = [
-    "base",
+    "factory",
     "ConvEncoderUMAP",
     "decorator",
     "definition",
     "ED_TCN",
     "FrameClassificationModel",
     "get",
-    "Model",
+    "ModelFactory",
     "model",
     "model_family",
     "ParametricUMAPModel",

src/vak/models/decorator.py

@@ -1,22 +1,25 @@
-"""Decorator that makes a model class,
+"""Decorator that makes a :class:`vak.models.ModelFactory`,
 given a definition of the model,
-and another class that represents a
+and a :class:`lightning.LightningModule` that represents a
 family of models that the new model belongs to.
 
-The function returns a newly-created subclass
-of the class representing the family of models.
-The subclass can then be instantiated
-and have all model methods.
+The function returns a new instance of :class:`vak.models.ModelFactory`,
+that can create new instances of the model with its
+:meth:`~:class:`vak.models.ModelFactory.from_config` and
+:meth:`~:class:`vak.models.ModelFactory.from_instances` methods.
 """
 
 from __future__ import annotations
 
-from typing import Type
+from typing import Type, TYPE_CHECKING
+
+import lightning
 
-from .base import Model
 from .definition import validate as validate_definition
 from .registry import register_model
 
+if TYPE_CHECKING:
+    from .factory import ModelFactory
 
 class ModelDefinitionValidationError(Exception):
     """Exception raised when validating a model
@@ -28,16 +31,16 @@ class ModelDefinitionValidationError(Exception):
     pass
 
 
-def model(family: Type[Model]):
-    """Decorator that makes a model class,
+def model(family: lightning.pytorch.LightningModule):
+    """Decorator that makes a :class:`vak.models.ModelFactory`,
     given a definition of the model,
-    and another class that represents a
+    and a :class:`lightning.LightningModule` that represents a
     family of models that the new model belongs to.
 
-    Returns a newly-created subclass
-    of the class representing the family of models.
-    The subclass can then be instantiated
-    and have all model methods.
+    The function returns a new instance of :class:`vak.models.ModelFactory`,
+    that can create new instances of the model with its
+    :meth:`~:class:`vak.models.ModelFactory.from_config` and
+    :meth:`~:class:`vak.models.ModelFactory.from_instances` methods.
 
     Parameters
     ----------
@@ -46,50 +49,40 @@ def model(family: Type[Model]):
         A class with all the class variables required
         by :func:`vak.models.definition.validate`.
         See docstring of that function for specification.
-    family : subclass of vak.models.Model
+        See also :class:`vak.models.definition.ModelDefinition`,
+        but note that it is not necessary to subclass
+        :class:`~vak.models.definition.ModelDefinition` to
+        define a model.
+    family : lightning.LightningModule
         The class representing the family of models
         that the new model will belong to.
         E.g., :class:`vak.models.FrameClassificationModel`.
+        Should be a subclass of :class:`lightning.LightningModule`
+        that was registered with the
+        :func:`vak.models.registry.model_family` decorator.
 
     Returns
     -------
-    model : type
-        A sub-class of ``model_family``,
-        with attribute ``definition``,
+    model_factory : vak.models.ModelFactory
+        An instance of :class:`~vak.models.ModelFactory`,
+        with attribute ``definition`` and ``family``,
         that will be used when making
-        new instances of the model.
+        new instances of the model by calling the
+        :meth:`~vak.models.ModelFactory.from_config` method
+        or the :meth:`~:class:`vak.models.ModelFactory.from_instances` method.
     """
 
-    def _model(definition: Type):
-        if not issubclass(family, Model):
-            raise TypeError(
-                "The ``family`` argument to the ``vak.models.model`` decorator"
-                "should be a subclass of ``vak.models.base.Model``,"
-                f"but the type was: {type(family)}, "
-                "which was not recognized as a subclass "
-                "of ``vak.models.base.Model``."
-            )
-
-        try:
-            validate_definition(definition)
-        except ValueError as err:
-            raise ModelDefinitionValidationError(
-                f"Validation failed for the following model definition:\n{definition}"
-            ) from err
-        except TypeError as err:
-            raise ModelDefinitionValidationError(
-                f"Validation failed for the following model definition:\n{definition}"
-            ) from err
-
-        attributes = dict(family.__dict__)
-        attributes.update({"definition": definition})
-        subclass_name = definition.__name__
-        subclass = type(subclass_name, (family,), attributes)
-        subclass.__module__ = definition.__module__
-
-        # finally, add model to registry
-        register_model(subclass)
-
-        return subclass
+    def _model(definition: Type) -> ModelFactory:
+        from .factory import ModelFactory  # avoid circular import
+
+        model_factory = ModelFactory(
+            definition,
+            family
+        )
+        model_factory.__name__ = definition.__name__
+        model_factory.__doc__ = definition.__doc__
+        model_factory.__module__ = definition.__module__
+        register_model(model_factory)
+        return model_factory
 
     return _model

src/vak/models/definition.py

@@ -27,10 +27,7 @@
 class ModelDefinition:
     """A class that represents the definition of a neural network model.
 
-    Note it is **not** necessary to sub-class this class;
-    it exists mainly for type-checking purposes.
-
-    A model definition is a class that has the following class variables:
+    A model definition is any class that has the following class variables:
 
     network: torch.nn.Module or dict
         Neural network.
@@ -48,6 +45,12 @@ class ModelDefinition:
         Used by ``vak.models.base.Model`` and its
         sub-classes that represent model families. E.g., those classes will do:
         ``network = self.definition.network(**self.definition.default_config['network'])``.
+
+    Note it is **not** necessary to sub-class this class;
+    it exists mainly for type-checking purposes.
+
+    For more detail, see :func:`vak.models.decorator.model`
+    and :class:`vak.models.ModelFactory`.
     """
 
     network: Union[torch.nn.Module, dict]
@@ -67,7 +70,7 @@ class ModelDefinition:
 }
 
 
-def validate(definition: Type[ModelDefinition]) -> Type[ModelDefinition]:
+def validate(definition: Type) -> Type:
     """Validate a model definition.
 
     A model definition is a class that has the following class variables:
@@ -124,8 +127,8 @@ def validate(definition: Type[ModelDefinition]) -> Type[ModelDefinition]:
     converting it into a sub-class ofhttps://peps.python.org/pep-0416/
     ``vak.models.Model``.
 
-    It's also used by ``vak.models.Model``
-    to validate a definition when initializing
+    It's also used by :class:`vak.models.ModelFactory`,
+    to validate a definition before building 
     a new model instance from the definition.
     """
     # need to set this default first