Tools

Transformers Models

onnxscript.tools.transformers_models.get_model_and_inputs(model: str, config: str, dynamic_shapes: bool, device: str = 'cpu', num_hidden_layers: int = 1, with_mask: bool = True, implementation: str = 'eager', dtype: str | None = None, warmup: int = 5, repeat: int = 10) tuple[Any, list[tuple[Tensor, ...]], dict | None][source]

Returns a model and a couple of dummy inputs.

Parameters:
  • model – model name, ‘phi’, ‘llama’, …

  • config – ‘small’, ‘medium’, ‘large’, …

  • dynamic_shapes – dynamic or static shapes

  • device – ‘cpu’ or ‘cuda’

  • num_hidden_layers – Number of hidden layers.

  • with_mask – One input or two inputs.

  • implementation – eager or sdpa

  • warmup – Number of inputs to generate.

  • repeat – Number of inputs to generate for repeat.

  • dtype – If specified, cast the model and the inputs into this type.

Returns:

model and list of inputs