Source code for onnxscript.tools.transformers_models.phi

# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
# pylint: disable=import-outside-toplevel
from __future__ import annotations

from typing import Any, Sequence

import torch

import onnxscript.tools.transformers_models


def _prepare_config_and_inputs(
    batch_size: int,
    seq_length: int,
    vocab_size: int,
    type_sequence_label_size: int = 2,
    type_vocab_size: int = 16,
    num_labels: int = 3,
    num_choices: int = 4,
    use_input_mask: bool = False,
    use_token_type_ids: bool = False,
    use_labels: bool = False,
) -> tuple[Any, ...]:
    input_ids = onnxscript.tools.transformers_models.ids_tensor(
        [batch_size, seq_length], vocab_size
    )

    input_mask = None
    if use_input_mask:
        input_mask = torch.tril(torch.ones(batch_size, seq_length))

    token_type_ids = None
    if use_token_type_ids:
        assert type_vocab_size > 0, "type_vocab_size is null"
        token_type_ids = onnxscript.tools.transformers_models.ids_tensor(
            [batch_size, seq_length], type_vocab_size
        )

    sequence_labels = None
    token_labels = None
    choice_labels = None
    if use_labels:
        assert type_sequence_label_size > 0, "type_sequence_label_size is null"
        assert num_labels > 0, "num_labels is null"
        assert num_choices > 0, "num_choices is null"
        sequence_labels = onnxscript.tools.transformers_models.ids_tensor(
            [batch_size], type_sequence_label_size
        )
        token_labels = onnxscript.tools.transformers_models.ids_tensor(
            [batch_size, seq_length], num_labels
        )
        choice_labels = onnxscript.tools.transformers_models.ids_tensor(
            [batch_size], num_choices
        )

    return (
        input_ids,
        token_type_ids,
        input_mask,
        sequence_labels,
        token_labels,
        choice_labels,
    )


def get_phi_model(
    input_dims: Sequence[tuple[int, int]] = ((13, 7), (14, 7), (15, 8)),
    hidden_size: int = 32,
    num_hidden_layers: int = 2,
    vocab_size: int = 99,
    intermediate_size: int = 16,
    max_position_embeddings: int = 512,
    num_attention_heads: int = 4,
    num_key_value_heads: int = 2,
    _attn_implementation: str = "eager",  # needed value to remove graph breaks
    with_mask: bool = True,
) -> tuple[Any, list[tuple[torch.Tensor, ...]], dict]:
    """
    Returns a model.
    See `PhiConfig
    <https://huggingface.co/docs/transformers/main/en/model_doc/phi#transformers.PhiConfig>`_.
    The parameters are chosen for a unit test configuration from `test_modeling_phi.py
    <https://github.com/huggingface/transformers/blob/main/tests/models/phi/test_modeling_phi.py>`_.
    """
    from transformers import PhiConfig
    from transformers.models.phi.modeling_phi import PhiModel

    dynamic_shapes = {0: {0: "batch", 1: "length"}}
    if with_mask:
        dynamic_shapes.update({1: {0: "batch", 1: "length"}})

    config = PhiConfig(
        hidden_size=hidden_size,
        num_hidden_layers=num_hidden_layers,
        vocab_size=vocab_size,
        intermediate_size=intermediate_size,
        max_position_embeddings=max_position_embeddings,
        num_attention_heads=num_attention_heads,
        num_key_value_heads=num_key_value_heads,
    )
    if _attn_implementation:
        config._attn_implementation = _attn_implementation  # pylint: disable=protected-access

    if with_mask:

        class PhiModelWrapper(torch.nn.Module):
            def __init__(self, config):
                super().__init__()
                self.model = PhiModel(config)

            def forward(self, input_ids, attention_mask):
                model_output = self.model(
                    input_ids, attention_mask=attention_mask, use_cache=False
                )
                return model_output.to_tuple()

        def generate_example_inputs(batch: int, seq: int, vocab_size: int):
            (
                input_ids,
                _,  # token_type_ids,
                input_mask,
                _,  # sequence_labels,
                _,  # token_labels,
                _,  # choice_labels,
            ) = _prepare_config_and_inputs(
                batch_size=batch,
                seq_length=seq,
                vocab_size=vocab_size,
                use_input_mask=True,
            )
            return input_ids, input_mask

        example_args_collection = []
        for b, s in input_dims:
            example_args_collection.append(generate_example_inputs(b, s, vocab_size))

        return PhiModelWrapper(config), example_args_collection, dynamic_shapes

    # no mask

    class PhiModelWrapperNoMask(torch.nn.Module):
        def __init__(self, config):
            super().__init__()
            self.model = PhiModel(config)

        def forward(self, input_ids):
            model_output = self.model(input_ids, use_cache=False)
            return model_output.to_tuple()

    def generate_example_inputs_no_mask(batch: int, seq: int, vocab_size: int):
        (
            input_ids,
            _,  # token_type_ids,
            _,  # input_mask,
            _,  # sequence_labels,
            _,  # token_labels,
            _,  # choice_labels,
        ) = _prepare_config_and_inputs(
            batch_size=batch,
            seq_length=seq,
            vocab_size=vocab_size,
            use_input_mask=True,
        )
        return (input_ids,)

    example_args_collection = []
    for b, s in input_dims:
        example_args_collection.append(generate_example_inputs_no_mask(b, s, vocab_size))

    return PhiModelWrapperNoMask(config), example_args_collection, dynamic_shapes


[docs] def get_phi_model_from_config( warmup: int = 5, repeat: int = 10, config: str = "small", num_hidden_layers: int = 1, implementation: str = "eager", dynamic_shapes: bool = False, with_mask: bool = True, ) -> tuple[Any, list[tuple[torch.Tensor, ...]], dict]: """ Returns a model Phi to test or benchmark. Args: warmup: Number of inputs to generate. repeat: Number of inputs to generate for repeat. config: small, medium or large num_hidden_layers: number of hidden layers implementation: eager or sdpa with_mask: One or two inputs. dynamic_shapes: dynamic shapes or not Returns: Model and list of inputs. """ if config == "small": conf_dict = dict( input_dims=onnxscript.tools.transformers_models.get_input_dims_for_llm( dynamic_shapes, warmup, repeat ), hidden_size=32, num_hidden_layers=num_hidden_layers, vocab_size=99, intermediate_size=16, max_position_embeddings=512, num_attention_heads=4, num_key_value_heads=2, _attn_implementation=implementation, with_mask=with_mask, ) elif config == "medium": conf_dict = dict( input_dims=onnxscript.tools.transformers_models.get_input_dims_for_llm( dynamic_shapes, warmup, repeat ), hidden_size=1024, num_hidden_layers=num_hidden_layers, vocab_size=1024, intermediate_size=1024, num_attention_heads=4, num_key_value_heads=4, max_position_embeddings=1024, _attn_implementation=implementation, with_mask=with_mask, ) elif config in ("large", "default"): conf_dict = dict( input_dims=onnxscript.tools.transformers_models.get_input_dims_for_llm( dynamic_shapes, warmup, repeat ), hidden_size=2048, num_hidden_layers=num_hidden_layers, vocab_size=51200, intermediate_size=8192, num_attention_heads=32, num_key_value_heads=None, max_position_embeddings=2048, _attn_implementation=implementation, with_mask=with_mask, ) else: raise ValueError(f"Unexpected configuration {config!r}.") return get_phi_model(**conf_dict) # type: ignore[arg-type]