onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -0,0 +1,388 @@
+import copy
+import inspect
+import os
+import pprint
+import time
+from typing import Any, Dict, Optional, Tuple
+import torch
+import transformers
+from ...helpers.config_helper import update_config, build_diff_config
+from ...tasks import reduce_model_config, random_input_kwargs
+from .hub_api import (
+    task_from_arch,
+    task_from_id,
+    get_pretrained_config,
+    download_code_modelid,
+    architecture_from_config,
+    find_package_source,
+)
+from .model_specific import HANDLED_MODELS, load_specific_model, instantiate_specific_model
+
+
+def _code_needing_rewriting(model: Any) -> Any:
+    from onnx_diagnostic.torch_export_patches.patch_module_helper import code_needing_rewriting
+
+    return code_needing_rewriting(model)
+
+
+def _preprocess_model_id(
+    model_id: str, subfolder: Optional[str], same_as_pretrained: bool, use_pretrained: bool
+) -> Tuple[str, Optional[str], bool, bool]:
+    if subfolder or "//" not in model_id:
+        return model_id, subfolder, same_as_pretrained, use_pretrained
+    spl = model_id.split("//")
+    if spl[-1] == "pretrained":
+        return _preprocess_model_id("//".join(spl[:-1]), "", True, True)
+    if spl[-1] in {"transformer", "vae"}:
+        # known subfolder
+        return "//".join(spl[:-1]), spl[-1], same_as_pretrained, use_pretrained
+    return model_id, subfolder, same_as_pretrained, use_pretrained
+
+
+def get_untrained_model_with_inputs(
+    model_id: str,
+    config: Optional[Any] = None,
+    task: Optional[str] = "",
+    inputs_kwargs: Optional[Dict[str, Any]] = None,
+    model_kwargs: Optional[Dict[str, Any]] = None,
+    verbose: int = 0,
+    dynamic_rope: Optional[bool] = None,
+    use_pretrained: bool = False,
+    same_as_pretrained: bool = False,
+    use_preinstalled: bool = True,
+    add_second_input: int = 1,
+    subfolder: Optional[str] = None,
+    use_only_preinstalled: bool = False,
+) -> Dict[str, Any]:
+    """
+    Gets a non initialized model similar to the original model
+    based on the model id given to the function.
+    The model size is reduced compare to the original model.
+    No weight is downloaded, only the configuration file sometimes.
+
+    :param model_id: model id, ex: :epkg:`arnir0/Tiny-LLM`
+    :param config: to overwrite the configuration
+    :param task: model task, can be overwritten, otherwise, it is automatically determined
+    :param input_kwargs: parameters sent to input generation
+    :param model_kwargs: to change the model generation
+    :param verbose: display found information
+    :param dynamic_rope: use dynamic rope (see :class:`transformers.LlamaConfig`)
+    :param same_as_pretrained: if True, do not change the default values
+        to get a smaller model
+    :param use_pretrained: download the pretrained weights as well
+    :param use_preinstalled: use preinstalled configurations
+    :param add_second_input: provides others inputs to check a model
+        supports different shapes
+    :param subfolder: subfolder to use for this model id
+    :param use_only_preinstalled: use only preinstalled version
+    :return: dictionary with a model, inputs, dynamic shapes, and the configuration,
+        some necessary rewriting as well
+
+    Example:
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        from onnx_diagnostic.helpers import string_type
+        from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
+
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM", verbose=1)
+
+        print("-- model size:", data['size'])
+        print("-- number of parameters:", data['n_weights'])
+        print("-- inputs:", string_type(data['inputs'], with_shape=True))
+        print("-- dynamic shapes:", pprint.pformat(data['dynamic_shapes']))
+        print("-- configuration:", pprint.pformat(data['configuration']))
+    """
+    if task == "":
+        task = None
+    assert not use_preinstalled or not use_only_preinstalled, (
+        f"model_id={model_id!r}, preinstalled model is only available "
+        f"if use_only_preinstalled is False."
+    )
+    model_id, subfolder, same_as_pretrained, use_pretrained = _preprocess_model_id(
+        model_id,
+        subfolder,
+        same_as_pretrained=same_as_pretrained,
+        use_pretrained=use_pretrained,
+    )
+    if verbose:
+        print(
+            f"[get_untrained_model_with_inputs] model_id={model_id!r}, subfolder={subfolder!r}"
+        )
+        if use_preinstalled:
+            print(f"[get_untrained_model_with_inputs] use preinstalled {model_id!r}")
+    if config is None:
+        config = get_pretrained_config(
+            model_id,
+            use_preinstalled=use_preinstalled,
+            use_only_preinstalled=use_only_preinstalled,
+            subfolder=subfolder,
+            **(model_kwargs or {}),
+        )
+
+    model, task_, mkwargs, diff_config = None, None, {}, None
+    if use_pretrained and same_as_pretrained:
+        if model_id in HANDLED_MODELS:
+            model, task_, config = load_specific_model(model_id, verbose=verbose)
+
+    if task is None:
+        task = task_
+    if model is None:
+        arch = architecture_from_config(config)
+        if task is None and arch is None:
+            task = task_from_id(model_id, subfolder=subfolder)
+        assert task is not None or arch is not None, (
+            f"Unable to determine the architecture for model {model_id!r}, "
+            f"archs={arch!r}, conf={config}"
+        )
+        if verbose:
+            print(f"[get_untrained_model_with_inputs] architecture={arch!r}")
+            print(f"[get_untrained_model_with_inputs] cls={config.__class__.__name__!r}")
+        if task is None:
+            task = task_from_arch(arch, model_id=model_id, subfolder=subfolder)
+        if verbose:
+            print(f"[get_untrained_model_with_inputs] task={task!r}")
+
+        # model kwagrs
+        if dynamic_rope is not None:
+            assert (
+                type(config) is not dict
+            ), f"Unable to set dynamic_rope if the configuration is a dictionary\n{config}"
+            assert hasattr(config, "rope_scaling"), f"Missing 'rope_scaling' in\n{config}"
+            config.rope_scaling = (
+                {"rope_type": "dynamic", "factor": 10.0} if dynamic_rope else None
+            )
+
+        # updating the configuration
+        config0 = copy.deepcopy(config)
+        mkwargs = reduce_model_config(config, task) if not same_as_pretrained else {}
+        if model_kwargs:
+            for k, v in model_kwargs.items():
+                if isinstance(v, dict):
+                    if k in mkwargs:
+                        mkwargs[k].update(v)
+                    else:
+                        mkwargs[k] = v
+                else:
+                    mkwargs[k] = v
+        if mkwargs:
+            update_config(config, mkwargs)
+        try:
+            diff_config = build_diff_config(config0, config)
+        except (ValueError, AttributeError, TypeError) as e:
+            diff_config = f"DIFF CONFIG ERROR {e}"
+        if verbose:
+            if diff_config:
+                print("[get_untrained_model_with_inputs] -- updated config")
+                pprint.pprint(diff_config)
+                print("[get_untrained_model_with_inputs] --")
+
+        # SDPA
+        if model_kwargs and "attn_implementation" in model_kwargs:
+            if hasattr(config, "_attn_implementation_autoset"):
+                config._attn_implementation_autoset = False
+            config._attn_implementation = model_kwargs["attn_implementation"]  # type: ignore[union-attr]
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] config._attn_implementation="
+                    f"{config._attn_implementation!r}"  # type: ignore[union-attr]
+                )
+        elif verbose:
+            print(
+                f"[get_untrained_model_with_inputs] default config._attn_implementation="
+                f"{getattr(config, '_attn_implementation', '?')!r}"  # type: ignore[union-attr]
+            )
+
+        if find_package_source(config) == "diffusers":
+            import diffusers
+
+            package_source = diffusers
+        else:
+            package_source = transformers
+
+        if verbose:
+            print(
+                f"[get_untrained_model_with_inputs] package_source={package_source.__name__} "
+                f"from {package_source.__file__}"
+            )
+        if use_pretrained:
+            begin = time.perf_counter()
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] pretrained model_id {model_id!r}, "
+                    f"subfolder={subfolder!r}"
+                )
+            model = transformers.AutoModel.from_pretrained(
+                model_id, subfolder=subfolder or "", trust_remote_code=True, **mkwargs
+            )
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] -- done(1) in "
+                    f"{time.perf_counter() - begin}s"
+                )
+        else:
+            begin = time.perf_counter()
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] instantiate model_id {model_id!r}, "
+                    f"subfolder={subfolder!r}"
+                )
+            if arch is not None:
+                try:
+                    cls_model = getattr(package_source, arch)
+                except AttributeError as e:
+                    # The code of the models is not in transformers but in the
+                    # repository of the model. We need to download it.
+                    pyfiles = download_code_modelid(model_id, verbose=verbose)
+                    if pyfiles:
+                        if "." in arch:
+                            cls_name = arch
+                        else:
+                            modeling = [_ for _ in pyfiles if "/modeling_" in _]
+                            assert len(modeling) == 1, (
+                                f"Unable to guess the main file implemented class "
+                                f"{arch!r} from {pyfiles}, found={modeling}."
+                            )
+                            last_name = os.path.splitext(os.path.split(modeling[0])[-1])[0]
+                            cls_name = f"{last_name}.{arch}"
+                        if verbose:
+                            print(
+                                f"[get_untrained_model_with_inputs] "
+                                f"custom code for {cls_name!r}"
+                            )
+                            print(
+                                f"[get_untrained_model_with_inputs] from folder "
+                                f"{os.path.split(pyfiles[0])[0]!r}"
+                            )
+                        cls_model = (
+                            transformers.dynamic_module_utils.get_class_from_dynamic_module(
+                                cls_name,
+                                pretrained_model_name_or_path=os.path.split(pyfiles[0])[0],
+                            )
+                        )
+                    else:
+                        raise AttributeError(
+                            f"Unable to find class 'tranformers.{arch}'. "
+                            f"The code needs to be downloaded, config="
+                            f"\n{pprint.pformat(config)}."
+                        ) from e
+            else:
+                assert same_as_pretrained and use_pretrained, (
+                    f"Model {model_id!r} cannot be built, the model cannot be built. "
+                    f"It must be downloaded. Use same_as_pretrained=True "
+                    f"and use_pretrained=True, arch={arch!r}, config={config}"
+                )
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] -- done(2) in "
+                    f"{time.perf_counter() - begin}s"
+                )
+
+            seed = int(os.environ.get("SEED", "17"))
+            torch.manual_seed(seed)
+
+            if verbose:
+                begin = time.perf_counter()
+                print(
+                    f"[get_untrained_model_with_inputs] "
+                    f"instantiate_specific_model {cls_model}"
+                )
+
+            model = instantiate_specific_model(cls_model, config)
+
+            if verbose:
+                print(
+                    f"[get_untrained_model_with_inputs] -- done(3) in "
+                    f"{time.perf_counter() - begin}s (model is {type(model)})"
+                )
+
+            if model is None:
+
+                if verbose:
+                    print(
+                        f"[get_untrained_model_with_inputs] "
+                        f"instantiate_specific_model(2) {cls_model}"
+                    )
+
+                try:
+                    if type(config) is dict:
+                        model = cls_model(**config)
+                    else:
+                        model = cls_model(config)
+                except RuntimeError as e:
+                    raise RuntimeError(
+                        f"Unable to instantiate class {cls_model.__name__} with\n{config}"
+                    ) from e
+
+                if verbose:
+                    print(
+                        f"[get_untrained_model_with_inputs] -- done(4) in "
+                        f"{time.perf_counter() - begin}s (model is {type(model)})"
+                    )
+
+    # input kwargs
+    seed = int(os.environ.get("SEED", "17")) + 1
+    torch.manual_seed(seed)
+    kwargs, fct = random_input_kwargs(config, task)  # type: ignore[arg-type]
+    if verbose:
+        print(f"[get_untrained_model_with_inputs] use fct={fct}")
+        if os.environ.get("PRINT_CONFIG") in (1, "1"):
+            print(f"-- input kwargs for task {task!r}")
+            pprint.pprint(kwargs)
+    if inputs_kwargs:
+        kwargs.update(inputs_kwargs)
+
+    # This line is important. Some models may produce different
+    # outputs even with the same inputs in training mode.
+    model.eval()  # type: ignore[union-attr]
+    res = fct(model, config, add_second_input=add_second_input, **kwargs)
+
+    res["input_kwargs"] = kwargs
+    res["model_kwargs"] = mkwargs
+    if diff_config is not None:
+        res["dump_info"] = dict(config_diff=diff_config)
+
+    sizes = compute_model_size(model)
+    res["model"] = model
+    res["configuration"] = config
+    res["size"] = sizes[0]
+    res["n_weights"] = sizes[1]
+    res["task"] = task
+
+    update = {}
+    for k, v in res.items():
+        if k.startswith(("inputs", "dynamic_shapes")) and isinstance(v, dict):
+            update[k] = filter_out_unexpected_inputs(model, v, verbose=verbose)
+    res.update(update)
+
+    rewrite = _code_needing_rewriting(model.__class__.__name__)
+    if rewrite:
+        res["rewrite"] = rewrite
+    return res
+
+
+def filter_out_unexpected_inputs(
+    model: torch.nn.Module, kwargs: Dict[str, Any], verbose: int = 0
+):
+    """
+    Removes input names in kwargs if no parameter names was found in ``model.forward``.
+    """
+    sig = inspect.signature(model.forward)
+    allowed = set(sig.parameters)
+    new_kwargs = {k: v for k, v in kwargs.items() if k in allowed}
+    diff = set(kwargs) - set(new_kwargs)
+    if diff and verbose:
+        print(f"[filter_out_unexpected_inputs] removed {diff}")
+    return new_kwargs
+
+
+def compute_model_size(model: torch.nn.Module) -> Tuple[int, int]:
+    """Returns the size of the models (weights only) and the number of the parameters."""
+    param_size = 0
+    nparams = 0
+    for param in model.parameters():
+        param_size += param.nelement() * param.element_size()
+        nparams += param.nelement()
+    return param_size, nparams

onnx_diagnostic/torch_models/hghub/model_specific.py

@@ -0,0 +1,76 @@
+from typing import Any, Dict, Tuple
+
+
+def instantiate_specific_model(cls_model: type, config: Any) -> object:
+    """
+    Instantiates some model requiring some specific code.
+    """
+    if cls_model.__name__ == "CosmosTransformer3DModel":
+        return instantiate_CosmosTransformer3DModel(cls_model, config)
+    return None
+
+
+def instantiate_CosmosTransformer3DModel(cls_model: type, config: Any) -> object:
+    kwargs = dict(
+        in_channels=config.in_channels,
+        out_channels=config.out_channels,
+        attention_head_dim=config.attention_head_dim,
+        mlp_ratio=config.mlp_ratio,
+        num_layers=config.num_layers,
+        text_embed_dim=config.text_embed_dim,
+        adaln_lora_dim=config.adaln_lora_dim,
+        max_size=config.max_size,
+        patch_size=config.patch_size,
+        rope_scale=config.rope_scale,
+        concat_padding_mask=config.concat_padding_mask,
+        extra_pos_embed_type=config.extra_pos_embed_type,
+    )
+    return cls_model(**kwargs)
+
+
+class SpecificConfig:
+    """Creates a specific configuration for the loaded model."""
+
+    def __init__(self, **kwargs):
+        self._atts = set(kwargs)
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {k: getattr(self, k) for k in self._atts if k != "_atts"}
+
+
+def load_specific_model(
+    model_id: str, verbose: int = 0, **kwargs
+) -> Tuple[Any, str, SpecificConfig]:
+    """
+    Some models do not have any generic to be loaded.
+    This functions
+
+    :param model_id: model id
+    :param verbose: verbosiy
+    :param kwargs: additional parameters
+    :return: the model, the task associated to it, a configuration
+    """
+    assert model_id in HANDLED_MODELS, (
+        f"Unable to load model_id={model_id!r}, "
+        f"no function is mapped to this id in {sorted(HANDLED_MODELS)}"
+    )
+    return HANDLED_MODELS[model_id](model_id, verbose=verbose, **kwargs)
+
+
+def _load_bingsu_adetailer(model_id: str, verbose: int = 0) -> Tuple[Any, str, SpecificConfig]:
+    """See `Bingsu/adetailer <https://huggingface.co/Bingsu/adetailer>`_."""
+    from huggingface_hub import hf_hub_download
+    from ultralytics import YOLO
+
+    path = hf_hub_download("Bingsu/adetailer", "face_yolov8n.pt")
+    model = YOLO(path)
+    return (
+        model,
+        "object-detection",
+        SpecificConfig(architecture=type(model), image_size=224, num_channels=3),
+    )
+
+
+HANDLED_MODELS = {"Bingsu/adetailer": _load_bingsu_adetailer}

onnx_diagnostic/torch_models/llms.py

@@ -0,0 +1,2 @@
+from .untrained.llm_phi2 import get_phi2
+from .untrained.llm_tiny_llm import get_tiny_llm

onnx_diagnostic/torch_models/untrained/llm_phi2.py

@@ -0,0 +1,113 @@
+from typing import Any, Dict
+import torch
+import transformers
+from ...helpers.cache_helper import make_dynamic_cache
+
+
+def get_phi2(
+    batch_size: int = 1,
+    sequence_length: int = 30,
+    sequence_length2: int = 3,
+    dynamic_rope: bool = False,
+    use_dim_not_dynamic: bool = False,
+    **kwargs,
+) -> Dict[str, Any]:
+    """
+    Gets a non initialized model similar to :epkg:`microsoft/phi-2`.
+
+    :param batch_size: batch size
+    :param sequence_length: sequence length
+    :param sequence_length2: new sequence length
+    :param dynamic_rope: use dynamic rope (see :class:`transformers.LlamaConfig`)
+    :param use_dim_not_dynamic: uses ``torch.export.Dim`` and not a string for the batch size,
+        the sequence length and the cache length
+    :param kwargs: to overwrite the configuration, example ``num_hidden_layers=1``
+    :return: dictionary
+
+    See :ref:`l-plot-tiny-llm-export-patched` for an example with a similar model.
+    """
+    config = {
+        "_name_or_path": "microsoft/phi-2",
+        "architectures": ["PhiForCausalLM"],
+        "attention_dropout": 0.0,
+        "bos_token_id": 50256,
+        "embd_pdrop": 0.0,
+        "eos_token_id": 50256,
+        "hidden_act": "gelu_new",
+        "hidden_size": 2560,
+        "initializer_range": 0.02,
+        "intermediate_size": 10240,
+        "layer_norm_eps": 1e-05,
+        "max_position_embeddings": 2048,
+        "model_type": "phi",
+        "num_attention_heads": 32,
+        "num_hidden_layers": 32,
+        "num_key_value_heads": 32,
+        "partial_rotary_factor": 0.4,
+        "qk_layernorm": False,
+        "resid_pdrop": 0.1,
+        "rope_scaling": {"rope_type": "dynamic", "factor": 10.0} if dynamic_rope else None,
+        "rope_theta": 10000.0,
+        "tie_word_embeddings": False,
+        "torch_dtype": "float16",
+        "transformers_version": "4.37.0",
+        "use_cache": True,
+        "vocab_size": 51200,
+    }
+    config.update(**kwargs)
+    conf = transformers.PhiConfig(**config)
+    model = transformers.PhiForCausalLM(conf)
+    model.eval()
+
+    # now the inputs
+    cache_last_dim = 80
+    max_token_id = config["vocab_size"] - 1
+    n_layers = config["num_hidden_layers"]
+    num_key_value_heads = config["num_key_value_heads"]
+
+    if use_dim_not_dynamic:
+        batch = torch.export.Dim("batch", min=1, max=1024)
+        seq_length = torch.export.Dim("seq_length", min=1, max=4096)
+        cache_length = torch.export.Dim("cache_length", min=1, max=4096)
+    else:
+        batch = "batch"
+        seq_length = "seq_length"
+        cache_length = "cache_length"
+
+    shapes = {
+        "input_ids": {0: batch, 1: seq_length},
+        "position_ids": {
+            0: batch,
+            1: torch.export.Dim.DYNAMIC,  # cache_length + seq_length
+        },
+        "attention_mask": {
+            0: batch,
+            1: torch.export.Dim.DYNAMIC,  # cache_length + seq_length
+        },
+        "past_key_values": [{0: batch, 2: cache_length} for _ in range(n_layers * 2)],
+    }
+    inputs = dict(
+        input_ids=torch.randint(0, max_token_id, (batch_size, sequence_length2)).to(
+            torch.int64
+        ),
+        attention_mask=torch.ones((batch_size, sequence_length + sequence_length2)).to(
+            torch.int64
+        ),
+        position_ids=torch.arange(sequence_length, sequence_length + sequence_length2)
+        .to(torch.int64)
+        .expand((batch_size, -1)),
+        past_key_values=make_dynamic_cache(
+            [
+                (
+                    torch.randn(
+                        batch_size, num_key_value_heads, sequence_length, cache_last_dim
+                    ),
+                    torch.randn(
+                        batch_size, num_key_value_heads, sequence_length, cache_last_dim
+                    ),
+                )
+                for i in range(n_layers)
+            ]
+        ),
+    )
+    return dict(inputs=inputs, model=model, dynamic_shapes=shapes, configuration=conf)

onnx_diagnostic/torch_models/untrained/llm_tiny_llm.py

@@ -0,0 +1,76 @@
+from typing import Any, Dict
+import transformers
+
+
+def get_tiny_llm(
+    batch_size: int = 2,
+    sequence_length: int = 30,
+    sequence_length2: int = 3,
+    dynamic_rope: bool = False,
+    use_static_cache: bool = False,
+    **kwargs,
+) -> Dict[str, Any]:
+    """
+    Gets a non initialized model similar to :epkg:`arnir0/Tiny-LLM`.
+
+    :param batch_size: batch size
+    :param sequence_length: sequence length
+    :param sequence_length2: new sequence length
+    :param dynamic_rope: use dynamic rope (see :class:`transformers.LlamaConfig`)
+    :param use_static_cache: use StaticCache instead of DynamicCache
+    :param kwargs: to overwrite the configuration, example ``num_hidden_layers=1``
+    :return: dictionary
+
+    See :ref:`l-plot-tiny-llm-export` or :ref:`l-plot-tiny-llm-export-patched` for examples.
+    """
+    from ...tasks.text_generation import get_inputs
+
+    config = {
+        "architectures": ["LlamaForCausalLM"],
+        "bos_token_id": 1,
+        "eos_token_id": 2,
+        "hidden_act": "silu",
+        "hidden_size": 192,
+        "initializer_range": 0.02,
+        "intermediate_size": 1024,
+        "max_position_embeddings": 1024,
+        "model_type": "llama",
+        "num_attention_heads": 2,
+        "num_hidden_layers": 1,
+        "num_key_value_heads": 1,
+        "pretraining_tp": 1,
+        "rms_norm_eps": 1e-05,
+        "rope_scaling": {"rope_type": "dynamic", "factor": 10.0} if dynamic_rope else None,
+        "tie_word_embeddings": False,
+        "torch_dtype": "float32",
+        "transformers_version": "4.31.0.dev0",
+        "use_cache": True,
+        "vocab_size": 32000,
+    }
+
+    config.update(**kwargs)
+    conf = transformers.LlamaConfig(**config)
+    if use_static_cache:
+        conf.cache_implementation = "static"
+    model = transformers.LlamaForCausalLM(conf)
+    model.eval()
+
+    res = get_inputs(
+        model,
+        conf,
+        dummy_max_token_id=config["vocab_size"],  # type: ignore[arg-type]
+        num_hidden_layers=config["num_hidden_layers"],  # type: ignore[arg-type]
+        batch_size=batch_size,
+        sequence_length=sequence_length,
+        sequence_length2=sequence_length2,
+        dynamic_rope=dynamic_rope,
+        num_key_value_heads=config["num_key_value_heads"],  # type: ignore[arg-type]
+        cls_cache="StaticCache" if use_static_cache else "DynamicCache",
+    )
+
+    return dict(
+        inputs=res["inputs"],
+        model=model,
+        dynamic_shapes=res["dynamic_shapes"],
+        configuration=conf,
+    )