onnx_diagnostic.torch_models.llms

onnx_diagnostic.torch_models.llms.get_phi2(batch_size: int = 1, sequence_length: int = 30, sequence_length2: int = 3, dynamic_rope: bool = False, **kwargs) → Dict[str, Any]

Gets a non initialized model similar to microsoft/phi-2.

Parameters:

• batch_size – batch size

• sequence_length – sequence length

• sequence_length2 – new sequence length

• dynamic_rope – use dynamic rope (see transformers.LlamaConfig)

• kwargs – to overwrite the configuration, example num_hidden_layers=1

Returns:

dictionary

See Export Tiny-LLM with patches for an example with a similar model.

onnx_diagnostic.torch_models.llms.get_tiny_llm(batch_size: int = 2, sequence_length: int = 30, sequence_length2: int = 3, dynamic_rope: bool = False, use_static_cache: bool = False, **kwargs) → Dict[str, Any]

Gets a non initialized model similar to arnir0/Tiny-LLM.

Parameters:

• batch_size – batch size

• sequence_length – sequence length

• sequence_length2 – new sequence length

• dynamic_rope – use dynamic rope (see transformers.LlamaConfig)

• use_static_cache – use StaticCache instead of DynamicCache

• kwargs – to overwrite the configuration, example num_hidden_layers=1

Returns:

dictionary

See Steel method forward to guess inputs and dynamic shapes (with Tiny-LLM) or Export Tiny-LLM with patches for examples.