evaluate.py

import torch
import numpy as np
import random
from glob import glob
import argparse
import json
import os
import sys

sys.setrecursionlimit(100000)
sys.path.append(os.path.normpath(os.path.dirname(os.path.realpath(__file__))))
sys.path.append(
    os.path.normpath(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../"))
)
from hashlib import md5
from utils import log
from progressbar import ProgressBar
from agent import Agent
from models.prover import Prover
import pdb

import xlrd

similar_record_file = (
    "/users/zhangliao/Downloads/Documents/research/Master-Thesis/result4.xlsx"
)


def read_similar_record():

    wb = xlrd.open_workbook(filename=similar_record_file)

    # print(wb.sheet_names())
    sheet = wb.sheet_by_name("Sheet3")
    target_goals = sheet.col_values(0)[1::]
    learn_from = sheet.col_values(1)[1::]
    return target_goals, learn_from


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("method", type=str)
    parser.add_argument("eval_id", type=str)
    parser.add_argument("--datapath", type=str, default="../data")
    parser.add_argument("--projs_split", type=str, default="../projs_split.json")
    parser.add_argument(
        "--split", choices=["train", "valid", "test"], type=str, default="test"
    )
    parser.add_argument("--file", type=str)
    parser.add_argument("--proof", type=str)
    parser.add_argument("--filter", type=str)
    parser.add_argument("--path", type=str)
    parser.add_argument("--output_dir", type=str, default="evaluation")
    parser.add_argument("--max_num_tactics", type=int, default=300)
    parser.add_argument("--timeout", type=int, default=600)
    parser.add_argument("--hammer_timeout", type=int, default=100)
    parser.add_argument("--depth_limit", type=int, default=50)
    parser.add_argument(
        "--beam_width", type=int, default=20
    )  # lots of timeout when >200
    parser.add_argument("--num_tactic_candidates", type=int, default=20)
    parser.add_argument(
        "--lens_norm", type=float, default=0.5, help="lengths normalization"
    )
    parser.add_argument("--tac_grammar", type=str, default="tactics.ebnf")
    parser.add_argument("--term_embedding_dim", type=int, default=256)
    parser.add_argument("--size_limit", type=int, default=50)
    parser.add_argument(
        "--embedding_dim",
        type=int,
        default=256,
        help="dimension of the grammar embeddings",
    )
    parser.add_argument(
        "--symbol_dim",
        type=int,
        default=256,
        help="dimension of the terminal/nonterminal symbol embeddings",
    )
    parser.add_argument(
        "--hidden_dim", type=int, default=256, help="dimension of the LSTM controller"
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--eval_similar", type=int)



#    parser.add_argument("--eval_test", type=int)

    opts = parser.parse_args()


    log(opts)
    opts.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if opts.device.type == "cpu":
        log("using CPU", "WARNING")

    torch.manual_seed(opts.seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(opts.seed)
    random.seed(opts.seed)

    if "ours" in opts.method:
        model = Prover(opts)
        log("loading model checkpoint from %s.." % opts.path)
        if opts.device.type == "cpu":
            checkpoint = torch.load(opts.path, map_location="cpu")
        else:
            checkpoint = torch.load(opts.path)
        model.load_state_dict(checkpoint["state_dict"])
        model.to(opts.device)
    else:
        model = None

    agent = Agent(model, None, None, opts)

    if opts.file:
        files = [opts.file]
    else:
        files = []
        projs = json.load(open(opts.projs_split))["projs_" + opts.split]
        for proj in projs:
            files.extend(
                glob(os.path.join(opts.datapath, "%s/**/*.json" % proj), recursive=True)
            )

    if opts.filter:
        files = [
            f
            for f in files
            if md5(f.encode("utf-8")).hexdigest().startswith(opts.filter)
        ]

    print(files)
    results = []
    bar = ProgressBar(max_value=len(files))
    if opts.eval_similar is not None:
        for i, f in enumerate(files):
            print("begin to evaluate Similar")
            target_goals, learn_from = read_similar_record()
            results.extend(agent.evaluate_similar(filename, target_goals, learn_from))
    else:
        for i, f in enumerate(files):
            print("file: ", f)
            # print('cuda memory allocated before file: ', torch.cuda.memory_allocated(opts.device), file=sys.stderr)
            results.extend(agent.evaluate(f, opts.proof))
            bar.update(i)

    oup_dir = os.path.join(opts.output_dir, opts.eval_id)
    if not os.path.exists(oup_dir):
        os.makedirs(oup_dir)
    if opts.filter is None and opts.file is None:
        oup_file = os.path.join(oup_dir, "results.json")
    elif opts.file is None:
        oup_file = os.path.join(oup_dir, "%s.json" % opts.filter)
    elif opts.proof is None:
        oup_file = os.path.join(
            oup_dir,
            "%s.json"
            % os.path.sep.join(opts.file.split(os.path.sep)[2:]).replace(
                os.path.sep, "-"
            ),
        )
    else:
        oup_file = os.path.join(
            oup_dir,
            "%s-%s.json"
            % (
                os.path.sep.join(opts.file.split(os.path.sep)[2:]).replace(
                    os.path.sep, "-"
                ),
                opts.proof,
            ),
        )
    opts = vars(opts)
    del opts["device"]
    json.dump({"options": opts, "results": results}, open(oup_file, "wt"))
    log("results saved to " + oup_file)