TokenSkip-Evaluation-代码解析

该函数主要两个功能：1.通过原始数据集和针对该数据集微调后的模型得到原始思维链输出-train 模式。2.评估指定模型在指定数据集上的表现，包括答题准确率，token 长度-test 模式

import os
import json
import torch
import random
import argparse
import numpy as np
from tqdm import tqdm
from time import time
from copy import deepcopy
from peft import PeftModel
from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest
from transformers import AutoTokenizer, AutoModelForCausalLM

from eval.utils import generate_completions
from data_processing.process_utils import *
from data_processing.answer_extraction import *
from eval.eval_script import *

# import time

def set_random_seed(seed):
    # 设置随机种子
    random.seed(seed)
    # 控制python 内置随机数生成器的随机种子
    os.environ['PYTHONHASHSEED'] = str(seed)
    # 控制 Python 内部哈希随机性（dict/set 顺序）
    np.random.seed(seed)
    # 控制NumPy 随机数种子
    torch.manual_seed(seed)
    # 控制CPU 上的 PyTorch 随机数
    torch.cuda.manual_seed(seed)
    # 控制 GPU上的 PyTorhch 随机数
    torch.backends.cudnn.deterministic = True
    # 强制 cuDNN 使用 确定性（deterministic） 算法。 也就是说： 同样的输入、同样的模型、同样的硬件 → 输出结果完全一致。
    # 这对 复现实验结果 非常重要。
    torch.backends.cudnn.benchmark = False
    # 关闭 cuDNN 的 自动算法优化（auto-tuner）。防止：1.每次运行选择的算法可能不同；2.输入形状稍有变化就会重新 benchmark；
    # 3.所以会导致结果不确定性。为了结果完全可复现，应关闭 benchmark（设为 False）。
def read_data(path):
    # 读取数据
    if path.endswith("json"):
        data = json.load(open(path, "r"))
    # json.load:一次性读取整个文本数据，该函数的好处能自动将其变为字典，并且支持嵌套，里面的字典也能识别
    # 。f.load 只是将其变为字符串。
    # example:
    # [{"id": 1, "text": "hello"}, {"id": 2, "text": "world"}]可以一次性读
    elif path.endswith("jsonl"):
        # json 文件只有一个对象（一般外层会用一个列表括号或字典括号包起来），len 为 1，jsonl 有多个对象，len 为 n
        data = []
        with open(path, "r") as file:
            for line in file:
                line = json.loads(line)
                # {"id": 1, "text": "hello"}\n{"id": 2, "text": "world"}得分行读
                data.append(line)
    else:
        raise NotImplementedError()
    return data

def infer(args, test_data, answer_extraction_fn):
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True)
    # 加载分词器
    # 一些模型（例如 Qwen、ChatGLM、Yi）必须加这个参数，否则会报错。
    prompts = []
    # test_data为包含了所有数据的列表，列表中各项皆为字典
    for example in test_data:
        prompt = ""
        for mess in example['messages']:
            if mess['role'] == 'user':
                if args.model_type == 'llama3':
                    # 构建对应 prompt
                    if args.compression_ratio < 1.0:
                        prompt += f"{tokenizer.bos_token}" + "<|start_header_id|>user<|end_header_id|>\n\nPlease reason step by step, and put your final answer within \\boxed{}.\n" + f"{mess['content']}\n{tokenizer.eos_token}{args.compression_ratio}{tokenizer.eos_token}{tokenizer.eos_token}<|start_header_id|>assistant<|end_header_id|>\n\n"
                    else:
                        prompt += f"{tokenizer.bos_token}" + "<|start_header_id|>user<|end_header_id|>\n\nPlease reason step by step, and put your final answer within \\boxed{}.\n" + f"{mess['content']}\n{tokenizer.eos_token}<|start_header_id|>assistant<|end_header_id|>\n\n"
                elif args.model_type == 'qwen':
                    if args.compression_ratio < 1.0:
                        prompt += "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nPlease reason step by step, and put your final answer within \\boxed{}.\n" + f"{mess['content']}<|eot_id|>{args.compression_ratio}<|eot_id|><|im_end|>\n<|im_start|>assistant\n"
                    else:
                        prompt += "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nPlease reason step by step, and put your final answer within \\boxed{}.\n" + f"{mess['content']}<|im_end|>\n<|im_start|>assistant\n"
                else:
                    raise NotImplementedError()
            elif mess['role'] == 'assistant':
                prompt += mess['content'].rstrip()
                # 拼接上 assiatant 的内容，但是在之前已经被清空了
                # rstrip:去掉 assistant 文本尾部的空格/换行，避免在后续再拼模板时出现多余空行或空格造成 token 差异。
            prompt = prompt.lstrip()
            # lstrip它会去掉字符串左边（开头）的所有空白字符，但右边的空白不会动。
        example['prompt'] = prompt
        # 放进字典的 prompt 里
        prompts.append(prompt)
        # 单独存放
    print("Loading model and tokenizer...")
    if args.use_vllm:
        # 是否多模态,暂时用不上，先不解析
        if args.use_adapter:
            # 是否微调过
            model = LLM(model=args.model_path, tokenizer=args.tokenizer_path, trust_remote_code=True, enable_lora=True, tensor_parallel_size=len(os.environ['CUDA_VISIBLE_DEVICES'].split(",")), max_model_len=16000)
        else:
            model = LLM(model=args.model_path, tokenizer=args.tokenizer_path, trust_remote_code=True, tensor_parallel_size=len(os.environ['CUDA_VISIBLE_DEVICES'].split(",")))
        # 加载模型
        eos_token = tokenizer.eos_token if tokenizer is not None and tokenizer.eos_token is not None else '</s>'
        stop_words = [eos_token]
        
        torch.cuda.synchronize()
        start_time = time()
        if args.use_adapter:
            outputs = model.generate(prompts, SamplingParams(temperature=args.temperature, top_p=1.0, max_tokens=args.max_new_tokens, n=1, stop=stop_words), lora_request=LoRARequest("sql_adapter", 1, args.adapter_path))
        else:
            outputs = model.generate(prompts, SamplingParams(temperature=args.temperature, top_p=1.0, max_tokens=args.max_new_tokens, n=1, stop=stop_words))
        torch.cuda.synchronize()
        total_time = time() - start_time
        
        outputs = sorted(outputs, key=lambda x: int(x.request_id)) # sort outputs by request_id
        outputs = [output.outputs[0].text for output in outputs]
    else:
        tokenizer = AutoTokenizer.from_pretrained(
            args.tokenizer_path,
            trust_remote_code=True,
        )
        model = AutoModelForCausalLM.from_pretrained(
            args.model_path,
            torch_dtype=torch.float16,
            trust_remote_code=True,
            device_map="auto",
        )

        if args.use_adapter:
            model = PeftModel.from_pretrained(model, args.adapter_path, device_map="auto")
            model = model.merge_and_unload()
            # 加载了适配器之后直接将参数和主模型合并，随后弃掉适配器

        # set padding side to left for batch generation
        tokenizer.padding_side = "left"
        # set pad token to eos token if pad token is not set (as is the case for llama models)
        if tokenizer.pad_token is None:
            # 空的时候才设置，qwen 并不空
            tokenizer.pad_token = tokenizer.eos_token
            tokenizer.pad_token_id = tokenizer.eos_token_id
        print(tokenizer.pad_token_id)
        # time.sleep(10)
        # eos？待会儿测试一下
        stop_id_sequences = []
        if tokenizer.eos_token_id is not None:
            stop_id_sequences = [[tokenizer.eos_token_id]]
        # 分词器有停止符，把停止符放进去
        do_sample = False if args.temperature == 0.0 else True
        # 是否做非确定采样，False 等于不做确定性采样，只做确定性采样
        torch.cuda.synchronize()
        start_time = time()
        outputs, _ = generate_completions(
            model=model,
            tokenizer=tokenizer,
            prompts=prompts,
            max_new_tokens=args.max_new_tokens,
            do_sample=do_sample,
            temperature=args.temperature,
            top_p=1.0,
            batch_size=args.eval_batch_size,
            stop_id_sequences=stop_id_sequences if stop_id_sequences else None,
            end_of_generation_id_sequence=[tokenizer.eos_token_id] if tokenizer.eos_token_id is not None else None
        )
        # 好像只要考虑新 tokens的最大生产数量就行
        torch.cuda.synchronize()
        total_time = time() - start_time
            
    model_outputs = outputs
    # 获得模型输出
    cot_lengths = []
    for model_completion in model_outputs:
        cot = model_completion.split('\n\nThe final answer is:')[0]
        # 忽略输出中 final answer 后面的部分，只提取前面的思维链
        cot_length = tokenizer(cot, return_tensors="pt")['input_ids'].shape[1]
        # 计算思维链 token 数量
        cot_lengths.append(cot_length)
        # 加入总思维链长度列表中
    predictions = [eval(answer_extraction_fn)(item['messages'][-2]['content'], output, task='cot') 
                   for item, output in tqdm(zip(test_data, model_outputs), desc="extract answer", 
                                            total=len(model_outputs))]
    # 提取出所有模型输出的答案
    assert len(model_outputs) > 0, f"{len(model_outputs)}"

    results = []
    for example, output, pred, cot_length in zip(test_data, model_outputs, predictions, cot_lengths):
        item = deepcopy(example)
        item.update({
            'model_output': output,
            'prediction': pred,
            'cot_length': cot_length,
        })
        results.append(item)
    # 把提取出的东西加进字典里，就可以退出啦
    return results, total_time


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=str, default="outputs/Qwen2.5-7B-Instruct/gsm8k/", help="default to `model_path`_predictions")
    # 输出评估结果或初始思维链数据存放位置
    parser.add_argument("--model-path", type=str, default="/your_model_path/Qwen2.5-7B-Instruct")
    # 模型位置，用以获取模型
    parser.add_argument("--tokenizer-path", type=str, default="/your_model_path/Qwen2.5-7B-Instruct")
    # 分词器位置，一般和模型放一起
    parser.add_argument("--adapter-path", type=str, default="/your_model_path/TokenSkip-Qwen2.5-7B-Instruct-GSM8K")
    # 适配器位置，这种要lora微调之后才有
    parser.add_argument("--model-size", type=str, choices=['3b', '7b', '8b', '13b', '14b', '33b', '34b', '70b'], default="7b")
    # 选择模型规模大小
    parser.add_argument("--model-type", type=str, choices=['llama3', 'qwen'], default="qwen")
    # 模型类型选择，目前只有 llama3和 qwen
    parser.add_argument("--use_adapter", action='store_true', default=False, help="whether to use LoRA")
    # 是否使用适配器，不用的话 lora 调了也是白调，因为 lora 是不会动原模型参数的，通过训练适配器来对原模型参数增加修正项。
    parser.add_argument("--compression_ratio", type=float, default=1.0, help="compression ratio for cot.")
    # 压缩率，在评估模式下，用来告诉模型需要压百分之多少，模型会根据比例来压出对应压缩率的思维链
    parser.add_argument("--benchmark", type=str, choices=['gsm8k', 'math'], default="gsm8k")
    # 数据集，说明用的数据集是哪个
    parser.add_argument("--data-type", type=str, choices=['train', 'test'], default="test")
    # 数据类型，train 的话就是 train 模式，用来让模型输出原始思维链用的，test 的话就是 test 模式，用来评估模型的思维链压缩能力
    parser.add_argument("--use_vllm", action='store_true', default=False, help="whether to use vllm")
    # 是否使用了多模态大模型
    parser.add_argument("--max_num_examples", type=int, default=100000000000000, help="maximum number of examples to evaluate.")
    # 最大样本数，超过这个数，多余的可不要了
    parser.add_argument("--max_new_tokens", type=int, default=512)
    # 限定了一句话下的最大token 数
    parser.add_argument("--eval_batch_size", type=int, default=16, help="batch size for evaluation.")
    # 输入batch 数
    parser.add_argument("--temperature", type=float, default=0.0)
    # 温度，改变概率分布的，目前不知道什么用
    parser.add_argument("--seed", type=int, default=42)
    # 随机种子
    args, unparsed_args = parser.parse_known_args()
    # 存储好上述超参数
    # os.environ['CUDA_VISIBLE_DEVICES'] = "1,2,5,6,7"

    if args.benchmark == 'math' and args.use_adapter:
        args.max_new_tokens = args.max_new_tokens * args.compression_ratio
    # 暂时没用到 math 数据集，不知道啥意思
    print(f"Evaluating {args.model_path}", flush=True)
    print(f"Max new tokens: {args.max_new_tokens}, eval batch size: {args.eval_batch_size}, temperature: {args.temperature}, seed: {args.seed}\n", flush=True)
    # print("end?")
    if args.use_adapter:
        print(f"Adapter path {args.adapter_path}, compression ratio: {args.compression_ratio}", flush=True)
    # 用适配器说明已经进入评估阶段，微调完了，评估阶段是肯定要压缩率这个参数的
    if args.use_adapter:
        args.output_dir = os.path.join(args.output_dir, f"{args.model_size}/", f"TokenSkip/", f"{args.compression_ratio}/")
        # 使用适配器说明是 test 模式，那必有压缩率，输出压缩数据到对应压缩率文件
    else:
        args.output_dir = os.path.join(args.output_dir, f"{args.model_size}/", f"Original/{args.data_type}/")
        # 没用说明是 train 模式，输出原始思维链到 Original 文件夹
    # 这个是建立输出路径的

    test_conf = read_data(f"configs/{args.benchmark}_{args.data_type}.json")
    # 找到对应数据集相应的参数文件
    # print("end??")
    for src, info in test_conf.items():
        # example:
        # src:gsm8k-train
        # 遍历参数文件内各项
        fname = os.path.join(args.output_dir, "test_data", "test.jsonl")
        # 创建测试样例对应文件名
        input_dir = os.path.dirname(fname)
        # 获取对应文件夹
        os.makedirs(input_dir, exist_ok=True)
        # 创建该文件夹
        metric_path = os.path.join(args.output_dir, "samples", "metrics.json")
        # 创建评估结果对应文件名
        if os.path.exists(metric_path) and read_data(metric_path)['n_samples'] > 0:
            continue
        # 如果显示已经评估过了，就跳过此次评估
        # 否则继续
        with open(fname, "w") as file:
            # 打开文件
            data = read_data(info['test_path'])
            # example：
            # "datasets/gsm8k/train.jsonl"
            # 读取数据路径
            for i, sample in enumerate(tqdm(data, desc=f'processing {src}')):
                # 开始从对应数据路径-读取数据，便利数据中的每个问题，思维链，及答案
                fn = eval(info['process_fn'])
                # 直接调取函数，eval 可以直接获取同名函数，
                # 这么写可以提升阅读便利性,需要调用不同函数的时候特别方便。
                sample['id'] = sample.get('id', f"{src}-{i}")
                #读取第 i 个问题的 id，没有的话就返回默认值，默认值就是第二个参数f"{src}-{i}"，并且给它加进去
                for j, item in enumerate(fn(sample)):
                    item['dataset'] = src
                    item['id'] = f"{src}-test-{i}-{j}"
                    assert 'answer' in item
                    print(json.dumps(item), file=file, flush=True)
                    # file:输出文件夹，打开文件夹，将 item 写入
                    # flush=True:强制立刻把内容写进文件（不走缓存）
                # 这里其实展示输入数据集的部分样例
            output_dir = os.path.join(args.output_dir, "samples")
            os.makedirs(output_dir, exist_ok=True)
            # 创建采样文件夹
        set_random_seed(args.seed)
        # 要固定随机种子了
        print("Loading data...")
        test_data = []
        with open(os.path.join(input_dir, f"test.jsonl")) as fin:
            # 打开直接处理好的 test.jsonl 文件夹
            for line in fin:
                # 枚举其中每个数据
                example = json.loads(line)
                # 读取
                messages = example['messages']
                # 取出 messages
                assert messages[-1]['role'] == 'assistant'
                # 最后一个角色不是assistant说明数据集有问题，报错
                example['reference'] = example.get('reference', '') or [mess['content'] for mess in messages if
                                                                        mess['role'] == 'assistant']
                # or:如果左边没东西，就返回右边的，左边有个默认值''，train 模式下可以直接返回右边，返回思维链
                for mess in messages:
                    if mess['role'] == 'assistant':
                        mess['content'] = ''
                # 原先 assiatant 里的思维链加答案进入了 reference, 随后把 assistant 的内容给清空
                example['messages'] = messages
                # 更改完成的 meesages 放回数据字典中
                test_data.append(example)
                # 塞入 test_data 里

        if args.max_num_examples and len(test_data) > args.max_num_examples:
            test_data = random.sample(test_data, args.max_num_examples)
        # 数据量超过最大限制，裁到上限
        if not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)
        # 没懂，这个文件夹还能不存在的
        results, total_time = infer(args, test_data, info['answer_extraction_fn'])
        # results 包含了所有输入输出，思维链长度等数据项的字典
        print("Finished inference...")

        os.environ['TOKENIZERS_PARALLELISM'] = "false"

        invalid_outputs = []
        labels = []
        for item in results:
            # 遍历所有条目
            if len(item['prediction']) == 0:
                # 没预测出来结果
                invalid_outputs.append({'prompt': item['prompt'], 'output':  item['model_output'], 'answer': item['prediction']})
                res = False
                extract_ans = None
            else:
                extract_ans = item['prediction']
                # 这玩意儿搞啥的没动，可能也是调试的
                res = eval_math(item)
                # 评估模型输出与标准答案是否匹配
                # res为 True 或 False的布尔型
            labels.append(res)

        for item, label in zip(results, labels):
            # 更新各条目中模型输出的正确性
            item['accuracy'] = label

        print("Calculating accuracy...")
        acc = 0
        for item in results:
            acc += item['accuracy']
            # 虽然是布尔值，但是在加法里算数字
        print("output acc = {:.5f}".format(acc / len(results) * 100), flush=True)
        # 计算正确率并输出
        avg_cot_length = sum(item['cot_length'] for item in results) / len(results)
        print("output avg_cot_length = {:.5f}".format(avg_cot_length), flush=True)
        # 输出平均思维链长度
        print("number of invalid outputs: {}".format(len(invalid_outputs)), flush=True)
        # 输出无效的模型输出
        pred_fname = "predictions.jsonl"
        for item in results:
            with open(os.path.join(output_dir, pred_fname), 'a+', encoding='utf-8') as fout:
                line = json.dumps(item, ensure_ascii=False)
                fout.write(line + '\n')
        # 写入预测文件
        metric_fname = "metrics.json"
        with open(os.path.join(output_dir, metric_fname), "w") as fout:
            json.dump({
                "n_samples": len(results),
                "accuracy": sum(item['accuracy'] for item in results) / len(results),
                "avg_cot_length": avg_cot_length,
                'sample_latency': total_time / len(test_data),
            }, fout, indent=4)
        # 将结果写入评估文件