inputs2

xadupre · xadupre · commit 9df1cdcf8395 · 2025-04-24T16:00:10.000+02:00
diff --git a/onnx_diagnostic/tasks/automatic_speech_recognition.py b/onnx_diagnostic/tasks/automatic_speech_recognition.py
@@ -33,6 +33,7 @@ def get_inputs(
     head_dim: int,
     batch_size: int = 2,
     sequence_length: int = 30,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -68,6 +69,7 @@ def get_inputs(
             use_cache:bool,return_dict:bool
         )
     """
+    assert not add_second_input, "add_second_input=True not yet implemented"
     batch = torch.export.Dim("batch", min=1, max=1024)
     seq_length = "seq_length"
 
diff --git a/onnx_diagnostic/tasks/feature_extraction.py b/onnx_diagnostic/tasks/feature_extraction.py
@@ -22,6 +22,7 @@ def get_inputs(
     batch_size: int,
     sequence_length: int,
     dummy_max_token_id: int,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -34,6 +35,7 @@ def get_inputs(
         token_type_ids:T7s1x13[0,0:A0.0],
         attention_mask:T7s1x13[1,1:A1.0])
     """
+    assert not add_second_input, "add_second_input=True not yet implemented"
     batch = torch.export.Dim("batch", min=1, max=1024)
     seq_length = "sequence_length"
     shapes = {
diff --git a/onnx_diagnostic/tasks/fill_mask.py b/onnx_diagnostic/tasks/fill_mask.py
@@ -22,6 +22,7 @@ def get_inputs(
     batch_size: int,
     sequence_length: int,
     dummy_max_token_id: int,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -34,6 +35,7 @@ def get_inputs(
         token_type_ids:T7s1x13[0,0:A0.0],
         attention_mask:T7s1x13[1,1:A1.0])
     """
+    assert not add_second_input, "add_second_input=True not yet implemented"
     batch = torch.export.Dim("batch", min=1, max=1024)
     seq_length = "sequence_length"
     shapes = {
diff --git a/onnx_diagnostic/tasks/image_classification.py b/onnx_diagnostic/tasks/image_classification.py
@@ -27,6 +27,7 @@ def get_inputs(
     input_channels: int,
     batch_size: int = 2,
     dynamic_rope: bool = False,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -40,6 +41,7 @@ def get_inputs(
     :param input_height: input height
     :return: dictionary
     """
+    assert not add_second_input, "add_second_input=True not yet implemented"
     assert isinstance(
         input_width, int
     ), f"Unexpected type for input_width {type(input_width)}{config}"
diff --git a/onnx_diagnostic/tasks/image_text_to_text.py b/onnx_diagnostic/tasks/image_text_to_text.py
@@ -32,6 +32,7 @@ def get_inputs(
     sequence_length2: int = 3,
     n_images: int = 2,
     dynamic_rope: bool = False,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
diff --git a/onnx_diagnostic/tasks/mixture_of_expert.py b/onnx_diagnostic/tasks/mixture_of_expert.py
@@ -41,6 +41,7 @@ def get_inputs(
     sequence_length2: int = 3,
     n_images: int = 2,
     dynamic_rope: bool = False,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -60,6 +61,7 @@ def get_inputs(
     :param dynamic_rope: use dynamic rope (see :class:`transformers.LlamaConfig`)
     :return: dictionary
     """
+    assert not add_second_input, "add_second_input=True not yet implemented"
     raise NotImplementedError(f"get_inputs not yet implemented for task {__TASK__!r}.")
 
 
diff --git a/onnx_diagnostic/tasks/sentence_similarity.py b/onnx_diagnostic/tasks/sentence_similarity.py
@@ -22,6 +22,7 @@ def get_inputs(
     batch_size: int,
     sequence_length: int,
     dummy_max_token_id: int,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -34,6 +35,7 @@ def get_inputs(
         token_type_ids:T7s1x13[0,0:A0.0],
         attention_mask:T7s1x13[1,1:A1.0])
     """
+    assert not add_second_input, "add_second_input=True not yet implemented"
     batch = torch.export.Dim("batch", min=1, max=1024)
     seq_length = "seq_length"
     shapes = {
diff --git a/onnx_diagnostic/tasks/text2text_generation.py b/onnx_diagnostic/tasks/text2text_generation.py
@@ -28,6 +28,7 @@ def get_inputs(
     batch_size: int = 2,
     sequence_length: int = 30,
     sequence_length2: int = 3,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
diff --git a/onnx_diagnostic/tasks/text_classification.py b/onnx_diagnostic/tasks/text_classification.py
@@ -22,6 +22,7 @@ def get_inputs(
     batch_size: int,
     sequence_length: int,
     dummy_max_token_id: int,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -34,6 +35,7 @@ def get_inputs(
         token_type_ids:T7s1x13[0,0:A0.0],
         attention_mask:T7s1x13[1,1:A1.0])
     """
+    assert not add_second_input, "add_second_input=True not yet implemented"
     batch = torch.export.Dim("batch", min=1, max=1024)
     seq_length = "seq_length"  # torch.export.Dim("sequence_length", min=1, max=1024)
     shapes = {
diff --git a/onnx_diagnostic/tasks/text_generation.py b/onnx_diagnostic/tasks/text_generation.py
@@ -71,6 +71,7 @@ def get_inputs(
     num_key_value_heads: Optional[int] = None,
     head_dim: Optional[int] = None,
     cls_cache: Optional[Union[type, str]] = None,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -88,6 +89,7 @@ def get_inputs(
         :class:`transformers.cache_utils.DynamicCache`
     :return: dictionary
     """
+    assert not add_second_input, "add_second_input=True not yet implemented"
     batch = torch.export.Dim("batch", min=1, max=1024)
     seq_length = "seq_length"  # torch.export.Dim("seq_length", min=1, max=4096)
     cache_length = "cache_length"  # torch.export.Dim("cache_length", min=1, max=4096)
@@ -192,7 +194,23 @@ def get_inputs(
             ]
         ),
     )
-    return dict(inputs=inputs, dynamic_shapes=shapes)
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            dummy_max_token_id=dummy_max_token_id,
+            num_hidden_layers=num_hidden_layers,
+            batch_size=batch_size + 1,
+            sequence_length=sequence_length + 1,
+            sequence_length2=sequence_length2 + 1,
+            dynamic_rope=dynamic_rope,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=head_dim,
+            cls_cache=cls_cache,
+            **kwargs,
+        )
+    return res
 
 
 def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
diff --git a/onnx_diagnostic/tasks/zero_shot_image_classification.py b/onnx_diagnostic/tasks/zero_shot_image_classification.py
@@ -34,6 +34,7 @@ def get_inputs(
     input_height: int = 224,
     input_channels: int = 3,
     batch_size_image=3,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -54,6 +55,7 @@ def get_inputs(
     # attention_mask:T7s2x7
     # pixel_values:T1s2x3x224x224
     """
+    assert not add_second_input, "add_second_input=True not yet implemented"
     assert isinstance(
         input_width, int
     ), f"Unexpected type for input_width {type(input_width)}{config}"
diff --git a/onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py b/onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py
@@ -3569,3 +3569,74 @@ def _ccached_tiiuae_falcon_mamba_tiny_dev():
             "vocab_size": 65024,
         }
     )
+
+
+def _ccached_facebook_bart_base():
+    "facebook/bart-base"
+    return transformers.BartConfig(
+        **{
+            "_name_or_path": "bart-base",
+            "activation_dropout": 0.1,
+            "activation_function": "gelu",
+            "add_bias_logits": false,
+            "add_final_layer_norm": false,
+            "architectures": ["BartModel"],
+            "attention_dropout": 0.1,
+            "bos_token_id": 0,
+            "classif_dropout": 0.1,
+            "classifier_dropout": 0.0,
+            "d_model": 768,
+            "decoder_attention_heads": 12,
+            "decoder_ffn_dim": 3072,
+            "decoder_layerdrop": 0.0,
+            "decoder_layers": 6,
+            "decoder_start_token_id": 2,
+            "dropout": 0.1,
+            "early_stopping": true,
+            "encoder_attention_heads": 12,
+            "encoder_ffn_dim": 3072,
+            "encoder_layerdrop": 0.0,
+            "encoder_layers": 6,
+            "eos_token_id": 2,
+            "forced_eos_token_id": 2,
+            "forced_bos_token_id": 0,
+            "gradient_checkpointing": false,
+            "id2label": {"0": "LABEL_0", "1": "LABEL_1", "2": "LABEL_2"},
+            "init_std": 0.02,
+            "is_encoder_decoder": true,
+            "label2id": {"LABEL_0": 0, "LABEL_1": 1, "LABEL_2": 2},
+            "max_position_embeddings": 1024,
+            "model_type": "bart",
+            "no_repeat_ngram_size": 3,
+            "normalize_before": false,
+            "normalize_embedding": true,
+            "num_beams": 4,
+            "num_hidden_layers": 6,
+            "pad_token_id": 1,
+            "scale_embedding": false,
+            "task_specific_params": {
+                "summarization": {
+                    "length_penalty": 1.0,
+                    "max_length": 128,
+                    "min_length": 12,
+                    "num_beams": 4,
+                },
+                "summarization_cnn": {
+                    "length_penalty": 2.0,
+                    "max_length": 142,
+                    "min_length": 56,
+                    "num_beams": 4,
+                },
+                "summarization_xsum": {
+                    "length_penalty": 1.0,
+                    "max_length": 62,
+                    "min_length": 11,
+                    "num_beams": 6,
+                },
+            },
+            "torch_dtype": "float32",
+            "transformers_version": "4.12.0.dev0",
+            "use_cache": true,
+            "vocab_size": 50265,
+        }
+    )
diff --git a/onnx_diagnostic/torch_models/hghub/model_inputs.py b/onnx_diagnostic/torch_models/hghub/model_inputs.py
@@ -17,6 +17,7 @@ def get_untrained_model_with_inputs(
     dynamic_rope: Optional[bool] = None,
     same_as_pretrained: bool = False,
     use_preinstalled: bool = True,
+    add_second_input: bool = False,
 ) -> Dict[str, Any]:
     """
     Gets a non initialized model similar to the original model
@@ -34,6 +35,8 @@ def get_untrained_model_with_inputs(
     :param same_as_pretrained: if True, do not change the default values
         to get a smaller model
     :param use_preinstalled: use preinstalled configurations
+    :param add_second_input: provides a second inputs to check a model
+        supports different shapes
     :return: dictionary with a model, inputs, dynamic shapes, and the configuration
 
     Example:
@@ -107,7 +110,7 @@ def get_untrained_model_with_inputs(
     # This line is important. Some models may produce different
     # outputs even with the same inputs in training mode.
     model.eval()
-    res = fct(model, config, **kwargs)
+    res = fct(model, config, add_second_input=add_second_input, **kwargs)
 
     res["input_kwargs"] = kwargs
     res["model_kwargs"] = mkwargs
