onnxruntime-directml 1.20.0__cp313-cp313-win_amd64.whl

onnxruntime/transformers/onnx_exporter.py

@@ -0,0 +1,717 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.  See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+import logging
+import os
+from pathlib import Path
+
+import numpy
+import torch
+from affinity_helper import AffinitySetting
+from benchmark_helper import OptimizerInfo, Precision, create_onnxruntime_session
+from huggingface_models import MODEL_CLASSES
+from quantize_helper import QuantizeHelper
+from torch_onnx_export_helper import torch_onnx_export
+from transformers import AutoConfig, AutoFeatureExtractor, AutoTokenizer, LxmertConfig, TransfoXLConfig
+
+from onnxruntime.transformers.models.gpt2.gpt2_helper import (
+    PRETRAINED_GPT2_MODELS,
+    GPT2ModelNoPastState,
+    TFGPT2ModelNoPastState,
+)
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+
+logger = logging.getLogger(__name__)
+
+# Workaround by replacing torch.triu using self-defined op
+# Since torch.triu cannot be exported to ONNX. See https://github.com/pytorch/pytorch/issues/32968
+torch_func = {"triu": torch.triu}
+
+
+def triu_onnx(x, diagonal=0, out=None):
+    assert out is None
+    assert len(x.shape) == 2 and x.size(0) == x.size(1)
+
+    torch_triu = torch_func["triu"]
+    template = torch_triu(torch.ones((1024, 1024), dtype=torch.uint8), diagonal)
+    mask = template[: x.size(0), : x.size(1)]
+    return torch.where(mask.bool(), x, torch.zeros_like(x))
+
+
+def replace_torch_functions():
+    torch.triu = triu_onnx
+
+
+def restore_torch_functions():
+    torch.triu = torch_func["triu"]
+
+
+def create_onnxruntime_input(vocab_size, batch_size, sequence_length, input_names, config, data_type=numpy.int64):
+    if config.model_type in ["vit", "swin"]:
+        input_ids = numpy.random.rand(batch_size, 3, config.image_size, config.image_size).astype(numpy.float32)
+        inputs = {"pixel_values": input_ids}
+        return inputs
+
+    input_ids = numpy.random.randint(low=0, high=vocab_size - 1, size=(batch_size, sequence_length), dtype=data_type)
+    inputs = {"input_ids": input_ids}
+
+    if "attention_mask" in input_names:
+        attention_mask = numpy.ones([batch_size, sequence_length], dtype=data_type)
+        inputs["attention_mask"] = attention_mask
+
+    if "token_type_ids" in input_names:
+        segment_ids = numpy.zeros([batch_size, sequence_length], dtype=data_type)
+        inputs["token_type_ids"] = segment_ids
+
+    if config.is_encoder_decoder:
+        inputs["decoder_input_ids"] = input_ids
+
+    if isinstance(config, LxmertConfig):
+        inputs["visual_feats"] = numpy.random.randn(1, 1, config.visual_feat_dim).astype(numpy.float32)
+        inputs["visual_pos"] = numpy.random.randn(1, 1, config.visual_pos_dim).astype(numpy.float32)
+    if isinstance(config, TransfoXLConfig):
+        inputs["tf_transfo_xl_model/transformer/pos_emb/einsum/Einsum/inputs_1:0"] = numpy.zeros(
+            [config.hidden_size], dtype=numpy.float32
+        )
+    return inputs
+
+
+def filter_inputs(inputs, input_names):
+    remaining_model_inputs = {}
+    for input_name in input_names:
+        if input_name in inputs:
+            remaining_model_inputs[input_name] = inputs[input_name]
+    return remaining_model_inputs
+
+
+def flatten(inputs):
+    return [[flatten(i) for i in inputs] if isinstance(inputs, (list, tuple)) else inputs]
+
+
+def update_flatten_list(inputs, res_list):
+    for i in inputs:
+        res_list.append(i) if not isinstance(i, (list, tuple)) else update_flatten_list(i, res_list)
+    return res_list
+
+
+def build_dynamic_axes(example_inputs, outputs_flatten):
+    sequence_length = example_inputs["input_ids"].shape[-1]
+
+    dynamic_axes = {key: {0: "batch_size", 1: "seq_len"} for key in example_inputs}
+
+    output_names = ["output_" + str(i + 1) for i in range(len(outputs_flatten))]
+    for i, output_name in enumerate(output_names):
+        dynamic_axes[output_name] = {0: "batch_size"}
+        dims = outputs_flatten[i].shape
+        for j, dim in enumerate(dims):
+            if dim == sequence_length:
+                dynamic_axes[output_name].update({j: "seq_len"})
+    return dynamic_axes, output_names
+
+
+def validate_onnx_model(
+    onnx_model_path,
+    example_inputs,
+    example_outputs_flatten,
+    use_gpu,
+    fp16,
+    output_names=None,
+):
+    test_session = create_onnxruntime_session(onnx_model_path, use_gpu, enable_all_optimization=False)
+    if test_session is None:
+        logger.error(f"{onnx_model_path} is an invalid ONNX model")
+        return False
+
+    logger.info(f"{onnx_model_path} is a valid ONNX model")
+
+    # Compare the inference result with PyTorch or Tensorflow
+    example_ort_inputs = {k: t.numpy() for k, t in example_inputs.items()}
+    example_ort_outputs = test_session.run(output_names, example_ort_inputs)
+    if len(example_outputs_flatten) != len(example_ort_outputs):
+        logger.error(
+            f"Number of output tensors expected {len(example_outputs_flatten)}, got {len(example_ort_outputs)}"
+        )
+        return False
+
+    for i in range(len(example_outputs_flatten)):
+        abs_diff = numpy.amax(numpy.abs(example_ort_outputs[i] - example_outputs_flatten[i].cpu().numpy()))
+        if abs_diff > 1e-4:
+            logger.info(f"Max absolute diff={abs_diff} for output tensor {i}")
+
+        rtol = 5e-02 if fp16 else 1e-4
+        atol = 1e-01 if fp16 else 1e-4
+        if not numpy.allclose(
+            example_ort_outputs[i],
+            example_outputs_flatten[i].cpu().numpy(),
+            rtol=rtol,
+            atol=atol,
+        ):
+            logger.error(f"Output tensor {i} is not close: rtol={rtol}, atol={atol}")
+            return False
+
+    logger.info(f"inference result of onnxruntime is validated on {onnx_model_path}")
+    return True
+
+
+def get_onnx_file_path(
+    onnx_dir: str,
+    model_name: str,
+    input_count: int,
+    optimized_by_script: bool,
+    use_gpu: bool,
+    precision: Precision,
+    optimized_by_onnxruntime: bool,
+    use_external_data: bool,
+):
+    from re import sub
+
+    normalized_model_name = sub(r"[^a-zA-Z0-9_]", "_", model_name)
+
+    if not optimized_by_script:
+        filename = f"{normalized_model_name}_{input_count}"
+    else:
+        device = "gpu" if use_gpu else "cpu"
+        filename = f"{normalized_model_name}_{input_count}_{precision}_{device}"
+
+    if optimized_by_onnxruntime:
+        filename += "_ort"
+
+    directory = onnx_dir
+    # ONNXRuntime will not write external data so the raw and optimized models shall be in same directory.
+    if use_external_data and not optimized_by_onnxruntime:
+        directory = os.path.join(onnx_dir, filename)
+        if not os.path.exists(directory):
+            os.makedirs(directory)
+
+    return os.path.join(directory, f"{filename}.onnx")
+
+
+def add_filename_suffix(file_path: str, suffix: str) -> str:
+    """
+    Append a suffix at the filename (before the extension).
+    Args:
+        path: pathlib.Path The actual path object we would like to add a suffix
+        suffix: The suffix to add
+    Returns: path with suffix appended at the end of the filename and before extension
+    """
+    path = Path(file_path)
+    return str(path.parent.joinpath(path.stem + suffix).with_suffix(path.suffix))
+
+
+def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwrite, model_fusion_statistics):
+    if overwrite or not os.path.exists(ort_model_path):
+        Path(ort_model_path).parent.mkdir(parents=True, exist_ok=True)
+        from optimizer import get_fusion_statistics, optimize_by_onnxruntime
+
+        # Use onnxruntime to optimize model, which will be saved to *_ort.onnx
+        _ = optimize_by_onnxruntime(
+            onnx_model_path,
+            use_gpu=use_gpu,
+            optimized_model_path=ort_model_path,
+            opt_level=99,
+        )
+        model_fusion_statistics[ort_model_path] = get_fusion_statistics(ort_model_path)
+    else:
+        logger.info(f"Skip optimization since model existed: {ort_model_path}")
+
+
+def optimize_onnx_model(
+    onnx_model_path,
+    optimized_model_path,
+    model_type,
+    num_attention_heads,
+    hidden_size,
+    use_gpu,
+    precision,
+    use_raw_attention_mask,
+    overwrite,
+    model_fusion_statistics,
+    use_external_data_format,
+    optimization_options=None,
+):
+    if overwrite or not os.path.exists(optimized_model_path):
+        Path(optimized_model_path).parent.mkdir(parents=True, exist_ok=True)
+
+        from fusion_options import FusionOptions
+        from optimizer import optimize_model
+
+        if optimization_options is None:
+            optimization_options = FusionOptions(model_type)
+        optimization_options.use_raw_attention_mask(use_raw_attention_mask)
+        if precision == Precision.FLOAT16:
+            optimization_options.enable_gelu_approximation = True
+        if precision == Precision.INT8:
+            optimization_options.enable_embed_layer_norm = False
+
+        # For swin models, the num_attention_heads is a list, which isn't supported yet, so set to 0 for now
+        if model_type == "swin":
+            num_attention_heads = 0
+            hidden_size = 0
+
+        # Use script to optimize model.
+        # Use opt_level <= 1 for models to be converted to fp16, because some fused op (like FusedGemm) has only fp32 and no fp16.
+        # It is better to be conservative so we use opt_level=0 here, in case MemcpyFromHost is added to the graph by OnnxRuntime.
+        opt_model = optimize_model(
+            onnx_model_path,
+            model_type,
+            num_heads=num_attention_heads,
+            hidden_size=hidden_size,
+            opt_level=0,
+            optimization_options=optimization_options,
+            use_gpu=use_gpu,
+            only_onnxruntime=False,
+        )
+        if model_type == "bert_keras" or model_type == "bert_tf":
+            opt_model.use_dynamic_axes()
+
+        model_fusion_statistics[optimized_model_path] = opt_model.get_fused_operator_statistics()
+
+        if precision == Precision.FLOAT16:
+            opt_model.convert_float_to_float16(keep_io_types=True)
+
+        opt_model.save_model_to_file(optimized_model_path, use_external_data_format)
+    else:
+        logger.info(f"Skip optimization since model existed: {optimized_model_path}")
+
+
+def modelclass_dispatcher(model_name, custom_model_class):
+    if custom_model_class is not None:
+        if custom_model_class in MODEL_CLASSES:
+            return custom_model_class
+        else:
+            raise Exception("Valid model class: " + " ".join(MODEL_CLASSES))
+
+    if model_name in PRETRAINED_GPT2_MODELS:
+        return "GPT2ModelNoPastState"
+
+    import re
+
+    if re.search("-squad$", model_name) is not None:
+        return "AutoModelForQuestionAnswering"
+    elif re.search("-mprc$", model_name) is not None:
+        return "AutoModelForSequenceClassification"
+    elif re.search("gpt2", model_name) is not None:
+        return "AutoModelWithLMHead"
+
+    return "AutoModel"
+
+
+def load_pretrained_model(model_name, config, cache_dir, custom_model_class, is_tf_model=False):
+    model_class_name = modelclass_dispatcher(model_name, custom_model_class)
+
+    if model_class_name == "GPT2ModelNoPastState":
+        if is_tf_model:
+            return TFGPT2ModelNoPastState.from_pretrained(model_name, config=config, cache_dir=cache_dir)
+        else:
+            return GPT2ModelNoPastState.from_pretrained(model_name, config=config, cache_dir=cache_dir)
+
+    if is_tf_model:
+        model_class_name = "TF" + model_class_name
+
+    transformers_module = __import__("transformers", fromlist=[model_class_name])
+    logger.info(f"Model class name: {model_class_name}")
+    model_class = getattr(transformers_module, model_class_name)
+
+    return model_class.from_pretrained(model_name, config=config, cache_dir=cache_dir)
+
+
+def load_pt_model(model_name, model_class, cache_dir, config_modifier):
+    config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir)
+    if hasattr(config, "return_dict"):
+        config.return_dict = False
+
+    config_modifier.modify(config)
+
+    model = load_pretrained_model(model_name, config=config, cache_dir=cache_dir, custom_model_class=model_class)
+
+    return config, model
+
+
+def load_tf_model(model_name, model_class, cache_dir, config_modifier):
+    config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir)
+
+    config_modifier.modify(config)
+    # Loading tf model from transformers limits the cpu affinity to {0} when KMP_AFFINITY is set
+    # Restore the affinity after model loading for expected ORT performance
+    affinity_setting = AffinitySetting()
+    affinity_setting.get_affinity()
+    model = load_pretrained_model(
+        model_name,
+        config=config,
+        cache_dir=cache_dir,
+        custom_model_class=model_class,
+        is_tf_model=True,
+    )
+    affinity_setting.set_affinity()
+
+    return config, model
+
+
+# For test only
+def load_pt_model_from_tf(model_name):
+    # Note that we could get pt model from tf, but model source and its structure in this case is different from directly using
+    # load_pt_model() and load_tf_model() even with the same name. Therefore it should not be used for comparing with them
+    from convert_tf_models_to_pytorch import tf2pt_pipeline
+
+    config, model = tf2pt_pipeline(model_name)
+
+    return config, model
+
+
+def validate_and_optimize_onnx(
+    model_name,
+    use_external_data_format,
+    model_type,
+    onnx_dir,
+    input_names,
+    use_gpu,
+    precision,
+    optimize_info,
+    validate_onnx,
+    use_raw_attention_mask,
+    overwrite,
+    config,
+    model_fusion_statistics,
+    onnx_model_path,
+    example_inputs,
+    example_outputs_flatten,
+    output_names,
+    fusion_options,
+):
+    is_valid_onnx_model = True
+    if validate_onnx:
+        is_valid_onnx_model = validate_onnx_model(
+            onnx_model_path,
+            example_inputs,
+            example_outputs_flatten,
+            use_gpu,
+            False,
+            output_names,
+        )
+    if optimize_info == OptimizerInfo.NOOPT:
+        return onnx_model_path, is_valid_onnx_model, config.vocab_size
+
+    if (
+        optimize_info == OptimizerInfo.BYSCRIPT or precision == Precision.FLOAT16 or precision == Precision.INT8
+    ):  # Use script (optimizer.py) to optimize
+        optimized_model_path = get_onnx_file_path(
+            onnx_dir,
+            model_name,
+            len(input_names),
+            True,
+            use_gpu,
+            precision,
+            False,
+            use_external_data_format,
+        )
+        optimize_onnx_model(
+            onnx_model_path,
+            optimized_model_path,
+            model_type,
+            config.num_attention_heads,
+            config.hidden_size,
+            use_gpu,
+            precision,
+            use_raw_attention_mask,
+            overwrite,
+            model_fusion_statistics,
+            use_external_data_format,
+            fusion_options,
+        )
+
+        onnx_model_path = optimized_model_path
+        if validate_onnx:
+            is_valid_onnx_model = validate_onnx_model(
+                onnx_model_path,
+                example_inputs,
+                example_outputs_flatten,
+                use_gpu,
+                precision == Precision.FLOAT16,
+                output_names,
+            )
+
+        if precision == Precision.INT8:
+            logger.info(f"Quantizing model: {onnx_model_path}")
+            QuantizeHelper.quantize_onnx_model(onnx_model_path, onnx_model_path, use_external_data_format)
+            logger.info(f"Finished quantizing model: {onnx_model_path}")
+
+    if optimize_info == OptimizerInfo.BYORT:  # Use OnnxRuntime to optimize
+        if is_valid_onnx_model:
+            ort_model_path = add_filename_suffix(onnx_model_path, "_ort")
+            optimize_onnx_model_by_ort(
+                onnx_model_path,
+                ort_model_path,
+                use_gpu,
+                overwrite,
+                model_fusion_statistics,
+            )
+
+    return (
+        onnx_model_path,
+        is_valid_onnx_model,
+        config.num_labels if model_type in ["vit", "swin"] else config.vocab_size,
+    )
+
+
+def export_onnx_model_from_pt(
+    model_name,
+    opset_version,
+    use_external_data_format,
+    model_type,
+    model_class,
+    config_modifier,
+    cache_dir,
+    onnx_dir,
+    input_names,
+    use_gpu,
+    precision,
+    optimizer_info,
+    validate_onnx,
+    use_raw_attention_mask,
+    overwrite,
+    model_fusion_statistics,
+    fusion_options,
+):
+    config, model = load_pt_model(model_name, model_class, cache_dir, config_modifier)
+    # config, model = load_pt_model_from_tf(model_name)
+    model.cpu()
+
+    example_inputs = None
+    max_input_size = None
+
+    if model_type in ["vit", "swin"]:
+        image_processor = AutoFeatureExtractor.from_pretrained(model_name, cache_dir=cache_dir)
+        data = numpy.random.randint(
+            low=0, high=256, size=config.image_size * config.image_size * 3, dtype=numpy.uint8
+        ).reshape(config.image_size, config.image_size, 3)
+
+        example_inputs = image_processor(data, return_tensors="pt")
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
+        max_input_size = tokenizer.max_model_input_sizes.get(model_name, 1024)
+        example_inputs = tokenizer.encode_plus("This is a sample input", return_tensors="pt")
+
+    example_inputs = filter_inputs(example_inputs, input_names)
+
+    example_outputs = model(**example_inputs)
+
+    assert isinstance(example_outputs, (list, tuple)), f"type of output is not list or tuple: {type(example_outputs)}"
+
+    # Flatten is needed for gpt2 and distilgpt2.
+    example_outputs_flatten = flatten(example_outputs)
+    example_outputs_flatten = update_flatten_list(example_outputs_flatten, [])
+
+    onnx_model_path = get_onnx_file_path(
+        onnx_dir,
+        model_name,
+        len(input_names),
+        False,
+        use_gpu,
+        precision,
+        False,
+        use_external_data_format,
+    )
+
+    if overwrite or not os.path.exists(onnx_model_path):
+        logger.info(f"Exporting ONNX model to {onnx_model_path}")
+        Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
+
+        dynamic_axes = None
+        output_names = None
+
+        if model_type in ["vit", "swin"]:
+            dynamic_axes, output_names = {key: {0: "pixel_values"} for key in example_inputs}, ["logits"]
+        else:
+            dynamic_axes, output_names = build_dynamic_axes(example_inputs, example_outputs_flatten)
+
+        replace_torch_functions()
+        torch_onnx_export(
+            model=model,
+            args=tuple(example_inputs.values()),
+            f=onnx_model_path,
+            input_names=list(example_inputs.keys()),
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
+            do_constant_folding=True,
+            opset_version=opset_version,
+            use_external_data_format=use_external_data_format,
+        )
+        restore_torch_functions()
+    else:
+        logger.info(f"Skip export since model existed: {onnx_model_path}")
+
+    onnx_model_file, is_valid_onnx_model, vocab_size = validate_and_optimize_onnx(
+        model_name,
+        use_external_data_format,
+        model_type,
+        onnx_dir,
+        input_names,
+        use_gpu,
+        precision,
+        optimizer_info,
+        validate_onnx,
+        use_raw_attention_mask,
+        overwrite,
+        config,
+        model_fusion_statistics,
+        onnx_model_path,
+        example_inputs,
+        example_outputs_flatten,
+        None,
+        fusion_options,
+    )
+
+    return onnx_model_file, is_valid_onnx_model, vocab_size, max_input_size
+
+
+def export_onnx_model_from_tf(
+    model_name,
+    opset_version,
+    use_external_data_format,
+    model_type,
+    model_class,
+    config_modifier,
+    cache_dir,
+    onnx_dir,
+    input_names,
+    use_gpu,
+    precision,
+    optimizer_info,
+    validate_onnx,
+    use_raw_attention_mask,
+    overwrite,
+    model_fusion_statistics,
+    fusion_options,
+):
+    # Use CPU to export
+    import tensorflow as tf
+
+    tf.config.set_visible_devices([], "GPU")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
+    # Fix "Using pad_token, but it is not set yet" error.
+    if tokenizer.pad_token is None:
+        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+    max_input_size = tokenizer.max_model_input_sizes.get(model_name, 1024)
+
+    config, model = load_tf_model(model_name, model_class, cache_dir, config_modifier)
+    model.resize_token_embeddings(len(tokenizer))
+
+    example_inputs = tokenizer.encode_plus(
+        "This is a sample input",
+        return_tensors="tf",
+        max_length=max_input_size,
+        padding="max_length",
+        truncation=True,
+    )
+    example_inputs = filter_inputs(example_inputs, input_names)
+
+    if config.is_encoder_decoder:
+        example_inputs["decoder_input_ids"] = tokenizer.encode_plus(
+            "This is a sample input",
+            return_tensors="tf",
+            max_length=max_input_size,
+            padding="max_length",
+            truncation=True,
+        ).input_ids
+    if model_name == "unc-nlp/lxmert-base-uncased":
+        example_inputs["visual_feats"] = tf.random.normal([1, 1, config.visual_feat_dim])
+        example_inputs["visual_pos"] = tf.random.normal([1, 1, config.visual_pos_dim])
+
+    try:
+        # Use no past state for these models
+        if config.use_cache:
+            config.use_cache = False
+    except Exception:
+        pass
+
+    example_outputs = model(example_inputs, training=False)
+    output_names = None
+
+    # For xlnet models, only compare the last_hidden_state output.
+    if model_name == "xlnet-base-cased" or model_name == "xlnet-large-cased":
+        output_names = ["last_hidden_state"]
+        example_outputs = example_outputs["last_hidden_state"]
+
+    # Flatten is needed for gpt2 and distilgpt2. Output name sorting is needed for tf2onnx outputs to match onnx outputs.
+    from tensorflow.python.util import nest
+
+    example_outputs_flatten = nest.flatten(example_outputs)
+
+    onnx_model_path = get_onnx_file_path(
+        onnx_dir,
+        model_name,
+        len(input_names),
+        False,
+        use_gpu,
+        precision,
+        False,
+        use_external_data_format,
+    )
+    tf_internal_model_path = onnx_model_path[:-5] if use_external_data_format else onnx_model_path
+
+    if overwrite or not os.path.exists(tf_internal_model_path):
+        logger.info(f"Exporting ONNX model to {onnx_model_path}")
+        if not use_external_data_format:
+            Path(tf_internal_model_path).parent.mkdir(parents=True, exist_ok=True)
+
+        import zipfile
+
+        import tf2onnx
+
+        tf2onnx.logging.set_level(tf2onnx.logging.ERROR)
+        specs = []
+        for name, value in example_inputs.items():
+            dims = [None] * len(value.shape)
+            specs.append(tf.TensorSpec(tuple(dims), value.dtype, name=name))
+        _, _ = tf2onnx.convert.from_keras(
+            model,
+            input_signature=tuple(specs),
+            opset=opset_version,
+            large_model=use_external_data_format,
+            output_path=tf_internal_model_path,
+        )
+        if use_external_data_format:
+            # need to unpack the zip for run_onnxruntime()
+            with zipfile.ZipFile(tf_internal_model_path, "r") as z:
+                z.extractall(os.path.dirname(tf_internal_model_path))
+            tf_internal_model_path = os.path.join(os.path.dirname(tf_internal_model_path), "__MODEL_PROTO.onnx")
+            if os.path.exists(onnx_model_path):
+                os.remove(onnx_model_path)
+            os.rename(tf_internal_model_path, onnx_model_path)
+
+    else:
+        logger.info(f"Skip export since model existed: {onnx_model_path}")
+
+    model_type = model_type + "_tf"
+    optimized_onnx_path, is_valid_onnx_model, vocab_size = validate_and_optimize_onnx(
+        model_name,
+        use_external_data_format,
+        model_type,
+        onnx_dir,
+        input_names,
+        use_gpu,
+        precision,
+        optimizer_info,
+        validate_onnx,
+        use_raw_attention_mask,
+        overwrite,
+        config,
+        model_fusion_statistics,
+        onnx_model_path,
+        example_inputs,
+        example_outputs_flatten,
+        output_names,
+        fusion_options,
+    )
+
+    return (
+        optimized_onnx_path,
+        is_valid_onnx_model,
+        vocab_size,
+        max_input_size,
+    )