@@ -818,30 +818,19 @@ def test_ovmodel_load_with_compressed_weights(self, model_cls, model_type, trust
818
818
self .assertEqual (model ._openvino_config .quantization_config .bits , 8 )
819
819
self .assertEqual (model ._openvino_config .dtype , "int8" )
820
820
821
- if model .export_feature .startswith ("text2text-generation" ):
822
- models = [model .encoder , model .decoder ]
823
- if model .decoder_with_past is not None :
824
- models .append (model .decoder_with_past )
825
- elif model .export_feature == "text-to-image" :
826
- models = [model .unet , model .vae_encoder , model .vae_decoder ]
827
- models .append (model .text_encoder if model_type in ["stable-diffusion" , "sana" ] else model .text_encoder_2 )
828
- elif model_type == "open-clip" :
829
- models = [model .text_model , model .visual_model ]
830
- elif model .export_feature == "image-text-to-text" :
831
- models = list (model .submodels .values ())
832
- else :
833
- models = [model ]
834
-
835
821
if model_type == "open-clip" :
836
822
pytest .skip (reason = "ticket 161043" )
837
823
elif model_type == "t5" :
838
824
pytest .skip (reason = "ticket 160958" )
839
825
else :
840
826
check_optimization_not_applicable_to_optimized_model (model , quantization_config = {"bits" : 8 })
841
827
828
+ submodels = (
829
+ [model .text_model , model .visual_model ] if model_type == "open-clip" else model .ov_submodels .values ()
830
+ )
842
831
expected_ov_int8 = _ARCHITECTURES_TO_EXPECTED_INT8 [model_type ]
843
832
expected_ov_int8 = [{"int8" : it } for it in expected_ov_int8 ]
844
- check_compression_state_per_model (self , models , expected_ov_int8 )
833
+ check_compression_state_per_model (self , submodels , expected_ov_int8 )
845
834
846
835
@parameterized .expand (SUPPORTED_ARCHITECTURES_WITH_HYBRID_QUANTIZATION )
847
836
def test_ovmodel_hybrid_quantization (self , model_cls , model_type , expected_fake_nodes , expected_int8_nodes ):
@@ -938,11 +927,7 @@ def test_ovmodel_4bit_auto_compression_with_config(
938
927
# TODO: Check that AWQ was actually applied
939
928
pass
940
929
941
- submodels = []
942
- if isinstance (model , OVModelForCausalLM ):
943
- submodels = [model .model ]
944
- elif isinstance (model , OVModelForVisualCausalLM ):
945
- submodels = list (model .submodels .values ())
930
+ submodels = list (model .ov_submodels .values ())
946
931
check_compression_state_per_model (self , submodels , expected_num_weight_nodes_per_model )
947
932
948
933
model .save_pretrained (tmp_dir )
@@ -976,21 +961,11 @@ def test_ovmodel_load_with_uncompressed_weights(self, model_cls, model_type, tru
976
961
model = model_cls .from_pretrained (
977
962
MODEL_NAMES [model_type ], export = True , load_in_8bit = False , trust_remote_code = trust_remote_code
978
963
)
979
- if model .export_feature .startswith ("text2text-generation" ):
980
- models = [model .encoder , model .decoder ]
981
- if model .decoder_with_past is not None :
982
- models .append (model .decoder_with_past )
983
- elif model .export_feature == "text-to-image" :
984
- models = [model .unet , model .vae_encoder , model .vae_decoder ]
985
- models .append (model .text_encoder if model_type in ["stable-diffusion" , "sana" ] else model .text_encoder_2 )
986
- elif model_type == "open-clip" :
987
- models = [model .text_model , model .visual_model ]
988
- elif model .export_feature == "image-text-to-text" :
989
- models = list (model .submodels .values ())
990
- else :
991
- models = [model ]
992
964
993
- for i , submodel in enumerate (models ):
965
+ submodels = (
966
+ [model .text_model , model .visual_model ] if model_type == "open-clip" else model .ov_submodels .values ()
967
+ )
968
+ for i , submodel in enumerate (submodels ):
994
969
ov_model = submodel if isinstance (submodel , ov .Model ) else submodel .model
995
970
_ , num_weight_nodes = get_num_quantized_nodes (ov_model )
996
971
self .assertEqual (0 , num_weight_nodes ["int8" ])
@@ -1106,12 +1081,7 @@ def test_ovmodel_4bit_dynamic_with_config(
1106
1081
self .assertEqual (model .ov_config ["DYNAMIC_QUANTIZATION_GROUP_SIZE" ], str (group_size ))
1107
1082
self .assertEqual (model .ov_config ["KV_CACHE_PRECISION" ], "u8" )
1108
1083
1109
- submodels = []
1110
- if isinstance (model , OVModelForCausalLM ):
1111
- submodels = [model .model ]
1112
- elif isinstance (model , OVModelForVisualCausalLM ):
1113
- submodels = list (model .submodels .values ())
1114
- check_compression_state_per_model (self , submodels , expected_num_weight_nodes_per_model )
1084
+ check_compression_state_per_model (self , model .ov_submodels .values (), expected_num_weight_nodes_per_model )
1115
1085
1116
1086
model .save_pretrained (tmp_dir )
1117
1087
openvino_config = OVConfig .from_pretrained (tmp_dir )
0 commit comments