add tests

Author: eaidova

optimum/exporters/openvino/model_patcher.py

@@ -3988,7 +3988,16 @@ def __init__(
 
         @functools.wraps(model.__orig_forward)
         def patched_forward(*args, **kwargs):
-            return model.model.forward(*args, **kwargs)
+            fwd_args = inspect.signature(model.__orig_forward).parameters
+            internal_fwd_args = inspect.signature(model.model.forward).parameters
+            inputs = {}
+            for arg, fwd_arg_name in zip(args, fwd_args):
+                if fwd_arg_name in internal_fwd_args:
+                    inputs[fwd_arg_name] = arg
+            for key, value in kwargs.items():
+                if key in internal_fwd_args:
+                    inputs[key] = value
+            return model.model.forward(**inputs)
 
         model.forward = patched_forward
         self._internal_patcher = internal_patcher

optimum/exporters/openvino/stateful.py

@@ -290,7 +290,6 @@ def patch_stateful_decoder(config: PretrainedConfig, ov_model: ov.Model):
             openvino model
     """
 
-    log.warn(ov_model)
     key_value_input_names = [
         key_name for key in ov_model.inputs for key_name in key.get_names() if "key_values" in key_name
     ]

optimum/intel/openvino/modeling_visual_language.py

@@ -349,13 +349,13 @@ def __init__(
         language_model: ov.Model,
         text_embeddings: ov.Model,
         vision_embeddings: ov.Model,
-        lm_head: Optional[ov.Model] = None,
         config: PretrainedConfig = None,
         device: str = "CPU",
         dynamic_shapes: bool = True,
         ov_config: Optional[Dict[str, str]] = None,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         quantization_config: Union[OVWeightQuantizationConfig, Dict] = None,
+        lm_head: Optional[ov.Model] = None,
         **kwargs,
     ):
         self.config = config
@@ -717,6 +717,9 @@ def components(self):
     def _submodel_names(self):
         model_names = ["lm_model", "text_embeddings_model", "vision_embeddings_model"]
         for part in self.additional_parts:
+            if part == "lm_head" and getattr(self, part + "_model", None) is not None:
+                model_names.append(part + "_model")
+                continue 
             if getattr(self, part, None) is not None:
                 model_names.append(part + "_model")
         return model_names
@@ -2472,6 +2475,7 @@ def generate_image(
         image_token_num_per_image: int = 576,
         img_size: int = 384,
         patch_size: int = 16,
+        generator=None
     ):
         from PIL import Image
 
@@ -2520,7 +2524,7 @@ def generate_image(
             logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond)
             probs = torch.softmax(logits / temperature, dim=-1)
 
-            next_token = torch.multinomial(probs, num_samples=1)
+            next_token = torch.multinomial(probs, num_samples=1) if generator is None else torch.multinomial(probs, num_samples=1, generator=generator)
             generated_tokens[:, i] = next_token.squeeze(dim=-1)
 
             next_token = torch.cat([next_token.unsqueeze(dim=1), next_token.unsqueeze(dim=1)], dim=1).view(-1)
@@ -2563,11 +2567,10 @@ def preprocess_inputs(
                 },
                 {"role": "<|Assistant|>", "content": ""},
             ]
-            prompt = None
+            prepare_inputs = processor(conversations=conversation, images=[image], force_batchify=True)
         else:
-            conversation = None
-            prompt = text
-        prepare_inputs = processor(prompt=prompt, conversations=conversation, images=[image], force_batchify=True)
+            tokenizer = tokenizer if tokenizer is not None else processor.tokenizer
+            prepare_inputs = tokenizer(text, return_tensors="pt")
         required_keys = ["input_ids", "pixel_values", "images_seq_mask", "images_emb_mask"]
         inputs = {}
         for key in required_keys:

tests/openvino/test_modeling.py

@@ -2114,10 +2114,10 @@ class OVModelForVisualCausalLMIntegrationTest(unittest.TestCase):
     if is_transformers_version(">=", "4.40.0"):
         SUPPORTED_ARCHITECTURES += ["llava_next", "nanollava"]
     if is_transformers_version(">=", "4.45.0"):
-        SUPPORTED_ARCHITECTURES += ["minicpmv", "internvl2", "phi3_v", "qwen2_vl", "maira2"]
-    TASK = "image-text-to-text"
-    REMOTE_CODE_MODELS = ["internvl2", "minicpmv", "nanollava", "phi3_v", "maira2"]
 
+        SUPPORTED_ARCHITECTURES += ["janus", "minicpmv", "internvl2", "phi3_v", "qwen2_vl", "maira2"]
+    TASK = "image-text-to-text"
+    REMOTE_CODE_MODELS = ["internvl2", "minicpmv", "nanollava", "phi3_v", "maira2", "janus"]
     IMAGE = Image.open(
         requests.get(
             TEST_IMAGE_URL,
@@ -2216,8 +2216,8 @@ def test_compare_to_transformers(self, model_arch):
         with torch.no_grad():
             transformers_outputs = transformers_model.generate(**transformers_inputs, generation_config=gen_config)
 
-        # original minicpmv, internvl always skip input tokens in generation results, while transformers based approach provide them
-        if model_arch in ["minicpmv", "internvl2"]:
+        # original minicpmv, internvl, janus always skip input tokens in generation results, while transformers based approach provide them
+        if model_arch in ["minicpmv", "internvl2", "janus"]:
             ov_outputs = ov_outputs[:, inputs["input_ids"].shape[1] :]
         self.assertTrue(
             torch.equal(ov_outputs, transformers_outputs),

tests/openvino/utils_tests.py

@@ -170,6 +170,7 @@
     "st-bert": "sentence-transformers/all-MiniLM-L6-v2",
     "st-mpnet": "sentence-transformers/all-mpnet-base-v2",
     "sana": "katuni4ka/tiny-random-sana",
+    "janus": "katuni4ka/tiny-random-janus"
 }