add cmd

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.3.0
 +++++
 
+* :pr:`30`: adds command to test a model id
 * :pr:`29`: adds helpers to measure the memory peak and run benchmark
   on different processes
 * :pr:`28`: adds command line to print out the configuration for a model id,

_doc/api/torch_models/index.rst

@@ -7,6 +7,7 @@ onnx_diagnostic.torch_models
 
     hghub/index
     llms
+    test_helper
 
 .. automodule:: onnx_diagnostic.torch_models
     :members:

_doc/api/torch_models/test_helper.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.torch_models.test_helper
+========================================
+
+.. automodule:: onnx_diagnostic.torch_models.test_helper
+    :members:
+    :no-undoc-members:

_unittests/ut_torch_models/test_test_helpers.py

@@ -0,0 +1,20 @@
+import copy
+import unittest
+from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.torch_models.test_helper import get_inputs_for_task
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_get_inputs_function_for_tasks
+
+
+class TestTestHelper(ExtTestCase):
+    def test_get_inputs_for_task(self):
+        fcts = get_get_inputs_function_for_tasks()
+        for task in self.subloop(sorted(fcts)):
+            data = get_inputs_for_task(task)
+            self.assertIsInstance(data, dict)
+            self.assertIn("inputs", data)
+            self.assertIn("dynamic_shapes", data)
+            copy.deepcopy(data["inputs"])
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_xrun_doc/test_command_lines.py

@@ -4,11 +4,12 @@
 from onnx_diagnostic.ext_test_case import ExtTestCase
 from onnx_diagnostic._command_lines_parser import (
     get_main_parser,
+    get_parser_config,
     get_parser_find,
     get_parser_lighten,
     get_parser_print,
     get_parser_unlighten,
-    get_parser_config,
+    get_parser_validate,
 )
 
 
@@ -55,6 +56,13 @@ def test_parser_config(self):
         text = st.getvalue()
         self.assertIn("mid", text)
 
+    def test_parser_validate(self):
+        st = StringIO()
+        with redirect_stdout(st):
+            get_parser_validate().print_help()
+        text = st.getvalue()
+        self.assertIn("mid", text)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_xrun_doc/test_command_lines_exe.py

@@ -34,6 +34,13 @@ def test_parser_config(self):
         text = st.getvalue()
         self.assertIn("LlamaForCausalLM", text)
 
+    def test_parser_validate(self):
+        st = StringIO()
+        with redirect_stdout(st):
+            main(["validate", "-t", "text-generation"])
+        text = st.getvalue()
+        self.assertIn("dynamic_shapes", text)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/_command_lines_parser.py

@@ -227,22 +227,20 @@ def _cmd_config(argv: List[Any]):
         print(f"task: {task_from_id(args.mid)}")
 
 
-def get_parser_inputs() -> ArgumentParser:
+def get_parser_validate() -> ArgumentParser:
     parser = ArgumentParser(
         prog="test",
         description=dedent(
             """
         Prints out dummy inputs for a particular task or a model id.
         """
         ),
-        epilog="If the model id is specified, one untrained "
-        "version of it is instantiated.",
+        epilog="If the model id is specified, one untrained version of it is instantiated.",
     )
     parser.add_argument(
         "-m",
         "--mid",
         type=str,
-        required=True,
         help="model id, usually <author>/<name>",
     )
     parser.add_argument(
@@ -274,6 +272,29 @@ def get_parser_inputs() -> ArgumentParser:
     return parser
 
 
+def _cmd_validate(argv: List[Any]):
+    from .helpers import string_type
+    from .torch_models.test_helper import get_inputs_for_task
+
+    parser = get_parser_validate()
+    args = parser.parse_args(argv[1:])
+    assert args.task or args.mid, "A model id or a task needs to be specified."
+    if not args.mid:
+        data = get_inputs_for_task(args.task)
+        if args.verbose:
+            print(f"task: {args.task}")
+        max_length = max(len(k) for k in data["inputs"]) + 1
+        print("-- inputs")
+        for k, v in data["inputs"].items():
+            print(f"  + {k.ljust(max_length)}: {string_type(v, with_shape=True)}")
+        print("-- dynamic_shapes")
+        for k, v in data["dynamic_shapes"].items():
+            vs = str(v).replace("<class 'onnx_diagnostic.torch_models.hghub.model_inputs.", "").replace("'>", "").replace("_DimHint(type=<_DimHintType.DYNAMIC: 3>", "DYNAMIC").replace("_DimHint(type=<_DimHintType.AUTO: 3>", "AUTO")
+            print(f"  + {k.ljust(max_length)}: {vs}")
+
+    # validate_model(args.input, verbose=args.verbose, watch=set(args.names.split(",")))
+
+
 def get_main_parser() -> ArgumentParser:
     parser = ArgumentParser(
         prog="onnx_diagnostic",
@@ -289,13 +310,13 @@ def get_main_parser() -> ArgumentParser:
         lighten    - makes an onnx model lighter by removing the weights,
         unlighten  - restores an onnx model produces by the previous experiment
         print      - prints the model on standard output
-        test       - tests a model
+        validate   - validate a model
         """
         ),
     )
     parser.add_argument(
         "cmd",
-        choices=["config", "find", "lighten", "print", "unlighten", "test"],
+        choices=["config", "find", "lighten", "print", "unlighten", "validate"],
         help="Selects a command.",
     )
     return parser
@@ -308,7 +329,7 @@ def main(argv: Optional[List[Any]] = None):
         print=_cmd_print,
         find=_cmd_find,
         config=_cmd_config,
-        text=_cmd_test,
+        validate=_cmd_validate,
     )
 
     if argv is None:
@@ -328,6 +349,7 @@ def main(argv: Optional[List[Any]] = None):
                 print=get_parser_print,
                 find=get_parser_find,
                 config=get_parser_config,
+                validate=get_parser_validate,
             )
             cmd = argv[0]
             if cmd not in parsers:

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -141,69 +141,92 @@ def _pick(config, *atts):
 
 
 def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
-    """Inputs kwargs"""
+    """
+    Inputs kwargs.
+
+    If the configuration is None, the function selects typical dimensions.
+    """
+    fcts = get_get_inputs_function_for_tasks()
+    assert task in fcts, f"Unsupported task {task!r}, supprted are {sorted(fcts)}"
     if task == "text-generation":
-        check_hasattr(
-            config,
-            "vocab_size",
-            "hidden_size",
-            "num_attention_heads",
-            ("num_key_value_heads", "num_attention_heads"),
-            "intermediate_size",
-            "hidden_size",
-        )
+        if config is not None:
+            check_hasattr(
+                config,
+                "vocab_size",
+                "hidden_size",
+                "num_attention_heads",
+                ("num_key_value_heads", "num_attention_heads"),
+                "intermediate_size",
+                "hidden_size",
+            )
         kwargs = dict(
             batch_size=2,
             sequence_length=30,
             sequence_length2=3,
-            head_dim=getattr(
-                config, "head_dim", config.hidden_size // config.num_attention_heads
+            head_dim=(
+                16
+                if config is None
+                else getattr(
+                    config, "head_dim", config.hidden_size // config.num_attention_heads
+                )
             ),
-            dummy_max_token_id=config.vocab_size - 1,
-            num_hidden_layers=config.num_hidden_layers,
-            num_key_value_heads=_pick(config, "num_key_value_heads", "num_attention_heads"),
-            intermediate_size=config.intermediate_size,
-            hidden_size=config.hidden_size,
+            dummy_max_token_id=31999 if config is None else (config.vocab_size - 1),
+            num_hidden_layers=4 if config is None else config.num_hidden_layers,
+            num_key_value_heads=(
+                24
+                if config is None
+                else _pick(config, "num_key_value_heads", "num_attention_heads")
+            ),
+            intermediate_size=1024 if config is None else config.intermediate_size,
+            hidden_size=512 if config is None else config.hidden_size,
         )
         fct = get_inputs_for_text_generation
     elif task == "text2text-generation":
-        check_hasattr(
-            config,
-            "vocab_size",
-            "hidden_size",
-            "num_attention_heads",
-            ("num_hidden_layers", "num_layers"),
-            ("n_positions", "d_model"),
-            (
-                "num_key_value_heads",
-                "num_heads",
-                ("decoder_attention_heads", "encoder_attention_heads"),
-            ),
-        )
+        if config is not None:
+            check_hasattr(
+                config,
+                "vocab_size",
+                "hidden_size",
+                "num_attention_heads",
+                ("num_hidden_layers", "num_layers"),
+                ("n_positions", "d_model"),
+                (
+                    "num_key_value_heads",
+                    "num_heads",
+                    ("decoder_attention_heads", "encoder_attention_heads"),
+                ),
+            )
         kwargs = dict(
             batch_size=2,
             sequence_length=30,
             sequence_length2=3,
-            head_dim=config.d_kv if hasattr(config, "d_kv") else 1,
-            dummy_max_token_id=config.vocab_size - 1,
-            num_hidden_layers=_pick(config, "num_hidden_layers", "num_layers"),
-            num_key_value_heads=_pick(
-                config,
-                "num_key_value_heads",
-                "num_heads",
-                (sum, "encoder_attention_heads", "decoder_attention_heads"),
+            head_dim=16 if config is None else (config.d_kv if hasattr(config, "d_kv") else 1),
+            dummy_max_token_id=31999 if config is None else config.vocab_size - 1,
+            num_hidden_layers=(
+                8 if config is None else _pick(config, "num_hidden_layers", "num_layers")
+            ),
+            num_key_value_heads=(
+                16
+                if config is None
+                else _pick(
+                    config,
+                    "num_key_value_heads",
+                    "num_heads",
+                    (sum, "encoder_attention_heads", "decoder_attention_heads"),
+                )
             ),
-            encoder_dim=_pick(config, "n_positions", "d_model"),
+            encoder_dim=512 if config is None else _pick(config, "n_positions", "d_model"),
         )
         fct = get_inputs_for_text2text_generation  # type: ignore
     elif task == "image-classification":
-        check_hasattr(config, "image_size", "num_channels")
-        if isinstance(config.image_size, int):
+        if config is not None:
+            check_hasattr(config, "image_size", "num_channels")
+        if config is None or isinstance(config.image_size, int):
             kwargs = dict(
                 batch_size=2,
-                input_width=config.image_size,
-                input_height=config.image_size,
-                input_channels=config.num_channels,
+                input_width=224 if config is None else config.image_size,
+                input_height=224 if config is None else config.image_size,
+                input_channels=3 if config is None else config.num_channels,
             )
         else:
             kwargs = dict(
@@ -214,32 +237,41 @@ def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callabl
             )
         fct = get_inputs_for_image_classification  # type: ignore
     elif task == "image-text-to-text":
-        check_hasattr(
-            config,
-            "vocab_size",
-            "hidden_size",
-            "num_attention_heads",
-            ("num_key_value_heads", "num_attention_heads"),
-            "intermediate_size",
-            "hidden_size",
-            "vision_config",
-        )
-        check_hasattr(config.vision_config, "image_size", "num_channels")
+        if config is not None:
+            check_hasattr(
+                config,
+                "vocab_size",
+                "hidden_size",
+                "num_attention_heads",
+                ("num_key_value_heads", "num_attention_heads"),
+                "intermediate_size",
+                "hidden_size",
+                "vision_config",
+            )
+            check_hasattr(config.vision_config, "image_size", "num_channels")
         kwargs = dict(
             batch_size=2,
             sequence_length=30,
             sequence_length2=3,
-            head_dim=getattr(
-                config, "head_dim", config.hidden_size // config.num_attention_heads
+            head_dim=(
+                16
+                if config is None
+                else getattr(
+                    config, "head_dim", config.hidden_size // config.num_attention_heads
+                )
+            ),
+            dummy_max_token_id=31999 if config is None else config.vocab_size - 1,
+            num_hidden_layers=4 if config is None else config.num_hidden_layers,
+            num_key_value_heads=(
+                8
+                if config is None
+                else _pick(config, "num_key_value_heads", "num_attention_heads")
             ),
-            dummy_max_token_id=config.vocab_size - 1,
-            num_hidden_layers=config.num_hidden_layers,
-            num_key_value_heads=_pick(config, "num_key_value_heads", "num_attention_heads"),
-            intermediate_size=config.intermediate_size,
-            hidden_size=config.hidden_size,
-            width=config.vision_config.image_size,
-            height=config.vision_config.image_size,
-            num_channels=config.vision_config.num_channels,
+            intermediate_size=1024 if config is None else config.intermediate_size,
+            hidden_size=512 if config is None else config.hidden_size,
+            width=224 if config is None else config.vision_config.image_size,
+            height=224 if config is None else config.vision_config.image_size,
+            num_channels=3 if config is None else config.vision_config.num_channels,
         )
         fct = get_inputs_for_image_text_to_text  # type: ignore
     else:
@@ -682,3 +714,13 @@ def get_inputs_for_text2text_generation(
         # encoder_outputs=torch.randn(batch_size, sequence_length2, encoder_dim),
     )
     return dict(inputs=inputs, dynamic_shapes=shapes)
+
+
+def get_get_inputs_function_for_tasks() -> Dict[str, Callable]:
+    """Returns all the function producing dummy inputs for every task."""
+    return {
+        "image-classification": get_inputs_for_image_classification,
+        "text-generation": get_inputs_for_text_generation,
+        "text2text-generation": get_inputs_for_text2text_generation,
+        "image-text-to-text": get_inputs_for_image_text_to_text,
+    }