Supports model options in command lines (#70)

* Support model options in command lines

* spell

Author: xadupre

CHANGELOGS.rst

@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.4.1
++++++
+
+* :pr:`70`: support models options in command lines
+
 0.4.0
 +++++
 

_doc/api/tasks/index.rst

@@ -1,6 +1,33 @@
 onnx_diagnostic.tasks
 =====================
 
+All submodules contains the three following functions:
+
+* ``reduce_model_config(config) -> kwargs``:
+  updates the configuration to get a smaller model more suitable
+  for unit tests
+* ``random_input_kwargs(config) -> kwargs, get_inputs``:
+  produces values ``get_inputs`` can take to generate dummy inputs
+  suitable for a model defined by its configuration
+* ``get_inputs(model, config, *args, **kwargs) -> dict(inputs=..., dynamic_shapes=...)``:
+  generates the dummy inputs and dynamic shapes for a specific model and configuration.
+
+For a specific task, you would write:
+
+.. code-block:: python
+
+    kwargs, get_inputs = random_input_kwargs(config)
+    dummies = get_inputs(model, config, **kwargs)
+
+Or:
+
+.. code-block:: python
+
+    from onnx_diagnostic.tasks import random_input_kwargs
+
+    kwargs, get_inputs = random_input_kwargs(config, task)  # "text-generation" for example
+    dummies = get_inputs(model, config, **kwargs)
+
 .. toctree::
     :maxdepth: 1
     :caption: modules

_doc/index.rst

@@ -173,6 +173,7 @@ Size of the package:
 Older versions
 ++++++++++++++
 
+* `0.4.1 <../v0.4.1/index.html>`_
 * `0.4.0 <../v0.4.0/index.html>`_
 * `0.3.0 <../v0.3.0/index.html>`_
 * `0.2.2 <../v0.2.2/index.html>`_

_unittests/ut_tasks/try_tasks.py

@@ -98,6 +98,32 @@ def test_text2text_generation(self):
             )
         print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
 
+    @never_test()
+    def test_text_generation_phi4(self):
+        # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k phi4
+
+        import torch
+        from transformers import RobertaTokenizer, T5ForConditionalGeneration
+
+        tokenizer = RobertaTokenizer.from_pretrained("microsoft/Phi-4-mini-instruct")
+        model = T5ForConditionalGeneration.from_pretrained("microsoft/Phi-4-mini-instruct")
+
+        text = "def greet(user): print(f'hello <extra_id_0>!')"
+        input_ids = tokenizer(text, return_tensors="pt").input_ids
+        mask = (
+            torch.tensor([1 for i in range(input_ids.shape[1])])
+            .to(torch.int64)
+            .reshape((1, -1))
+        )
+
+        # simply generate a single sequence
+        print()
+        with steal_forward(model):
+            generated_ids = model.generate(
+                decoder_input_ids=input_ids, attention_mask=mask, max_length=100
+            )
+        print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
+
     @never_test()
     def test_imagetext2text_generation(self):
         # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k etext2t

_unittests/ut_torch_models/test_hghub_api.py

@@ -74,7 +74,14 @@ def test_task_from_id_long(self):
     def test_get_pretrained_config(self):
         conf = get_pretrained_config("microsoft/phi-2")
         self.assertNotEmpty(conf)
-        print(conf)
+
+    @requires_transformers("4.50")
+    @requires_torch("2.7")
+    @hide_stdout()
+    def test_get_pretrained_config_options(self):
+        conf = get_pretrained_config("microsoft/phi-2", num_key_value_heads=16)
+        self.assertNotEmpty(conf)
+        self.assertEqual(conf.num_key_value_heads, 16)
 
     @requires_transformers("4.50")
     @requires_torch("2.7")

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.4.0"
+__version__ = "0.4.1"
 __author__ = "Xavier Dupré"

onnx_diagnostic/_command_lines_parser.py

@@ -214,6 +214,22 @@ def get_parser_config() -> ArgumentParser:
         action=BooleanOptionalAction,
         help="displays the task as well",
     )
+    parser.add_argument(
+        "-c",
+        "--cached",
+        default=True,
+        action=BooleanOptionalAction,
+        help="uses cached configuration, only available for some of them, "
+        "mostly for unit test purposes",
+    )
+    parser.add_argument(
+        "--mop",
+        metavar="KEY=VALUE",
+        nargs="*",
+        help="Additional model options, use to change some parameters of the model, "
+        "example: --mop attn_implementation=eager",
+        action=_ParseDict,
+    )
     return parser
 
 
@@ -222,7 +238,11 @@ def _cmd_config(argv: List[Any]):
 
     parser = get_parser_config()
     args = parser.parse_args(argv[1:])
-    print(get_pretrained_config(args.mid))
+    conf = get_pretrained_config(args.mid, **(args.mop or {}))
+    print(conf)
+    for k, v in sorted(conf.__dict__.items()):
+        if "_implementation" in k:
+            print(f"config.{k}={v!r}")
     if args.task:
         print("------")
         print(f"task: {task_from_id(args.mid)}")
@@ -238,6 +258,19 @@ def __call__(self, parser, namespace, values, option_string=None):
                 key = split_items[0].strip()  # we remove blanks around keys, as is logical
                 value = split_items[1]
 
+                if value in ("True", "true", "False", "false"):
+                    d[key] = bool(value)
+                    continue
+                try:
+                    d[key] = int(value)
+                    continue
+                except (TypeError, ValueError):
+                    pass
+                try:
+                    d[key] = float(value)
+                    continue
+                except (TypeError, ValueError):
+                    pass
                 d[key] = value
 
         setattr(namespace, self.dest, d)
@@ -321,6 +354,14 @@ def get_parser_validate() -> ArgumentParser:
         "inputs use to export, example: --iop cls_cache=SlidingWindowCache",
         action=_ParseDict,
     )
+    parser.add_argument(
+        "--mop",
+        metavar="KEY=VALUE",
+        nargs="*",
+        help="Additional model options, use to change some parameters of the model, "
+        "example: --mop attn_implementation=eager",
+        action=_ParseDict,
+    )
     return parser
 
 
@@ -371,6 +412,7 @@ def _cmd_validate(argv: List[Any]):
             drop_inputs=None if not args.drop else args.drop.split(","),
             ortfusiontype=args.ortfusiontype,
             input_options=args.iop,
+            model_options=args.mop,
         )
         print("")
         print("-- summary --")

onnx_diagnostic/helpers/config_helper.py

@@ -28,13 +28,18 @@ def check_hasattr(config: Any, *args: Union[str, Tuple[Any, ...]]):
 def update_config(config: Any, mkwargs: Dict[str, Any]):
     """Updates a configuration with different values."""
     for k, v in mkwargs.items():
+        if k == "attn_implementation":
+            config._attn_implementation = v
+            if getattr(config, "_attn_implementation_autoset", False):
+                config._attn_implementation_autoset = False
+            continue
         if isinstance(v, dict):
             assert hasattr(
                 config, k
             ), f"missing attribute {k!r} in config={config}, cannot update it with {v}"
             update_config(getattr(config, k), v)
-        else:
-            setattr(config, k, v)
+            continue
+        setattr(config, k, v)
 
 
 def _pick(config, *atts):

onnx_diagnostic/tasks/__init__.py

@@ -33,7 +33,7 @@ def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
     """Reduces a model size."""
     tasks = {mod.__TASK__: mod.reduce_model_config for mod in __TASKS__}
     assert task in tasks, f"Task {task!r} not found in {sorted(tasks)}"
-    return tasks[task](config, task)
+    return tasks[task](config)
 
 
 def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
@@ -45,4 +45,4 @@ def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callabl
     """
     tasks = {mod.__TASK__: mod.random_input_kwargs for mod in __TASKS__}
     assert task in tasks, f"Task {task!r} not found in {sorted(tasks)}"
-    return tasks[task](config, task)
+    return tasks[task](config)

onnx_diagnostic/tasks/automatic_speech_recognition.py

@@ -7,7 +7,7 @@
 __TASK__ = "automatic-speech-recognition"
 
 
-def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_decoder_layers"):
@@ -129,7 +129,7 @@ def get_inputs(
     return dict(inputs=inputs, dynamic_shapes=shapes)
 
 
-def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
+def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
     """
     Inputs kwargs.