recipe metadata updates (#394)

* metadata update path

* updates from review

Author: bfineran

src/sparseml/keras/optim/manager.py

@@ -25,8 +25,7 @@
 from sparseml.keras.optim.modifier import Modifier, ScheduledModifier
 from sparseml.keras.utils.compat import keras
 from sparseml.keras.utils.logger import KerasLogger
-from sparseml.optim import BaseManager
-from sparseml.utils import load_recipe_yaml_str
+from sparseml.optim import BaseManager, load_recipe_yaml_str
 from sparsezoo.objects import Recipe
 
 
@@ -42,6 +41,7 @@ class ScheduledModifierManager(BaseManager, Modifier):
     def from_yaml(
         file_path: Union[str, Recipe],
         add_modifiers: List[Modifier] = None,
+        **recipe_variables,
     ):
         """
         Convenience function used to create the manager of multiple modifiers from a
@@ -55,9 +55,11 @@ def from_yaml(
             'zoo:model/stub/path?recipe_type=transfer'
         :param add_modifiers: additional modifiers that should be added to the
             returned manager alongside the ones loaded from the recipe file
+        :param recipe_variables: additional variable values to override the recipe
+            with (i.e. num_epochs, init_lr)
         :return: ScheduledModifierManager() created from the recipe file
         """
-        yaml_str = load_recipe_yaml_str(file_path)
+        yaml_str = load_recipe_yaml_str(file_path, **recipe_variables)
         modifiers = Modifier.load_list(yaml_str)
         if add_modifiers:
             modifiers.extend(add_modifiers)

src/sparseml/optim/helpers.py

@@ -17,20 +17,103 @@
 """
 
 
+import logging
 import re
 from typing import Any, Dict, Tuple, Union
 
 import yaml
 
 from sparseml.utils import UnknownVariableException, restricted_eval
+from sparsezoo import Zoo
+from sparsezoo.objects import Recipe
 
 
 __all__ = [
+    "load_recipe_yaml_str",
     "load_recipe_yaml_str_no_classes",
     "rewrite_recipe_yaml_string_with_classes",
+    "update_recipe_variables",
     "evaluate_recipe_yaml_str_equations",
 ]
 
+_LOGGER = logging.getLogger(__name__)
+
+
+def load_recipe_yaml_str(
+    file_path: Union[str, Recipe],
+    **variable_overrides,
+) -> str:
+    """
+    Loads a YAML recipe file to a string or
+    extracts recipe from YAML front matter in a sparsezoo markdown recipe card.
+    Recipes can also be provided as SparseZoo model stubs or Recipe
+    objects.
+
+    YAML front matter: https://jekyllrb.com/docs/front-matter/
+
+    :param file_path: file path to recipe YAML file or markdown recipe card or
+        stub to a SparseZoo model whose recipe will be downloaded and loaded.
+        SparseZoo stubs should be preceded by 'zoo:', and can contain an optional
+        '?recipe_type=<type>' parameter or include a `/<type>` subpath. Can also
+        be a SparseZoo Recipe object. i.e. '/path/to/local/recipe.yaml',
+        'zoo:model/stub/path', 'zoo:model/stub/path?recipe_type=transfer_learn',
+        'zoo:model/stub/path/transfer_learn'
+    :param variable_overrides: dict of variable values to replace
+        in the loaded yaml string. Default is None
+    :return: the recipe YAML configuration loaded as a string
+    """
+    if isinstance(file_path, Recipe):
+        # download and unwrap Recipe object
+        file_path = file_path.downloaded_path()
+
+    if not isinstance(file_path, str):
+        raise ValueError(f"file_path must be a str, given {type(file_path)}")
+
+    if file_path.startswith("zoo:"):
+        # download from zoo stub
+        recipe = Zoo.download_recipe_from_stub(file_path)
+        file_path = recipe.downloaded_path()
+
+    # load the yaml string
+    if "\n" in file_path or "\r" in file_path:
+        # treat as raw yaml passed in
+        yaml_str = file_path
+        extension = "unknown"
+    else:
+        # load yaml from file_path
+        extension = file_path.lower().split(".")[-1]
+        if extension not in ["md", "yaml"]:
+            raise ValueError(
+                "Unsupported file extension for recipe. Excepted '.md' or '.yaml'. "
+                f"Received {file_path}"
+            )
+        with open(file_path, "r") as yaml_file:
+            yaml_str = yaml_file.read()
+
+    if extension == "md" or extension == "unknown":
+        # extract YAML front matter from markdown recipe card
+        # adapted from
+        # https://github.com/jonbeebe/frontmatter/blob/master/frontmatter
+        yaml_delim = r"(?:---|\+\+\+)"
+        yaml = r"(.*?)"
+        re_pattern = r"^\s*" + yaml_delim + yaml + yaml_delim
+        regex = re.compile(re_pattern, re.S | re.M)
+        result = regex.search(yaml_str)
+
+        if result:
+            yaml_str = result.group(1)
+        elif extension == "md":
+            # fail if we know whe should have extracted front matter out
+            raise RuntimeError(
+                "Could not extract YAML front matter from recipe card:"
+                " {}".format(file_path)
+            )
+
+    if variable_overrides:
+        update_recipe_variables(yaml_str, variable_overrides)
+
+    return yaml_str
+
 
 def load_recipe_yaml_str_no_classes(recipe_yaml_str: str) -> str:
     """
@@ -57,6 +140,30 @@ def rewrite_recipe_yaml_string_with_classes(recipe_contianer: Any) -> str:
     return pattern.sub(r"!\g<class_name>", updated_yaml_str)
 
 
+def update_recipe_variables(recipe_yaml_str: str, variables: Dict[str, Any]) -> str:
+    """
+    :param recipe_yaml_str: YAML string of a SparseML recipe
+    :param variables: variables dictionary to update recipe top level variables with
+    :return: given recipe with variables updated
+    """
+
+    container = load_recipe_yaml_str_no_classes(recipe_yaml_str)
+    if not isinstance(container, dict):
+        # yaml string does not create a dict, return original string
+        return recipe_yaml_str
+
+    for key in variables:
+        if key not in container:
+            _LOGGER.warning(
+                f"updating recipe variable {key} but {key} is not currently "
+                "set in existing recipe. This change may have no impact on the recipe "
+                "modifiers"
+            )
+
+    container.update(variables)
+    return rewrite_recipe_yaml_string_with_classes(container)
+
+
 def evaluate_recipe_yaml_str_equations(recipe_yaml_str: str) -> str:
     """
     :param recipe_yaml_str: YAML string of a SparseML recipe

src/sparseml/pytorch/optim/manager.py

@@ -24,10 +24,9 @@
 from torch.nn import Module
 from torch.optim.optimizer import Optimizer
 
-from sparseml.optim import BaseManager
+from sparseml.optim import BaseManager, load_recipe_yaml_str
 from sparseml.pytorch.optim.modifier import Modifier, ScheduledModifier
 from sparseml.pytorch.utils import BaseLogger
-from sparseml.utils import load_recipe_yaml_str
 from sparsezoo.objects import Recipe
 
 
@@ -248,6 +247,7 @@ class ScheduledModifierManager(BaseManager, Modifier):
     def from_yaml(
         file_path: Union[str, Recipe],
         add_modifiers: List[Modifier] = None,
+        **recipe_variables,
     ):
         """
         Convenience function used to create the manager of multiple modifiers from a
@@ -261,9 +261,11 @@ def from_yaml(
             'zoo:model/stub/path?recipe_type=transfer'
         :param add_modifiers: additional modifiers that should be added to the
             returned manager alongside the ones loaded from the recipe file
+        :param recipe_variables: additional variable values to override the recipe
+            with (i.e. num_epochs, init_lr)
         :return: ScheduledModifierManager() created from the recipe file
         """
-        yaml_str = load_recipe_yaml_str(file_path)
+        yaml_str = load_recipe_yaml_str(file_path, **recipe_variables)
         modifiers = Modifier.load_list(yaml_str)
 
         if add_modifiers:

src/sparseml/tensorflow_v1/optim/manager.py

@@ -21,10 +21,9 @@
 import itertools
 from typing import Any, Callable, Dict, List, Tuple, Union
 
-from sparseml.optim import BaseManager, BaseScheduled
+from sparseml.optim import BaseManager, BaseScheduled, load_recipe_yaml_str
 from sparseml.tensorflow_v1.optim.modifier import NM_RECAL, Modifier, ScheduledModifier
 from sparseml.tensorflow_v1.utils import tf_compat
-from sparseml.utils import load_recipe_yaml_str
 from sparsezoo.objects import Recipe
 
 
@@ -78,6 +77,7 @@ class ScheduledModifierManager(BaseManager, Modifier):
     def from_yaml(
         file_path: Union[str, Recipe],
         add_modifiers: List[Modifier] = None,
+        **recipe_variables,
     ):
         """
         Convenience function used to create the manager of multiple modifiers from a
@@ -91,9 +91,11 @@ def from_yaml(
             'zoo:model/stub/path?recipe_type=transfer'
         :param add_modifiers: additional modifiers that should be added to the
             returned manager alongside the ones loaded from the recipe file
+        :param recipe_variables: additional variable values to override the recipe
+            with (i.e. num_epochs, init_lr)
         :return: ScheduledModifierManager() created from the recipe file
         """
-        yaml_str = load_recipe_yaml_str(file_path)
+        yaml_str = load_recipe_yaml_str(file_path, **recipe_variables)
         modifiers = Modifier.load_list(yaml_str)
         if add_modifiers:
             modifiers.extend(add_modifiers)

src/sparseml/utils/helpers.py

@@ -21,16 +21,13 @@
 import fnmatch
 import logging
 import os
-import re
 import sys
 from collections import OrderedDict
 from typing import Any, Callable, Dict, Iterable, List, Tuple, Union
 from urllib.parse import urlparse
 
 import numpy
 
-from sparsezoo import Zoo
-from sparsezoo.objects import Recipe
 from sparsezoo.utils import load_numpy_list
 
 
@@ -59,7 +56,6 @@
     "NumpyArrayBatcher",
     "tensor_export",
     "tensors_export",
-    "load_recipe_yaml_str",
     "parse_optimization_str",
 ]
 
@@ -768,74 +764,6 @@ def _tensors_export_batch(
     )
 
 
-def load_recipe_yaml_str(file_path: Union[str, Recipe]) -> str:
-    """
-    Loads a YAML recipe file to a string or
-    extracts recipe from YAML front matter in a sparsezoo markdown recipe card.
-    Recipes can also be provided as SparseZoo model stubs or Recipe
-    objects.
-
-    YAML front matter: https://jekyllrb.com/docs/front-matter/
-
-    :param file_path: file path to recipe YAML file or markdown recipe card or
-        stub to a SparseZoo model whose recipe will be downloaded and loaded.
-        SparseZoo stubs should be preceded by 'zoo:', and can contain an optional
-        '?recipe_type=<type>' parameter or include a `/<type>` subpath. Can also
-        be a SparseZoo Recipe object. i.e. '/path/to/local/recipe.yaml',
-        'zoo:model/stub/path', 'zoo:model/stub/path?recipe_type=transfer_learn',
-        'zoo:model/stub/path/transfer_learn'
-    :return: the recipe YAML configuration loaded as a string
-    """
-    if isinstance(file_path, Recipe):
-        # download and unwrap Recipe object
-        file_path = file_path.downloaded_path()
-
-    if not isinstance(file_path, str):
-        raise ValueError(f"file_path must be a str, given {type(file_path)}")
-
-    if file_path.startswith("zoo:"):
-        # download from zoo stub
-        recipe = Zoo.download_recipe_from_stub(file_path)
-        file_path = recipe.downloaded_path()
-
-    # load the yaml string
-    if "\n" in file_path or "\r" in file_path:
-        # treat as raw yaml passed in
-        yaml_str = file_path
-        extension = "unknown"
-    else:
-        # load yaml from file_path
-        extension = file_path.lower().split(".")[-1]
-        if extension not in ["md", "yaml"]:
-            raise ValueError(
-                "Unsupported file extension for recipe. Excepted '.md' or '.yaml'. "
-                "Received {}".format(file_path)
-            )
-        with open(file_path, "r") as yaml_file:
-            yaml_str = yaml_file.read()
-
-    if extension == "md" or extension == "unknown":
-        # extract YAML front matter from markdown recipe card
-        # adapted from
-        # https://github.com/jonbeebe/frontmatter/blob/master/frontmatter
-        yaml_delim = r"(?:---|\+\+\+)"
-        yaml = r"(.*?)"
-        re_pattern = r"^\s*" + yaml_delim + yaml + yaml_delim
-        regex = re.compile(re_pattern, re.S | re.M)
-        result = regex.search(yaml_str)
-
-        if result:
-            yaml_str = result.group(1)
-        elif extension == "md":
-            # fail if we know whe should have extracted front matter out
-            raise RuntimeError(
-                "Could not extract YAML front matter from recipe card:"
-                " {}".format(file_path)
-            )
-
-    return yaml_str
-
-
 def parse_optimization_str(optim_full_name: str) -> Tuple[str, str, Any]:
     """
     :param optim_full_name: A name of a pretrained model optimization. i.e.