如何在Python代码中直接调用Huggingface训练脚本的main()函数并传参运行

你目前尝试的直接传字典调用script.main({'arg1' : 1, 'arg2' : 2})的写法无法运行，原因是原脚本的main函数没有定义入参，它内部默认直接读取全局的命令行参数列表解析，既接收不到你传入的字典，也不支持字典格式的参数输入。

方案1：修改原脚本（最稳定，无副作用） #

可选1.1：支持列表形式传参（和bash执行逻辑完全对齐） #

HuggingFace的HfArgumentParser自带的parse_args_into_dataclasses方法本身支持传入自定义参数列表，只要修改main函数增加可选入参即可：

# script.py修改后的main函数
def main(custom_args: list[str] | None = None):
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    # 有自定义参数就用自定义的，没有就默认读取bash传入的命令行参数
    model_args, data_args, training_args = parser.parse_args_into_dataclasses(args=custom_args)
    # 剩余原有逻辑完全不变
    (...)

调用时直接把bash命令的参数拆成字符串列表传入即可：

import script

script.main(["--arg1", "1", "--arg2", "2"])

可选1.2：支持字典形式传参 #

如果想直接用字典传参，可以利用HfArgumentParser的parse_dict方法修改main函数：

# script.py修改后的main函数
def main(custom_arg_dict: dict | None = None):
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if custom_arg_dict:
        model_args, data_args, training_args = parser.parse_dict(custom_arg_dict)
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    # 剩余原有逻辑完全不变
    (...)

修改完成后就可以用你预期的字典方式调用：

import script

script.main({'arg1' : 1, 'arg2' : 2})

方案2：不修改原脚本（适合无权限改动原脚本的场景） #

可以通过临时修改Python全局的sys.argv参数列表实现传参，调用完成后再恢复原值避免影响其他逻辑：

import sys
import script

# 备份原始argv
ori_argv = sys.argv.copy()
# 替换为目标参数，第一个元素为脚本名无需改动，后面追加参数
sys.argv = [sys.argv[0], "--arg1", "1", "--arg2", "2"]
# 执行main
script.main()
# 恢复原始argv
sys.argv = ori_argv

注意：该方案是修改全局变量，多线程场景下可能出现参数冲突，仅推荐单线程场景使用。

补充说明 #

Bash中传入的命令行参数并不直接等价于字典，它传入Python后首先会被解析为sys.argv的字符串列表，是HfArgumentParser将列表按照参数规则转换成了对应的数据类实例，两种格式需要通过HfArgumentParser自带的方法做转换才可以通用。

内容的提问来源于stack exchange，提问作者Vincent