Skip to content

Commit ad62a15

Browse files
committed
add test for baichuan 13b
1 parent 4c8cb0e commit ad62a15

File tree

4 files changed

+26
-15
lines changed

4 files changed

+26
-15
lines changed

optimum/exporters/openvino/convert.py

+9-10
Original file line numberDiff line numberDiff line change
@@ -388,16 +388,15 @@ def ts_patched_forward(*args, **kwargs):
388388
out_tensor.get_tensor().set_names({output_names[idx]})
389389

390390
for idx, inp_tensor in enumerate(ov_model.inputs):
391-
if idx < len(ordered_input_names):
392-
input_name = ordered_input_names[idx]
393-
inp_tensor.get_tensor().set_names({input_name})
394-
inp_data = flatten_inputs[idx]
395-
static_shape = PartialShape(inp_data.shape)
396-
dims = inputs.get(input_name, [])
397-
for dim in dims:
398-
static_shape[dim] = -1
399-
inp_tensor.get_node().set_partial_shape(static_shape)
400-
inp_tensor.get_node().set_element_type(get_element_type(inp_data.cpu().numpy().dtype))
391+
input_name = ordered_input_names[idx]
392+
inp_tensor.get_tensor().set_names({input_name})
393+
inp_data = flatten_inputs[idx]
394+
static_shape = PartialShape(inp_data.shape)
395+
dims = inputs.get(input_name, [])
396+
for dim in dims:
397+
static_shape[dim] = -1
398+
inp_tensor.get_node().set_partial_shape(static_shape)
399+
inp_tensor.get_node().set_element_type(get_element_type(inp_data.cpu().numpy().dtype))
401400
ov_model.validate_nodes_and_infer_types()
402401

403402
if stateful:

optimum/exporters/openvino/model_patcher.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -329,9 +329,9 @@ def _llama_gemma_update_causal_mask(self, attention_mask, input_tensor, cache_po
329329
offset = 0
330330
mask_shape = attention_mask.shape
331331
mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype
332-
causal_mask[: mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]] = (
333-
mask_slice
334-
)
332+
causal_mask[
333+
: mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]
334+
] = mask_slice
335335

336336
if (
337337
self.config._attn_implementation == "sdpa"

tests/openvino/test_modeling.py

+12-1
Original file line numberDiff line numberDiff line change
@@ -510,6 +510,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
510510
SUPPORTED_ARCHITECTURES = (
511511
"bart",
512512
"baichuan2",
513+
"baichuan2-13b",
513514
"gpt_bigcode",
514515
"blenderbot",
515516
"blenderbot-small",
@@ -540,7 +541,17 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase):
540541
"falcon",
541542
)
542543
GENERATION_LENGTH = 100
543-
REMOTE_CODE_MODELS = ("chatglm", "minicpm", "baichuan2", "jais", "qwen", "internlm2", "olmo", "orion")
544+
REMOTE_CODE_MODELS = (
545+
"chatglm",
546+
"minicpm",
547+
"baichuan2",
548+
"baichuan2-13b",
549+
"jais",
550+
"qwen",
551+
"internlm2",
552+
"olmo",
553+
"orion",
554+
)
544555

545556
@parameterized.expand(SUPPORTED_ARCHITECTURES)
546557
def test_compare_to_transformers(self, model_arch):

tests/openvino/utils_tests.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
"bert": "hf-internal-testing/tiny-random-bert",
2525
"bart": "hf-internal-testing/tiny-random-bart",
2626
"baichuan2": "katuni4ka/tiny-random-baichuan2",
27+
"baichuan2-13b": "katuni4ka/tiny-random-baichuan2-13b",
2728
"bigbird_pegasus": "hf-internal-testing/tiny-random-bigbird_pegasus",
2829
"blenderbot-small": "hf-internal-testing/tiny-random-BlenderbotModel",
2930
"blenderbot": "hf-internal-testing/tiny-random-BlenderbotModel",
@@ -80,7 +81,7 @@
8081
"poolformer": "hf-internal-testing/tiny-random-PoolFormerModel",
8182
"qwen": "katuni4ka/tiny-random-qwen",
8283
"qwen2": "Qwen/Qwen1.5-0.5B",
83-
"resnet": "hf-internal-testing/tiny-random-resnet",
84+
"resnet": "hf-internal-testing/tiny-rantiny-random-baichuan2-13bdom-resnet",
8485
"roberta": "hf-internal-testing/tiny-random-roberta",
8586
"roformer": "hf-internal-testing/tiny-random-roformer",
8687
"segformer": "hf-internal-testing/tiny-random-SegformerModel",

0 commit comments

Comments
 (0)