onnx_diagnostic.tasks.feature_extraction

onnx_diagnostic.tasks.feature_extraction.get_inputs(model: Module, config: Any | None, batch_size: int, sequence_length: int, dummy_max_token_id: int, add_second_input: bool = False, **kwargs)

Generates inputs for task feature-extraction. Example:

input_ids:T7s1x13[101,72654:A16789.23076923077],
token_type_ids:T7s1x13[0,0:A0.0],
attention_mask:T7s1x13[1,1:A1.0])

onnx_diagnostic.tasks.feature_extraction.random_input_kwargs(config: Any) → Tuple[Dict[str, Any], Callable]

Inputs kwargs.

If the configuration is None, the function selects typical dimensions.

onnx_diagnostic.tasks.feature_extraction.reduce_model_config(config: Any) → Dict[str, Any]

Reduces a model size.