math_eval.py

import random
import os
import argparse
import time
from vllm import LLM, SamplingParams
from datetime import datetime
from tqdm import tqdm

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

from evaluate import evaluate
from utils import set_seed, load_jsonl, save_jsonl, construct_prompt
from parser import *
from data_loader import load_data
from model_utils import load_hf_lm_and_tokenizer, generate_completions


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_names", default="gsm8k,math", type=str)
    parser.add_argument("--data_dir", default="./data", type=str)
    parser.add_argument("--model_name_or_path", default="gpt-4", type=str)
    parser.add_argument("--output_dir", default="./output", type=str)
    parser.add_argument("--prompt_type", default="tool-integrated", type=str)
    parser.add_argument("--split", default="test", type=str)
    parser.add_argument("--num_test_sample", default=-1, type=int)  # -1 for full data
    parser.add_argument("--seed", default=0, type=int)
    parser.add_argument("--start", default=0, type=int)
    parser.add_argument("--end", default=-1, type=int)
    parser.add_argument("--temperature", default=0, type=float)
    parser.add_argument("--n_sampling", default=1, type=int)
    parser.add_argument("--top_p", default=1, type=float)
    parser.add_argument("--max_tokens_per_call", default=None, type=int)
    parser.add_argument("--shuffle", action="store_true")
    parser.add_argument("--use_vllm", action="store_true")
    parser.add_argument("--save_outputs", action="store_true")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--use_safetensors", action="store_true")
    parser.add_argument("--num_shots", type=int, default=0)
    parser.add_argument(
        "--apply_chat_template",
        action="store_true",
        help="Apply chat template to prompt.",
    )
    parser.add_argument("--pipeline_parallel_size", type=int, default=1)
    parser.add_argument(
        "--adapt_few_shot",
        action="store_true",
        help="Few shot for multiple-choice questions, zero shot for others.",
    )
    parser.add_argument(
        "--pangu_think_mode",
        choices=["slow", "fast", "auto"],
        default="slow",
    )
    args = parser.parse_args()
    args.top_p = (
        1 if args.temperature == 0 else args.top_p
    )  # top_p must be 1 when using greedy sampling (vllm)
    return args


def prepare_data(data_name, args):
    examples = load_data(data_name, args.split, args.data_dir)

    # sample `num_test_sample` from dataset
    if args.num_test_sample > 0:
        # examples = random.sample(examples, min(args.num_test_sample, len(examples)))
        examples = examples[: args.num_test_sample]

    # shuffle
    if args.shuffle:
        random.seed(datetime.now().timestamp())
        random.shuffle(examples)

    # select start and end
    examples = examples[args.start : len(examples) if args.end == -1 else args.end]

    # get out_file name
    dt_string = datetime.now().strftime("%m-%d_%H-%M")
    model_name = "/".join(args.model_name_or_path.split("/")[-2:])
    out_file_prefix = f"{args.split}_{args.prompt_type}_{args.num_test_sample}_seed{args.seed}_t{args.temperature}_top-p{args.top_p}"

    if args.prompt_type == "pangu":
        out_file_prefix += f"_{args.pangu_think_mode}"

    output_dir = args.output_dir

    out_file = f"{output_dir}/{data_name}/{out_file_prefix}.jsonl"
    os.makedirs(f"{output_dir}/{data_name}", exist_ok=True)

    # load all processed samples
    processed_samples = []
    if not args.overwrite:
        processed_files = [
            f
            for f in os.listdir(f"{output_dir}/{data_name}/")
            if f.endswith(".jsonl") and f.startswith(out_file_prefix)
        ]
        for f in processed_files:
            processed_samples.extend(list(load_jsonl(f"{output_dir}/{data_name}/{f}")))

    # dedepulicate
    processed_samples = {sample["idx"]: sample for sample in processed_samples}
    processed_idxs = list(processed_samples.keys())
    processed_samples = list(processed_samples.values())
    examples = [example for example in examples if example["idx"] not in processed_idxs]
    return examples, processed_samples, out_file


def setup(args):
    # import json
    # model_tokenizers = json.load(open("Qwen/model_tokenizer.json", "r"))
    # model_tokenizers[args.output_dir.split("/")[-1]] = args.model_name_or_path
    # json.dump(model_tokenizers, open("Qwen/model_tokenizer.json", "w"))

    # load model
    if torch.cuda.is_available():
        num_gpus = torch.cuda.device_count()
        print(f"Using {num_gpus} GPU")
    elif torch.npu.is_available():
        num_gpus = torch.npu.device_count()
        print(f"Using {num_gpus} NPU")
    else:
        raise ValueError("No GPU or NPU available.")

    data_list = args.data_names.split(",")
    need_eval_data_list = []
    if not args.overwrite:
        for data_name in data_list:
            out_file_prefix = f"{args.split}_{args.prompt_type}_{args.num_test_sample}_seed{args.seed}_t{args.temperature}_top-p{args.top_p}"
            out_file = f"{args.output_dir}/{data_name}/{out_file_prefix}.jsonl"
            out_metric_json = out_file.replace(".jsonl", "_metrics.json")

            if os.path.exists(out_metric_json):
                print(f"Skipping {data_name} because {out_metric_json} already exists.")
                continue
            else:
                need_eval_data_list.append(data_name)

        if len(need_eval_data_list) == 0:
            print("All datasets already evaluated. Exiting.")
            exit(0)
        data_list = need_eval_data_list

    if args.use_vllm:
        if "pangu" in args.model_name_or_path.lower():
            # if torch.cuda.is_available():
            #     max_num_batched_tokens=32768
            # else:
            #     max_num_batched_tokens=4096

            vllm_kwargs = dict(
                max_num_seqs=32,
                max_model_len=131072,
                # max_num_batched_tokens=max_num_batched_tokens,
                tokenizer_mode="slow",
                dtype="bfloat16",
                distributed_executor_backend="mp",
                gpu_memory_utilization=0.95,
                enable_prefix_caching=False,
                enable_chunked_prefill=False,
                hf_overrides={"max_position_embeddings": 131072},
            )
        else:
            vllm_kwargs = dict(
                dtype="auto",
            )

            if args.prompt_type == "qwen25-math-cot":
                vllm_kwargs["hf_overrides"] = {
                    "max_position_embeddings": 131072,
                    # "rope_theta": 1000000,
                    # "rope_scaling": {
                    #     "rope_type": "yarn",
                    #     "factor": 32.0,
                    #     "original_max_position_embeddings": 4096,
                    # },
                }

        llm = LLM(
            model=args.model_name_or_path,
            tensor_parallel_size=num_gpus // args.pipeline_parallel_size,
            pipeline_parallel_size=args.pipeline_parallel_size,
            # gpu_memory_utilization=0.8,
            trust_remote_code=True,
            **vllm_kwargs,
        )
        # tokenizer = None
        # if args.apply_chat_template:
        #     tokenizer = AutoTokenizer.from_pretrained(
        #         args.model_name_or_path, trust_remote_code=True
        #     )
        tokenizer = llm.get_tokenizer()
    else:
        llm, tokenizer = load_hf_lm_and_tokenizer(
            model_name_or_path=args.model_name_or_path,
            load_in_half=True,
            use_fast_tokenizer=True,
            use_safetensors=args.use_safetensors,
        )

    # infer & eval
    results = []
    for data_name in data_list:
        results.append(main(llm, tokenizer, data_name, args))

    # add "avg" result to data_list and results
    data_list.append("avg")
    results.append(
        {
            "acc": sum([result["acc"] for result in results]) / len(results),
        }
    )

    # print all results
    pad = max([len(data_name) for data_name in data_list])
    print("\t".join(data_name.ljust(pad, " ") for data_name in data_list))
    print("\t".join([f"{result['acc']:.1f}".ljust(pad, " ") for result in results]))


def is_multi_choice(answer):
    for c in answer:
        if c not in ["A", "B", "C", "D", "E"]:
            return False
    return True


def main(llm, tokenizer, data_name, args):
    examples, processed_samples, out_file = prepare_data(data_name, args)
    print(f"data: {data_name}, remain samples: {len(examples)}, out_file: {out_file}")

    if len(examples) > 0:
        print(examples[0])

    samples = []
    for example in tqdm(examples, total=len(examples)):
        idx = example["idx"]

        # parse question and answer
        example["question"] = parse_question(example, data_name)
        if example["question"] == "":
            continue
        gt_cot, gt_ans = parse_ground_truth(example, data_name)
        example["gt_ans"] = gt_ans

        if args.apply_chat_template:
            question = example["question"].strip()
            if args.prompt_type == "pangu":
                match args.pangu_think_mode:
                    case "slow":
                        pass
                    case "fast":
                        question = question + " /no_think"
                    case "auto":
                        question = question + " /auto_think"
                    case _:
                        raise ValueError(
                            f"Invalid pangu_think_mode {args.pangu_think_mode}"
                        )

            full_prompt = tokenizer.apply_chat_template(
                [{"role": "user", "content": question}],
                tokenize=False,
                add_generation_prompt=True,
            )
        else:
            full_prompt = construct_prompt(example, data_name, args)

        if idx == args.start:
            print(full_prompt)

        sample = {
            "idx": idx,
            # "id": example["id"],
            "question": example["question"],
            "gt_cot": gt_cot,
            "gt": gt_ans,
            "prompt": full_prompt,
        }

        # add remain fields
        for key in [
            "level",
            "type",
            "unit",
            "solution_type",
            "choices",
            "solution",
            "ques_type",
            "ans_type",
            "answer_type",
            "dataset",
            "subfield",
            "filed",
            "theorem",
            "answer",
        ]:
            if key in example:
                sample[key] = example[key]
        samples.append(sample)

    # repeat n times
    input_prompts = [
        sample["prompt"] for sample in samples for _ in range(args.n_sampling)
    ]
    current_prompts = [(i, prompt) for i, prompt in enumerate(input_prompts)]
    end_prompts = []

    # start inference
    # measure time use
    start_time = time.time()

    prompts = [item[1] for item in current_prompts]
    if args.use_vllm:
        outputs = llm.generate(
            prompts,
            SamplingParams(
                temperature=args.temperature,
                top_p=args.top_p,
                max_tokens=args.max_tokens_per_call,
                n=1,
            ),
        )

        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:
        outputs = generate_completions(
            model=llm,
            tokenizer=tokenizer,
            prompts=prompts,
            max_new_tokens=args.max_tokens_per_call,
            batch_size=16,
        )

    assert len(outputs) == len(current_prompts)

    for (i, query), output in zip(current_prompts, outputs):
        end_prompts.append((i, output))

    end_prompts = sorted(end_prompts, key=lambda x: x[0])
    codes = [code for _, code in end_prompts]

    results = []
    token_lengths = []
    for code in codes:
        if not code or code == "error":
            prediction = None
        else:
            prediction = extract_answer(code, data_name)
            prediction = strip_string(
                prediction, skip_unit=data_name in STRIP_EXCEPTIONS
            )
        results.append(prediction)
        token_lengths.append(len(tokenizer.tokenize(code)))

    time_use = time.time() - start_time

    all_samples = []
    for i, sample in enumerate(samples):
        code = codes[i * args.n_sampling : (i + 1) * args.n_sampling]
        preds = results[i * args.n_sampling : (i + 1) * args.n_sampling]
        preds = [val if (val is not None) else "None" for val in preds]
        token_len = token_lengths[i * args.n_sampling : (i + 1) * args.n_sampling]

        for j in range(len(preds)):
            if sample["gt"] in ["A", "B", "C", "D", "E"] and preds[j] not in [
                "A",
                "B",
                "C",
                "D",
                "E",
            ]:
                preds[j] = choice_answer_clean(code[j])

            elif is_multi_choice(sample["gt"]) and not is_multi_choice(preds[j]):
                preds[j] = "".join(
                    [c for c in preds[j] if c in ["A", "B", "C", "D", "E"]]
                )

        # sample.pop("prompt")
        sample.update({"code": code, "pred": preds, "token_len": token_len})
        all_samples.append(sample)

    # add processed samples
    all_samples.extend(processed_samples)
    all_samples, result_json = evaluate(
        samples=all_samples,
        data_name=data_name,
        prompt_type=args.prompt_type,
        execute=True,
    )

    # save outputs
    if len(processed_samples) < len(all_samples) and args.save_outputs:
        save_jsonl(all_samples, out_file)

    result_json["time_use_in_second"] = time_use
    result_json["time_use_in_minite"] = (
        f"{int(time_use // 60)}:{int(time_use % 60):02d}"
    )

    with open(out_file.replace(".jsonl", f"_metrics.json"), "w") as f:
        json.dump(result_json, f, indent=4)
    return result_json


if __name__ == "__main__":
    args = parse_args()
    set_seed(args.seed)
    setup(args)