Add tests

Author: nikita-savelyevv

optimum/intel/openvino/quantization.py

@@ -737,6 +737,12 @@ def _prepare_visual_causal_lm_dataset(self, config: OVWeightQuantizationConfig,
                 "You have entered a string value for dataset. You can only choose between"
                 f"{list(PREDEFINED_VISUAL_LM_DATASETS.keys())}, but the {dataset_name} was found"
             )
+        if config.processor is None:
+            raise ValueError(
+                "`processor` must be specified in order to run data-aware weight compression. "
+                "Please provide it as a model id, or a path to a directory containing all the required "
+                "configuration files."
+            )
 
         dataset_metadata = PREDEFINED_VISUAL_LM_DATASETS[dataset_name]
         dataset = datasets.load_dataset(dataset_metadata["name"], split=dataset_metadata["split"]).shuffle(seed=0)

tests/openvino/test_quantization.py

@@ -63,6 +63,7 @@
     OVWeightQuantizationConfig,
     OVDynamicQuantizationConfig,
     OVModelOpenCLIPForZeroShotImageClassification,
+    OVModelForVisualCausalLM,
 )
 from optimum.intel.openvino.configuration import (
     OVQuantizationMethod,
@@ -298,6 +299,7 @@ class OVWeightCompressionTest(unittest.TestCase):
         (OVStableDiffusionPipeline, "stable-diffusion"),
         (OVStableDiffusionXLPipeline, "stable-diffusion-xl"),
         (OVModelOpenCLIPForZeroShotImageClassification, "open-clip"),
+        (OVModelForVisualCausalLM, "llava"),
     )
 
     SUPPORTED_ARCHITECTURES_WITH_HYBRID_QUANTIZATION = (
@@ -439,6 +441,9 @@ def test_ovmodel_load_with_compressed_weights(self, model_cls, model_type):
             models.append(model.text_encoder if model_type == "stable-diffusion" else model.text_encoder_2)
         elif model_type == "open-clip":
             models = [model.text_model, model.visual_model]
+        elif model.export_feature == "image-text-to-text":
+            models = [model.lm_model, model.vision_embeddings_model, model.text_embeddings_model]
+            models += [getattr(model, part) for part in model.additional_parts]
         else:
             models = [model]
 
@@ -528,16 +533,16 @@ def test_ovmodel_4bit_auto_compression_with_config(
                 # TODO: Check that AWQ was actually applied
                 pass
 
-            tokenizer = AutoTokenizer.from_pretrained(model_id)
-            if tokenizer.pad_token is None:
-                tokenizer.pad_token = tokenizer.eos_token
+            ov_model = model
+            if model_cls == OVModelForVisualCausalLM:
+                ov_model = model.lm_model
 
-            _, num_weight_nodes = get_num_quantized_nodes(model)
+            _, num_weight_nodes = get_num_quantized_nodes(ov_model)
             expected_num_weight_nodes.update({k: 0 for k in set(num_weight_nodes) - set(expected_num_weight_nodes)})
             self.assertEqual(expected_num_weight_nodes, num_weight_nodes)
             model.save_pretrained(tmp_dir)
 
-            wc_rt_info = model.model.get_rt_info()["nncf"]["weight_compression"]
+            wc_rt_info = ov_model.get_rt_info()["nncf"]["weight_compression"]
             self.assertEqual(quantization_config.quant_method.lower() == "awq", wc_rt_info["awq"].value == "True")
             self.assertEqual(
                 quantization_config.scale_estimation or False, wc_rt_info["scale_estimation"].value == "True"
@@ -568,6 +573,9 @@ def test_ovmodel_load_with_uncompressed_weights(self, model_cls, model_type):
             models.append(model.text_encoder if model_type == "stable-diffusion" else model.text_encoder_2)
         elif model_type == "open-clip":
             models = [model.text_model, model.visual_model]
+        elif model.export_feature == "image-text-to-text":
+            models = [model.lm_model, model.vision_embeddings_model, model.text_embeddings_model]
+            models += [getattr(model, part) for part in model.additional_parts]
         else:
             models = [model]
 

tests/openvino/utils_tests.py

@@ -13,6 +13,7 @@
 #  limitations under the License.
 
 import numpy as np
+import openvino as ov
 import torch
 
 
@@ -170,18 +171,20 @@
     "stable-diffusion-xl": (366, 34, 42, 66),
     "stable-diffusion-xl-refiner": (366, 34, 42, 66),
     "open-clip": (20, 28),
+    "llava": (30, 18, 2),
 }
 
 
-def get_num_quantized_nodes(ov_model):
+def get_num_quantized_nodes(model):
     num_fake_quantize = 0
     num_weight_nodes = {
         "int8": 0,
         "int4": 0,
         "f4e2m1": 0,
         "f8e8m0": 0,
     }
-    for elem in ov_model.model.get_ops():
+    ov_model = model if isinstance(model, ov.Model) else model.model
+    for elem in ov_model.get_ops():
         if "FakeQuantize" in elem.name:
             num_fake_quantize += 1
         for i in range(elem.get_output_size()):