add test for baichuan 13b

Author: eaidova

optimum/exporters/openvino/convert.py

@@ -388,16 +388,15 @@ def ts_patched_forward(*args, **kwargs):
                 out_tensor.get_tensor().set_names({output_names[idx]})
 
         for idx, inp_tensor in enumerate(ov_model.inputs):
-            if idx < len(ordered_input_names):
-                input_name = ordered_input_names[idx]
-                inp_tensor.get_tensor().set_names({input_name})
-                inp_data = flatten_inputs[idx]
-                static_shape = PartialShape(inp_data.shape)
-                dims = inputs.get(input_name, [])
-                for dim in dims:
-                    static_shape[dim] = -1
-                inp_tensor.get_node().set_partial_shape(static_shape)
-                inp_tensor.get_node().set_element_type(get_element_type(inp_data.cpu().numpy().dtype))
+            input_name = ordered_input_names[idx]
+            inp_tensor.get_tensor().set_names({input_name})
+            inp_data = flatten_inputs[idx]
+            static_shape = PartialShape(inp_data.shape)
+            dims = inputs.get(input_name, [])
+            for dim in dims:
+                static_shape[dim] = -1
+            inp_tensor.get_node().set_partial_shape(static_shape)
+            inp_tensor.get_node().set_element_type(get_element_type(inp_data.cpu().numpy().dtype))
         ov_model.validate_nodes_and_infer_types()
 
         if stateful:

optimum/exporters/openvino/model_patcher.py

@@ -329,9 +329,9 @@ def _llama_gemma_update_causal_mask(self, attention_mask, input_tensor, cache_po
                 offset = 0
             mask_shape = attention_mask.shape
             mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype
-            causal_mask[: mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]] = (
-                mask_slice
-            )
+            causal_mask[
+                : mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]
+            ] = mask_slice
 
     if (
         self.config._attn_implementation == "sdpa"

tests/openvino/test_modeling.py

@@ -510,6 +510,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES = (
         "bart",
         "baichuan2",
+        "baichuan2-13b",
         "gpt_bigcode",
         "blenderbot",
         "blenderbot-small",
@@ -540,7 +541,17 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
         "falcon",
     )
     GENERATION_LENGTH = 100
-    REMOTE_CODE_MODELS = ("chatglm", "minicpm", "baichuan2", "jais", "qwen", "internlm2", "olmo", "orion")
+    REMOTE_CODE_MODELS = (
+        "chatglm",
+        "minicpm",
+        "baichuan2",
+        "baichuan2-13b",
+        "jais",
+        "qwen",
+        "internlm2",
+        "olmo",
+        "orion",
+    )
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
     def test_compare_to_transformers(self, model_arch):

tests/openvino/utils_tests.py

@@ -24,6 +24,7 @@
     "bert": "hf-internal-testing/tiny-random-bert",
     "bart": "hf-internal-testing/tiny-random-bart",
     "baichuan2": "katuni4ka/tiny-random-baichuan2",
+    "baichuan2-13b": "katuni4ka/tiny-random-baichuan2-13b",
     "bigbird_pegasus": "hf-internal-testing/tiny-random-bigbird_pegasus",
     "blenderbot-small": "hf-internal-testing/tiny-random-BlenderbotModel",
     "blenderbot": "hf-internal-testing/tiny-random-BlenderbotModel",
@@ -80,7 +81,7 @@
     "poolformer": "hf-internal-testing/tiny-random-PoolFormerModel",
     "qwen": "katuni4ka/tiny-random-qwen",
     "qwen2": "Qwen/Qwen1.5-0.5B",
-    "resnet": "hf-internal-testing/tiny-random-resnet",
+    "resnet": "hf-internal-testing/tiny-rantiny-random-baichuan2-13bdom-resnet",
     "roberta": "hf-internal-testing/tiny-random-roberta",
     "roformer": "hf-internal-testing/tiny-random-roformer",
     "segformer": "hf-internal-testing/tiny-random-SegformerModel",