evaluate.py

"""
模型评估脚本
计算困惑度和使用lm-evaluation-harness进行下游任务评估
"""
import os
import argparse
import math
import torch
from tqdm import tqdm
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from datasets import load_dataset


def calculate_perplexity(
    model_path: str,
    dataset_name: str = "wikitext",
    dataset_config: str = "wikitext-2-raw-v1",
    split: str = "test",
    stride: int = 512,
    max_length: int = 1024,
    device: str = "cuda"
):
    """
    计算模型在指定数据集上的困惑度

    Args:
        model_path: 模型路径
        dataset_name: 数据集名称
        dataset_config: 数据集配置
        split: 数据集划分
        stride: 滑动窗口步长
        max_length: 最大序列长度
        device: 设备
    """
    print(f"Loading model from {model_path}...")
    model = GPT2LMHeadModel.from_pretrained(model_path)
    tokenizer = GPT2Tokenizer.from_pretrained(model_path)

    model = model.to(device)
    model.eval()

    print(f"Loading dataset {dataset_name}/{dataset_config}...")
    dataset = load_dataset(dataset_name, dataset_config, split=split)

    # 将所有文本拼接
    text = "\n\n".join(dataset["text"])
    encodings = tokenizer(text, return_tensors="pt")

    seq_len = encodings.input_ids.size(1)
    print(f"Total tokens: {seq_len:,}")

    nlls = []
    prev_end_loc = 0

    for begin_loc in tqdm(range(0, seq_len, stride), desc="Calculating PPL"):
        end_loc = min(begin_loc + max_length, seq_len)
        trg_len = end_loc - prev_end_loc

        input_ids = encodings.input_ids[:, begin_loc:end_loc].to(device)
        target_ids = input_ids.clone()
        target_ids[:, :-trg_len] = -100

        with torch.no_grad():
            outputs = model(input_ids, labels=target_ids)
            neg_log_likelihood = outputs.loss * trg_len

        nlls.append(neg_log_likelihood)

        prev_end_loc = end_loc
        if end_loc == seq_len:
            break

    ppl = torch.exp(torch.stack(nlls).sum() / end_loc)
    return ppl.item()


def calculate_perplexity_on_file(
    model_path: str,
    data_file: str,
    block_size: int = 1024,
    device: str = "cuda"
):
    """
    计算模型在二进制数据文件上的困惑度
    """
    import numpy as np

    print(f"Loading model from {model_path}...")
    model = GPT2LMHeadModel.from_pretrained(model_path)
    model = model.to(device)
    model.eval()

    print(f"Loading data from {data_file}...")
    data = np.memmap(data_file, dtype=np.uint16, mode='r')
    num_tokens = len(data)
    print(f"Total tokens: {num_tokens:,}")

    total_loss = 0
    total_tokens = 0
    num_batches = (num_tokens - 1) // block_size

    for i in tqdm(range(0, num_batches), desc="Calculating PPL"):
        start = i * block_size
        end = start + block_size + 1

        if end > num_tokens:
            break

        chunk = torch.from_numpy(data[start:end].astype(np.int64)).unsqueeze(0).to(device)
        input_ids = chunk[:, :-1]
        labels = chunk[:, 1:]

        with torch.no_grad():
            outputs = model(input_ids, labels=labels)
            loss = outputs.loss

        total_loss += loss.item() * block_size
        total_tokens += block_size

    avg_loss = total_loss / total_tokens
    ppl = math.exp(avg_loss)
    return ppl


def run_lm_eval(
    model_path: str,
    tasks: str = "lambada_openai,hellaswag,piqa,winogrande,arc_easy,arc_challenge",
    num_fewshot: int = 0,
    batch_size: int = 4,
    device: str = "cuda:0"
):
    """
    使用lm-evaluation-harness进行评估

    需要先安装: pip install lm-eval
    """
    import subprocess

    cmd = [
        "lm_eval",
        "--model", "hf",
        "--model_args", f"pretrained={model_path}",
        "--tasks", tasks,
        "--device", device,
        "--batch_size", str(batch_size),
    ]

    if num_fewshot > 0:
        cmd.extend(["--num_fewshot", str(num_fewshot)])

    print(f"Running: {' '.join(cmd)}")
    subprocess.run(cmd)


def compare_models(
    model_paths: list,
    model_names: list,
    tasks: str = "lambada_openai,hellaswag,piqa,winogrande",
    device: str = "cuda"
):
    """
    对比多个模型的性能
    """
    results = {}

    for path, name in zip(model_paths, model_names):
        print(f"\n{'='*50}")
        print(f"Evaluating: {name}")
        print(f"{'='*50}")

        # 计算WikiText-2困惑度
        ppl = calculate_perplexity(path, device=device)
        print(f"WikiText-2 PPL: {ppl:.2f}")

        results[name] = {"wikitext2_ppl": ppl}

    # 打印对比表格
    print("\n" + "=" * 60)
    print("Results Comparison")
    print("=" * 60)
    print(f"{'Model':<30} {'WikiText-2 PPL':<15}")
    print("-" * 60)
    for name, metrics in results.items():
        print(f"{name:<30} {metrics['wikitext2_ppl']:<15.2f}")

    return results


def main():
    parser = argparse.ArgumentParser(description="Evaluate GPT-2 models")

    parser.add_argument("--model_path", type=str, required=True,
                        help="Path to the model")
    parser.add_argument("--eval_type", type=str, default="ppl",
                        choices=["ppl", "lm_eval", "both"],
                        help="Evaluation type")

    # PPL评估参数
    parser.add_argument("--dataset", type=str, default="wikitext",
                        help="Dataset for PPL evaluation")
    parser.add_argument("--dataset_config", type=str, default="wikitext-2-raw-v1",
                        help="Dataset config")
    parser.add_argument("--data_file", type=str, default=None,
                        help="Binary data file for PPL evaluation")

    # lm-eval参数
    parser.add_argument("--tasks", type=str,
                        default="lambada_openai,hellaswag,piqa,winogrande,arc_easy,arc_challenge",
                        help="Tasks for lm-eval")
    parser.add_argument("--num_fewshot", type=int, default=0,
                        help="Number of few-shot examples")
    parser.add_argument("--batch_size", type=int, default=4,
                        help="Batch size")

    parser.add_argument("--device", type=str, default="cuda",
                        help="Device")

    args = parser.parse_args()

    if args.eval_type in ["ppl", "both"]:
        print("\n" + "=" * 50)
        print("Perplexity Evaluation")
        print("=" * 50)

        if args.data_file:
            ppl = calculate_perplexity_on_file(
                args.model_path,
                args.data_file,
                device=args.device
            )
            print(f"PPL on {args.data_file}: {ppl:.2f}")
        else:
            ppl = calculate_perplexity(
                args.model_path,
                args.dataset,
                args.dataset_config,
                device=args.device
            )
            print(f"PPL on {args.dataset}/{args.dataset_config}: {ppl:.2f}")

    if args.eval_type in ["lm_eval", "both"]:
        print("\n" + "=" * 50)
        print("LM Evaluation Harness")
        print("=" * 50)

        run_lm_eval(
            args.model_path,
            args.tasks,
            args.num_fewshot,
            args.batch_size,
            args.device
        )


if __name__ == "__main__":
    main()