fix zero shot

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.4.0
 +++++
 
+* :pr:`55`: add support for text-classification
 * :pr:`54`: add support for fill-mask, refactoring
 * :pr:`52`: add support for zero-shot-image-classification
 * :pr:`50`: add support for onnxruntime fusion

onnx_diagnostic/tasks/zero_shot_image_classification.py

@@ -62,9 +62,9 @@ def get_inputs(
     ), f"Unexpected type for input_height {type(input_height)}{config}"
 
     batch = torch.export.Dim("batch", min=1, max=1024)
-    seq_length = torch.export.Dim("seq_length", min=1, max=4096)
+    seq_length = "seq_length"  # torch.export.Dim("seq_length", min=1, max=4096)
     shapes = {
-        "inputs_ids": {0: batch, 1: seq_length},
+        "input_ids": {0: batch, 1: seq_length},
         "attention_mask": {0: batch, 1: seq_length},
         "pixel_values": {
             0: torch.export.Dim("batch_img", min=1, max=1024),

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -105,6 +105,7 @@ def get_untrained_model_with_inputs(
     # outputs even with the same inputs in training mode.
     model.eval()
     res = fct(model, config, **kwargs)
+
     res["input_kwargs"] = kwargs
     res["model_kwargs"] = mkwargs
 
@@ -118,19 +119,24 @@ def get_untrained_model_with_inputs(
     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)
+            update[k] = filter_out_unexpected_inputs(model, v, verbose=verbose)
     res.update(update)
     return res
 
 
-def filter_out_unexpected_inputs(model: torch.nn.Module, kwargs: Dict[str, Any]):
+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)
-    kwargs = {k: v for k, v in kwargs.items() if k in allowed}
-    return kwargs
+    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]: