• STaR 方法代码开源，这里给出一个中文代码解读地址：repo
• 入口点：iteration_train.py；
• 关键代码：device_train.py, device_inference.py, and create_finetune_tfrecords.py；
• 基于 JAX、RAY，在 Google TPU 上实现；

入口点：iteration_train.py

if __name__ == "__main__":
    args = parse_args()
    print(args)
    task = args.task                                                                    # 选择数据集/任务：论文中有 CommonsenseQA、GSM8K
    experiment_name = "_".join(sys.argv[1:])                                            # 实验参数以_分割，拼接在一起命名
    experiment_name = ''.join(ch for ch in experiment_name if ch.isalnum() or ch == "_")# 确保 name 只有字母、数字、下划线（符合文件命名格式）
    if args.no_prompt:
        eval_seq = 128 + args.gen_length
    os.makedirs(f"configs/{experiment_name}", exist_ok=True)
    shutil.copy(f"configs/qa_base.json", f"configs/{experiment_name}/base.json")        # 复制一份实验配置模版
    prev_config = f"configs/{experiment_name}/base.json"                                # 实验配置模版的路径（后续代码会修改这个复制文件）
    new_json = make_first_config()

    os.makedirs(f'data/{experiment_name}', exist_ok=True)
    os.makedirs(f'{task}/{experiment_name}', exist_ok=True)
    os.makedirs(f'result_logs/', exist_ok=True)
    with open(f"result_logs/{experiment_name}.txt", "a+") as f:
        print("================================", file=f)                               # 类似 f.write
        print(args, file=f)
    for cur_iter in range(1, args.n_iters):                                             # 论文中的外循环迭代次数，重复多少次 STaR 微调方法
        exp_iteration = f"{experiment_name}_{cur_iter}"
        gen_train() # Generate the training set
        train_set = gen_records() # Create the tfrecords from the data                  # "{experiment_name}/{exp_iteration}.index"
        config_name = gen_config(train_set) # Create the new configuration file         # 核心是修改 total_steps
        train_model() # Train the new model
        eval_model() # Evaluate the new model
        prev_config = config_name  # Prepare for next iteration
        if args.copy_n > 0:
            copy_files()                                                                # [TODO] 复制上次外循环的一些配置文件，暂时不知道有啥用

parse_args() 标准的解析命令行参数，但是这里代码参数非常多。论文中，对一些技术细节写的比较模糊或者看不明白，这里需要结合代码分析。、

启动命令参数 parse_args()

• 说明：对于 bool 参数，在启动命令中带 --bool_params 或者不带这个参数即可提现，不用具体赋值

参数	取值范围	默认值	说明
--no_prompt	bool	true	eval时是否移出prompts （不用few-shot prompting，训练默认都是用的，对比实验不用）
--base_epochs	float	1.0	第一次 iter 的 epoch
--add_epochs	float	0.2	不同 iter 中需要 add 的 epoch
--few_shot_train	bool	false	是否使用 few-shot 训练
--steady_grow	bool	false	是否使用固定数量的 epoch
--start_steps	float	40.0	第一次外循环的步数（不同外循环步数可能不同）
--exponential_grow	bool	false	是否使用指数增长
--add_steps	float	20.0	steady_grow 配对参数，每次迭代中增加的步数
--grow_steps	float	1.2	exponential_grow 配对参数，每次迭代中按比例增长
--p_rationalization	float	1.0	使用合理化的错误样本比例
--p_show_hint_save	float	0.0	保存合理化提示的比例 [TODO]
--rationalize	bool	false	是否使用合理化
--start_iter	int	1	起始迭代数
--n_iters	int	64	外部循环迭代的最大次数 （论文中的外循环，使用多少次 STaR 微调）
--copy_n	int	0	每次迭代中需要复制的文件数
--n_train_samples	int	10000	训练样本数
--gradient_accumulation_steps	int	8	梯度累积的步数 Batch size
--task	str	“commonsenseqa”	运行的任务类型 ，论文中有 CommonsenseQA、GSM8K 两个数据集
--direct	bool	false	是否直接预测（不使用scratchpad）
--gen_length	int	96	生成输出的长度
--sequence_count	int	10	每个batch的平均序列数量
--base_model_location	str	“gs://checkpoint-bucket/step_383500/”	微调模型的检查点路径
--dry_run	bool	false	是否进行快速运行以可视化输出
--skip_eval	bool	false	是否跳过评估（例如算术任务）

训练epoch、step是否随着外循环迭代而增长？

epoch 控制参数：
step 控制参数：steady_grow、exponential_grow 或者都不选。三选一。选了 steady_grow、exponential_grow 分别还有一个配对的配置参数：add_steps、grow_steps（比例）。不选的话根据下面计算步数：

# Count data points
        total_count = 0
        for cur_file in sorted(os.listdir(record_folder(cur_iter - 1)), key=lambda x: int(x.split('.')[0].split("_")[-1])):
            with open(f"{record_folder(cur_iter - 1)}/{cur_file}", encoding='utf-8') as train_file:
                train_file_text = train_file.read()
                total_count += len(train_file_text.split("\n\n"))
                print(len(train_file_text.split("\n\n")))
        train_epochs = args.base_epochs + args.add_epochs * (cur_iter - 1)
        cur_steps = int(total_count * train_epochs // (args.gradient_accumulation_steps * args.sequence_count))
        return cur_steps

配置文件

qa_base.json

configs/qa_base.json 是实验的基础配置文件，运行实验会复制这个 template 然后不断修改这里的 value。

{
    "layers": 28,
    "d_model": 4096,
    "n_heads": 16,
    "n_vocab": 50400,
    "norm": "layernorm",
    "pe": "rotary",
    "pe_rotary_dims": 64,
    "seq": 1536, // 模型上下文窗口长度
    "cores_per_replica": 8, // device_inference 中用到，模型并行的参数，模型要分散到多个cores上来进行模型的计算
    "per_replica_batch": 1,	// device_inference 中用到，数据并行的参数，数据并行中每个模块并行的batch大小
    "gradient_accumulation_steps": 8, // 始终是 args.gradient_accumulation_steps
    "warmup_steps": 100,
    "anneal_steps": 300000,
    "lr": 1e-06,
    "end_lr": 1e-06,
    "weight_decay": 0.0,
    "total_steps": 383500,	   // 来自 get_n_steps()，有三种配置模式，见上面
    "tpu_size": 8,
    "p_rationalization": 1.0, // 始终是 args.p_rationalization
    "bucket": "checkpoint-bucket",			// 模型 ckpt 存储桶名
    "model_dir": "full_qa_4",				// 模型存储路径
    "train_set": "qa_train_4.index",
    "val_set": {
      "index": "qa.val.index"
    },
    "eval_harness_tasks": [
      "lambada",
      "piqa",
      "hellaswag",
      "winogrande",
      "mathqa",
      "pubmedqa"
    ],
    "val_batches": 100,
    "val_every": 10000,
    "ckpt_every": 10000,
    "keep_every": 10000,
    "name": "slow_grow_full_epoch_0",			// 这里会不断修改为 "{experiment_name}_0"
    "wandb_project": "full_6",	// wandb是一个日志服务，这里是日志记录的所属项目
    "comment": "",
    "target_save_folder": "commonsenseqa/iterative_full/iterative_full_0", // 文件存储所在文件夹路径
    "target_save": "commonsenseqa/slow_grow_full_epoch/slow_grow_full_epoch_0/slow_grow_full_epoch_0.txt" // 文件存储位置：文件和 name 同名，target_save_folder+name+".txt"
  }

训练核心代码

外层调用：iteration_train.py

调用侧代码（iteration_train.py）：

# main:
    for cur_iter in range(1, args.n_iters):                                             # 论文中的外循环迭代次数，重复多少次 STaR 微调方法
        exp_iteration = f"{experiment_name}_{cur_iter}"
        gen_train() # Generate the training set （第一次不执行）
        train_set = gen_records() # Create the tfrecords from the data                  # "{experiment_name}/{exp_iteration}.index"
        config_name = gen_config(train_set) # Create the new configuration file         # 核心是修改 total_steps
        train_model() # Train the new model

在训练前，需要先生成训练数据集（rationale generation）。核心是：gen_train()，然后通过 train_model() 开始微调模型。

def gen_records():
    gen_cmd = f'python3 create_finetune_tfrecords.py {record_folder(cur_iter - 1)} {record_folder(cur_iter - 1)}'
    print(f"Creating records for finetuning {cur_iter}: {gen_cmd}")
    if not args.dry_run and (cur_iter >= args.start_iter):
        os.system(gen_cmd)
    train_set = f"{experiment_name}/{exp_iteration}.index"
    with open(f"data/{train_set}", "w") as new_data_file:
        new_data_file.write(f"{record_folder(cur_iter - 1)}.tfrecords")
    return train_set

def train_model():
    model_cmd = f"python3 device_train.py --config {config_name} --tune-model-path={args.base_model_location}"
    print(f"Train model {cur_iter}: {model_cmd}")
    if not args.dry_run and (cur_iter >= args.start_iter):
        os.system(model_cmd)

rationale generation 代码 gen_train：device_inference.py

device_inference.py

参数	取值范围	默认值	说明
--config	str	None	配置文件路径
--direct	bool	false	是否直接预测（不使用scratchpad）
--rationalize	bool	false	是否使用合理化
--no_prompt	bool	false	eval时是否移出prompts （不用few-shot prompting，训练默认都是用的，对比实验不用）
--few_shot_train	bool	false	训练时是否移除few-shot-prompts
--show_hint_prompt	bool	false	是否需要提示提示
--split	str	“dev”	split的数据集（train,dev) gen_train里是–split=train，eval_model 里是 dev
--dataset_mode	str	“cqa”	使用的数据集(注意cqa在另一个文件默认值是全写，有代码做了兼容，这里默认值不能改，必须是cqa)
--n_train_samples	int	3000	训练样本数量
--gen_length	int	96	生成长度
--eval_batch_size	int	8	评估时的批量大小
--p_show_hint_save	float	0.0	保存合理化提示的比例
--ckpt_step	int	-1	要评估的检查点，-1表示最终检查点
--eval_seq	int	-1	序列长度，-1表示使用参数文件中的配置 （seq是模型上下文tokens最大长度）

此时传入的参数是：

• prev_config：用的上次迭代的配置，因为这里用上一次学习好的模型来生成数据集；
• gen_length 输出长度；

    if args.no_prompt:
        eval_seq = 128 + args.gen_length

如果按默认值，这里gen_length是128+96=224

• p_show_hint_save：合理化相关的参数
• n_train_samples：训练样本，默认是 10000（论文里始终保持这个数）

def gen_train():
    train_cmd = f"python3 device_inference.py --config={prev_config} --split=train --gen_length={args.gen_length} --p_show_hint_save={args.p_show_hint_save} "
    if task != "commonsenseqa":
        train_cmd += f" --dataset_mode={task} "
    if args.rationalize:
        train_cmd += " --rationalize "
    if args.few_shot_train:
        train_cmd += " --few_shot_train "
    if cur_iter > 1 and args.no_prompt:
        train_cmd += f" --no_prompt --eval_seq {eval_seq} "
    train_cmd += f" --n_train_samples={args.n_train_samples} "
    train_cmd += f" >> result_logs/{experiment_name}.txt"
    print(f"Generating training set {cur_iter} using model {cur_iter - 1}: {train_cmd}")
    if not args.dry_run and (cur_iter >= args.start_iter):
        if (cur_iter == 1) and os.path.exists(record_folder(0) + f"/{experiment_name}_0.txt"):
            print("First file cached") # 第一次不执行
        else:
            os.system(train_cmd)

注意：第一次运行 gen_train 的时候不执行，需要先微调后才执行合理化。

接下来分析 device_inference.py 中的代码：

if __name__ == "__main__":
    # 参数解析
    args = parse_args()
    print(args)
    split = args.split                              # 'dev'
    params = json.load(smart_open(args.config))     # smart_open 是一个用于打开文件的函数，支持多种文件格式和存储后端，本地文件，aws s3，gcs 等等

    # 初始化 wandb
    project = params.get("wandb_project", "mesh-transformer-jax")               # 日志服务所属的项目，随便什么值，这里不重要
    experiment_details = params["name"].split("_")
    wandb_name = "_".join(experiment_details[:-1])
    wandb_iteration = int(experiment_details[-1])
    wandb.init(project=project, name=wandb_name, config=params, resume=True)    # resume=True: 表示如果有相同名称的实验已经存在，则恢复该实验的状态，而不是创建一个新的实验。

    # 根据配置加载不同的 prompt 设置
    prompts_file = "prompts.txt" if not args.direct else "prompts_direct.txt"   # 默认不带 direct，即用带 few-shot 和 rationales 的 prompt
    prompts_file = f"{args.dataset_mode}/{prompts_file}"                        
    if args.no_prompt:
        commonsense_prompts = []
    else:
        with basic_open(prompts_file) as prompts:
            commonsense_prompts = prompts.read().split("\n\n")
    prompts_hint_file = "prompts_answer_key.txt" if not args.direct else "prompts_direct_answer_key.txt"
    prompts_hint_file = f"{args.dataset_mode}/{prompts_hint_file}"
    if args.no_prompt and not args.show_hint_prompt:
        commonsense_prompts_hint = []
    else:
        with basic_open(prompts_hint_file) as prompts:
            commonsense_prompts_hint = prompts.read().split("\n\n")

    # 参数设置
    per_replica_batch = params["per_replica_batch"]                             # 数据并行参数：1
    cores_per_replica = params["cores_per_replica"]                             # 模型并行参数：模型并行中的每个 replica 的核心数，默认是 8
    target_save = params["target_save"] if split != "dev" else f'{args.dataset_mode}/new_dev.txt'
    seq = params["seq"] if args.eval_seq == -1 else args.eval_seq
    hint_seq = seq
    set_opt(params)

    mesh_shape = (jax.device_count() // cores_per_replica, cores_per_replica)   # (replica 数量，每个 replica 的核心数)
    devices = np.array(jax.devices()).reshape(mesh_shape)                       # 为每个 replica 划分 cores，形成一个资源分配矩阵
    ckpt_path = get_ckpt_path(params, args.ckpt_step)                           # 默认用最新的 ckpt
    with jax.experimental.maps.mesh(devices, ('dp', 'mp')):                     # 并行策略的维度：dp，数据并行，mp，模型并行
        network = load_model(params, ckpt_path, devices, mesh_shape)

        dataset = get_dataset(args)
        dataset_keys = set([datakey for datakey, _ in dataset])

        total_batch = per_replica_batch * jax.device_count() // cores_per_replica * args.eval_batch_size    # 数据并行侧，一次性输入的数据 batch 大小
        gen_params = {"top_p": np.ones(total_batch) * 0.9, "temp": np.ones(total_batch) * 0.01}             # top_p: 控制生成文本的多样性的一种采样策略, Nucleus Sampling; temp: 温度参数，用于控制生成文本的随机性。温度越高，生成的文本越随机；温度越低，生成的文本越确定。

        accurate_count = eval_examples(dataset, commonsense_prompts, commonsense_prompts_hint, direct=args.direct)
        for cur_key, cur_counts in accurate_count.items():
            print(f"{split}, {cur_key}, {get_score(cur_counts)}")
            wandb.log({f"{split}_{cur_key}_accuracy": get_score(cur_counts), "iteration": wandb_iteration})

• 最开始，参数解析，注意一方面参数来自于外层调用传入的（前文分析了），另一部分来自配置文件 json；
• 初始化 wandb：Weights & Biases（通常简称为 WandB）是一个用于机器学习实验管理和可视化的工具。它提供了一系列功能，帮助研究人员和开发者更好地跟踪、管理和可视化他们的机器学习实验。
• 然后是根据配置加载不同的 prompt 设置
  • arg.direct：不用带 rationales 的 prompt，默认是用；
  • 加载不带合理化（但有rationales或者无rationales的配置）/ 或者不使用 few-shot；
  • 加载带合理化（hint）的 prompt （且带有 rationales）；
• 然后是从config读一些配置：注意数据集分 train、dev

# seq 是模型上下文窗口长度，input tokens 不能超过这个
seq = params["seq"] if args.eval_seq == -1 else args.eval_seq
hint_seq = seq

    "cores_per_replica": 8, // device_inference 中用到，模型并行的参数，模型要分散到多个cores上来进行模型的计算
    "per_replica_batch": 1,	// device_inference 中用到，数据并行的参数，数据并行中每个模块并行的batch大小

• replica 指的应该是大模型并行的其中一个部分。per_replica_batch 是数据并行的参数。cores_per_replica 是每个 replia 分配的核心数，是模型并行的参数，模型要分散到多个cores上来进行模型的计算。
  • 数据并行：数据并行是将训练数据分割成多个小批次，并在多个设备上并行处理这些小批次。每个设备都有一个完整的模型副本，计算梯度后再进行参数更新。
  • 模型并行：模型并行是将一个模型的不同部分分布在多个计算设备上。适用于模型非常大，以至于单个设备无法容纳整个模型的情况。

    mesh_shape = (jax.device_count() // cores_per_replica, cores_per_replica)   # (replica 数量，每个 replica 的核心数)
    devices = np.array(jax.devices()).reshape(mesh_shape)                       # 为每个 replica 划分 cores，形成一个资源分配矩阵
    ckpt_path = get_ckpt_path(params, args.ckpt_step)                           # 默认用最新的 ckpt
    with jax.experimental.maps.mesh(devices, ('dp', 'mp')):                     # 并行策略的维度：dp，数据并行，mp，模型并行

注意：eval_batch_size 主要是 cache 样本，样本缓存到这个数，才执行（减少模型io开销）。

eval_examples

def eval_examples(data_examples, few_shot_prompts, few_shot_prompts_hint, direct=False):
    accurate_count = {}
    tokenizer = transformers.GPT2TokenizerFast.from_pretrained('gpt2')

    main_examples, hint_examples = [], []
    pbar = tqdm(data_examples, smoothing=0)
    for data_example in pbar:   # 逐个遍历：而单个样本的执行和合理化样本的执行都是 cache 到一个 batch 再执行
        main_examples.append(data_example)
        if len(main_examples) == args.eval_batch_size:  # 默认值 8
            successful_examples = eval_batch(           # 评估
                main_examples, few_shot_prompts, seq, tokenizer,
                args.gen_length, gen_params, accurate_count, target_save, direct=direct
            )
            for example_idx, example in enumerate(main_examples):
                if (example_idx not in successful_examples) and (random.random() < params.get('p_rationalization', 1.)): # p_rationalization 默认值是 1
                    hint_examples.append(example)   # 如果回答失败，加入 hint 合理化样本中
            main_examples = [] # 清空队列

        if args.rationalize and len(hint_examples) >= args.eval_batch_size: # 合理化
            cur_hint_examples = hint_examples[:args.eval_batch_size]
            cur_hint_examples = [                                           # hint 样本修改 key
                (hint_example_key + "_r", hint_example) for hint_example_key, hint_example in cur_hint_examples
            ]
            eval_batch(                                                     # 评估
                cur_hint_examples, few_shot_prompts_hint, hint_seq, tokenizer,
                args.gen_length, gen_params, accurate_count, target_save, hint=True, direct=direct  # 开启 hint 合理化
            )
            hint_examples = hint_examples[args.eval_batch_size:]            # 清空当前合理化的样本
        pbar.set_description(f"{split} " + ", ".join([
            f"{cur_key}: {get_score(cur_counts):0.4f}" for cur_key, cur_counts in accurate_count.items()
        ]))
    return accurate_count

eval_batch

def eval_batch(examples, few_shot_prompts, seq, tok, gen_length, gen_params, accuracy, target_save, hint=False, direct=False):
    batch = examples_to_batch(examples, few_shot_prompts, seq, tok, hint=hint, direct=direct, p_show_hint_save=args.p_show_hint_save)   # 把example批处理成合适的prompt
    output = network.generate(batch["padded_batch"], batch["lengths"], gen_length, gen_params)    # 实际上执行输出的代码
    return eval_output(                                                                           # 评估输出结果，记录回答正确的样本
        output, batch["answers"], batch["base_context"], batch["classes"], accuracy, target_save, tok, direct=direct
    )

def examples_to_batch(data_examples, few_shot_prompts, seq, tokenizer, hint=False, direct=False, p_show_hint_save=0.1):
    batch = {
        "base_context": [],
        "initial_batch": [],
        "lengths": [],
        "padded_batch": [],
        "answers": [],
        "classes": []                                   # 分类
    }
    for data_class, data_example in data_examples:
        batch['classes'].append(data_class)
        # Context, without the few-shot prompt
        hintless_base_context = question_to_context(data_example, hint=False, dataset_mode=args.dataset_mode, direct=direct)    # 不带 hint
        base_context = question_to_context(data_example, hint=hint, dataset_mode=args.dataset_mode, direct=direct)
        if args.dataset_mode == "arithmetic":
            few_shot_prompts = base_context.split("\n\n")[:-1]
            base_context = base_context.split("\n\n")[-1]
            hintless_base_context = hintless_base_context.split("\n\n")[-1]

        if random.random() < p_show_hint_save:  # 默认是 0
            hintless_base_context = base_context

        # We always want to act as if no hint was given
        if args.few_shot_train:
            if args.dataset_mode == "arithmetic":
                raise NotImplementedError
            else:
                save_context = "\n\n".join(commonsense_prompts) + "\n\n"
                save_context += hintless_base_context
                batch['base_context'].append(save_context)
        else:
            batch['base_context'].append(hintless_base_context)

        # Input tokens
        if args.no_prompt:
            context = ""
        else:
            context = "\n\n".join(few_shot_prompts) + "\n\n"            # 最终prompt部分 1：默认带 few-shot

        context += base_context                                         # 最终prompt部分 2：当前问题（可能带有合理化）
        tokens = tokenizer.encode(context)                              # tokenizer
        batch['initial_batch'].append(tokens)
        # Input lengths
        batch['lengths'].append(len(tokens))
        # Padded tokens
        provided_ctx = len(tokens)
        pad_amount = max(seq - provided_ctx, 0)                         # seq 是最大窗口长度，如果不够这个长度需要 pad
        if provided_ctx > seq:
            tokens = tokens[-seq:]                                      # 如果超出，需要截断
        batch['padded_batch'].append(np.pad(tokens, ((pad_amount, 0),)).astype(np.uint32))

        # Answer
        if args.dataset_mode == "arithmetic":
            if len(data_example.split("\n")) >= 3:
                target = data_example.split("\n")[-3]
            else:
                target = "invalid"
        elif args.dataset_mode == "cqa":
            target = data_example['answerKey']
        elif args.dataset_mode == "gsm":
            target = data_example['answer'].split("#### ")[-1]
        batch['answers'].append(target)
    batch["lengths"] = np.asarray(batch["lengths"], dtype=np.uint32)
    batch["padded_batch"] = np.array(batch["padded_batch"])
    return batch

def question_to_context(data_example, hint=False, dataset_mode='cqa', direct=False):
    """"
    将问题转为 prompt

    - hint: 是否开启合理化
    """
    if dataset_mode == 'cqa':
        context = f"Q: {data_example['question']['stem']}\nAnswer Choices:\n"
        for choice in data_example['question']['choices']:
            if hint and (choice['label'].lower() == data_example['answerKey'].lower()):
                context += f"({choice['label'].lower()}) {choice['text']} (CORRECT)\n"
            else:
                context += f"({choice['label'].lower()}) {choice['text']}\n"
        context += "A:"
    elif dataset_mode == 'gsm':
        context = f"Q: {data_example['question']}"
        if hint:
            chosen_hint = data_example['answer']                # gsm 竟然直接把答案作为 hint
            context += f" ({chosen_hint})"
        context += "\nA:"
    elif dataset_mode == "arithmetic":
        context = ""
        for example_split, next_example_split in zip(data_example.split('Target:')[:-1], data_example.split('Target:')[1:]):
            if direct and "</scratch>" in example_split:
                context += example_split.split("</scratch>")[-1]
            else:
                context += example_split
            context += "Target:"
            if hint:
                context += " " + next_example_split.split("\n")[-5]
    return context

eval_output

def eval_output(output, answers, context, example_classes, accuracy, target_save, tokenizer, show=False, direct=False, endoftext="<|endoftext|>"):
    """
    评估输出结果，统计准确率，并将成功的示例保存到指定文件中。

    参数:
    - output (list): 模型的输出结果。
    - answers (list): 正确答案列表。
    - context (list): 上下文列表。
    - example_classes (list): 示例类别列表。
    - accuracy (dict): 用于统计准确率的字典。
    - target_save (str): 成功示例保存的文件路径。
    - tokenizer (transformers.PreTrainedTokenizer): 用于处理文本的分词器。
    - show (bool, optional): 是否打印成功示例到控制台。默认为 False。
    - direct (bool, optional): 是否使用直接预测，跳过scratchpad。默认为 False。
    - endoftext (str, optional): 用于标记文本结束的字符串。默认为 "<|endoftext|>"。

    返回:
    - list: 成功示例的索引列表。
    """
    successful_examples = []
    enum_outputs = enumerate(output[1][0][:, :, 0])
    for (idx, o), target, cur_base_context, example_class in zip(enum_outputs, answers, context, example_classes):
        cur_output = tokenizer.decode(o)
        output_numbers = cur_output.split('\n')
        if example_class not in accuracy:
            accuracy[example_class] = {'accurate': 0, 'total': 0}
        accuracy[example_class]['total'] += 1
        if len(output_numbers) == 0:
            continue
        try:
            if args.dataset_mode == "cqa":
                output_numbers = output_numbers[0]
                if "<|endoftext|>" in output_numbers:
                    output_numbers = output_numbers.split("<|endoftext|>")[0]
                output_prediction = output_numbers[-3]                                  # 选项
            elif args.dataset_mode == "gsm":
                output_prediction = ""
                for line_idx, line in enumerate(output_numbers):
                    if "####" in line:
                        output_numbers = "\n".join(output_numbers[:line_idx + 1])
                        if "<|endoftext|>" in output_numbers:
                            output_numbers = output_numbers.split("<|endoftext|>")[0]
                        output_prediction = output_numbers.split("####")[-1].strip()
                        break
            elif args.dataset_mode == "arithmetic":
                if len(output_numbers) == 0:
                    continue
                elif "<|endoftext|>" in output_numbers:
                    prediction_index = output_numbers.index("<|endoftext|>") - 1
                elif "</scratch>" in output_numbers:
                    prediction_index = output_numbers.index("</scratch>") + 1
                    if prediction_index == len(output_numbers):
                        continue
                else:
                    if direct and len(output_numbers) > 1:
                        prediction_index = 1
                    else:
                        prediction_index = 0
                output_prediction = output_numbers[prediction_index]                      # 计算结果

            if "<|endoftext|>" in output_prediction:
                output_prediction = output_prediction.split("<|endoftext|>")[0]

            correct = output_prediction.lower() == target.lower()                         # 判断输出是否和目标一致
            if correct:
                accuracy[example_class]['accurate'] += 1                                  # 回答正确，计数++
                with basic_open(target_save, 'a+') as new_train_f:
                    if args.dataset_mode == "cqa" or args.dataset_mode == "gsm":
                        new_example = cur_base_context + output_numbers + endoftext       # 正确回答的样本作为新的训练样本
                    elif args.dataset_mode == "arithmetic":
                        if args.few_shot_train:
                            raise NotImplementedError
                        joined_output = "\n".join(output_numbers[:prediction_index + 1])
                        if "<|endoftext|>" in joined_output:
                            joined_output = joined_output.split("<|endoftext|>")[0]
                        new_example = cur_base_context + joined_output + endoftext       # 正确回答的样本作为新的训练样本
                    if show:
                        print(new_example)
                    print(new_example, file=new_train_f, end="")                         # 把回答正确的样本写入文件中
                successful_examples.append(idx)
        except IndexError:
            pass
    return successful_examples

合理化部分代码总结

结合代码以及论文解读[NeurlPS 2022] STaR: Self-Taught Reasoner Bootstrapping Reasoning With Reasoning 现在重新来理解论文。

论文基本思路是，先给出few-shot，让模型参考few-shot在回答answer前带上rationales，如果回答不正确，就加上hint回答，最终把回答正确的样本留下进行下一轮微调。

在具体代码实现上，首先在 eval_examples 中，对样本做了个 batch 级别的 cache，每满8个，才执行对应的推理（回答）。这里维护了两个cache 队列，一个是回答正确的队列，一个是直接回答失败的队列（因此，用合理化修改了原始prompt）。两个队列分别满8时分别执行重新的回答操作，具体是在 eval_batch 中实现。先通过 examples_to_batch 对 batch 样本批量处理prompt，比如加上few-shot template 等等（或者加上hint）。然后批量推理。然后通过eval_output评估是否回答正确。如果没有回答正确，那么加入hint的样本中。所有回答正确的样本都会保存作为下一次微调的数据集【注意，对于合理化的样本，保存的问题不带hint】。

所以，根据这个实现，再回答阅读论文中的问题：

注意：这个标里的细节。文字部分说“Note the final STaR model is trained on 78.2% of the training dataset with rationale generation, and an additional 8.5% from rationalization”，而表格里不带合理化的STaR准确率只有68.8%，这里78.2%和68.8%有个差值！这里要怎么理解：因为带有合理化后，fine tune，导致模型处理hard问题的能力提升，所以在之后的实验中，部分问题不需要合理化就可以解出，所以涨了近10个点。