diff --git a/onnx_diagnostic/tasks/automatic_speech_recognition.py b/onnx_diagnostic/tasks/automatic_speech_recognition.py index b6da7e7a..f94997cb 100644 --- a/onnx_diagnostic/tasks/automatic_speech_recognition.py +++ b/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 diff --git a/onnx_diagnostic/tasks/feature_extraction.py b/onnx_diagnostic/tasks/feature_extraction.py index cda396bd..8c65ec08 100644 --- a/onnx_diagnostic/tasks/feature_extraction.py +++ b/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 diff --git a/onnx_diagnostic/tasks/fill_mask.py b/onnx_diagnostic/tasks/fill_mask.py index 167993d2..ba116ee6 100644 --- a/onnx_diagnostic/tasks/fill_mask.py +++ b/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) diff --git a/onnx_diagnostic/tasks/image_classification.py b/onnx_diagnostic/tasks/image_classification.py index 3a993399..ac7dab7f 100644 --- a/onnx_diagnostic/tasks/image_classification.py +++ b/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) ) diff --git a/onnx_diagnostic/tasks/image_text_to_text.py b/onnx_diagnostic/tasks/image_text_to_text.py index b36e6036..1c538980 100644 --- a/onnx_diagnostic/tasks/image_text_to_text.py +++ b/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"): diff --git a/onnx_diagnostic/tasks/mask_generation.py b/onnx_diagnostic/tasks/mask_generation.py index 23641880..4d91c110 100644 --- a/onnx_diagnostic/tasks/mask_generation.py +++ b/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) diff --git a/onnx_diagnostic/tasks/mixture_of_expert.py b/onnx_diagnostic/tasks/mixture_of_expert.py index 1376ade2..c3006594 100644 --- a/onnx_diagnostic/tasks/mixture_of_expert.py +++ b/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( diff --git a/onnx_diagnostic/tasks/object_detection.py b/onnx_diagnostic/tasks/object_detection.py index e85e6355..b8006642 100644 --- a/onnx_diagnostic/tasks/object_detection.py +++ b/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) ) diff --git a/onnx_diagnostic/tasks/sentence_similarity.py b/onnx_diagnostic/tasks/sentence_similarity.py index 5c7b7b04..471dd845 100644 --- a/onnx_diagnostic/tasks/sentence_similarity.py +++ b/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) diff --git a/onnx_diagnostic/tasks/summarization.py b/onnx_diagnostic/tasks/summarization.py index 4384f29d..679e285c 100644 --- a/onnx_diagnostic/tasks/summarization.py +++ b/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 diff --git a/onnx_diagnostic/tasks/text2text_generation.py b/onnx_diagnostic/tasks/text2text_generation.py index 9eb038f5..e2fe4b2a 100644 --- a/onnx_diagnostic/tasks/text2text_generation.py +++ b/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 diff --git a/onnx_diagnostic/tasks/text_classification.py b/onnx_diagnostic/tasks/text_classification.py index 14866f7c..55684110 100644 --- a/onnx_diagnostic/tasks/text_classification.py +++ b/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) diff --git a/onnx_diagnostic/tasks/text_generation.py b/onnx_diagnostic/tasks/text_generation.py index 601f729b..fa872bdb 100644 --- a/onnx_diagnostic/tasks/text_generation.py +++ b/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")