import argparse
import json
import logging
from lm_eval import evaluator, tasks

logging.getLogger("openai").setLevel(logging.WARNING)

# python Evaluation.py --model gpt2 --model_args pretrained=gpt2-xl --device cuda:2 --tasks piqa --output_path /mnt/home/joel/evolver/result/piqa.json

# wikitext (ppl)
# lambada (acc) (ppl)
# hellaswag (acc)
# winogrande (acc)

# mathqa (acc)
# piqa (acc)

def evaluate_lm(args, model):
    args.model = model
    args.model_args = "pretrained=" + args.model_name_or_path
    args.num_fewshot = 0
    args.device = "cuda:" + args.CUDA_VISIBLE_DEVICES
    args.limit = None
    args.description_dict_path = None
    args.no_cache = False
    args.check_integrity = False
    
    if args.dataset == "all_tasks":
        task_names = tasks.ALL_TASKS
    else:
        task_names = args.dataset.split(",")
    description_dict = {}
    if args.description_dict_path:
        with open(args.description_dict_path, 'r') as f:
            description_dict = json.load(f)

    results = evaluator.simple_evaluate(
        model=args.model,
        model_args=args.model_args,
        tasks=task_names,
        num_fewshot=args.num_fewshot,
        batch_size=args.eval_batch_size,
        device=args.device,
        no_cache=args.no_cache,
        limit=args.limit,
        description_dict=description_dict,
        check_integrity=args.check_integrity
    )

    dumped = json.dumps(results, indent=2)
    
    print(dumped)

    if args.output_log:
        with open(args.output_log, "w") as f:
            f.write(dumped)

    print(evaluator.make_table(results))
