#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------
"""
Convert a Bert large model exported by Keras2Onnx to a tiny model for test purpose.
The input model is generated like the following (need install keras2onnx from source):

    import numpy
    import keras2onnx
    from transformers import (TFBertForQuestionAnswering, BertTokenizer)

    tokenizer = BertTokenizer.from_pretrained'bert-large-uncased-whole-word-masking-finetuned-squad', do_lower_case=True, cache_dir=cache_dir)
    model = TFBertForQuestionAnswering.from_pretrained(model_name_or_path, cache_dir=cache_dir)

    question, text = "What is ONNX Runtime?", "ONNX Runtime is a performance-focused inference engine for ONNX models."
    inputs = tokenizer.encode_plus(question, text, add_special_tokens=True, return_tensors='tf')

    output_model_path =  os.path.join(output_dir, 'keras_{}.onnx'.format(model_name_or_path))
    if not os.path.exists(output_model_path):
        model.predict(inputs)
        onnx_model = keras2onnx.convert_keras(model, model.name)
        keras2onnx.save_model(onnx_model, output_model_path)
"""

import onnx
import onnx.utils
import sys
import argparse
import numpy as np
from onnx import ModelProto, TensorProto, numpy_helper
from onnxruntime_tools.transformers.onnx_model import OnnxModel
import os
import onnxruntime
import random
from pathlib import Path
import timeit

DICT_SIZE = 20
SEQ_LEN = 7
""" This class creates a tiny bert model for test purpose. """


class TinyBertOnnxModel(OnnxModel):
    def __init__(self, model, verbose):
        super(TinyBertOnnxModel, self).__init__(model, verbose)
        self.resize_model()

    def resize_weight(self, initializer_name, target_shape):
        weight = self.get_initializer(initializer_name)
        w = numpy_helper.to_array(weight)

        target_w = w
        if len(target_shape) == 1:
            target_w = w[:target_shape[0]]
        elif len(target_shape) == 2:
            target_w = w[:target_shape[0], :target_shape[1]]
        elif len(target_shape) == 3:
            target_w = w[:target_shape[0], :target_shape[1], :target_shape[2]]
        else:
            print("at most 3 dimensions")

        tensor = onnx.helper.make_tensor(name=initializer_name + '_resize',
                                         data_type=TensorProto.FLOAT,
                                         dims=target_shape,
                                         vals=target_w.flatten().tolist())

        return tensor

    def resize_model(self):
        graph = self.model.graph
        initializers = graph.initializer

        # parameters of input base model.
        old_parameters = {
            "seq_len": 26,
            "hidden_size": 1024,
            "num_heads": 16,
            "size_per_head": 64,
            "word_dict_size": [28996, 30522],  # list of supported dictionary size.
            "max_word_position": 512
        }

        # parameters of output tiny model.
        new_parameters = {
            "seq_len": SEQ_LEN,
            "hidden_size": 8,
            "num_heads": 2,
            "size_per_head": 4,
            "word_dict_size": DICT_SIZE,
            "max_word_position": 10
        }

        for input in graph.input:
            if (input.type.tensor_type.shape.dim[1].dim_value == old_parameters["seq_len"]):
                print("input", input.name, input.type.tensor_type.shape)
                input.type.tensor_type.shape.dim[1].dim_value = new_parameters["seq_len"]
                print("=>", input.type.tensor_type.shape)

        reshapes = {}
        for initializer in initializers:
            tensor = numpy_helper.to_array(initializer)
            dtype = np.float32 if initializer.data_type == 1 else np.int32
            if len(tensor.shape) == 1 and tensor.shape[0] == 1:
                if tensor == old_parameters["num_heads"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["num_heads"], "=>[", new_parameters["num_heads"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["num_heads"]], dtype=dtype),
                                                initializer.name))
                elif tensor == old_parameters["seq_len"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["seq_len"], "=>[", new_parameters["seq_len"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["seq_len"]], dtype=dtype), initializer.name))
                elif tensor == old_parameters["size_per_head"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["size_per_head"], "=>[", new_parameters["size_per_head"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["size_per_head"]], dtype=dtype),
                                                initializer.name))
                elif tensor == old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["hidden_size"], "=>[", new_parameters["hidden_size"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["hidden_size"]], dtype=dtype),
                                                initializer.name))
                elif tensor == 4 * old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          4 * old_parameters["hidden_size"], "=>[", 4 * new_parameters["hidden_size"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([4 * new_parameters["hidden_size"]], dtype=dtype),
                                                initializer.name))
            elif len(tensor.shape) == 0:
                if tensor == old_parameters["num_heads"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["num_heads"], "=>", new_parameters["num_heads"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["num_heads"], dtype=dtype), initializer.name))
                elif tensor == old_parameters["seq_len"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["seq_len"], "=>", new_parameters["seq_len"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["seq_len"], dtype=dtype), initializer.name))
                elif tensor == old_parameters["size_per_head"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["size_per_head"], "=>", new_parameters["size_per_head"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["size_per_head"], dtype=dtype),
                                                initializer.name))
                elif tensor == old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["hidden_size"], "=>", new_parameters["hidden_size"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["hidden_size"], dtype=dtype),
                                                initializer.name))
                elif tensor == 4 * old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          4 * old_parameters["hidden_size"], "=>", 4 * new_parameters["hidden_size"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(4 * new_parameters["hidden_size"], dtype=dtype),
                                                initializer.name))
                elif tensor == 1.0 / np.sqrt(old_parameters["size_per_head"]):
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          1.0 / np.sqrt(old_parameters["size_per_head"]), "=>",
                          1.0 / np.sqrt(new_parameters["size_per_head"]))
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(1.0 / np.sqrt(new_parameters["size_per_head"]), dtype=dtype),
                                                initializer.name))

            new_shape = []
            shape_changed = False
            for dim in tensor.shape:
                if (dim == old_parameters["hidden_size"]):
                    new_shape.append(new_parameters["hidden_size"])
                    shape_changed = True
                elif (dim == 4 * old_parameters["hidden_size"]):
                    new_shape.append(4 * new_parameters["hidden_size"])
                    shape_changed = True
                elif (dim in old_parameters["word_dict_size"]):
                    new_shape.append(new_parameters["word_dict_size"])
                    shape_changed = True
                elif (dim == old_parameters["max_word_position"]):
                    new_shape.append(new_parameters["max_word_position"])
                    shape_changed = True
                else:
                    new_shape.append(dim)
            if shape_changed:
                reshapes[initializer.name] = new_shape
                print("initializer", initializer.name, tensor.shape, "=>", new_shape)

        for initializer_name in reshapes:
            self.replace_input_of_all_nodes(initializer_name, initializer_name + '_resize')
            tensor = self.resize_weight(initializer_name, reshapes[initializer_name])
            self.model.graph.initializer.extend([tensor])

        self.use_dynamic_axes()

    def use_dynamic_axes(self, dynamic_batch_dim='batch_size', seq_len=7):
        """
        Update input and output shape to use dynamic axes.
        """
        dynamic_batch_inputs = {}
        for input in self.model.graph.input:
            dim_proto = input.type.tensor_type.shape.dim[0]
            dim_proto.dim_param = dynamic_batch_dim
            dim_proto = input.type.tensor_type.shape.dim[1]
            dim_proto.dim_value = seq_len

        for output in self.model.graph.output:
            dim_proto = output.type.tensor_type.shape.dim[0]
            dim_proto.dim_param = dynamic_batch_dim
            dim_proto = output.type.tensor_type.shape.dim[1]
            dim_proto.dim_value = seq_len


def generate_test_data(onnx_file,
                       output_path,
                       batch_size,
                       sequence_length,
                       use_cpu=True,
                       input_tensor_only=False,
                       dictionary_size=DICT_SIZE,
                       test_cases=3):
    input_data_type = np.int32
    for test_case in range(test_cases):
        input_1 = np.random.randint(dictionary_size, size=(batch_size, sequence_length), dtype=input_data_type)
        tensor_1 = numpy_helper.from_array(input_1, 'input_ids')

        actual_seq_len = random.randint(sequence_length - 3, sequence_length)
        input_2 = np.zeros((batch_size, sequence_length), dtype=input_data_type)
        temp = np.ones((batch_size, actual_seq_len), dtype=input_data_type)
        input_2[:temp.shape[0], :temp.shape[1]] = temp
        tensor_2 = numpy_helper.from_array(input_2, 'attention_mask')

        input_3 = np.zeros((batch_size, sequence_length), dtype=input_data_type)
        tensor_3 = numpy_helper.from_array(input_3, 'token_type_ids')

        path = os.path.join(output_path, 'test_data_set_' + str(test_case))
        try:
            os.mkdir(path)
        except OSError:
            print("Creation of the directory %s failed" % path)
        else:
            print("Successfully created the directory %s " % path)

        if input_tensor_only:
            return

        sess_options = onnxruntime.SessionOptions()
        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
        sess = onnxruntime.InferenceSession(onnx_file, sess_options, providers=['CPUExecutionProvider'])

        input1_name = sess.get_inputs()[0].name
        output_names = [output.name for output in sess.get_outputs()]
        inputs = {'input_ids': input_1, 'attention_mask': input_2, 'token_type_ids': input_3}
        print("inputs", inputs)
        result = sess.run(output_names, inputs)

        with open(os.path.join(path, 'input_{}.pb'.format(0)), 'wb') as f:
            f.write(tensor_1.SerializeToString())
        with open(os.path.join(path, 'input_{}.pb'.format(1)), 'wb') as f:
            f.write(tensor_2.SerializeToString())
        with open(os.path.join(path, 'input_{}.pb'.format(2)), 'wb') as f:
            f.write(tensor_3.SerializeToString())

        for i, output_name in enumerate(output_names):
            tensor_result = numpy_helper.from_array(
                np.asarray(result[i]).reshape((batch_size, sequence_length)), output_names[i])
            with open(os.path.join(path, 'output_{}.pb'.format(i)), 'wb') as f:
                f.write(tensor_result.SerializeToString())

        start_time = timeit.default_timer()

        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED

        path_prefix = onnx_file[:-5]  #remove .onnx suffix
        if use_cpu:
            sess_options.optimized_model_filepath = path_prefix + "_optimized_cpu.onnx"
        else:
            sess_options.optimized_model_filepath = path_prefix + "_optimized_gpu.onnx"

        session = onnxruntime.InferenceSession(onnx_file, sess_options=sess_options,
                                               providers=onnxruntime.get_available_providers())
        if use_cpu:
            session.set_providers(['CPUExecutionProvider'])  # use cpu
        else:
            if 'CUDAExecutionProvider' not in session.get_providers():
                print("Warning: GPU not found")
                continue
        outputs = session.run(None, inputs)
        evalTime = timeit.default_timer() - start_time
        if outputs[0].tolist() != result[0].tolist():
            print("Error: not same result after optimization. use_cpu={}, no_opt_output={}, opt_output={}".format(
                use_cpu, result[0].tolist(), outputs[1].tolist()))
        print("** Evaluation done in total {} secs".format(evalTime))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', required=True, type=str)
    parser.add_argument('--output', required=True, type=str)
    parser.add_argument('--float16', required=False, action='store_true')
    parser.set_defaults(float16=False)
    args = parser.parse_args()

    model = ModelProto()
    with open(args.input, "rb") as f:
        model.ParseFromString(f.read())

    bert_model = TinyBertOnnxModel(model, False)

    if args.float16:
        bert_model.convert_model_float32_to_float16()

    bert_model.update_graph()
    bert_model.remove_unused_constant()

    print("opset version", bert_model.model.opset_import[0].version)

    with open(args.output, "wb") as out:
        out.write(bert_model.model.SerializeToString())

    p = Path(args.output)
    data_path = p.parent

    batch_size = 1
    sequence_length = SEQ_LEN

    generate_test_data(args.output, data_path, batch_size, sequence_length, use_cpu=not args.float16)


if __name__ == "__main__":
    main()

#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------
"""
Convert a Bert large model exported by Keras2Onnx to a tiny model for test purpose.
The input model is generated like the following (need install keras2onnx from source):

    import numpy
    import keras2onnx
    from transformers import (TFBertForQuestionAnswering, BertTokenizer)

    tokenizer = BertTokenizer.from_pretrained'bert-large-uncased-whole-word-masking-finetuned-squad', do_lower_case=True, cache_dir=cache_dir)
    model = TFBertForQuestionAnswering.from_pretrained(model_name_or_path, cache_dir=cache_dir)

    question, text = "What is ONNX Runtime?", "ONNX Runtime is a performance-focused inference engine for ONNX models."
    inputs = tokenizer.encode_plus(question, text, add_special_tokens=True, return_tensors='tf')

    output_model_path =  os.path.join(output_dir, 'keras_{}.onnx'.format(model_name_or_path))
    if not os.path.exists(output_model_path):
        model.predict(inputs)
        onnx_model = keras2onnx.convert_keras(model, model.name)
        keras2onnx.save_model(onnx_model, output_model_path)
"""

import onnx
import onnx.utils
import sys
import argparse
import numpy as np
from onnx import ModelProto, TensorProto, numpy_helper
from onnxruntime_tools.transformers.onnx_model import OnnxModel
import os
import onnxruntime
import random
from pathlib import Path
import timeit

DICT_SIZE = 20
SEQ_LEN = 7
""" This class creates a tiny bert model for test purpose. """


class TinyBertOnnxModel(OnnxModel):
    def __init__(self, model, verbose):
        super(TinyBertOnnxModel, self).__init__(model, verbose)
        self.resize_model()

    def resize_weight(self, initializer_name, target_shape):
        weight = self.get_initializer(initializer_name)
        w = numpy_helper.to_array(weight)

        target_w = w
        if len(target_shape) == 1:
            target_w = w[:target_shape[0]]
        elif len(target_shape) == 2:
            target_w = w[:target_shape[0], :target_shape[1]]
        elif len(target_shape) == 3:
            target_w = w[:target_shape[0], :target_shape[1], :target_shape[2]]
        else:
            print("at most 3 dimensions")

        tensor = onnx.helper.make_tensor(name=initializer_name + '_resize',
                                         data_type=TensorProto.FLOAT,
                                         dims=target_shape,
                                         vals=target_w.flatten().tolist())

        return tensor

    def resize_model(self):
        graph = self.model.graph
        initializers = graph.initializer

        # parameters of input base model.
        old_parameters = {
            "seq_len": 26,
            "hidden_size": 1024,
            "num_heads": 16,
            "size_per_head": 64,
            "word_dict_size": [28996, 30522],  # list of supported dictionary size.
            "max_word_position": 512
        }

        # parameters of output tiny model.
        new_parameters = {
            "seq_len": SEQ_LEN,
            "hidden_size": 8,
            "num_heads": 2,
            "size_per_head": 4,
            "word_dict_size": DICT_SIZE,
            "max_word_position": 10
        }

        for input in graph.input:
            if (input.type.tensor_type.shape.dim[1].dim_value == old_parameters["seq_len"]):
                print("input", input.name, input.type.tensor_type.shape)
                input.type.tensor_type.shape.dim[1].dim_value = new_parameters["seq_len"]
                print("=>", input.type.tensor_type.shape)

        reshapes = {}
        for initializer in initializers:
            tensor = numpy_helper.to_array(initializer)
            dtype = np.float32 if initializer.data_type == 1 else np.int32
            if len(tensor.shape) == 1 and tensor.shape[0] == 1:
                if tensor == old_parameters["num_heads"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["num_heads"], "=>[", new_parameters["num_heads"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["num_heads"]], dtype=dtype),
                                                initializer.name))
                elif tensor == old_parameters["seq_len"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["seq_len"], "=>[", new_parameters["seq_len"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["seq_len"]], dtype=dtype), initializer.name))
                elif tensor == old_parameters["size_per_head"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["size_per_head"], "=>[", new_parameters["size_per_head"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["size_per_head"]], dtype=dtype),
                                                initializer.name))
                elif tensor == old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["hidden_size"], "=>[", new_parameters["hidden_size"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["hidden_size"]], dtype=dtype),
                                                initializer.name))
                elif tensor == 4 * old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          4 * old_parameters["hidden_size"], "=>[", 4 * new_parameters["hidden_size"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([4 * new_parameters["hidden_size"]], dtype=dtype),
                                                initializer.name))
            elif len(tensor.shape) == 0:
                if tensor == old_parameters["num_heads"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["num_heads"], "=>", new_parameters["num_heads"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["num_heads"], dtype=dtype), initializer.name))
                elif tensor == old_parameters["seq_len"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["seq_len"], "=>", new_parameters["seq_len"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["seq_len"], dtype=dtype), initializer.name))
                elif tensor == old_parameters["size_per_head"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["size_per_head"], "=>", new_parameters["size_per_head"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["size_per_head"], dtype=dtype),
                                                initializer.name))
                elif tensor == old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["hidden_size"], "=>", new_parameters["hidden_size"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["hidden_size"], dtype=dtype),
                                                initializer.name))
                elif tensor == 4 * old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          4 * old_parameters["hidden_size"], "=>", 4 * new_parameters["hidden_size"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(4 * new_parameters["hidden_size"], dtype=dtype),
                                                initializer.name))
                elif tensor == 1.0 / np.sqrt(old_parameters["size_per_head"]):
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          1.0 / np.sqrt(old_parameters["size_per_head"]), "=>",
                          1.0 / np.sqrt(new_parameters["size_per_head"]))
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(1.0 / np.sqrt(new_parameters["size_per_head"]), dtype=dtype),
                                                initializer.name))

            new_shape = []
            shape_changed = False
            for dim in tensor.shape:
                if (dim == old_parameters["hidden_size"]):
                    new_shape.append(new_parameters["hidden_size"])
                    shape_changed = True
                elif (dim == 4 * old_parameters["hidden_size"]):
                    new_shape.append(4 * new_parameters["hidden_size"])
                    shape_changed = True
                elif (dim in old_parameters["word_dict_size"]):
                    new_shape.append(new_parameters["word_dict_size"])
                    shape_changed = True
                elif (dim == old_parameters["max_word_position"]):
                    new_shape.append(new_parameters["max_word_position"])
                    shape_changed = True
                else:
                    new_shape.append(dim)
            if shape_changed:
                reshapes[initializer.name] = new_shape
                print("initializer", initializer.name, tensor.shape, "=>", new_shape)

        for initializer_name in reshapes:
            self.replace_input_of_all_nodes(initializer_name, initializer_name + '_resize')
            tensor = self.resize_weight(initializer_name, reshapes[initializer_name])
            self.model.graph.initializer.extend([tensor])

        self.use_dynamic_axes()

    def use_dynamic_axes(self, dynamic_batch_dim='batch_size', seq_len=7):
        """
        Update input and output shape to use dynamic axes.
        """
        dynamic_batch_inputs = {}
        for input in self.model.graph.input:
            dim_proto = input.type.tensor_type.shape.dim[0]
            dim_proto.dim_param = dynamic_batch_dim
            dim_proto = input.type.tensor_type.shape.dim[1]
            dim_proto.dim_value = seq_len

        for output in self.model.graph.output:
            dim_proto = output.type.tensor_type.shape.dim[0]
            dim_proto.dim_param = dynamic_batch_dim
            dim_proto = output.type.tensor_type.shape.dim[1]
            dim_proto.dim_value = seq_len


def generate_test_data(onnx_file,
                       output_path,
                       batch_size,
                       sequence_length,
                       use_cpu=True,
                       input_tensor_only=False,
                       dictionary_size=DICT_SIZE,
                       test_cases=3):
    input_data_type = np.int32
    for test_case in range(test_cases):
        input_1 = np.random.randint(dictionary_size, size=(batch_size, sequence_length), dtype=input_data_type)
        tensor_1 = numpy_helper.from_array(input_1, 'input_ids')

        actual_seq_len = random.randint(sequence_length - 3, sequence_length)
        input_2 = np.zeros((batch_size, sequence_length), dtype=input_data_type)
        temp = np.ones((batch_size, actual_seq_len), dtype=input_data_type)
        input_2[:temp.shape[0], :temp.shape[1]] = temp
        tensor_2 = numpy_helper.from_array(input_2, 'attention_mask')

        input_3 = np.zeros((batch_size, sequence_length), dtype=input_data_type)
        tensor_3 = numpy_helper.from_array(input_3, 'token_type_ids')

        path = os.path.join(output_path, 'test_data_set_' + str(test_case))
        try:
            os.mkdir(path)
        except OSError:
            print("Creation of the directory %s failed" % path)
        else:
            print("Successfully created the directory %s " % path)

        if input_tensor_only:
            return

        sess_options = onnxruntime.SessionOptions()
        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
        sess = onnxruntime.InferenceSession(onnx_file, sess_options, providers=['CPUExecutionProvider'])

        input1_name = sess.get_inputs()[0].name
        output_names = [output.name for output in sess.get_outputs()]
        inputs = {'input_ids': input_1, 'attention_mask': input_2, 'token_type_ids': input_3}
        print("inputs", inputs)
        result = sess.run(output_names, inputs)

        with open(os.path.join(path, 'input_{}.pb'.format(0)), 'wb') as f:
            f.write(tensor_1.SerializeToString())
        with open(os.path.join(path, 'input_{}.pb'.format(1)), 'wb') as f:
            f.write(tensor_2.SerializeToString())
        with open(os.path.join(path, 'input_{}.pb'.format(2)), 'wb') as f:
            f.write(tensor_3.SerializeToString())

        for i, output_name in enumerate(output_names):
            tensor_result = numpy_helper.from_array(
                np.asarray(result[i]).reshape((batch_size, sequence_length)), output_names[i])
            with open(os.path.join(path, 'output_{}.pb'.format(i)), 'wb') as f:
                f.write(tensor_result.SerializeToString())

        start_time = timeit.default_timer()

        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED

        path_prefix = onnx_file[:-5]  #remove .onnx suffix
        if use_cpu:
            sess_options.optimized_model_filepath = path_prefix + "_optimized_cpu.onnx"
        else:
            sess_options.optimized_model_filepath = path_prefix + "_optimized_gpu.onnx"

        session = onnxruntime.InferenceSession(onnx_file, sess_options=sess_options,
                                               providers=onnxruntime.get_available_providers())
        if use_cpu:
            session.set_providers(['CPUExecutionProvider'])  # use cpu
        else:
            if 'CUDAExecutionProvider' not in session.get_providers():
                print("Warning: GPU not found")
                continue
        outputs = session.run(None, inputs)
        evalTime = timeit.default_timer() - start_time
        if outputs[0].tolist() != result[0].tolist():
            print("Error: not same result after optimization. use_cpu={}, no_opt_output={}, opt_output={}".format(
                use_cpu, result[0].tolist(), outputs[1].tolist()))
        print("** Evaluation done in total {} secs".format(evalTime))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', required=True, type=str)
    parser.add_argument('--output', required=True, type=str)
    parser.add_argument('--float16', required=False, action='store_true')
    parser.set_defaults(float16=False)
    args = parser.parse_args()

    model = ModelProto()
    with open(args.input, "rb") as f:
        model.ParseFromString(f.read())

    bert_model = TinyBertOnnxModel(model, False)

    if args.float16:
        bert_model.convert_model_float32_to_float16()

    bert_model.update_graph()
    bert_model.remove_unused_constant()

    print("opset version", bert_model.model.opset_import[0].version)

    with open(args.output, "wb") as out:
        out.write(bert_model.model.SerializeToString())

    p = Path(args.output)
    data_path = p.parent

    batch_size = 1
    sequence_length = SEQ_LEN

    generate_test_data(args.output, data_path, batch_size, sequence_length, use_cpu=not args.float16)


if __name__ == "__main__":
    main()

#------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation.  All rights reserved. # Licensed under the MIT License. #-------------------------------------------------------------------------- Convert a Bert large model exported by Keras2Onnx to a tiny model for test purpose. The input model is generated like the following (need install keras2onnx from source):      import numpy     import keras2onnx     from transformers import (TFBertForQuestionAnswering, BertTokenizer)      tokenizer = BertTokenizer.from_pretrained'bert-large-uncased-whole-word-masking-finetuned-squad', do_lower_case=True, cache_dir=cache_dir)     model = TFBertForQuestionAnswering.from_pretrained(model_name_or_path, cache_dir=cache_dir)      question, text = "What is ONNX Runtime?", "ONNX Runtime is a performance-focused inference engine for ONNX models."     inputs = tokenizer.encode_plus(question, text, add_special_tokens=True, return_tensors='tf')      output_model_path =  os.path.join(output_dir, 'keras_{}.onnx'.format(model_name_or_path))     if not os.path.exists(output_model_path):         model.predict(inputs)         onnx_model = keras2onnx.convert_keras(model, model.name)         keras2onnx.save_model(onnx_model, output_model_path) This class creates a tiny bert model for test purpose. # parameters of input base model. # list of supported dictionary size. # parameters of output tiny model. Update input and output shape to use dynamic axes. #remove .onnx suffix # use cpu