def get_model_io_config(model_name: str, task: str, model: "PreTrainedModel", **kwargs) -> Optional[dict]:
    """Get the input/output config for the model_name and task."""
    # just log a debug message if io_config is not found
    # this is not a critical error and the caller may not need the io_config
    export_config = get_export_config(model_name, task, **kwargs)
    if not export_config:
        return None

    if is_peft_model(model):
        # if pytorch_model is PeftModel, we need to get the base model
        # otherwise, the model forward has signature (*args, **kwargs)
        model = model.get_base_model()

    inputs = export_config.ordered_inputs(model)
    input_names = list(inputs.keys())
    output_names = list(export_config.outputs.keys())
    dynamic_axes = dict(chain(inputs.items(), export_config.outputs.items()))
    # optimum has the total sequence length as "past_sequence_length  + 1" but that is not always the case
    # change it to "past_sequence_length + sequence_length" if past is used
    for value in dynamic_axes.values():
        for axis, axis_name in value.items():
            if axis_name == "past_sequence_length + 1":
                value[axis] = "past_sequence_length + sequence_length"
    # dynamic_shapes should follow input order and format
    dynamic_shapes = _unflatten_past_key_values_with_check(inputs)
    return {
        "input_names": input_names,
        "output_names": output_names,
        "dynamic_axes": dynamic_axes,
        "dynamic_shapes": dynamic_shapes,
    }
Get the input/output config for the model_name and task.
olive/common/hf/model_io.py