#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    MATH-dataset evaluation entry-point.
#
#    Differences from `test.py`:
#      1. Loads the MATH test split (`problem`, `solution`, `answer` fields).
#      2. Answers are LaTeX expressions (e.g. \\frac{14}{3}, \\sqrt{2}, \\boxed{...}),
#         not plain numbers. We extract the predicted answer using the same
#         "The answer is: <ans>" template the model is trained on (see `train.py`,
#         where every example's target ends with `f"The answer is: {answer}"`).
#      3. Equivalence checking goes through Hendrycks/Minerva-style LaTeX
#         normalisation, with numeric and SymPy fallbacks.

import logging
import math
import re
import os
from dataclasses import dataclass, field
from typing import Dict, Optional, Sequence

import torch
import transformers
from torch.nn import functional as F
import json
import yaml
import argparse

from peft import PeftModel, LoraConfig, TaskType, get_peft_model
from peft import PeftModel
from datasets import load_dataset, concatenate_datasets
from accelerate.utils import set_seed
from src.model import (
    CODI_CT,
    ModelArguments,
    DataArguments,
    TrainingArguments,
    load_safetensors_state_dict,
)

import numpy as np

do_print = True

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)


# ---------------------------------------------------------------------------
# I/O helpers
# ---------------------------------------------------------------------------
def read_json(file_path):
    """Read a JSON file and return the parsed Python object."""
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            return json.load(file)
    except Exception as e:
        print(f"读取JSON文件时出错: {e}")
        return None


def write_json(data, file_path):
    """Dump a Python object to JSON."""
    try:
        with open(file_path, 'w', encoding='utf-8') as file:
            json.dump(data, file, ensure_ascii=False, indent=4)
    except Exception as e:
        print(f"写入JSON文件时出错: {e}")


# ---------------------------------------------------------------------------
# MATH answer extraction & equivalence
# ---------------------------------------------------------------------------
ANSWER_PROMPT = "The answer is:"


def extract_math_answer(full_output: str) -> str:
    """
    Extract the final answer from a model output, following the training format.

    Training (see train.py preprocess): every target ends with
        "The answer is: {answer}<eos>"
    where {answer} is either a plain number (GSM8K) or a LaTeX expression (MATH).

    Strategy (most-specific to least-specific):
      1. Split on "The answer is:" and take the trailing chunk.
      2. Fall back to \\boxed{...} if present anywhere in the output.
      3. Otherwise return the full output (the equivalence check will handle it).
    """
    if full_output is None:
        return ""
    text = full_output

    if ANSWER_PROMPT in text:
        ans = text.split(ANSWER_PROMPT)[-1]
    elif "The answer is" in text:  # tolerate missing colon
        ans = text.split("The answer is")[-1].lstrip(":")
    else:
        boxed = _last_boxed_only_string(text)
        ans = boxed if boxed is not None else text

    ans = ans.strip()
    # Drop a trailing period / eos artefacts
    ans = ans.rstrip(".").strip()
    # Drop surrounding inline-math $ ... $
    if ans.startswith("$") and ans.endswith("$") and len(ans) >= 2:
        ans = ans[1:-1].strip()
    # If the answer is wrapped in \\boxed{...}, peel it
    boxed_inner = _last_boxed_only_string(ans)
    if boxed_inner is not None:
        ans = boxed_inner.strip()
    return ans


def _last_boxed_only_string(s: str):
    """Return the inner content of the last \\boxed{...} (brace-balanced)."""
    if s is None or "\\boxed" not in s:
        # Some models emit \fbox instead
        if s is None or "\\fbox" not in s:
            return None
        idx = s.rfind("\\fbox")
        prefix_len = len("\\fbox")
    else:
        idx = s.rfind("\\boxed")
        prefix_len = len("\\boxed")

    i = idx + prefix_len
    # skip optional whitespace
    while i < len(s) and s[i] == " ":
        i += 1
    if i >= len(s) or s[i] != "{":
        return None
    depth = 0
    start = i
    for j in range(i, len(s)):
        if s[j] == "{":
            depth += 1
        elif s[j] == "}":
            depth -= 1
            if depth == 0:
                return s[start + 1:j]
    return None


# ---- Hendrycks/Minerva-style LaTeX normalisation -----------------------------
def _fix_fracs(string: str) -> str:
    """Normalise `\\frac ab` and `\\frac{a}b` to `\\frac{a}{b}`."""
    substrs = string.split("\\frac")
    new_str = substrs[0]
    for sub in substrs[1:]:
        new_str += "\\frac"
        if sub.startswith("{"):
            new_str += sub
        else:
            try:
                a = sub[0]
                b = sub[1]
                rest = sub[2:] if len(sub) > 2 else ""
                new_str += "{" + a + "}{" + b + "}" + rest
            except IndexError:
                return string
    return new_str


def _fix_sqrt(string: str) -> str:
    """Turn `\\sqrt3` into `\\sqrt{3}`."""
    if "\\sqrt" not in string:
        return string
    splits = string.split("\\sqrt")
    new = splits[0]
    for sp in splits[1:]:
        if sp and sp[0] != "{":
            new += "\\sqrt{" + sp[0] + "}" + sp[1:]
        else:
            new += "\\sqrt" + sp
    return new


def _fix_a_slash_b(string: str) -> str:
    """`a/b` -> `\\frac{a}{b}` (only for simple integer cases)."""
    if len(string.split("/")) != 2:
        return string
    a, b = string.split("/")
    try:
        ai = int(a)
        bi = int(b)
        if string == f"{ai}/{bi}":
            return f"\\frac{{{ai}}}{{{bi}}}"
        return string
    except (ValueError, TypeError):
        return string


def _remove_right_units(string: str) -> str:
    """Drop trailing `\\text{...}` units (e.g. `12\\text{ cm}` -> `12`)."""
    if "\\text{" in string:
        splits = string.split("\\text{")
        return splits[0].rstrip()
    return string


def _strip_string(string: str) -> str:
    """Canonicalise a LaTeX answer string for equivalence comparison."""
    if string is None:
        return ""
    s = str(string)

    # newlines, \\!, double backslashes
    s = s.replace("\n", "")
    s = s.replace("\\!", "")
    s = s.replace("\\\\", "\\")

    # tfrac/dfrac -> frac
    s = s.replace("tfrac", "frac").replace("dfrac", "frac")
    # \left / \right
    s = s.replace("\\left", "").replace("\\right", "")
    # degree marks
    s = s.replace("^{\\circ}", "").replace("^\\circ", "")
    # currency / percent
    s = s.replace("\\$", "")
    s = s.replace("$", "")
    s = s.replace("\\%", "")
    s = s.replace("\\,", "")
    s = s.replace("\\ ", "")
    # text wrappers
    s = re.sub(r"\\text\s*\{\s*([^{}]*?)\s*\}", r"\1", s)
    s = re.sub(r"\\mbox\s*\{\s*([^{}]*?)\s*\}", r"\1", s)
    s = re.sub(r"\\mathrm\s*\{\s*([^{}]*?)\s*\}", r"\1", s)
    s = re.sub(r"\\mathbf\s*\{\s*([^{}]*?)\s*\}", r"\1", s)
    # boxed
    boxed = _last_boxed_only_string(s)
    if boxed is not None:
        s = boxed
    # remove trailing units `\\text{ ... }`
    s = _remove_right_units(s)
    # remove all whitespace
    s = s.replace(" ", "")
    # ".5" -> "0.5"
    s = re.sub(r"(^|[^0-9])\.([0-9])", r"\g<1>0.\2", s)
    # "5." -> "5"
    s = re.sub(r"([0-9])\.(?=$|[^0-9])", r"\1", s)
    # 1,000 -> 1000
    s = re.sub(r"(\d),(\d{3})", r"\1\2", s)

    s = _fix_sqrt(s)
    s = _fix_fracs(s)
    s = _fix_a_slash_b(s)

    # X = ans  -> ans
    if "=" in s:
        parts = s.split("=")
        if len(parts) == 2 and parts[1] != "":
            s = parts[1]

    # strip surrounding parens / braces
    while len(s) >= 2 and ((s[0] == "{" and s[-1] == "}") or (s[0] == "(" and s[-1] == ")")):
        s = s[1:-1]

    return s.strip()


def _to_float(s: str):
    try:
        return float(s)
    except (TypeError, ValueError):
        return None


def is_math_equiv(pred: str, gold: str) -> bool:
    """Best-effort equivalence between a predicted and a gold MATH answer."""
    if pred is None or gold is None:
        return False
    p = _strip_string(pred)
    g = _strip_string(gold)
    if p == g and p != "":
        return True

    # numeric comparison
    pf, gf = _to_float(p), _to_float(g)
    if pf is not None and gf is not None:
        if math.isfinite(pf) and math.isfinite(gf) and abs(pf - gf) < 1e-6:
            return True

    # SymPy fallback (LaTeX -> expression). Best-effort; failures are non-fatal.
    try:
        from sympy import simplify, nsimplify, Rational, sympify
        try:
            from sympy.parsing.latex import parse_latex
        except Exception:
            parse_latex = None

        def _parse(x):
            if parse_latex is not None:
                try:
                    return parse_latex(x)
                except Exception:
                    pass
            try:
                return sympify(x.replace("^", "**"))
            except Exception:
                return None

        pe, ge = _parse(p), _parse(g)
        if pe is not None and ge is not None:
            try:
                if simplify(pe - ge) == 0:
                    return True
            except Exception:
                pass
    except Exception:
        pass

    return False


# ---------------------------------------------------------------------------
# Evaluation loop
# ---------------------------------------------------------------------------
def evaluation(model_args, data_args, training_args):
    if model_args.lora_init:
        task_type = TaskType.CAUSAL_LM
        if any(name in model_args.model_name_or_path.lower() for name in ["llama", "mistral", "falcon", "qwen"]):
            target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"]
        elif any(name in model_args.model_name_or_path.lower() for name in ["phi"]):
            target_modules = ["q_proj", "k_proj", "v_proj", "dense", "fc1", "fc2"]
        elif any(name in model_args.model_name_or_path.lower() for name in ["gpt2"]):
            target_modules = ["c_attn", "c_proj", 'c_fc']
        else:
            raise ValueError(f"Only support LLAMA, Mistral, Falcon, Phi-2, but got {model_args.model_name_or_path}.")
        lora_config = LoraConfig(
            task_type=task_type,
            inference_mode=False,
            r=model_args.lora_r,
            lora_alpha=model_args.lora_alpha,
            lora_dropout=0.1,
            target_modules=target_modules,
            init_lora_weights=True,
        )
    else:
        raise NotImplementedError

    model = CODI_CT(model_args, training_args, lora_config)
    state_dict = load_safetensors_state_dict(model_args.ckpt_dir)

    model.load_state_dict(state_dict, strict=False)
    model.codi.tie_weights()

    tokenizer_path = model_args.model_name_or_path
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        tokenizer_path,
        token=model_args.token,
        model_max_length=training_args.model_max_length,
        padding_side="left",
        use_fast=False,
    )

    if tokenizer.pad_token_id is None:
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
        tokenizer.pad_token_id = model.pad_token_id
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids('[PAD]')

    device = "cuda"
    model = model.to('cuda')
    model.to(torch.bfloat16)

    ######################
    #      dataset       #
    ######################
    logging.warning("Downloading MATH test data")
    # MATH uses (problem, solution, answer). Training (see train.py) wires:
    #   question = problem.strip() + "\n"
    #   cot      = solution.strip() + "\n"
    #   target   = "The answer is: {answer}"
    # so at test time we feed `problem` as the question.
    math_test_path = data_args.data_path if getattr(data_args, "data_path", None) else \
        '/data/oss_bucket_0/liqi/dataset/MATH/test.json'
    test_set = read_json(math_test_path)
    if test_set is None:
        raise FileNotFoundError(f"Could not load MATH test set from {math_test_path}")

    question_name = "problem"
    answer_name = "answer"

    logging.warning("Formatting inputs...")
    question = [f"{example[question_name].strip().replace('  ', ' ')}" for example in test_set]
    answer = [str(example[answer_name]) for example in test_set]  # keep raw LaTeX strings

    logging.warning("Tokenizing inputs...")
    eval_step = math.ceil(len(question) / data_args.batch_size)
    logging.warning(
        f"Total example: {len(question)} | eval batch size: {data_args.batch_size} "
        f"eval steps: {eval_step}"
    )

    question_data = []
    for i in range(eval_step):
        if i < eval_step - 1:
            batch = tokenizer(
                question[i * data_args.batch_size:(i + 1) * data_args.batch_size],
                return_tensors="pt",
                padding="longest",
                truncation=True,
                max_length=761,
            )
        else:
            batch = tokenizer(
                question[i * data_args.batch_size:],
                return_tensors="pt",
                padding="longest",
                truncation=True,
                max_length=761,
            )

        if training_args.remove_eos:
            bot_tensor = torch.tensor([model.bot_id], dtype=torch.long).expand(batch["input_ids"].size(0), 1)
        else:
            bot_tensor = torch.tensor(
                [tokenizer.eos_token_id, model.bot_id], dtype=torch.long
            ).expand(batch["input_ids"].size(0), 2)
        batch["input_ids"] = torch.cat((batch["input_ids"], bot_tensor), dim=1)
        batch["attention_mask"] = torch.cat((batch["attention_mask"], torch.ones_like(bot_tensor)), dim=1)
        batch['input_len'] = len(batch['input_ids'][0])
        question_data.append(batch.to(device))

    model.eval()
    gen_kwargs = {
        "max_new_tokens": 256,
        "temperature": 0.1,
        "top_k": 40,
        "top_p": 0.95,
        "do_sample": True,
    }

    ans_pred_list = []
    set_seed(42)
    len_cot = []
    full_out = []
    model.eval()
    embed_layer = model.get_embd(model.codi, model.model_name)

    for step, batch in enumerate(question_data):
        batch_size = batch["input_ids"].size(0)
        with torch.no_grad():
            past_key_values = None
            outputs = model.codi(
                input_ids=batch["input_ids"],
                use_cache=True,
                output_hidden_states=True,
                past_key_values=past_key_values,
                attention_mask=batch["attention_mask"],
            )
            past_key_values = outputs.past_key_values

            if training_args.enable_top_k:
                logits_last = outputs.logits[:, -1, :]
                latent_embd = model._compute_topk_weighted_embedding(logits_last, embed_layer)
            else:
                latent_embd = outputs.hidden_states[-1][:, -1, :].unsqueeze(1)

            if training_args.use_prj:
                latent_embd = model.prj(latent_embd)

            inf_latent_iterations = training_args.inf_latent_iterations

            for _ in range(inf_latent_iterations):
                outputs = model.codi(
                    inputs_embeds=latent_embd,
                    use_cache=True,
                    output_hidden_states=True,
                    past_key_values=past_key_values,
                )
                past_key_values = outputs.past_key_values

                if training_args.enable_top_k:
                    logits_last = outputs.logits[:, -1, :]
                    latent_embd = model._compute_topk_weighted_embedding(logits_last, embed_layer)
                else:
                    latent_embd = outputs.hidden_states[-1][:, -1, :].unsqueeze(1)

                if training_args.use_prj:
                    latent_embd = model.prj(latent_embd)

            if training_args.remove_eos:
                eot_emb = model.get_embd(model.codi, model.model_name)(
                    torch.tensor([model.eot_id], dtype=torch.long, device='cuda')
                ).unsqueeze(0).to(device)
            else:
                eot_emb = model.get_embd(model.codi, model.model_name)(
                    torch.tensor([model.eot_id, tokenizer.eos_token_id], dtype=torch.long, device='cuda')
                ).unsqueeze(0).to(device)

            eot_emb = eot_emb.expand(batch["input_ids"].size(0), -1, -1)
            output = eot_emb

            seq_len = 0
            finished = torch.zeros(batch_size, dtype=torch.bool, device="cuda")
            pred_tokens = [[] for _ in range(batch_size)]
            for i in range(gen_kwargs["max_new_tokens"]):
                seq_len += 1
                out = model.codi(
                    inputs_embeds=output,
                    output_hidden_states=False,
                    attention_mask=None,
                    use_cache=True,
                    output_attentions=False,
                    past_key_values=past_key_values,
                )
                past_key_values = out.past_key_values
                logits = out.logits[:, -1, :model.codi.config.vocab_size - 1]

                if training_args.greedy:
                    next_token_ids = torch.argmax(logits, dim=-1).squeeze(-1)
                else:
                    logits /= gen_kwargs["temperature"]
                    if gen_kwargs["top_k"] > 1:
                        top_k_values, _ = torch.topk(logits, gen_kwargs["top_k"], dim=-1)
                        min_top_k_value = top_k_values[:, -1].unsqueeze(-1)
                        logits[logits < min_top_k_value] = -float("inf")

                    if gen_kwargs["top_p"] < 1.0:
                        sorted_logit, sorted_indices = torch.sort(logits, descending=True, dim=-1)
                        cumulative_probs = torch.cumsum(F.softmax(sorted_logit, dim=-1), dim=-1)

                        sorted_indices_to_remove = cumulative_probs > gen_kwargs["top_p"]
                        if sorted_indices_to_remove.any():
                            sorted_indices_to_remove = sorted_indices_to_remove.roll(1, dims=-1)
                            sorted_indices_to_remove[:, 0] = False

                        for b in range(logits.size(0)):
                            logits[b, sorted_indices[b, sorted_indices_to_remove[b]]] = -float("inf")

                    probs = F.softmax(logits, dim=-1)
                    next_token_ids = torch.multinomial(probs, num_samples=1).squeeze(-1)

                for b in range(batch_size):
                    if not finished[b]:
                        pred_tokens[b].append(next_token_ids[b].item())
                        if next_token_ids[b] == tokenizer.eos_token_id:
                            finished[b] = True

                if finished.all():
                    break

                output = model.get_embd(model.codi, model.model_name)(next_token_ids).unsqueeze(1).to(device)

            for mini_step, pred_token in enumerate(pred_tokens):
                len_cot.append(len(pred_token))
                decoded_pred = tokenizer.decode(pred_token, skip_special_tokens=True)
                full_out.append(decoded_pred)
                pred_ans = extract_math_answer(decoded_pred)
                gold_ans = answer[step * data_args.batch_size + mini_step]
                if do_print:
                    print(f"Question {step * data_args.batch_size + mini_step} Starts...")
                    # print(f"Q: {question[step * data_args.batch_size + mini_step]}")
                    # print(decoded_pred)
                    # print(f"Question {step * data_args.batch_size + mini_step} Ends")
                    # print(f"Prediction={pred_ans}; Groundtruth={gold_ans}; "
                    #       f"Match={is_math_equiv(pred_ans, gold_ans)}")
                    # print("")
                ans_pred_list.append(pred_ans)

    save_dir = training_args.output_dir
    ckpt_name = model_args.ckpt_dir.split('/')[-1]
    save_dir = f"{save_dir}/{ckpt_name}"
    os.makedirs(save_dir, exist_ok=True)

    save_list = []
    for gt, pred, out in zip(test_set, ans_pred_list, full_out):
        save_list.append({
            "question": gt[question_name],
            "ground truth": gt[answer_name],
            "full output": out,
            "prediction": pred,
            "match": is_math_equiv(pred, str(gt[answer_name])),
        })

    accuracy = compute_accuracy(answer, ans_pred_list)
    write_json(
        {
            "accuracy": 100 * accuracy,
            "average length of COT": sum(len_cot) / len(len_cot),
            "output": save_list,
        },
        f"{save_dir}/answer.json",
    )

    print(f"adapter: {model_args.adapter_name_or_path} | MATH test accuracy: {100 * accuracy:.2f}% | ")
    print(f"average length of COT: {sum(len_cot) / len(len_cot)}")

    del model
    del question_data
    import gc
    gc.collect()
    torch.cuda.empty_cache()

    return 100 * accuracy


def compute_accuracy(gold: list, pred: list) -> float:
    """MATH accuracy via LaTeX-aware equivalence."""
    if len(gold) == 0:
        return 0.0
    correct = 0
    for p, g in zip(pred, gold):
        if is_math_equiv(p, g):
            correct += 1
    return correct / len(gold)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser(description="CODI MATH Testing with YAML config support")
    arg_parser.add_argument("--config_file", type=str, default=None, help="Path to YAML config file")
    args, remaining_args = arg_parser.parse_known_args()

    if args.config_file:
        print(f"Loading configuration from YAML file: {args.config_file}")
        with open(args.config_file, 'r') as f:
            config = yaml.safe_load(f)

        cmd_args = []
        for key, value in config.items():
            if value is None:
                continue
            if isinstance(value, bool):
                cmd_args.extend([f"--{key}", "True" if value else "False"])
            else:
                cmd_args.extend([f"--{key}", str(value)])

        parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
        model_args, data_args, training_args = parser.parse_args_into_dataclasses(args=cmd_args)
    else:
        print("Loading configuration from command line arguments")
        parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # ---------------------------------------------------------------
    # Checkpoint list: leave empty to evaluate the single ckpt_dir
    # supplied via config / command line.
    # ---------------------------------------------------------------
    # ckpt_list = []
    # ckpt_list = ['990', '979', '968', '957', '946', '935', '924', '913', '902', '891', '880', '869', '858', '847', '836']
    # ckpt_list = ['990', '968', '946', '924', '902', '880', '858', '836', '814', '792', '770', '748', '726', '704', '682']
    # ckpt_list = ['990', '968', '946', '924', '902', '880', '858', '836', '814', '792', '770', '748', '726', '704', '682', '660', '638', '616']
    # ckpt_list = ['6300', '6279', '6258', '6237', '6216', '6195', '6174', '6153', '6132', '6111', '6090', '6069', '6048', '6027', '6006', '5985', '5964', '5943']
    # ckpt_list = ['4400', '4378', '4356', '4334', '4312', '4290', '4268', '4246', '4224', '4202', '4180']
    # ckpt_list = ['6300', '6279', '6258', '6237', '6216', '6195', '6174', '6153', '6132', '6111', '6090', '6069', '6048', '6027', '6006']
    # ckpt_list = ['1100', '1089', '1078', '1067', '1056', '1045', '1034', '1023', '1012', '1001', '990', '979', '968', '957', '946']
    # ckpt_list = ['9982', '9936', '9890', '9844', '9798', '9752', '9706', '9660', '9614', '9568']
    # ckpt_list = ['6300', '6258', '6216', '6174', '6132', '6090', '6048', '6006', '5964', '5922']
    ckpt_list = ['6300', '6279', '6258', '6237', '6216', '6195', '6174', '6153', '6132', '6111', '6090', '6069', '6048', '6027', '6006', '5985', '5964', '5943', '5922']
    # ckpt_list = ['6000', '5980', '5960', '5940', '5920', '5900', '5880', '5860', '5840', '5820', '5800', '5780', '5760', '5740', '5720']
    if not ckpt_list:
        accu_list = []
        for _ in range(training_args.inf_num_iterations):
            accu = evaluation(model_args, data_args, training_args)
            accu_list.append(accu)
            torch.cuda.empty_cache()
        avg_accu = sum(accu_list) / len(accu_list)
        print(f"Average accuracy over {training_args.inf_num_iterations} sampling: {avg_accu:.2f}%")
        write_json({"avg_accuracy": avg_accu, "all_runs": accu_list},
                   f"{training_args.output_dir}/math_eval_summary.json")
    else:
        all_ckpt_results = {}
        base_dir = model_args.ckpt_dir
        for ckpt_id in ckpt_list:
            print(f"\n{'=' * 60}")
            print(f"Evaluating checkpoint: {ckpt_id}")
            print(f"{'=' * 60}")
            accu_list = []
            for _ in range(training_args.inf_num_iterations):
                model_args.ckpt_dir = os.path.join(base_dir, f"checkpoint-{ckpt_id}")
                accu = evaluation(model_args, data_args, training_args)
                accu_list.append(accu)
                torch.cuda.empty_cache()
            avg_accu = sum(accu_list) / len(accu_list)
            all_ckpt_results[ckpt_id] = avg_accu
            print(f"Checkpoint {ckpt_id} | Average accuracy over "
                  f"{training_args.inf_num_iterations} sampling: {avg_accu:.2f}%")

        write_json(all_ckpt_results, f"{training_args.output_dir}/all_ckpt_results.json")
        print(f"\n{'=' * 60}")
        print("Summary of all checkpoints:")
        for ckpt_dir, avg_accu in all_ckpt_results.items():
            print(f"  {ckpt_dir}: {avg_accu:.2f}%")