Update num_hidden_layers to 4

onnx_diagnostic/tasks/automatic_speech_recognition.py

@@ -15,7 +15,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     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)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     update_config(config, kwargs)
     return kwargs
 

onnx_diagnostic/tasks/feature_extraction.py

@@ -9,7 +9,7 @@
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_hidden_layers")
-    kwargs = dict(num_hidden_layers=min(config.num_hidden_layers, 2))
+    kwargs = dict(num_hidden_layers=min(config.num_hidden_layers, 4))
     update_config(config, kwargs)
     return kwargs
 

onnx_diagnostic/tasks/fill_mask.py

@@ -9,7 +9,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 2),
+        num_hidden_layers=min(config.num_hidden_layers, 4),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)

onnx_diagnostic/tasks/image_classification.py

@@ -17,7 +17,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
-            min(config.num_hidden_layers, 2)
+            min(config.num_hidden_layers, 4)
             if hasattr(config, "num_hidden_layers")
             else len(config.hidden_sizes)
         )

onnx_diagnostic/tasks/image_text_to_text.py

@@ -10,7 +10,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     if hasattr(config, "mm_tokens_per_image"):
         config.mm_tokens_per_image = min(config.mm_tokens_per_image, 2)
     if hasattr(config, "vision_config"):

onnx_diagnostic/tasks/mask_generation.py

@@ -9,7 +9,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     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)

onnx_diagnostic/tasks/mixture_of_expert.py

@@ -11,7 +11,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     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)
     if hasattr(config, "audio_processor") and hasattr(

onnx_diagnostic/tasks/object_detection.py

@@ -10,7 +10,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
-            min(config.num_hidden_layers, 2)
+            min(config.num_hidden_layers, 4)
             if hasattr(config, "num_hidden_layers")
             else len(config.hidden_sizes)
         )

onnx_diagnostic/tasks/sentence_similarity.py

@@ -9,7 +9,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 2),
+        num_hidden_layers=min(config.num_hidden_layers, 4),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)

onnx_diagnostic/tasks/summarization.py

@@ -12,7 +12,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if hasattr(config, "num_decoder_layers"):
         config.num_decoder_layers = min(config.num_decoder_layers, 2)
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     update_config(config, kwargs)
     return kwargs
 

onnx_diagnostic/tasks/text2text_generation.py

@@ -12,7 +12,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if hasattr(config, "num_decoder_layers"):
         config.num_decoder_layers = min(config.num_decoder_layers, 2)
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, 4)
     update_config(config, kwargs)
     return kwargs
 

onnx_diagnostic/tasks/text_classification.py

@@ -9,7 +9,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 2),
+        num_hidden_layers=min(config.num_hidden_layers, 4),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)

onnx_diagnostic/tasks/text_generation.py

@@ -25,7 +25,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if config.__class__.__name__ == "FalconMambaConfig":
         check_hasattr(config, "conv_kernel", "state_size", "intermediate_size")  # 4 and 8
         kwargs = dict(
-            num_hidden_layers=min(config.num_hidden_layers, 2),
+            num_hidden_layers=min(config.num_hidden_layers, 4),
             intermediate_size=256 if config is None else min(512, config.intermediate_size),
             hidden_size=512 if config is None else min(512, config.hidden_size),
             cls_cache="MambaCache",
@@ -37,7 +37,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
             head_dim=getattr(
                 config, "head_dim", config.hidden_size // config.num_attention_heads
             ),
-            num_hidden_layers=min(config.num_hidden_layers, 2),
+            num_hidden_layers=min(config.num_hidden_layers, 4),
             num_key_value_heads=(
                 config.num_key_value_heads
                 if hasattr(config, "num_key_value_heads")