reference/code/dMoE/evaluations/simple-evals/simple_evals.py

import json
import argparse
import pandas as pd
from . import common
from .browsecomp_eval import BrowseCompEval
from .drop_eval import DropEval
from .gpqa_eval import GPQAEval
from .humaneval_eval import HumanEval
from .math_eval import MathEval
from .mgsm_eval import MGSMEval
from .mmlu_eval import MMLUEval
from .simpleqa_eval import SimpleQAEval
from .sampler.chat_completion_sampler import (
    OPENAI_SYSTEM_MESSAGE_API,
    OPENAI_SYSTEM_MESSAGE_CHATGPT,
    ChatCompletionSampler,
)
from .sampler.o_chat_completion_sampler import OChatCompletionSampler
from .sampler.responses_sampler import ResponsesSampler
from .sampler.claude_sampler import ClaudeCompletionSampler, CLAUDE_SYSTEM_MESSAGE_LMSYS

from .sampler.opensource_sampler import LLaDASampler
from .sampler.llama_sampler import LLaMASampler

import torch
from torch.distributed import scatter_object_list, gather_object
import torch.distributed as dist

def init_distributed():
    from datetime import datetime, timedelta
    dist.init_process_group(
        backend="nccl", timeout=timedelta(hours=24)
    )
    rank = dist.get_rank()
    world = dist.get_world_size()
    torch.cuda.set_device(rank)
    return rank, world

def main():
    parser = argparse.ArgumentParser(
        description="Run sampling and evaluations using different samplers and evaluations."
    )

    parser.add_argument(
        "--task", type=str, required=True, choices=["mmlu", "math", "gpqa", "humaneval"]
    )
    parser.add_argument(
        "--list-models", action="store_true", help="List available models"
    )
    parser.add_argument("--model", type=str, help="Select a model by name")
    parser.add_argument("--debug", action="store_true", help="Run in debug mode")
    parser.add_argument(
        "--examples", type=int, help="Number of examples to use (overrides default)"
    )
    # Generation Config
    parser.add_argument(
        "--remasking", type=str, default="low_confidence", help="Remasking strategy"
    )
    parser.add_argument(
        "--steps", type=int, default=128, help="Number of sampling steps"
    )
    parser.add_argument(
        "--max-len", type=int, default=128, help="Length of generated sequence"
    )
    parser.add_argument(
        "--block", type=int, default=64, help="Block length for sampling"
    )

    # for Cache
    parser.add_argument("--kv-cache-masked", action="store_true")
    parser.add_argument("--kv-cache-decoded", action="store_true")
    
    parser.add_argument("--cache-steps", type=int, default=None)
    parser.add_argument("--window-size", type=int, default=0)

    args = parser.parse_args()

    if torch.cuda.device_count() > 1:
        rank, world = init_distributed()
    else:
        world, rank = 1, 0

    models = {
        "llada": LLaDASampler(
            model_name="GSAI-ML/LLaDA-8B-Instruct",
            generation_kwargs={
                'steps': args.steps, 
                'gen_length': args.max_len, 
                'block_length': args.block, 
                'temperature': 0., 
                'cfg_scale': 0., 
                'remasking': args.remasking,
                'enable_cache': args.kv_cache_masked or args.kv_cache_decoded,
                'cache_reloading_step': args.cache_steps,
                'window_size': args.window_size
            },
            kv_cache_masked=args.kv_cache_masked,
            kv_cache_decoded=args.kv_cache_decoded
        ),
        "llama": LLaMASampler(
            model_name = "meta-llama/Llama-3.1-8B-Instruct",
            generation_kwargs={
                'max_new_tokens': args.max_len
            }
        )
    }

    if args.list_models:
        print("Available models:")
        for model_name in models.keys():
            print(f" - {model_name}")
        return

    if args.model:
        if args.model not in models:
            print(f"Error: Model '{args.model}' not found.")
            return
        models = {args.model: models[args.model]}

    grading_sampler = ChatCompletionSampler(model="gpt-4o")
    equality_checker = ChatCompletionSampler(model="gpt-4-turbo-preview")
    # ^^^ used for fuzzy matching, just for math

    def get_evals(eval_name, debug_mode):
        num_examples = (
            args.examples if args.examples is not None else (5 if debug_mode else None)
        )
        # Set num_examples = None to reproduce full evals
        match eval_name:
            case "mmlu":
                return MMLUEval(num_examples=10 if debug_mode else num_examples)
            case "math":
                return MathEval(
                    equality_checker=equality_checker,
                    num_examples=num_examples,
                    n_repeats=1 if debug_mode else 10,
                )
            case "gpqa":
                return GPQAEval(
                    n_repeats=1 if debug_mode else 1, num_examples=num_examples,
                )
            case "mgsm":
                return MGSMEval(num_examples_per_lang=10 if debug_mode else 250)
            case "drop":
                return DropEval(
                    num_examples=10 if debug_mode else num_examples,
                    train_samples_per_prompt=3,
                )
            case "humaneval":
                return HumanEval(num_examples=10 if debug_mode else num_examples)
            case "simpleqa":
                return SimpleQAEval(
                    grader_model=grading_sampler,
                    num_examples=10 if debug_mode else num_examples,
                )
            case "browsecomp":
                return BrowseCompEval(
                    grader_model=grading_sampler,
                    num_examples=10 if debug_mode else num_examples,
                )
            case _:
                raise Exception(f"Unrecognized eval type: {eval_name}")

    evals = {
        eval_name: get_evals(eval_name, args.debug)
        for eval_name in [args.task]
    }
    debug_suffix = "_DEBUG" if args.debug else ""
    print(debug_suffix)
    mergekey2resultpath = {}
    for model_name, sampler in models.items():
        sampler.init_model()
        for eval_name, eval_obj in evals.items():
            result = eval_obj(sampler, rank = rank, world = world)

            # ^^^ Gather from different gpus:
            sampler._free_memory()
            print("Start aggregate results")
            if world > 1:
                gathered = [None] * world
                dist.all_gather_object(gathered, result)   
                flat = [o for sub in gathered for o in sub]  # 展平
                result = flat

            result = common.aggregate_results(result)

            # ^^^ how to use a sampler
            if (world > 1 and rank ==0) or world == 1:
                file_stem = f"{eval_name}_{model_name}"
                file_stem += f"_{args.remasking}_steps_{args.steps}_len_{args.max_len}_block_{args.block}"
                file_stem += f"_decoded_{args.kv_cache_decoded}_masked_{args.kv_cache_masked}_cache{args.cache_steps}_window{args.window_size}"
                report_filename = f"simple_eval_results/{file_stem}{debug_suffix}.html"
                print(f"Writing report to {report_filename}")
                with open(report_filename, "w") as fh:
                    fh.write(common.make_report(result))
                metrics = result.metrics | {"score": result.score}
                print(metrics)
                result_filename = f"simple_eval_results/{file_stem}{debug_suffix}.json"
                with open(result_filename, "w") as f:
                    f.write(json.dumps(metrics, indent=2))
                print(f"Writing results to {result_filename}")
                mergekey2resultpath[f"{file_stem}"] = result_filename

            if world > 1:
                dist.barrier()
    
    if (world > 1 and rank ==0) or world == 1:
        merge_metrics = []
        for eval_model_name, result_filename in mergekey2resultpath.items():
            try:
                result = json.load(open(result_filename, "r+"))
            except Exception as e:
                print(e, result_filename)
                continue
            result = result.get("f1_score", result.get("score", None))
            eval_name = eval_model_name[: eval_model_name.find("_")]
            model_name = eval_model_name[eval_model_name.find("_") + 1 :]
            merge_metrics.append(
                {"eval_name": eval_name, "model_name": model_name, "metric": result}
            )
        merge_metrics_df = pd.DataFrame(merge_metrics).pivot(
            index=["model_name"], columns="eval_name"
        )
        print("\nAll results: ")
        print(merge_metrics_df.to_markdown())
        return merge_metrics


if __name__ == "__main__":
    main()