Repository: Project-MONAI/MONAI
File: monai/bundle/scripts.py (Python)

Code before:
        key_in_ckpt="",
        use_trace=False,
        input_shape=None,
        converter_kwargs={},
    )

    parser = ConfigParser()

    parser.read_config(f=config_file_)
    if meta_file_ is not None:
        parser.read_meta(f=meta_file_)

    # the rest key-values in the _args are to override config content
    for k, v in _args.items():
        parser[k] = v

    # The convert_to_onnx must have an `inputs` as input, no matter what the `use_trace` is.
    # If the `input_shape` is not provided, will try to parse it from the parser to generate a random `inputs`.
    if not input_shape_:
        input_shape_ = _get_fake_input_shape(parser=parser)

    inputs_ = [torch.rand(input_shape_)]

    converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_})

    def save_onnx(onnx_obj: Any, filename_prefix_or_stream: str | IO[Any], **kwargs: Any) -> None:
        onnx.save(onnx_obj, filename_prefix_or_stream)

    _export(
        convert_to_onnx,
        save_onnx,
        parser,
        net_id=net_id_,
        filepath=filepath_,
        ckpt_file=ckpt_file_,
        config_file=config_file_,
        key_in_ckpt=key_in_ckpt_,
        **converter_kwargs_,
    )


def ckpt_export(
    net_id: str | None = None,
    filepath: PathLike | None = None,
    ckpt_file: str | None = None,
    meta_file: str | Sequence[str] | None = None,
    config_file: str | Sequence[str] | None = None,
    key_in_ckpt: str | None = None,
    use_trace: bool | None = None,
    input_shape: Sequence[int] | None = None,
    args_file: str | None = None,

Review comment (refactor):
In this version, the `_export` function only supports exporting as files. I am not sure why the `| IO[Any]` is added here. If this is for exporting to IO, please also add this to the `_export` function and implement the IO exporting in it. Otherwise, I think we should remove this type hints.

Thanks  
Bin

Code after:
    net = converter(model=net, **kwargs)

    extra_files: dict = {}
    for i in ensure_tuple(config_file):
        # split the filename and directory
        filename = os.path.basename(i)
        # remove extension
        filename, _ = os.path.splitext(filename)
        # because all files are stored as JSON their name parts without extension must be unique
        if filename in extra_files:
            raise ValueError(f"Filename part '{filename}' is given multiple times in config file list.")
        # the file may be JSON or YAML but will get loaded and dumped out again as JSON
        extra_files[filename] = json.dumps(ConfigParser.load_config_file(i)).encode()

    # add .json extension to all extra files which are always encoded as JSON
    extra_files = {k + ".json": v for k, v in extra_files.items()}

    meta_values = parser.get().pop("_meta_", None)
    saver(net, filepath, meta_values=meta_values, more_extra_files=extra_files)

    logger.info(f"exported to file: {filepath}.")


def onnx_export(
    net_id: str | None = None,
    filepath: PathLike | None = None,
    ckpt_file: str | None = None,
    meta_file: str | Sequence[str] | None = None,
    config_file: str | Sequence[str] | None = None,
    key_in_ckpt: str | None = None,
    use_trace: bool | None = None,
    input_shape: Sequence[int] | None = None,
    args_file: str | None = None,
    converter_kwargs: Mapping | None = None,
    **override: Any,
) -> None:
    """
    Export the model checkpoint to an onnx model.

    Typical usage examples:

    .. code-block:: bash

        python -m monai.bundle onnx_export network --filepath <export path> --ckpt_file <checkpoint path> ...

    Args:
        net_id: ID name of the network component in the config, it must be `torch.nn.Module`.
        filepath: filepath where the onnx model is saved to.
        ckpt_file: filepath of the model checkpoint to load.
        meta_file: filepath of the metadata file, if it is a list of file paths, the content of them will be merged.
        config_file: filepath of the config file that contains extract network information,