experimental_experiment.torch_models.llm_model_helper

experimental_experiment.torch_models.llm_model_helper.get_ai21_jamba_15_mini(inputs_as_tuple: bool = False, input_cache: bool = True, batch_size: int = 1, common_dynamic_shapes: bool = False, **kwargs) → Dict[str, Any]

Gets a non initialized model.

Parameters:

• inputs_as_tuple – returns dummy inputs as a dictionary or not

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

• batch_size – batch size

• common_dynamic_shapes – if True returns dynamic shapes as well

Returns:

dictionary

See ai21labs/AI21-Jamba-1.5-Mini/config.json.

experimental_experiment.torch_models.llm_model_helper.get_all_mini_ml_l6_v1(inputs_as_tuple: bool = False, input_cache: bool = True, batch_size: int = 1, common_dynamic_shapes: bool = False, **kwargs) → Dict[str, Any]

Gets a non initialized model.

Parameters:

• inputs_as_tuple – returns dummy inputs as a dictionary or not

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

• batch_size – batch size

• common_dynamic_shapes – if True returns dynamic shapes as well

Returns:

dicionary

See all-MiniLM-L6-v1.

Model forward signature:

def forward(
    self,
    input_ids: Optional[torch.Tensor] = None,
    attention_mask: Optional[torch.Tensor] = None,
    token_type_ids: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.Tensor] = None,
    head_mask: Optional[torch.Tensor] = None,
    inputs_embeds: Optional[torch.Tensor] = None,
    encoder_hidden_states: Optional[torch.Tensor] = None,
    encoder_attention_mask: Optional[torch.Tensor] = None,
    past_key_values: Optional[List[torch.FloatTensor]] = None,
    use_cache: Optional[bool] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    return_dict: Optional[bool] = None,
) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:

experimental_experiment.torch_models.llm_model_helper.get_falcon_mamba_7b(batch_size: int = 2, input_cache: bool = True, inputs_as_tuple: bool = False, common_dynamic_shapes: bool = False, **kwargs) → Tuple[Any, Tuple[Any, ...] | Dict[str, Any]]

Gets a non initialized model.

Parameters:

• inputs_as_tuple – returns dummy inputs as a dictionary or not

• batch_size – batch size

• input_cache – generate data for this iteration with or without cache

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

• common_dynamic_shapes – if True returns dynamic shapes as well

Returns:

dictionary

See flacon-mamba-7b/config.json.

experimental_experiment.torch_models.llm_model_helper.get_llama32_9b_vision(inputs_as_tuple: bool = False, common_dynamic_shapes: bool = False, **kwargs) → Tuple[Any, Tuple[Any, ...] | Dict[str, Any]]

Gets a non initialized model.

Parameters:

• inputs_as_tuple – returns dummy inputs as a dictionary or not

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

• common_dynamic_shapes – if True returns dynamic shapes as well

Returns:

model, inputs

See MLlama.

experimental_experiment.torch_models.llm_model_helper.get_phi2(inputs_as_tuple: bool = False, input_cache: bool = True, batch_size: int = 1, common_dynamic_shapes: bool = False, **kwargs) → Dict[str, Any]

Gets a non initialized model.

Parameters:

• inputs_as_tuple – returns dummy inputs as a dictionary or not

• input_cache – generate data for this iteration with or without cache

• batch_size – batch size

• common_dynamic_shapes – if True returns dynamic shapes as well

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

Returns:

dict

See Phi-2/config.json.

experimental_experiment.torch_models.llm_model_helper.get_phi35_mini_instruct(inputs_as_tuple: bool = False, input_cache: bool = True, batch_size: int = 1, common_dynamic_shapes: bool = False, **kwargs) → Dict[str, Any]

Gets a non initialized model.

Parameters:

• inputs_as_tuple – returns dummy inputs as a dictionary or not

• batch_size – batch size

• input_cache – generate data for this iteration with or without cache

• common_dynamic_shapes – if True returns dynamic shapes as well

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

Returns:

dictionary

See Phi-3.5-mini-instruct/config.json.

experimental_experiment.torch_models.llm_model_helper.get_phi35_vision_instruct(inputs_as_tuple: bool = False, batch_size: int = 2, input_kind: LLMInputKind = LLMInputKind.input_ids, device: str = 'cpu', common_dynamic_shapes: bool = False, **kwargs) → Tuple[Any, Tuple[Any, ...] | Dict[str, Any], Any | None]

Gets a non initialized model.

Parameters:

• batch_size – batch size to use

• inputs_as_tuple – returns dummy inputs as a dictionary or not

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

• common_dynamic_shapes – if True returns dynamic shapes as well

Returns:

model, inputs, dynamic shapes

experimental_experiment.torch_models.llm_model_helper.get_phi3_vision_128k_instruct(inputs_as_tuple: bool = False, batch_size: int = 2, input_kind: LLMInputKind = LLMInputKind.input_ids, device: str = 'cpu', common_dynamic_shapes: bool = False, **kwargs) → Tuple[Any, Tuple[Any, ...] | Dict[str, Any], Any | None]

Gets a non initialized model.

Parameters:

• batch_size – batch size to use

• inputs_as_tuple – returns dummy inputs as a dictionary or not

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

• common_dynamic_shapes – if True returns dynamic shapes as well

Returns:

model, inputs, dynamic shapes

See Phi-3-vision-128k-instruct/config.json.

experimental_experiment.torch_models.llm_model_helper.get_smollm_1_7b(batch_size: int = 2, input_cache: bool = True, inputs_as_tuple: bool = False, common_dynamic_shapes: bool = False, **kwargs) → Tuple[Any, Tuple[Any, ...] | Dict[str, Any]]

Gets a non initialized model.

Parameters:

• inputs_as_tuple – returns dummy inputs as a dictionary or not

• batch_size – batch size

• input_cache – generate data for this iteration with or without cache

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

• common_dynamic_shapes – if True returns dynamic shapes as well

Returns:

dictionary

See SmolLM-1.7B.