refactoring

xadupre · xadupre · commit a79484285759 · 2025-04-15T10:20:33.000+02:00
diff --git a/_doc/api/helpers/config_helper.rst b/_doc/api/helpers/config_helper.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.helpers.config_helper
+=====================================
+
+.. automodule:: onnx_diagnostic.helpers.config_helper
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/helpers/index.rst b/_doc/api/helpers/index.rst
@@ -9,6 +9,7 @@ onnx_diagnostic.helpers
     args_helper
     bench_run
     cache_helper
+    config_helper
     helper
     memory_peak
     onnx_helper
diff --git a/_doc/api/tasks/automatic_speech_recognition.rst b/_doc/api/tasks/automatic_speech_recognition.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.automatic_speech_generation
+==================================================
+
+.. automodule:: onnx_diagnostic.tasks.automatic_speech_generation
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/tasks/fill_mask.rst b/_doc/api/tasks/fill_mask.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.fill_mask
+================================
+
+.. automodule:: onnx_diagnostic.tasks.fill_mask
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/tasks/image_classification.rst b/_doc/api/tasks/image_classification.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.image_classification
+===========================================
+
+.. automodule:: onnx_diagnostic.tasks.image_classification
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/tasks/image_text_to_text.rst b/_doc/api/tasks/image_text_to_text.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.image_text_to_text
+=========================================
+
+.. automodule:: onnx_diagnostic.tasks.image_text_to_text
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/tasks/index.rst b/_doc/api/tasks/index.rst
@@ -0,0 +1,26 @@
+onnx_diagnostic.tasks
+=====================
+
+.. toctree::
+    :maxdepth: 1
+    :caption: modules
+
+    automatic_speech_generation
+    fill_mask
+    image_classification
+    image_text_to_text      
+    text_generation
+    text2text_generation
+    zero_short_image_classification
+
+reduce_model_config
++++++++++++++++++++
+
+.. autofunction:: onnx_diagnostic.tasks.reduce_model_config
+    
+Other functions
++++++++++++++++
+
+.. automodule:: onnx_diagnostic.tasks
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/tasks/text2text_generation.rst b/_doc/api/tasks/text2text_generation.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.text2text_generation
+===========================================
+
+.. automodule:: onnx_diagnostic.tasks.text2text_generation
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/tasks/text_generation.rst b/_doc/api/tasks/text_generation.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.text_generation
+======================================
+
+.. automodule:: onnx_diagnostic.tasks.text_generation
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/tasks/zero_shot_image_classification.rst b/_doc/api/tasks/zero_shot_image_classification.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.zero_short_image_classification
+======================================================
+
+.. automodule:: onnx_diagnostic.tasks.zero_short_image_classification
+    :members:
+    :no-undoc-members:
diff --git a/onnx_diagnostic/helpers/config_helper.py b/onnx_diagnostic/helpers/config_helper.py
@@ -0,0 +1,44 @@
+from typing import Any, Dict, Tuple, Union
+
+
+def check_hasattr(config: Any, *args: Union[str, Tuple[Any, ...]]):
+    """
+    Checks the confiugation has all the attributes in ``args``.
+    Raises an exception otherwise.
+    """
+    for a in args:
+        assert isinstance(a, (str, tuple)), f"unexpected type {type(a)} in {args!r}"
+        if isinstance(a, str):
+            assert (isinstance(config, dict) and a in config) or hasattr(
+                config, a
+            ), f"Missing attribute {a!r} in\n{config}"
+        elif isinstance(a, tuple):
+            assert any(
+                (isinstance(name, str) and hasattr(config, name))
+                or all(hasattr(config, _) for _ in name)
+                for name in a
+            ), f"All attributes in {a!r} are missing from\n{config}"
+
+
+def update_config(config: Any, mkwargs: Dict[str, Any]):
+    """Updates a configuration with different values."""
+    for k, v in mkwargs.items():
+        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)
+
+
+def _pick(config, *atts):
+    """Returns the first value found in the configuration."""
+    for a in atts:
+        if isinstance(a, str):
+            if hasattr(config, a):
+                return getattr(config, a)
+        elif isinstance(a, tuple):
+            if all(hasattr(config, _) for _ in a[1:]):
+                return a[0]([getattr(config, _) for _ in a[1:]])
+    raise AssertionError(f"Unable to find any of these {atts!r} in {config}")
diff --git a/onnx_diagnostic/tasks/__init__.py b/onnx_diagnostic/tasks/__init__.py
@@ -0,0 +1,25 @@
+from typing import Any, Dict
+from . import (
+    automatic_speech_recognition,
+    image_classification,
+    image_text_to_text,
+    text_generation,
+    text2text_generation,
+    zero_shot_image_classification,
+)
+
+__TASKS__ = [
+    automatic_speech_recognition,
+    image_classification,
+    image_text_to_text,
+    text_generation,
+    text2text_generation,
+    zero_shot_image_classification,
+]
+
+
+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)
diff --git a/onnx_diagnostic/tasks/automatic_speech_recognition.py b/onnx_diagnostic/tasks/automatic_speech_recognition.py
@@ -0,0 +1,17 @@
+from typing import Any, Dict
+from ..helpers.config_helper import update_config
+
+__TASK__ = "automatic_speech_recognition"
+
+
+def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+    """Reduces a model size."""
+    kwargs = {}
+    if hasattr(config, "num_decoder_layers"):
+        config.num_decoder_layers = min(config.num_decoder_layers, 2)
+    if hasattr(config, "decoder_layers"):
+        config.decoder_layers = min(config.decoder_layers, 2)
+    if hasattr(config, "num_hidden_layers"):
+        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+    update_config(config, kwargs)
+    return kwargs
diff --git a/onnx_diagnostic/tasks/image_classification.py b/onnx_diagnostic/tasks/image_classification.py
@@ -0,0 +1,18 @@
+from typing import Any, Dict
+from ..helpers.config_helper import update_config, check_hasattr
+
+__TASK__ = "image-classification"
+
+
+def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+    """Reduces a model size."""
+    check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
+    kwargs = dict(
+        num_hidden_layers=(
+            min(config.num_hidden_layers, 2)
+            if hasattr(config, "num_hidden_layers")
+            else len(config.hidden_sizes)
+        )
+    )
+    update_config(config, kwargs)
+    return kwargs
diff --git a/onnx_diagnostic/tasks/image_text_to_text.py b/onnx_diagnostic/tasks/image_text_to_text.py
@@ -0,0 +1,15 @@
+from typing import Any, Dict
+from ..helpers.config_helper import update_config
+
+__TASK__ = "image-text-to-text"
+
+
+def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+    """Reduces a model size."""
+    kwargs = {}
+    if hasattr(config, "num_hidden_layers"):
+        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+    if hasattr(config, "vision_config") and hasattr(config.vision_config, "num_hidden_layers"):
+        config.vision_config.num_hidden_layers = min(config.vision_config.num_hidden_layers, 2)
+    update_config(config, kwargs)
+    return kwargs
diff --git a/onnx_diagnostic/tasks/text2text_generation.py b/onnx_diagnostic/tasks/text2text_generation.py
@@ -0,0 +1,17 @@
+from typing import Any, Dict
+from ..helpers.config_helper import update_config, check_hasattr
+
+__TASK__ = "text-to-text-generation"
+
+
+def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+    """Reduces a model size."""
+    kwargs = {}
+    if hasattr(config, "num_decoder_layers"):
+        check_hasattr(config.num_decoder_layers, "num_decoder_layers")
+        config.num_decoder_layers = min(config.num_decoder_layers, 2)
+    if hasattr(config, "num_hidden_layers"):
+        check_hasattr(config.num_decoder_layers, "num_hidden_layers")
+        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+    update_config(config, kwargs)
+    return kwargs
diff --git a/onnx_diagnostic/tasks/text_generation.py b/onnx_diagnostic/tasks/text_generation.py
@@ -0,0 +1,37 @@
+from typing import Any, Dict
+from ..helpers.config_helper import update_config, check_hasattr
+
+__TASK__ = "text-generation"
+
+
+def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+    """Reduces a model size."""
+    check_hasattr(
+        config,
+        ("head_dim", ("hidden_size", "num_attention_heads")),
+        "num_hidden_layers",
+        ("num_key_value_heads", "num_attention_heads"),
+        "intermediate_size",
+        "hidden_size",
+    )
+    kwargs = dict(
+        head_dim=getattr(config, "head_dim", config.hidden_size // config.num_attention_heads),
+        num_hidden_layers=min(config.num_hidden_layers, 2),
+        num_key_value_heads=(
+            config.num_key_value_heads
+            if hasattr(config, "num_key_value_heads")
+            else config.num_attention_heads
+        ),
+        intermediate_size=(
+            min(config.intermediate_size, 24576 // 4)
+            if config.intermediate_size % 4 == 0
+            else config.intermediate_size
+        ),
+        hidden_size=(
+            min(config.hidden_size, 3072 // 4)
+            if config.hidden_size % 4 == 0
+            else config.hidden_size
+        ),
+    )
+    update_config(config, kwargs)
+    return kwargs
diff --git a/onnx_diagnostic/tasks/zero_shot_image_classification.py b/onnx_diagnostic/tasks/zero_shot_image_classification.py
@@ -0,0 +1,23 @@
+from typing import Any, Dict
+from ..helpers.config_helper import update_config, check_hasattr
+
+__TASK__ = "zero-shot-image-generation"
+
+
+def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+    """Reduces a model size."""
+    check_hasattr(config, "vision_config", "text_config")
+    check_hasattr(config.vision_config, "num_hidden_layers", "num_attention_heads")
+    check_hasattr(config.text_config, "num_hidden_layers", "num_attention_heads")
+    kwargs = dict(
+        vision_config=dict(
+            num_hidden_layers=min(2, config.vision_config.num_hidden_layers),
+            num_attention_heads=min(2, config.vision_config.num_attention_heads),
+        ),
+        text_config=dict(
+            num_hidden_layers=min(2, config.text_config.num_hidden_layers),
+            num_attention_heads=min(2, config.text_config.num_attention_heads),
+        ),
+    )
+    update_config(config, kwargs)
+    return kwargs
diff --git a/onnx_diagnostic/torch_models/hghub/hub_data.py b/onnx_diagnostic/torch_models/hghub/hub_data.py
@@ -13,6 +13,7 @@
     ASTModel,feature-extraction
     AlbertModel,feature-extraction
     BeitForImageClassification,image-classification
+    BertForMaskedLM,fill-mask
     BigBirdModel,feature-extraction
     BlenderbotModel,feature-extraction
     BloomModel,feature-extraction
diff --git a/onnx_diagnostic/torch_models/hghub/model_inputs.py b/onnx_diagnostic/torch_models/hghub/model_inputs.py
