Upload 379 files

610d9c2 verified about 3 hours ago

22.8 kB

# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------

# It is a tool to generate test data for a bert model.
# The test data can be used by onnxruntime_perf_test tool to evaluate the inference latency.

import argparse
import os
import random
from pathlib import Path

import numpy as np
from onnx import ModelProto, TensorProto, numpy_helper
from onnx_model import OnnxModel


def fake_input_ids_data(
    input_ids: TensorProto, batch_size: int, sequence_length: int, dictionary_size: int
) -> np.ndarray:
    """Create input tensor based on the graph input of input_ids
    Args:
        input_ids (TensorProto): graph input of the input_ids input tensor
        batch_size (int): batch size
        sequence_length (int): sequence length
        dictionary_size (int): vocabulary size of dictionary
    Returns:
        np.ndarray: the input tensor created
    """
    assert input_ids.type.tensor_type.elem_type in [
        TensorProto.FLOAT,
        TensorProto.INT32,
        TensorProto.INT64,
    ]

    data = np.random.randint(dictionary_size, size=(batch_size, sequence_length), dtype=np.int32)

    if input_ids.type.tensor_type.elem_type == TensorProto.FLOAT:
        data = np.float32(data)
    elif input_ids.type.tensor_type.elem_type == TensorProto.INT64:
        data = np.int64(data)

    return data


def fake_segment_ids_data(segment_ids: TensorProto, batch_size: int, sequence_length: int) -> np.ndarray:
    """Create input tensor based on the graph input of segment_ids
    Args:
        segment_ids (TensorProto): graph input of the token_type_ids input tensor
        batch_size (int): batch size
        sequence_length (int): sequence length
    Returns:
        np.ndarray: the input tensor created
    """
    assert segment_ids.type.tensor_type.elem_type in [
        TensorProto.FLOAT,
        TensorProto.INT32,
        TensorProto.INT64,
    ]

    data = np.zeros((batch_size, sequence_length), dtype=np.int32)

    if segment_ids.type.tensor_type.elem_type == TensorProto.FLOAT:
        data = np.float32(data)
    elif segment_ids.type.tensor_type.elem_type == TensorProto.INT64:
        data = np.int64(data)

    return data


def get_random_length(max_sequence_length: int, average_sequence_length: int):
    assert average_sequence_length >= 1 and average_sequence_length <= max_sequence_length

    # For uniform distribution, we find proper lower and upper bounds so that the average is in the middle.
    if 2 * average_sequence_length > max_sequence_length:
        return random.randint(2 * average_sequence_length - max_sequence_length, max_sequence_length)
    else:
        return random.randint(1, 2 * average_sequence_length - 1)


def fake_input_mask_data(
    input_mask: TensorProto,
    batch_size: int,
    sequence_length: int,
    average_sequence_length: int,
    random_sequence_length: bool,
    mask_type: int = 2,
) -> np.ndarray:
    """Create input tensor based on the graph input of segment_ids.
    Args:
        input_mask (TensorProto): graph input of the attention mask input tensor
        batch_size (int): batch size
        sequence_length (int): sequence length
        average_sequence_length (int): average sequence length excluding paddings
        random_sequence_length (bool): whether use uniform random number for sequence length
        mask_type (int): mask type - 1: mask index (sequence length excluding paddings). Shape is (batch_size).
                                     2: 2D attention mask. Shape is (batch_size, sequence_length).
                                     3: key len, cumulated lengths of query and key. Shape is (3 * batch_size + 2).
    Returns:
        np.ndarray: the input tensor created
    """

    assert input_mask.type.tensor_type.elem_type in [
        TensorProto.FLOAT,
        TensorProto.INT32,
        TensorProto.INT64,
    ]

    if mask_type == 1:  # sequence length excluding paddings
        data = np.ones((batch_size), dtype=np.int32)
        if random_sequence_length:
            for i in range(batch_size):
                data[i] = get_random_length(sequence_length, average_sequence_length)
        else:
            for i in range(batch_size):
                data[i] = average_sequence_length
    elif mask_type == 2:  # 2D attention mask
        data = np.zeros((batch_size, sequence_length), dtype=np.int32)
        if random_sequence_length:
            for i in range(batch_size):
                actual_seq_len = get_random_length(sequence_length, average_sequence_length)
                for j in range(actual_seq_len):
                    data[i, j] = 1
        else:
            temp = np.ones((batch_size, average_sequence_length), dtype=np.int32)
            data[: temp.shape[0], : temp.shape[1]] = temp
    else:
        assert mask_type == 3
        data = np.zeros((batch_size * 3 + 2), dtype=np.int32)
        if random_sequence_length:
            for i in range(batch_size):
                data[i] = get_random_length(sequence_length, average_sequence_length)

            for i in range(batch_size + 1):
                data[batch_size + i] = data[batch_size + i - 1] + data[i - 1] if i > 0 else 0
                data[2 * batch_size + 1 + i] = data[batch_size + i - 1] + data[i - 1] if i > 0 else 0
        else:
            for i in range(batch_size):
                data[i] = average_sequence_length
            for i in range(batch_size + 1):
                data[batch_size + i] = i * average_sequence_length
                data[2 * batch_size + 1 + i] = i * average_sequence_length

    if input_mask.type.tensor_type.elem_type == TensorProto.FLOAT:
        data = np.float32(data)
    elif input_mask.type.tensor_type.elem_type == TensorProto.INT64:
        data = np.int64(data)

    return data


def output_test_data(directory: str, inputs: dict[str, np.ndarray]):
    """Output input tensors of test data to a directory
    Args:
        directory (str): path of a directory
        inputs (Dict[str, np.ndarray]): map from input name to value
    """
    if not os.path.exists(directory):
        try:
            os.mkdir(directory)
        except OSError:
            print(f"Creation of the directory {directory} failed")
        else:
            print(f"Successfully created the directory {directory} ")
    else:
        print(f"Warning: directory {directory} existed. Files will be overwritten.")

    for index, (name, data) in enumerate(inputs.items()):
        tensor = numpy_helper.from_array(data, name)
        with open(os.path.join(directory, f"input_{index}.pb"), "wb") as file:
            file.write(tensor.SerializeToString())


def fake_test_data(
    batch_size: int,
    sequence_length: int,
    test_cases: int,
    dictionary_size: int,
    verbose: bool,
    random_seed: int,
    input_ids: TensorProto,
    segment_ids: TensorProto,
    input_mask: TensorProto,
    average_sequence_length: int,
    random_sequence_length: bool,
    mask_type: int,
):
    """Create given number of input data for testing
    Args:
        batch_size (int): batch size
        sequence_length (int): sequence length
        test_cases (int): number of test cases
        dictionary_size (int): vocabulary size of dictionary for input_ids
        verbose (bool): print more information or not
        random_seed (int): random seed
        input_ids (TensorProto): graph input of input IDs
        segment_ids (TensorProto): graph input of token type IDs
        input_mask (TensorProto): graph input of attention mask
        average_sequence_length (int): average sequence length excluding paddings
        random_sequence_length (bool): whether use uniform random number for sequence length
        mask_type (int): mask type 1 is mask index; 2 is 2D mask; 3 is key len, cumulated lengths of query and key
    Returns:
        List[Dict[str,numpy.ndarray]]: list of test cases, where each test case is a dictionary
                                       with input name as key and a tensor as value
    """
    assert input_ids is not None

    np.random.seed(random_seed)
    random.seed(random_seed)

    all_inputs = []
    for _test_case in range(test_cases):
        input_1 = fake_input_ids_data(input_ids, batch_size, sequence_length, dictionary_size)
        inputs = {input_ids.name: input_1}

        if segment_ids:
            inputs[segment_ids.name] = fake_segment_ids_data(segment_ids, batch_size, sequence_length)

        if input_mask:
            inputs[input_mask.name] = fake_input_mask_data(
                input_mask, batch_size, sequence_length, average_sequence_length, random_sequence_length, mask_type
            )

        if verbose and len(all_inputs) == 0:
            print("Example inputs", inputs)
        all_inputs.append(inputs)
    return all_inputs


def generate_test_data(
    batch_size: int,
    sequence_length: int,
    test_cases: int,
    seed: int,
    verbose: bool,
    input_ids: TensorProto,
    segment_ids: TensorProto,
    input_mask: TensorProto,
    average_sequence_length: int,
    random_sequence_length: bool,
    mask_type: int,
    dictionary_size: int = 10000,
):
    """Create given number of input data for testing
    Args:
        batch_size (int): batch size
        sequence_length (int): sequence length
        test_cases (int): number of test cases
        seed (int): random seed
        verbose (bool): print more information or not
        input_ids (TensorProto): graph input of input IDs
        segment_ids (TensorProto): graph input of token type IDs
        input_mask (TensorProto): graph input of attention mask
        average_sequence_length (int): average sequence length excluding paddings
        random_sequence_length (bool): whether use uniform random number for sequence length
        mask_type (int): mask type 1 is mask index; 2 is 2D mask; 3 is key len, cumulated lengths of query and key
    Returns:
        List[Dict[str,numpy.ndarray]]: list of test cases, where each test case is a dictionary
                                       with input name as key and a tensor as value
    """
    all_inputs = fake_test_data(
        batch_size,
        sequence_length,
        test_cases,
        dictionary_size,
        verbose,
        seed,
        input_ids,
        segment_ids,
        input_mask,
        average_sequence_length,
        random_sequence_length,
        mask_type,
    )
    if len(all_inputs) != test_cases:
        print("Failed to create test data for test.")
    return all_inputs


def get_graph_input_from_embed_node(onnx_model, embed_node, input_index):
    if input_index >= len(embed_node.input):
        return None

    input = embed_node.input[input_index]
    graph_input = onnx_model.find_graph_input(input)
    if graph_input is None:
        parent_node = onnx_model.get_parent(embed_node, input_index)
        if parent_node is not None and parent_node.op_type == "Cast":
            graph_input = onnx_model.find_graph_input(parent_node.input[0])
    return graph_input


def find_bert_inputs(
    onnx_model: OnnxModel,
    input_ids_name: str | None = None,
    segment_ids_name: str | None = None,
    input_mask_name: str | None = None,
) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None]:
    """Find graph inputs for BERT model.
    First, we will deduce inputs from EmbedLayerNormalization node.
    If not found, we will guess the meaning of graph inputs based on naming.
    Args:
        onnx_model (OnnxModel): onnx model object
        input_ids_name (str, optional): Name of graph input for input IDs. Defaults to None.
        segment_ids_name (str, optional): Name of graph input for segment IDs. Defaults to None.
        input_mask_name (str, optional): Name of graph input for attention mask. Defaults to None.
    Raises:
        ValueError: Graph does not have input named of input_ids_name or segment_ids_name or input_mask_name
        ValueError: Expected graph input number does not match with specified input_ids_name, segment_ids_name
                    and input_mask_name
    Returns:
        Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: input tensors of input_ids,
                                                                                 segment_ids and input_mask
    """

    graph_inputs = onnx_model.get_graph_inputs_excluding_initializers()

    if input_ids_name is not None:
        input_ids = onnx_model.find_graph_input(input_ids_name)
        if input_ids is None:
            raise ValueError(f"Graph does not have input named {input_ids_name}")

        segment_ids = None
        if segment_ids_name:
            segment_ids = onnx_model.find_graph_input(segment_ids_name)
            if segment_ids is None:
                raise ValueError(f"Graph does not have input named {segment_ids_name}")

        input_mask = None
        if input_mask_name:
            input_mask = onnx_model.find_graph_input(input_mask_name)
            if input_mask is None:
                raise ValueError(f"Graph does not have input named {input_mask_name}")

        expected_inputs = 1 + (1 if segment_ids else 0) + (1 if input_mask else 0)
        if len(graph_inputs) != expected_inputs:
            raise ValueError(f"Expect the graph to have {expected_inputs} inputs. Got {len(graph_inputs)}")

        return input_ids, segment_ids, input_mask

    if len(graph_inputs) != 3:
        raise ValueError(f"Expect the graph to have 3 inputs. Got {len(graph_inputs)}")

    embed_nodes = onnx_model.get_nodes_by_op_type("EmbedLayerNormalization")
    if len(embed_nodes) == 1:
        embed_node = embed_nodes[0]
        input_ids = get_graph_input_from_embed_node(onnx_model, embed_node, 0)
        segment_ids = get_graph_input_from_embed_node(onnx_model, embed_node, 1)
        input_mask = get_graph_input_from_embed_node(onnx_model, embed_node, 7)

        if input_mask is None:
            for input in graph_inputs:
                input_name_lower = input.name.lower()
                if "mask" in input_name_lower:
                    input_mask = input
        if input_mask is None:
            raise ValueError("Failed to find attention mask input")

        return input_ids, segment_ids, input_mask

    # Try guess the inputs based on naming.
    input_ids = None
    segment_ids = None
    input_mask = None
    for input in graph_inputs:
        input_name_lower = input.name.lower()
        if "mask" in input_name_lower:  # matches input with name like "attention_mask" or "input_mask"
            input_mask = input
        elif (
            "token" in input_name_lower or "segment" in input_name_lower
        ):  # matches input with name like "segment_ids" or "token_type_ids"
            segment_ids = input
        else:
            input_ids = input

    if input_ids and segment_ids and input_mask:
        return input_ids, segment_ids, input_mask

    raise ValueError("Fail to assign 3 inputs. You might try rename the graph inputs.")


def get_bert_inputs(
    onnx_file: str,
    input_ids_name: str | None = None,
    segment_ids_name: str | None = None,
    input_mask_name: str | None = None,
) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None]:
    """Find graph inputs for BERT model.
    First, we will deduce inputs from EmbedLayerNormalization node.
    If not found, we will guess the meaning of graph inputs based on naming.
    Args:
        onnx_file (str): onnx model path
        input_ids_name (str, optional): Name of graph input for input IDs. Defaults to None.
        segment_ids_name (str, optional): Name of graph input for segment IDs. Defaults to None.
        input_mask_name (str, optional): Name of graph input for attention mask. Defaults to None.
    Returns:
        Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: input tensors of input_ids,
                                                                                 segment_ids and input_mask
    """
    model = ModelProto()
    with open(onnx_file, "rb") as file:
        model.ParseFromString(file.read())

    onnx_model = OnnxModel(model)
    return find_bert_inputs(onnx_model, input_ids_name, segment_ids_name, input_mask_name)


def parse_arguments():
    parser = argparse.ArgumentParser()

    parser.add_argument("--model", required=True, type=str, help="bert onnx model path.")

    parser.add_argument(
        "--output_dir",
        required=False,
        type=str,
        default=None,
        help="output test data path. Default is current directory.",
    )

    parser.add_argument("--batch_size", required=False, type=int, default=1, help="batch size of input")

    parser.add_argument(
        "--sequence_length",
        required=False,
        type=int,
        default=128,
        help="maximum sequence length of input",
    )

    parser.add_argument(
        "--input_ids_name",
        required=False,
        type=str,
        default=None,
        help="input name for input ids",
    )
    parser.add_argument(
        "--segment_ids_name",
        required=False,
        type=str,
        default=None,
        help="input name for segment ids",
    )
    parser.add_argument(
        "--input_mask_name",
        required=False,
        type=str,
        default=None,
        help="input name for attention mask",
    )

    parser.add_argument(
        "--samples",
        required=False,
        type=int,
        default=1,
        help="number of test cases to be generated",
    )

    parser.add_argument("--seed", required=False, type=int, default=3, help="random seed")

    parser.add_argument(
        "--verbose",
        required=False,
        action="store_true",
        help="print verbose information",
    )
    parser.set_defaults(verbose=False)

    parser.add_argument(
        "--only_input_tensors",
        required=False,
        action="store_true",
        help="only save input tensors and no output tensors",
    )
    parser.set_defaults(only_input_tensors=False)

    parser.add_argument(
        "-a",
        "--average_sequence_length",
        default=-1,
        type=int,
        help="average sequence length excluding padding",
    )

    parser.add_argument(
        "-r",
        "--random_sequence_length",
        required=False,
        action="store_true",
        help="use uniform random instead of fixed sequence length",
    )
    parser.set_defaults(random_sequence_length=False)

    parser.add_argument(
        "--mask_type",
        required=False,
        type=int,
        default=2,
        help="mask type: (1: mask index, 2: raw 2D mask, 3: key lengths, cumulated lengths of query and key)",
    )

    args = parser.parse_args()
    return args


def create_and_save_test_data(
    model: str,
    output_dir: str,
    batch_size: int,
    sequence_length: int,
    test_cases: int,
    seed: int,
    verbose: bool,
    input_ids_name: str | None,
    segment_ids_name: str | None,
    input_mask_name: str | None,
    only_input_tensors: bool,
    average_sequence_length: int,
    random_sequence_length: bool,
    mask_type: int,
):
    """Create test data for a model, and save test data to a directory.
    Args:
        model (str): path of ONNX bert model
        output_dir (str): output directory
        batch_size (int): batch size
        sequence_length (int): sequence length
        test_cases (int): number of test cases
        seed (int): random seed
        verbose (bool): whether print more information
        input_ids_name (str): graph input name of input_ids
        segment_ids_name (str): graph input name of segment_ids
        input_mask_name (str): graph input name of input_mask
        only_input_tensors (bool): only save input tensors,
        average_sequence_length (int): average sequence length excluding paddings
        random_sequence_length (bool): whether use uniform random number for sequence length
        mask_type(int): mask type
    """
    input_ids, segment_ids, input_mask = get_bert_inputs(model, input_ids_name, segment_ids_name, input_mask_name)

    all_inputs = generate_test_data(
        batch_size,
        sequence_length,
        test_cases,
        seed,
        verbose,
        input_ids,
        segment_ids,
        input_mask,
        average_sequence_length,
        random_sequence_length,
        mask_type,
    )

    for i, inputs in enumerate(all_inputs):
        directory = os.path.join(output_dir, "test_data_set_" + str(i))
        output_test_data(directory, inputs)

    if only_input_tensors:
        return

    import onnxruntime  # noqa: PLC0415

    providers = (
        ["CUDAExecutionProvider", "CPUExecutionProvider"]
        if "CUDAExecutionProvider" in onnxruntime.get_available_providers()
        else ["CPUExecutionProvider"]
    )
    session = onnxruntime.InferenceSession(model, providers=providers)
    output_names = [output.name for output in session.get_outputs()]

    for i, inputs in enumerate(all_inputs):
        directory = os.path.join(output_dir, "test_data_set_" + str(i))
        result = session.run(output_names, inputs)
        for i, output_name in enumerate(output_names):  # noqa: PLW2901
            tensor_result = numpy_helper.from_array(np.asarray(result[i]), output_name)
            with open(os.path.join(directory, f"output_{i}.pb"), "wb") as file:
                file.write(tensor_result.SerializeToString())


def main():
    args = parse_arguments()

    if args.average_sequence_length <= 0:
        args.average_sequence_length = args.sequence_length

    output_dir = args.output_dir
    if output_dir is None:
        # Default output directory is a sub-directory under the directory of model.
        p = Path(args.model)
        output_dir = os.path.join(p.parent, f"batch_{args.batch_size}_seq_{args.sequence_length}")

    if output_dir is not None:
        # create the output directory if not existed
        path = Path(output_dir)
        path.mkdir(parents=True, exist_ok=True)
    else:
        print("Directory existed. test data files will be overwritten.")

    create_and_save_test_data(
        args.model,
        output_dir,
        args.batch_size,
        args.sequence_length,
        args.samples,
        args.seed,
        args.verbose,
        args.input_ids_name,
        args.segment_ids_name,
        args.input_mask_name,
        args.only_input_tensors,
        args.average_sequence_length,
        args.random_sequence_length,
        args.mask_type,
    )

    print("Test data is saved to directory:", output_dir)


if __name__ == "__main__":
    main()

	# -------------------------------------------------------------------------
	# Copyright (c) Microsoft Corporation. All rights reserved.
	# Licensed under the MIT License.
	# --------------------------------------------------------------------------

	# It is a tool to generate test data for a bert model.
	# The test data can be used by onnxruntime_perf_test tool to evaluate the inference latency.

	import argparse
	import os
	import random
	from pathlib import Path

	import numpy as np
	from onnx import ModelProto, TensorProto, numpy_helper
	from onnx_model import OnnxModel


	def fake_input_ids_data(
	input_ids: TensorProto, batch_size: int, sequence_length: int, dictionary_size: int
	) -> np.ndarray:
	"""Create input tensor based on the graph input of input_ids

	Args:
	input_ids (TensorProto): graph input of the input_ids input tensor
	batch_size (int): batch size
	sequence_length (int): sequence length
	dictionary_size (int): vocabulary size of dictionary

	Returns:
	np.ndarray: the input tensor created
	"""
	assert input_ids.type.tensor_type.elem_type in [
	TensorProto.FLOAT,
	TensorProto.INT32,
	TensorProto.INT64,
	]

	data = np.random.randint(dictionary_size, size=(batch_size, sequence_length), dtype=np.int32)

	if input_ids.type.tensor_type.elem_type == TensorProto.FLOAT:
	data = np.float32(data)
	elif input_ids.type.tensor_type.elem_type == TensorProto.INT64:
	data = np.int64(data)

	return data


	def fake_segment_ids_data(segment_ids: TensorProto, batch_size: int, sequence_length: int) -> np.ndarray:
	"""Create input tensor based on the graph input of segment_ids

	Args:
	segment_ids (TensorProto): graph input of the token_type_ids input tensor
	batch_size (int): batch size
	sequence_length (int): sequence length

	Returns:
	np.ndarray: the input tensor created
	"""
	assert segment_ids.type.tensor_type.elem_type in [
	TensorProto.FLOAT,
	TensorProto.INT32,
	TensorProto.INT64,
	]

	data = np.zeros((batch_size, sequence_length), dtype=np.int32)

	if segment_ids.type.tensor_type.elem_type == TensorProto.FLOAT:
	data = np.float32(data)
	elif segment_ids.type.tensor_type.elem_type == TensorProto.INT64:
	data = np.int64(data)

	return data


	def get_random_length(max_sequence_length: int, average_sequence_length: int):
	assert average_sequence_length >= 1 and average_sequence_length <= max_sequence_length

	# For uniform distribution, we find proper lower and upper bounds so that the average is in the middle.
	if 2 * average_sequence_length > max_sequence_length:
	return random.randint(2 * average_sequence_length - max_sequence_length, max_sequence_length)
	else:
	return random.randint(1, 2 * average_sequence_length - 1)


	def fake_input_mask_data(
	input_mask: TensorProto,
	batch_size: int,
	sequence_length: int,
	average_sequence_length: int,
	random_sequence_length: bool,
	mask_type: int = 2,
	) -> np.ndarray:
	"""Create input tensor based on the graph input of segment_ids.

	Args:
	input_mask (TensorProto): graph input of the attention mask input tensor
	batch_size (int): batch size
	sequence_length (int): sequence length
	average_sequence_length (int): average sequence length excluding paddings
	random_sequence_length (bool): whether use uniform random number for sequence length
	mask_type (int): mask type - 1: mask index (sequence length excluding paddings). Shape is (batch_size).
	2: 2D attention mask. Shape is (batch_size, sequence_length).
	3: key len, cumulated lengths of query and key. Shape is (3 * batch_size + 2).

	Returns:
	np.ndarray: the input tensor created
	"""

	assert input_mask.type.tensor_type.elem_type in [
	TensorProto.FLOAT,
	TensorProto.INT32,
	TensorProto.INT64,
	]

	if mask_type == 1: # sequence length excluding paddings
	data = np.ones((batch_size), dtype=np.int32)
	if random_sequence_length:
	for i in range(batch_size):
	data[i] = get_random_length(sequence_length, average_sequence_length)
	else:
	for i in range(batch_size):
	data[i] = average_sequence_length
	elif mask_type == 2: # 2D attention mask
	data = np.zeros((batch_size, sequence_length), dtype=np.int32)
	if random_sequence_length:
	for i in range(batch_size):
	actual_seq_len = get_random_length(sequence_length, average_sequence_length)
	for j in range(actual_seq_len):
	data[i, j] = 1
	else:
	temp = np.ones((batch_size, average_sequence_length), dtype=np.int32)
	data[: temp.shape[0], : temp.shape[1]] = temp
	else:
	assert mask_type == 3
	data = np.zeros((batch_size * 3 + 2), dtype=np.int32)
	if random_sequence_length:
	for i in range(batch_size):
	data[i] = get_random_length(sequence_length, average_sequence_length)

	for i in range(batch_size + 1):
	data[batch_size + i] = data[batch_size + i - 1] + data[i - 1] if i > 0 else 0
	data[2 * batch_size + 1 + i] = data[batch_size + i - 1] + data[i - 1] if i > 0 else 0
	else:
	for i in range(batch_size):
	data[i] = average_sequence_length
	for i in range(batch_size + 1):
	data[batch_size + i] = i * average_sequence_length
	data[2 * batch_size + 1 + i] = i * average_sequence_length

	if input_mask.type.tensor_type.elem_type == TensorProto.FLOAT:
	data = np.float32(data)
	elif input_mask.type.tensor_type.elem_type == TensorProto.INT64:
	data = np.int64(data)

	return data


	def output_test_data(directory: str, inputs: dict[str, np.ndarray]):
	"""Output input tensors of test data to a directory

	Args:
	directory (str): path of a directory
	inputs (Dict[str, np.ndarray]): map from input name to value
	"""
	if not os.path.exists(directory):
	try:
	os.mkdir(directory)
	except OSError:
	print(f"Creation of the directory {directory} failed")
	else:
	print(f"Successfully created the directory {directory} ")
	else:
	print(f"Warning: directory {directory} existed. Files will be overwritten.")

	for index, (name, data) in enumerate(inputs.items()):
	tensor = numpy_helper.from_array(data, name)
	with open(os.path.join(directory, f"input_{index}.pb"), "wb") as file:
	file.write(tensor.SerializeToString())


	def fake_test_data(
	batch_size: int,
	sequence_length: int,
	test_cases: int,
	dictionary_size: int,
	verbose: bool,
	random_seed: int,
	input_ids: TensorProto,
	segment_ids: TensorProto,
	input_mask: TensorProto,
	average_sequence_length: int,
	random_sequence_length: bool,
	mask_type: int,
	):
	"""Create given number of input data for testing

	Args:
	batch_size (int): batch size
	sequence_length (int): sequence length
	test_cases (int): number of test cases
	dictionary_size (int): vocabulary size of dictionary for input_ids
	verbose (bool): print more information or not
	random_seed (int): random seed
	input_ids (TensorProto): graph input of input IDs
	segment_ids (TensorProto): graph input of token type IDs
	input_mask (TensorProto): graph input of attention mask
	average_sequence_length (int): average sequence length excluding paddings
	random_sequence_length (bool): whether use uniform random number for sequence length
	mask_type (int): mask type 1 is mask index; 2 is 2D mask; 3 is key len, cumulated lengths of query and key

	Returns:
	List[Dict[str,numpy.ndarray]]: list of test cases, where each test case is a dictionary
	with input name as key and a tensor as value
	"""
	assert input_ids is not None

	np.random.seed(random_seed)
	random.seed(random_seed)

	all_inputs = []
	for _test_case in range(test_cases):
	input_1 = fake_input_ids_data(input_ids, batch_size, sequence_length, dictionary_size)
	inputs = {input_ids.name: input_1}

	if segment_ids:
	inputs[segment_ids.name] = fake_segment_ids_data(segment_ids, batch_size, sequence_length)

	if input_mask:
	inputs[input_mask.name] = fake_input_mask_data(
	input_mask, batch_size, sequence_length, average_sequence_length, random_sequence_length, mask_type
	)

	if verbose and len(all_inputs) == 0:
	print("Example inputs", inputs)
	all_inputs.append(inputs)
	return all_inputs


	def generate_test_data(
	batch_size: int,
	sequence_length: int,
	test_cases: int,
	seed: int,
	verbose: bool,
	input_ids: TensorProto,
	segment_ids: TensorProto,
	input_mask: TensorProto,
	average_sequence_length: int,
	random_sequence_length: bool,
	mask_type: int,
	dictionary_size: int = 10000,
	):
	"""Create given number of input data for testing

	Args:
	batch_size (int): batch size
	sequence_length (int): sequence length
	test_cases (int): number of test cases
	seed (int): random seed
	verbose (bool): print more information or not
	input_ids (TensorProto): graph input of input IDs
	segment_ids (TensorProto): graph input of token type IDs
	input_mask (TensorProto): graph input of attention mask
	average_sequence_length (int): average sequence length excluding paddings
	random_sequence_length (bool): whether use uniform random number for sequence length
	mask_type (int): mask type 1 is mask index; 2 is 2D mask; 3 is key len, cumulated lengths of query and key

	Returns:
	List[Dict[str,numpy.ndarray]]: list of test cases, where each test case is a dictionary
	with input name as key and a tensor as value
	"""
	all_inputs = fake_test_data(
	batch_size,
	sequence_length,
	test_cases,
	dictionary_size,
	verbose,
	seed,
	input_ids,
	segment_ids,
	input_mask,
	average_sequence_length,
	random_sequence_length,
	mask_type,
	)
	if len(all_inputs) != test_cases:
	print("Failed to create test data for test.")
	return all_inputs


	def get_graph_input_from_embed_node(onnx_model, embed_node, input_index):
	if input_index >= len(embed_node.input):
	return None

	input = embed_node.input[input_index]
	graph_input = onnx_model.find_graph_input(input)
	if graph_input is None:
	parent_node = onnx_model.get_parent(embed_node, input_index)
	if parent_node is not None and parent_node.op_type == "Cast":
	graph_input = onnx_model.find_graph_input(parent_node.input[0])
	return graph_input


	def find_bert_inputs(
	onnx_model: OnnxModel,
	input_ids_name: str \| None = None,
	segment_ids_name: str \| None = None,
	input_mask_name: str \| None = None,
	) -> tuple[np.ndarray \| None, np.ndarray \| None, np.ndarray \| None]:
	"""Find graph inputs for BERT model.
	First, we will deduce inputs from EmbedLayerNormalization node.
	If not found, we will guess the meaning of graph inputs based on naming.

	Args:
	onnx_model (OnnxModel): onnx model object
	input_ids_name (str, optional): Name of graph input for input IDs. Defaults to None.
	segment_ids_name (str, optional): Name of graph input for segment IDs. Defaults to None.
	input_mask_name (str, optional): Name of graph input for attention mask. Defaults to None.

	Raises:
	ValueError: Graph does not have input named of input_ids_name or segment_ids_name or input_mask_name
	ValueError: Expected graph input number does not match with specified input_ids_name, segment_ids_name
	and input_mask_name

	Returns:
	Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: input tensors of input_ids,
	segment_ids and input_mask
	"""

	graph_inputs = onnx_model.get_graph_inputs_excluding_initializers()

	if input_ids_name is not None:
	input_ids = onnx_model.find_graph_input(input_ids_name)
	if input_ids is None:
	raise ValueError(f"Graph does not have input named {input_ids_name}")

	segment_ids = None
	if segment_ids_name:
	segment_ids = onnx_model.find_graph_input(segment_ids_name)
	if segment_ids is None:
	raise ValueError(f"Graph does not have input named {segment_ids_name}")

	input_mask = None
	if input_mask_name:
	input_mask = onnx_model.find_graph_input(input_mask_name)
	if input_mask is None:
	raise ValueError(f"Graph does not have input named {input_mask_name}")

	expected_inputs = 1 + (1 if segment_ids else 0) + (1 if input_mask else 0)
	if len(graph_inputs) != expected_inputs:
	raise ValueError(f"Expect the graph to have {expected_inputs} inputs. Got {len(graph_inputs)}")

	return input_ids, segment_ids, input_mask

	if len(graph_inputs) != 3:
	raise ValueError(f"Expect the graph to have 3 inputs. Got {len(graph_inputs)}")

	embed_nodes = onnx_model.get_nodes_by_op_type("EmbedLayerNormalization")
	if len(embed_nodes) == 1:
	embed_node = embed_nodes[0]
	input_ids = get_graph_input_from_embed_node(onnx_model, embed_node, 0)
	segment_ids = get_graph_input_from_embed_node(onnx_model, embed_node, 1)
	input_mask = get_graph_input_from_embed_node(onnx_model, embed_node, 7)

	if input_mask is None:
	for input in graph_inputs:
	input_name_lower = input.name.lower()
	if "mask" in input_name_lower:
	input_mask = input
	if input_mask is None:
	raise ValueError("Failed to find attention mask input")

	return input_ids, segment_ids, input_mask

	# Try guess the inputs based on naming.
	input_ids = None
	segment_ids = None
	input_mask = None
	for input in graph_inputs:
	input_name_lower = input.name.lower()
	if "mask" in input_name_lower: # matches input with name like "attention_mask" or "input_mask"
	input_mask = input
	elif (
	"token" in input_name_lower or "segment" in input_name_lower
	): # matches input with name like "segment_ids" or "token_type_ids"
	segment_ids = input
	else:
	input_ids = input

	if input_ids and segment_ids and input_mask:
	return input_ids, segment_ids, input_mask

	raise ValueError("Fail to assign 3 inputs. You might try rename the graph inputs.")


	def get_bert_inputs(
	onnx_file: str,
	input_ids_name: str \| None = None,
	segment_ids_name: str \| None = None,
	input_mask_name: str \| None = None,
	) -> tuple[np.ndarray \| None, np.ndarray \| None, np.ndarray \| None]:
	"""Find graph inputs for BERT model.
	First, we will deduce inputs from EmbedLayerNormalization node.
	If not found, we will guess the meaning of graph inputs based on naming.

	Args:
	onnx_file (str): onnx model path
	input_ids_name (str, optional): Name of graph input for input IDs. Defaults to None.
	segment_ids_name (str, optional): Name of graph input for segment IDs. Defaults to None.
	input_mask_name (str, optional): Name of graph input for attention mask. Defaults to None.

	Returns:
	Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: input tensors of input_ids,
	segment_ids and input_mask
	"""
	model = ModelProto()
	with open(onnx_file, "rb") as file:
	model.ParseFromString(file.read())

	onnx_model = OnnxModel(model)
	return find_bert_inputs(onnx_model, input_ids_name, segment_ids_name, input_mask_name)


	def parse_arguments():
	parser = argparse.ArgumentParser()

	parser.add_argument("--model", required=True, type=str, help="bert onnx model path.")

	parser.add_argument(
	"--output_dir",
	required=False,
	type=str,
	default=None,
	help="output test data path. Default is current directory.",
	)

	parser.add_argument("--batch_size", required=False, type=int, default=1, help="batch size of input")

	parser.add_argument(
	"--sequence_length",
	required=False,
	type=int,
	default=128,
	help="maximum sequence length of input",
	)

	parser.add_argument(
	"--input_ids_name",
	required=False,
	type=str,
	default=None,
	help="input name for input ids",
	)
	parser.add_argument(
	"--segment_ids_name",
	required=False,
	type=str,
	default=None,
	help="input name for segment ids",
	)
	parser.add_argument(
	"--input_mask_name",
	required=False,
	type=str,
	default=None,
	help="input name for attention mask",
	)

	parser.add_argument(
	"--samples",
	required=False,
	type=int,
	default=1,
	help="number of test cases to be generated",
	)

	parser.add_argument("--seed", required=False, type=int, default=3, help="random seed")

	parser.add_argument(
	"--verbose",
	required=False,
	action="store_true",
	help="print verbose information",
	)
	parser.set_defaults(verbose=False)

	parser.add_argument(
	"--only_input_tensors",
	required=False,
	action="store_true",
	help="only save input tensors and no output tensors",
	)
	parser.set_defaults(only_input_tensors=False)

	parser.add_argument(
	"-a",
	"--average_sequence_length",
	default=-1,
	type=int,
	help="average sequence length excluding padding",
	)

	parser.add_argument(
	"-r",
	"--random_sequence_length",
	required=False,
	action="store_true",
	help="use uniform random instead of fixed sequence length",
	)
	parser.set_defaults(random_sequence_length=False)

	parser.add_argument(
	"--mask_type",
	required=False,
	type=int,
	default=2,
	help="mask type: (1: mask index, 2: raw 2D mask, 3: key lengths, cumulated lengths of query and key)",
	)

	args = parser.parse_args()
	return args


	def create_and_save_test_data(
	model: str,
	output_dir: str,
	batch_size: int,
	sequence_length: int,
	test_cases: int,
	seed: int,
	verbose: bool,
	input_ids_name: str \| None,
	segment_ids_name: str \| None,
	input_mask_name: str \| None,
	only_input_tensors: bool,
	average_sequence_length: int,
	random_sequence_length: bool,
	mask_type: int,
	):
	"""Create test data for a model, and save test data to a directory.

	Args:
	model (str): path of ONNX bert model
	output_dir (str): output directory
	batch_size (int): batch size
	sequence_length (int): sequence length
	test_cases (int): number of test cases
	seed (int): random seed
	verbose (bool): whether print more information
	input_ids_name (str): graph input name of input_ids
	segment_ids_name (str): graph input name of segment_ids
	input_mask_name (str): graph input name of input_mask
	only_input_tensors (bool): only save input tensors,
	average_sequence_length (int): average sequence length excluding paddings
	random_sequence_length (bool): whether use uniform random number for sequence length
	mask_type(int): mask type
	"""
	input_ids, segment_ids, input_mask = get_bert_inputs(model, input_ids_name, segment_ids_name, input_mask_name)

	all_inputs = generate_test_data(
	batch_size,
	sequence_length,
	test_cases,
	seed,
	verbose,
	input_ids,
	segment_ids,
	input_mask,
	average_sequence_length,
	random_sequence_length,
	mask_type,
	)

	for i, inputs in enumerate(all_inputs):
	directory = os.path.join(output_dir, "test_data_set_" + str(i))
	output_test_data(directory, inputs)

	if only_input_tensors:
	return

	import onnxruntime # noqa: PLC0415

	providers = (
	["CUDAExecutionProvider", "CPUExecutionProvider"]
	if "CUDAExecutionProvider" in onnxruntime.get_available_providers()
	else ["CPUExecutionProvider"]
	)
	session = onnxruntime.InferenceSession(model, providers=providers)
	output_names = [output.name for output in session.get_outputs()]

	for i, inputs in enumerate(all_inputs):
	directory = os.path.join(output_dir, "test_data_set_" + str(i))
	result = session.run(output_names, inputs)
	for i, output_name in enumerate(output_names): # noqa: PLW2901
	tensor_result = numpy_helper.from_array(np.asarray(result[i]), output_name)
	with open(os.path.join(directory, f"output_{i}.pb"), "wb") as file:
	file.write(tensor_result.SerializeToString())


	def main():
	args = parse_arguments()

	if args.average_sequence_length <= 0:
	args.average_sequence_length = args.sequence_length

	output_dir = args.output_dir
	if output_dir is None:
	# Default output directory is a sub-directory under the directory of model.
	p = Path(args.model)
	output_dir = os.path.join(p.parent, f"batch_{args.batch_size}_seq_{args.sequence_length}")

	if output_dir is not None:
	# create the output directory if not existed
	path = Path(output_dir)
	path.mkdir(parents=True, exist_ok=True)
	else:
	print("Directory existed. test data files will be overwritten.")

	create_and_save_test_data(
	args.model,
	output_dir,
	args.batch_size,
	args.sequence_length,
	args.samples,
	args.seed,
	args.verbose,
	args.input_ids_name,
	args.segment_ids_name,
	args.input_mask_name,
	args.only_input_tensors,
	args.average_sequence_length,
	args.random_sequence_length,
	args.mask_type,
	)

	print("Test data is saved to directory:", output_dir)


	if __name__ == "__main__":
	main()