onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/models/longformer/generate_test_data.py

@@ -0,0 +1,347 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+# Generate test data for a longformer model, so that we can use onnxruntime_perf_test.exe to evaluate the inference latency.
+
+import argparse
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+from bert_test_data import fake_input_ids_data, fake_input_mask_data, output_test_data
+from onnx import ModelProto, TensorProto
+from onnx_model import OnnxModel
+
+
+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. If not specified, .",
+    )
+
+    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(
+        "-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(
+        "--global_tokens",
+        required=False,
+        type=int,
+        default=10,
+        help="number of global tokens",
+    )
+
+    parser.add_argument(
+        "--input_ids_name",
+        required=False,
+        type=str,
+        default=None,
+        help="input name for input ids",
+    )
+
+    parser.add_argument(
+        "--input_mask_name",
+        required=False,
+        type=str,
+        default=None,
+        help="input name for attention mask",
+    )
+
+    parser.add_argument(
+        "--global_mask_name",
+        required=False,
+        type=str,
+        default=None,
+        help="input name for global 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)
+
+    args = parser.parse_args()
+    return args
+
+
+def get_longformer_inputs(onnx_file, input_ids_name=None, input_mask_name=None, global_mask_name=None):
+    """
+    Get graph inputs for longformer model.
+    """
+    model = ModelProto()
+    with open(onnx_file, "rb") as f:
+        model.ParseFromString(f.read())
+
+    onnx_model = OnnxModel(model)
+    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}")
+
+        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}")
+
+        global_mask = None
+        if global_mask_name:
+            global_mask = onnx_model.find_graph_input(global_mask_name)
+            if global_mask is None:
+                raise ValueError(f"Graph does not have input named {global_mask_name}")
+
+        expected_inputs = 1 + (1 if input_mask else 0) + (1 if global_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, input_mask, global_mask
+
+    if len(graph_inputs) != 3:
+        raise ValueError(f"Expect the graph to have 3 inputs. Got {len(graph_inputs)}")
+
+    # Try guess the inputs based on naming.
+    input_ids = None
+    input_mask = None
+    global_mask = None
+    for input in graph_inputs:
+        input_name_lower = input.name.lower()
+        if "global" in input_name_lower:
+            global_mask = input
+        elif "mask" in input_name_lower:
+            input_mask = input
+        else:
+            input_ids = input
+
+    if input_ids and input_mask and global_mask:
+        return input_ids, input_mask, global_mask
+
+    raise ValueError("Fail to assign 3 inputs. You might try rename the graph inputs.")
+
+
+def fake_global_mask_data(global_mask, batch_size, sequence_length, num_global_tokens):
+    """
+    Fake data based on the graph input of segment_ids.
+    Args:
+        segment_ids (TensorProto): graph input of input tensor.
+    Returns:
+        data (np.array): the data for input tensor
+    """
+    data_type = global_mask.type.tensor_type.elem_type
+    assert data_type in [TensorProto.FLOAT, TensorProto.INT32, TensorProto.INT64]
+
+    if num_global_tokens > 0:
+        assert num_global_tokens <= sequence_length
+        data = np.zeros((batch_size, sequence_length), dtype=np.int32)
+        temp = np.ones((batch_size, num_global_tokens), dtype=np.int32)
+        data[: temp.shape[0], : temp.shape[1]] = temp
+    else:
+        data = np.zeros((batch_size, sequence_length), dtype=np.int32)
+
+    if data_type == TensorProto.FLOAT:
+        data = np.float32(data)
+    elif data_type == TensorProto.INT64:
+        data = np.int64(data)
+
+    return data
+
+
+def fake_test_data(
+    batch_size,
+    sequence_length,
+    test_cases,
+    dictionary_size,
+    verbose,
+    random_seed,
+    input_ids,
+    input_mask,
+    global_mask,
+    num_global_tokens,
+    average_sequence_length,
+    random_sequence_length,
+):
+    """
+    Generate fake input data for test.
+    """
+    assert input_ids is not None
+
+    np.random.seed(random_seed)
+    random.seed(random_seed)
+
+    all_inputs = []
+    for _ 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 input_mask:
+            inputs[input_mask.name] = fake_input_mask_data(
+                input_mask, batch_size, sequence_length, average_sequence_length, random_sequence_length
+            )
+
+        if global_mask:
+            inputs[global_mask.name] = fake_global_mask_data(
+                global_mask, batch_size, sequence_length, num_global_tokens
+            )
+
+        if verbose and len(all_inputs) == 0:
+            print("Example inputs", inputs)
+        all_inputs.append(inputs)
+
+    return all_inputs
+
+
+def generate_test_data(
+    batch_size,
+    sequence_length,
+    test_cases,
+    seed,
+    verbose,
+    input_ids,
+    input_mask,
+    global_mask,
+    num_global_tokens,
+    average_sequence_length,
+    random_sequence_length,
+):
+    dictionary_size = 10000
+    all_inputs = fake_test_data(
+        batch_size,
+        sequence_length,
+        test_cases,
+        dictionary_size,
+        verbose,
+        seed,
+        input_ids,
+        input_mask,
+        global_mask,
+        num_global_tokens,
+        average_sequence_length,
+        random_sequence_length,
+    )
+    if len(all_inputs) != test_cases:
+        print("Failed to create test data for test.")
+    return all_inputs
+
+
+def create_longformer_test_data(
+    model,
+    output_dir,
+    batch_size,
+    sequence_length,
+    test_cases,
+    seed,
+    verbose,
+    input_ids_name,
+    input_mask_name,
+    global_mask_name,
+    num_global_tokens,
+    average_sequence_length,
+    random_sequence_length,
+):
+    input_ids, input_mask, global_mask = get_longformer_inputs(model, input_ids_name, input_mask_name, global_mask_name)
+    all_inputs = generate_test_data(
+        batch_size,
+        sequence_length,
+        test_cases,
+        seed,
+        verbose,
+        input_ids,
+        input_mask,
+        global_mask,
+        num_global_tokens,
+        average_sequence_length,
+        random_sequence_length,
+    )
+
+    for i, inputs in enumerate(all_inputs):
+        output_test_data(output_dir, i, inputs)
+
+
+def main():
+    args = parse_arguments()
+
+    output_dir = args.output_dir
+    if output_dir is None:
+        # Default output directory is a sub-directory under the directory of model.
+        output_dir = os.path.join(
+            Path(args.model).parent,
+            f"b{args.batch_size}_s{args.sequence_length}_g{args.global_tokens}",
+        )
+
+    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.")
+
+    if args.average_sequence_length <= 0:
+        args.average_sequence_length = args.sequence_length
+
+    create_longformer_test_data(
+        args.model,
+        output_dir,
+        args.batch_size,
+        args.sequence_length,
+        args.samples,
+        args.seed,
+        args.verbose,
+        args.input_ids_name,
+        args.input_mask_name,
+        args.global_mask_name,
+        args.global_tokens,
+        args.average_sequence_length,
+    )
+
+    print("Test data is saved to directory:", output_dir)
+
+
+if __name__ == "__main__":
+    main()

onnxruntime/transformers/models/longformer/longformer_helper.py

@@ -0,0 +1,76 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.  See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+# This script helps creating dummy inputs for Longformer model.
+
+import logging
+
+import numpy
+import torch
+
+logger = logging.getLogger(__name__)
+
+PRETRAINED_LONGFORMER_MODELS = {
+    "longformer-base-4096": "allenai/longformer-base-4096",
+    "longformer-large-4096": "allenai/longformer-large-4096",
+    "longformer-random-tiny": "patrickvonplaten/longformer-random-tiny",  # A tiny model for debugging
+}
+
+
+class LongformerInputs:
+    def __init__(self, input_ids, attention_mask, global_attention_mask):
+        self.input_ids: torch.LongTensor = input_ids
+        self.attention_mask: torch.FloatTensor | torch.HalfTensor = attention_mask
+        self.global_attention_mask: torch.FloatTensor | torch.HalfTensor = global_attention_mask
+
+    def to_list(self) -> list:
+        return [v for v in [self.input_ids, self.attention_mask, self.global_attention_mask] if v is not None]
+
+    def to_tuple(self) -> tuple:
+        return tuple(v for v in self.to_list())
+
+    def get_ort_inputs(self) -> dict:
+        return {
+            "input_ids": numpy.ascontiguousarray(self.input_ids.cpu().numpy()),
+            "attention_mask": numpy.ascontiguousarray(self.attention_mask.cpu().numpy()),
+            "global_attention_mask": numpy.ascontiguousarray(self.global_attention_mask.cpu().numpy()),
+        }
+
+
+class LongformerHelper:
+    """A helper class for Longformer model conversion, inference and verification."""
+
+    @staticmethod
+    def get_dummy_inputs(
+        batch_size: int,
+        sequence_length: int,
+        num_global_tokens: int,
+        device: torch.device,
+        vocab_size: int = 100,
+    ) -> LongformerInputs:
+        """Create random inputs for Longformer model.
+        Returns torch tensors of input_ids, attention_mask and global_attention_mask tensors.
+        """
+
+        input_ids = torch.randint(
+            low=0,
+            high=vocab_size - 1,
+            size=(batch_size, sequence_length),
+            dtype=torch.long,
+            device=device,
+        )
+        attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device)
+        global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=device)
+        global_token_index = list(range(num_global_tokens))
+        global_attention_mask[:, global_token_index] = 1
+        return LongformerInputs(input_ids, attention_mask, global_attention_mask)
+
+    @staticmethod
+    def get_output_shapes(batch_size: int, sequence_length: int, hidden_size: int) -> dict[str, list[int]]:
+        """Returns a dictionary with output name as key, and shape as value."""
+        return {
+            "last_state": [batch_size, sequence_length, hidden_size],
+            "pooler": [batch_size, sequence_length],
+        }

onnxruntime/transformers/models/phi2/__init__.py

@@ -0,0 +1,12 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import os
+import sys
+
+sys.path.append(os.path.dirname(__file__))
+
+transformers_dir = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+if transformers_dir not in sys.path:
+    sys.path.append(transformers_dir)