[Bugfix] recipe_template cli

src/sparseml/pytorch/recipe_template/cli.py

@@ -14,15 +14,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-Usage: sparseml.pytorch.recipe_template [OPTIONS]
+Usage: sparseml.recipe_template [OPTIONS]
 
   Utility to create a recipe template based on specified options
 
   Example for using sparseml.recipe_template:
 
-       `sparseml.pytorch.recipe_template --pruning true --quantization true`
+       `sparseml.recipe_template --pruning true --quantization true`
 
-       `sparseml.pytorch.recipe_template --quantization vnni --lr constant `
+       `sparseml.recipe_template --quantization vnni --lr constant `
 
 Options:
   --version                       Show the version and exit.  [default: False]
@@ -100,11 +100,17 @@ def main(**kwargs):
 
     Example for using sparseml.pytorch.recipe_template:
 
-         `sparseml.pytorch.recipe_template --pruning true --quantization true`
+         `sparseml.recipe_template --pruning true --quantization true`
 
-         `sparseml.pytorch.recipe_template --quantization vnni --lr constant `
+         `sparseml.recipe_template --quantization true --target vnni --lr constant `
     """
     _LOGGER.debug(f"{kwargs}")
+    if kwargs.get("quantization") == "true":
+        target = kwargs.get("target")
+        if target == "vnni":
+            kwargs["mask_type"] = "block4"
+        elif target == "tensorrt":
+            kwargs["mask_type"] = "2:4"
     template = recipe_template(**kwargs)
     print(f"Template:\n{template}")
 

tests/sparseml/pytorch/recipe_template/test_end_to_end.py

@@ -12,9 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
 
+from click.testing import CliRunner
 from sparseml.pytorch import recipe_template
+from sparseml.pytorch.recipe_template.cli import main
 
 
 def test_function_entrypoint():
     recipe_template()
+
+
+@pytest.mark.parametrize(
+    "command",
+    [
+        ["--pruning", "true", "--quantization", "true"],
+        ["--quantization", "true", "--target", "vnni", "--lr", "constant"],
+    ],
+)
+def test_docstring_cli_examples(command, tmp_path):
+    runner = CliRunner()
+    command.extend(["--file_name", str(tmp_path / "temp.md")])
+    result = runner.invoke(main, command)
+    assert result.exit_code == 0