大语言模型 RLHF(一)——ChatGLM代码逐行解读

作者：小蓝xlanll | 2024-02-25 00:36:15

踩

chatglm

介绍：
为方便学习，对ChatGlm的代码做了逐行解读，这里主要是main方法，里面核心的部分如数据的解析，评估过程定义，模型推理训练的框架流程等。后续会针对ChatGLM核心优化代码做个解读，

代码获取：
完整代码可以参考下面git链接里的ptuning/main.py里获取
GitHub - Pillars-Creation/ChatGLM-LoRA: ChatGLM-6B添加了LoRA实现，以及部分核心代码的逐行讲解 ,实例部分是做了个新闻短标题的生成

训练评估过程：


    if training_args.do_train:  # 如果需要进行训练
        if "train" not in raw_datasets:  # 如果原始数据集中没有训练集
            raise ValueError("--do_train requires a train dataset")  # 抛出异常，提示需要提供训练集
        train_dataset = raw_datasets["train"]  # 获取训练集
        if data_args.max_train_samples is not None:  # 如果设置了最大训练样本数
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)  # 计算最大训练样本数
            train_dataset = train_dataset.select(range(max_train_samples))  # 选择最大训练样本数的子集
        with training_args.main_process_first(desc="train dataset map pre-processing"):  # 在主进程中进行训练集预处理
            train_dataset = train_dataset.map(  # 对训练集进行映射操作
                preprocess_function_train,  # 预处理函数
                batched=True,  # 是否对数据进行批处理
                num_proc=data_args.preprocessing_num_workers,  # 预处理使用的进程数
                remove_columns=column_names,  # 需要移除的列名
                load_from_cache_file=not data_args.overwrite_cache,  # 是否从缓存文件中加载数据
                desc="Running tokenizer on train dataset",  # 显示的描述信息
            )
        print_dataset_example(train_dataset[0])  # 打印训练集的第一个样本

评估指标定义过程：


        score_dict = {
            "rouge-1": [],
            "rouge-2": [],
            "rouge-l": [],
            "bleu-4": []
        }
        # 定义 score_dict，包含 rouge-1、rouge-2、rouge-l 和 bleu-4 四个指标
 
        for pred, label in zip(decoded_preds, decoded_labels):
            # 遍历解码后的 preds 和 labels
 
            hypothesis = list(jieba.cut(pred))
            reference = list(jieba.cut(label))
            # 使用 jieba 对预测值和真实值进行分词
 
            rouge = Rouge()
            scores = rouge.get_scores(' '.join(hypothesis), ' '.join(reference))
            result = scores[0]
            # 使用 Rouge 计算 ROUGE 指标，ROUGE（Recall-Oriented Understudy for Gisting Evaluation）
            # 是一种用于自动评估文本摘要和机器翻译的指标。它通过比较生成的摘要或翻译与参考摘要或翻译之间的重叠来计算得分。
            # ROUGE 指标包括 ROUGE-1、ROUGE-2 和 ROUGE-L 等，
            # 其中 ROUGE-1 表示单个词的重叠，ROUGE-2 表示两个词的重叠，ROUGE-L 表示最长公共子序列的重叠。
            # ROUGE 指标的取值范围为 0 到 1，值越大表示生成的摘要或翻译与参考摘要或翻译之间的重叠越多，即越好。
            # 在使用 Rouge 计算 ROUGE 指标时，rouge.get_scores() 方法返回一个包含多个指标的列表，每个指标都是一个字典，
            # 包含 precision、recall 和 f-measure 三个值。因此，scores[0] 表示第一个指标的字典，
            # 其中包含 precision、recall 和 f-measure 三个值。在这里，我们默认使用 ROUGE-1 指标，
            # 因此 scores[0] 表示 ROUGE-1 指标的字典。

具体main代码如下：

完整代码可以参考下面git链接里的ptuning/main.py里获取

GitHub - Pillars-Creation/ChatGLM-LoRA: ChatGLM-6B添加了LoRA实现，以及部分核心代码的逐行讲解 ,实例部分是做了个新闻短标题的生成


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--do_train", action="store_true",default="true")
    parser.add_argument("--train_file", type=str, default="train.json")
 
    # 解析参数
    args = parser.parse_args()
 
    # 将参数赋值给 sys.argv
    sys.argv = [sys.argv[0]] + [f"--{k}={v}" for k, v in vars(args).items()]
 
    # 解析命令行参数
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # 如果命令行参数只有一个，并且是一个 JSON 文件的路径，则解析该文件以获取参数
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        # 否则，解析命令行参数
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
    # 设置日志记录的格式和级别
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
 
    if training_args.should_log:
        # 如果需要记录日志，则设置日志级别为 info
        transformers.utils.logging.set_verbosity_info()
        log_level = training_args.get_process_log_level()
        logger.setLevel(log_level)
        # 获取日志级别，并设置 logger 的日志级别
        
        transformers.utils.logging.set_verbosity(log_level)
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
        # 设置 transformers 的日志级别，并启用默认处理程序和显式格式
 
        # Log on each process the small summary:
        logger.warning(
            f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
            + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
        )
        logger.info(f"Training/evaluation parameters {training_args}")
        # 打印进程、设备、GPU 数量、分布式训练和 16 位训练等信息，并打印训练/评估参数
 
        # Set seed before initializing model.
        set_seed(training_args.seed)
        # 设置随机种子
 
        # Load dataset
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
            extension = data_args.train_file.split(".")[-1]
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
            extension = data_args.validation_file.split(".")[-1]
        if data_args.test_file is not None:
            data_files["test"] = data_args.test_file
            extension = data_args.test_file.split(".")[-1]
        # 加载数据集文件
 
        raw_datasets = load_dataset(
            extension,
            data_files=data_files,
            cache_dir=model_args.cache_dir,
            use_auth_token=True if model_args.use_auth_token else None,
        )
        # 加载数据集，并设置缓存目录和身份验证令牌
 
    # Load pretrained model and tokenizer
    config = AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=True)
    config.pre_seq_len = model_args.pre_seq_len
    config.prefix_projection = model_args.prefix_projection
    # 从预训练模型中加载配置文件，并设置 pre_seq_len 和 prefix_projection
 
    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, trust_remote_code=True)
    # 从预训练模型中加载 tokenizer
 
    if model_args.ptuning_checkpoint is not None:
        # 如果设置了 ptuning_checkpoint
 
        # Evaluation
        # Loading extra state dict of prefix encoder
        model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True)
        # 从预训练模型中加载模型
 
        prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin"))
        new_prefix_state_dict = {}
        for k, v in prefix_state_dict.items():
            if k.startswith("transformer.prefix_encoder."):
                new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
        model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
        # 加载额外的 prefix encoder 的状态字典
 
    else:
        # Finetune
        model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True)
        # 从预训练模型中加载模型
 
    if model_args.quantization_bit is not None:
        print(f"Quantized to {model_args.quantization_bit} bit")
        model = model.quantize(model_args.quantization_bit).to(device)
        # 如果设置了 quantization_bit，将模型量化为指定的位数
 
    if model_args.pre_seq_len is not None:
        # 如果设置了 pre_seq_len
 
        # P-tuning v2
        model = model.half()
        model.transformer.prefix_encoder.float()
        # 将模型和 prefix encoder 转换为半精度浮点数
 
    else:
        # Finetune
        model = model.float()
        # 将模型转换为浮点数
 
    prefix = data_args.source_prefix if data_args.source_prefix is not None else ""
    # 如果设置了 source_prefix，将其赋值给 prefix，否则将其设置
 
    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    if training_args.do_train:
        column_names = raw_datasets["train"].column_names
    elif training_args.do_eval:
        column_names = raw_datasets["validation"].column_names
    elif training_args.do_predict:
        column_names = raw_datasets["test"].column_names
    else:
        logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
        return
 
    # Get the column names for input/target.
    prompt_column = data_args.prompt_column
    response_column = data_args.response_column
    history_column = data_args.history_column
    
    # Temporarily set max_target_length for training.
    max_target_length = data_args.max_target_length
 
    def preprocess_function_eval(examples):
        # 定义预处理函数，输入为 examples，输出为 model_inputs
 
        inputs, targets = [], []
        # 定义 inputs 和 targets 列表
 
        for i in range(len(examples[prompt_column])):
            # 遍历 examples 中的每个样本
 
            if examples[prompt_column][i] and examples[response_column][i]:
                # 如果 prompt 和 response 都不为空
 
                query = examples[prompt_column][i]
                # 将 prompt 赋值给 query
 
                if history_column is None or len(examples[history_column][i]) == 0:
                    prompt = query
                else:
                    prompt = ""
                    history = examples[history_column][i]
                    for turn_idx, (old_query, response) in enumerate(history):
                        prompt += "[Round {}]\n问：{}\n答：{}\n".format(turn_idx, old_query, response)
                    prompt += "[Round {}]\n问：{}\n答：".format(len(history), query)
                # 如果有 history_column，将历史对话和当前的 prompt 拼接起来，否则直接使用 prompt
 
                inputs.append(prompt)
                targets.append(examples[response_column][i])
                # 将 prompt 和 response 分别添加到 inputs 和 targets 中
 
        inputs = [prefix + inp for inp in inputs]
        # 将 prefix 和 inputs 中的每个元素拼接起来
 
        model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, truncation=True, padding=True)
        # 使用 tokenizer 对 inputs 进行编码，并进行截断和填充
 
        labels = tokenizer(text_target=targets, max_length=max_target_length, truncation=True)
        # 使用 tokenizer 对 targets 进行编码，并进行截断
 
        if data_args.ignore_pad_token_for_loss:
            labels["input_ids"] = [
                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
            ]
        # 如果 ignore_pad_token_for_loss 为 True，将 labels 中的 pad_token_id 替换为 -100
 
        model_inputs["labels"] = labels["input_ids"]
        # 将 labels 中的 input_ids 赋值给 model_inputs 中的 labels
 
        return model_inputs
        # 返回 model_inputs
 
    def preprocess_function_train(examples):
        # 定义预处理函数，输入为 examples，输出为 model_inputs
 
        max_seq_length = data_args.max_source_length + data_args.max_target_length
        # 计算最大序列长度，即 prompt 和 response 的最大长度之和
 
        model_inputs = {
            "input_ids": [],
            "labels": [],
        }
        # 定义 model_inputs，包含 input_ids 和 labels 两个字段
 
        for i in range(len(examples[prompt_column])):
            # 遍历 examples 中的每个样本
 
            if examples[prompt_column][i] and examples[response_column][i]:
                # 如果 prompt 和 response 都不为空
 
                query, answer = examples[prompt_column][i], examples[response_column][i]
                # 将 prompt 和 response 分别赋值给 query 和 answer
 
                if history_column is None:
                    prompt = query
                else:
                    prompt = ""
                    history = examples[history_column][i]
                    for turn_idx, (old_query, response) in enumerate(history):
                        prompt += "[Round {}]\n问：{}\n答：{}\n".format(turn_idx, old_query, response)
                    prompt += "[Round {}]\n问：{}\n答：".format(len(history), query)
                # 如果有 history_column，将历史对话和当前的 prompt 拼接起来，否则直接使用 prompt
 
                prompt = prefix + prompt
                # 将 prefix 和 prompt 拼接起来
 
                a_ids = tokenizer.encode(text=prompt, add_special_tokens=False)
                b_ids = tokenizer.encode(text=answer, add_special_tokens=False)
                # 使用 tokenizer 对 prompt 和 response 进行编码
 
                if len(a_ids) > data_args.max_source_length - 1:
                    a_ids = a_ids[: data_args.max_source_length - 1]
                # 如果 prompt 的长度超过了最大长度减一，截断 prompt
 
                if len(b_ids) > data_args.max_target_length - 2:
                    b_ids = b_ids[: data_args.max_target_length - 2]
                # 如果 response 的长度超过了最大长度减二，截断 response
 
                input_ids = tokenizer.build_inputs_with_special_tokens(a_ids, b_ids)
                # 将 prompt 和 response 拼接起来，并加上特殊的 token
 
                context_length = input_ids.index(tokenizer.bos_token_id)
                mask_position = context_length - 1
                labels = [-100] * context_length + input_ids[mask_position + 1:]
                # 构造 labels，其中 context_length 是 bos_token_id 的位置，mask_position 是下一个 token 的位置
 
                pad_len = max_seq_length - len(input_ids)
                input_ids = input_ids + [tokenizer.pad_token_id] * pad_len
                labels = labels + [tokenizer.pad_token_id] * pad_len
                # 将 input_ids 和 labels 补齐到最大长度
 
                if data_args.ignore_pad_token_for_loss:
                    labels = [(l if l != tokenizer.pad_token_id else -100) for l in labels]
                # 如果 ignore_pad_token_for_loss 为 True，将 labels 中的 pad_token_id 替换为 -100
 
                model_inputs["input_ids"].append(input_ids)
                model_inputs["labels"].append(labels)
                # 将 input_ids 和 labels 添加到 model_inputs 中
 
        return model_inputs
        # 返回 model_inputs
    
    def print_dataset_example(example):
        print("input_ids",example["input_ids"])
        print("inputs", tokenizer.decode(example["input_ids"]))
        print("label_ids", example["labels"])
        print("labels", tokenizer.decode(example["labels"]))
 
    if training_args.do_train:  # 如果需要进行训练
        if "train" not in raw_datasets:  # 如果原始数据集中没有训练集
            raise ValueError("--do_train requires a train dataset")  # 抛出异常，提示需要提供训练集
        train_dataset = raw_datasets["train"]  # 获取训练集
        if data_args.max_train_samples is not None:  # 如果设置了最大训练样本数
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)  # 计算最大训练样本数
            train_dataset = train_dataset.select(range(max_train_samples))  # 选择最大训练样本数的子集
        with training_args.main_process_first(desc="train dataset map pre-processing"):  # 在主进程中进行训练集预处理
            train_dataset = train_dataset.map(  # 对训练集进行映射操作
                preprocess_function_train,  # 预处理函数
                batched=True,  # 是否对数据进行批处理
                num_proc=data_args.preprocessing_num_workers,  # 预处理使用的进程数
                remove_columns=column_names,  # 需要移除的列名
                load_from_cache_file=not data_args.overwrite_cache,  # 是否从缓存文件中加载数据
                desc="Running tokenizer on train dataset",  # 显示的描述信息
            )
        print_dataset_example(train_dataset[0])  # 打印训练集的第一个样本
 
    if training_args.do_eval:  # 如果需要进行评估
        max_target_length = data_args.val_max_target_length  # 获取最大目标长度
        if "validation" not in raw_datasets:  # 如果原始数据集中没有验证集
            raise ValueError("--do_eval requires a validation dataset")  # 抛出异常，提示需要提供验证集
        eval_dataset = raw_datasets["validation"]  # 获取验证集
        if data_args.max_eval_samples is not None:  # 如果设置了最大评估样本数
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)  # 计算最大评估样本数
            eval_dataset = eval_dataset.select(range(max_eval_samples))  # 选择最大评估样本数的子集
        with training_args.main_process_first(desc="validation dataset map pre-processing"):  # 在主进程中进行验证集预处理
            eval_dataset = eval_dataset.map(  # 对验证集进行映射操作
                preprocess_function_eval,  # 预处理函数
                batched=True,  # 是否对数据进行批处理
                num_proc=data_args.preprocessing_num_workers,  # 预处理使用的进程数
                remove_columns=column_names,  # 需要移除的列名
                load_from_cache_file=not data_args.overwrite_cache,  # 是否从缓存文件中加载数据
                desc="Running tokenizer on validation dataset",  # 显示的描述信息
            )
        print_dataset_example(eval_dataset[0])  # 打印验证集的第一个样本
 
    if training_args.do_predict:
        # 如果设置了 do_predict 标志
 
        max_target_length = data_args.val_max_target_length
        # 设置最大目标长度为 val_max_target_length
 
        if "test" not in raw_datasets:
            raise ValueError("--do_predict requires a test dataset")
        # 如果 raw_datasets 中没有 test 数据集，抛出 ValueError 异常
 
        predict_dataset = raw_datasets["test"]
        # 将 test 数据集赋值给 predict_dataset
 
        if data_args.max_predict_samples is not None:
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))
        # 如果设置了 max_predict_samples，将 predict_dataset 截取到指定长度
 
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            predict_dataset = predict_dataset.map(
                preprocess_function_eval,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                remove_columns=column_names,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on prediction dataset",
            )
        # 使用 preprocess_function_eval 对 predict_dataset 进行预处理
 
        print_dataset_example(predict_dataset[0])
        # 打印 predict_dataset 的第一个样本
 
    # Data collator
    label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
    # 如果 ignore_pad_token_for_loss 为 True，将 label_pad_token_id 设置为 -100，否则设置为 tokenizer.pad_token_id
 
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=None,
        padding=False
    )
 
    # 定义 data_collator，用于将数据转换为模型所需的格式
 
    # Metric
    def compute_metrics(eval_preds):
        # 定义评估指标函数，输入为 eval_preds，输出为 score_dict
 
        preds, labels = eval_preds
        # 将 preds 和 labels 赋值给 preds 和 labels
 
        if isinstance(preds, tuple):
            preds = preds[0]
        # 如果 preds 是元组类型，将其转换为列表类型
 
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        # 使用 tokenizer 对 preds 进行解码
 
        if data_args.ignore_pad_token_for_loss:
            # 如果 ignore_pad_token_for_loss 为 True，将 labels 中的 -100 替换为 pad_token_id
            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
 
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
        # 使用 tokenizer 对 labels 进行解码
 
        score_dict = {
            "rouge-1": [],
            "rouge-2": [],
            "rouge-l": [],
            "bleu-4": []
        }
        # 定义 score_dict，包含 rouge-1、rouge-2、rouge-l 和 bleu-4 四个指标
 
        for pred, label in zip(decoded_preds, decoded_labels):
            # 遍历解码后的 preds 和 labels
 
            hypothesis = list(jieba.cut(pred))
            reference = list(jieba.cut(label))
            # 使用 jieba 对预测值和真实值进行分词
 
            rouge = Rouge()
            scores = rouge.get_scores(' '.join(hypothesis), ' '.join(reference))
            result = scores[0]
            # 使用 Rouge 计算 ROUGE 指标，ROUGE（Recall-Oriented Understudy for Gisting Evaluation）
            # 是一种用于自动评估文本摘要和机器翻译的指标。它通过比较生成的摘要或翻译与参考摘要或翻译之间的重叠来计算得分。
            # ROUGE 指标包括 ROUGE-1、ROUGE-2 和 ROUGE-L 等，
            # 其中 ROUGE-1 表示单个词的重叠，ROUGE-2 表示两个词的重叠，ROUGE-L 表示最长公共子序列的重叠。
            # ROUGE 指标的取值范围为 0 到 1，值越大表示生成的摘要或翻译与参考摘要或翻译之间的重叠越多，即越好。
            # 在使用 Rouge 计算 ROUGE 指标时，rouge.get_scores() 方法返回一个包含多个指标的列表，每个指标都是一个字典，
            # 包含 precision、recall 和 f-measure 三个值。因此，scores[0] 表示第一个指标的字典，
            # 其中包含 precision、recall 和 f-measure 三个值。在这里，我们默认使用 ROUGE-1 指标，
            # 因此 scores[0] 表示 ROUGE-1 指标的字典。
 
            for k, v in result.items():
                score_dict[k].append(round(v["f"] * 100, 4))
            # 将 ROUGE 指标添加到 score_dict 中
 
            bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
            score_dict["bleu-4"].append(round(bleu_score * 100, 4))
            # 计算 BLEU 指标，并将其添加到 score_dict 中
 
        for k, v in score_dict.items():
            score_dict[k] = float(np.mean(v))
        # 计算每个指标的平均值，并将其转换为浮点数类型
 
        return score_dict
        # 返回 score_dict
 
    # Override the decoding parameters of Seq2SeqTrainer
    training_args.generation_max_length = (
        training_args.generation_max_length
        if training_args.generation_max_length is not None
        else data_args.val_max_target_length
    )
    training_args.generation_num_beams = (
        data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
    )
    # 如果设置了 generation_max_length，将其赋值给 training_args.generation_max_length，否则将 val_max_target_length 赋值给其
    # 如果设置了 num_beams，将其赋值给 training_args.generation_num_beams，否则将其保持不变
 
    # Initialize our Trainer
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if training_args.predict_with_generate else None,
        save_prefixencoder=model_args.pre_seq_len is not None
    )
    # 初始化 Seq2SeqTrainer，包括模型、参数、训练集、验证集、tokenizer、data_collator、compute_metrics 和 save_prefixencoder
 
    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        # elif last_checkpoint is not None:
        #     checkpoint = last_checkpoint
 
        # 启用梯度检查点，减少显存使用，提高训练效率
        model.gradient_checkpointing_enable()
        # 启用输入梯度，使模型在训练时计算输入的梯度，提高训练效果
        model.enable_input_require_grads()
 
        # 开始训练模型，resume_from_checkpoint 表示是否从检查点恢复训练
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        # trainer.save_model()  # Saves the tokenizer too for easy upload
 
        # 将训练指标保存到 metrics 变量中
        metrics = train_result.metrics
        # 判断是否设置了 max_train_samples 参数，如果设置了，则将 train_samples 设置为 max_train_samples 和训练数据集大小的较小值，否则将其设置为训练数据集的大小
        max_train_samples = (
            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))
 
        # 将训练指标记录到日志中
        trainer.log_metrics("train", metrics)
        # 将训练指标保存到文件中
        trainer.save_metrics("train", metrics)
        # 保存训练状态，包括模型和优化器的参数
        trainer.save_state()
 
    # Evaluation
    results = {}
    # 定义一个空字典 results
 
    max_seq_length = data_args.max_source_length + data_args.max_target_length + 1
    # 计算最大序列长度
 
    if training_args.do_eval:
        # 如果设置了 do_eval 为 True
 
        logger.info("*** Evaluate ***")
        # 打印日志信息
 
        metrics = trainer.evaluate(metric_key_prefix="eval", do_sample=True, top_p=0.7, max_length=max_seq_length,
                                   temperature=0.95)
        # 使用 evaluate 方法评估模型，并返回评估指标
 
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
        # 计算评估样本数，并将其添加到 metrics 中
 
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
        # 打印评估指标，并将其保存到文件中
 
    if training_args.do_predict:
        # 如果设置了 do_predict 为 True
 
        logger.info("*** Predict ***")
        # 打印日志信息
 
        predict_results = trainer.predict(predict_dataset, metric_key_prefix="predict", max_length=max_seq_length,
                                          do_sample=True, top_p=0.7, temperature=0.95)
        # 使用 predict 方法预测模型，并返回预测结果
 
        metrics = predict_results.metrics
        max_predict_samples = (
            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
        )
        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))
        # 计算预测样本数，并将其添加到 metrics 中
 
        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)
        # 打印预测指标，并将其保存到文件中
 
        if trainer.is_world_process_zero():
            # 如果是主进程
 
            if training_args.predict_with_generate:
                # 如果设置了 predict_with_generate 为 True
 
                predictions = tokenizer.batch_decode(
                    predict_results.predictions, skip_special_tokens=True, clean_up_tokenization_spaces=True
                )
                # 使用 tokenizer 对预测值进行解码，并去除特殊标记和空格
 
                predictions = [pred.strip() for pred in predictions]
                # 去除预测值中的空格
 
                labels = tokenizer.batch_decode(
                    predict_results.label_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
                )
                # 使用 tokenizer 对真实值进行解码，并去除特殊标记和空格
 
                labels = [label.strip() for label in labels]
                # 去除真实值中的空格
 
                output_prediction_file = os.path.join(training_args.output_dir, "generated_predictions.txt")
                with open(output_prediction_file, "w", encoding="utf-8") as writer:
                    for p, l in zip(predictions, labels):
                        res = json.dumps({"labels": l, "predict": p}, ensure_ascii=False)
                        writer.write(f"{res}\n")
                # 将预测值和真实值写入文件中
 
    return results
    # 返回结果字典 results
 
 
def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()
 
 
if __name__ == "__main__":
    main()

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/小蓝xlanll/article/detail/137232